QK-Clip: Pushing Muon Further Along the Path of Scale-Up
Four months ago, we released Moonlight, validating the effectiveness of the Muon optimizer on a 16B MoE model. In Moonlight, we confirmed the necessity of adding Weight Decay to Muon, and also proposed the trick of transferring Adam hyperparameters via Update RMS alignment, which allowed Muon to be quickly applied to LLM training. However, when we tried to scale Muon further to models with over a hundred billion parameters, we ran into a new obstacle — MaxLogit explosion.
To address this problem, we propose a simple but extremely effective new method, which we call "QK-Clip". This method looks at and resolves the MaxLogit phenomenon from a very fundamental angle, without any loss in model performance, and it has become one of the key training techniques behind our latest trillion-parameter model, "Kimi K2".
Problem Description
Let's first briefly introduce the phenomenon of MaxLogit explosion. Recall the definition of Attention:
\begin{equation}\boldsymbol{O} = softmax(\boldsymbol{Q}\boldsymbol{K}^{\top})\boldsymbol{V}\end{equation}more
Here we've omitted the scaling factor $1/\sqrt{d}$, since it can always be absorbed into the definition of $\boldsymbol{Q},\boldsymbol{K}$. The "Logit" in "MaxLogit explosion" refers to the Attention matrix before the Softmax, i.e. $\boldsymbol{Q}\boldsymbol{K}^{\top}$, and MaxLogit refers to the maximum value across all these Logits, which we denote as
\begin{equation}S_{\max} = \max_{i,j}\, \boldsymbol{q}_i\cdot \boldsymbol{k}_j\end{equation}
Here, $\max$ is actually also taken over the batch_size dimension, ultimately yielding a scalar. MaxLogit explosion refers to the phenomenon where $S_{\max}$ keeps climbing as training progresses, growing linearly or even superlinearly, with no sign of stabilizing over a fairly long period.
MaxLogit explosion phenomenon
MaxLogit is essentially an outlier-value indicator, and its explosion means that outliers have grown beyond a controllable range. Specifically, we have
\begin{equation}|\boldsymbol{q}_i\cdot \boldsymbol{k}_j| \leq \Vert\boldsymbol{q}_i\Vert \Vert\boldsymbol{k}_j\Vert = \Vert\boldsymbol{x}_i\boldsymbol{W}_q\Vert \Vert\boldsymbol{x}_j\boldsymbol{W}_k\Vert \leq \Vert\boldsymbol{x}_i\Vert \Vert\boldsymbol{x}_j\Vert \Vert\boldsymbol{W}_q\Vert \Vert\boldsymbol{W}_k\Vert\label{eq:kexi}\end{equation}
Since $\boldsymbol{x}$ typically has RMSNorm applied to it, $\Vert\boldsymbol{x}_i\Vert \Vert\boldsymbol{x}_j\Vert$ generally won't explode, so MaxLogit explosion implies that the spectral norm $\Vert\boldsymbol{W}_q\Vert,\Vert\boldsymbol{W}_k\Vert$ is at risk of heading toward infinity, which is clearly bad news.
Since even very large values become smaller than 1 after Softmax, in relatively lucky cases this phenomenon won't have overly severe consequences — at worst it just wastes an Attention Head. But in worse cases, it may cause Grad Spikes or even training collapse. Therefore, to be safe, MaxLogit explosion should be avoided as much as possible.
Prior Attempts
In Muon Sequel: Why Did We Choose to Try Muon? we briefly analyzed that Weight Decay can, to some extent, prevent MaxLogit explosion, so the probability of MaxLogit explosion occurring in small models is quite low. Even for a 16B model like Moonlight, MaxLogit rises to at most 120 before automatically coming back down.
Moonlight's MaxLogit automatically comes back down
In other words, MaxLogit explosion appears more often in models with very large parameter counts — the bigger the model, the more instability factors during training, and the harder it becomes for Weight Decay alone to stabilize training. At that point, increasing Weight Decay can indeed strengthen control, but it also brings a noticeable performance cost, so that path is not viable. Another relatively direct idea is to directly add $\text{softcap}$ to the Logit:
\begin{equation}\boldsymbol{O} = softmax(\text{softcap}(\boldsymbol{Q}\boldsymbol{K}^{\top};\tau))\boldsymbol{V}\end{equation}
where $\text{softcap}(x;\tau) = \tau\tanh(x/\tau)$, as introduced by Google's Gemma2. Because of the boundedness of $\tanh$, $\text{softcap}$ can naturally guarantee that the Logit after $\text{softcap}$ is bounded, but it cannot guarantee that the Logit before $\text{softcap}$ is bounded (we've verified this empirically), so $\text{softcap}$ merely converts one problem into another — it doesn't actually solve the problem.
Perhaps Google itself realized this, since in the later Gemma3 they stopped using $\text{softcap}$ and switched to "QK-Norm" instead:
\begin{equation}\boldsymbol{O} = softmax(\tilde{\boldsymbol{Q}}\tilde{\boldsymbol{K}}{}^{\top})\boldsymbol{V},\quad \begin{aligned} \tilde{\boldsymbol{Q}}=&\,\text{RMSNorm}(\boldsymbol{Q}) \\ \tilde{\boldsymbol{K}}=&\,\text{RMSNorm}(\boldsymbol{K}) \end{aligned}\end{equation}
QK-Norm is indeed a very effective method for suppressing MaxLogit. However, it only applies to MHA, GQA, and the like — it doesn't apply to MLA, because QK-Norm requires materializing $\boldsymbol{Q},\boldsymbol{K}$, but for MLA, the $\boldsymbol{Q},\boldsymbol{K}$ during training differs from the one during decoding (as shown in the equation below). During decoding we cannot fully materialize the $\boldsymbol{K}$ from the training phase — in other words, QK-Norm cannot be done during decoding.
$$\require{cancel}\begin{array}{c|c} \text{training/Prefill} & \text{Decoding} \\ \\ \begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}_t^{(h)}\right] \\[10pt] \boldsymbol{o}_t^{(s)} = \frac{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)\boldsymbol{v}_i^{(s)}}{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)} \\[15pt] \boldsymbol{q}_i^{(s)} = \left[\boldsymbol{x}_i\boldsymbol{W}_{qc}^{(s)},\boldsymbol{x}_i\boldsymbol{W}_{qr}^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\right]\in\mathbb{R}^{d_k + d_r}\\ \boldsymbol{k}_i^{(s)} = \left[\boldsymbol{c}_i\boldsymbol{W}_{kc}^{(s)},\boldsymbol{x}_i\boldsymbol{W}_{kr}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\right]\in\mathbb{R}^{d_k + d_r} \\ \boldsymbol{v}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_v},\quad\boldsymbol{c}_i = \boldsymbol{x}_i \boldsymbol{W}_c\in\mathbb{R}^{d_c} \end{gathered} & \begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}\boldsymbol{W}_v^{(1)}, \boldsymbol{o}_t^{(2)}\boldsymbol{W}_v^{(2)}, \cdots, \boldsymbol{o}_t^{(h)}\boldsymbol{W}_v^{(h)}\right] \\[10pt] \boldsymbol{o}_t^{(s)} = \frac{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}}{}^{\top}\right)\boldsymbol{v}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}} }{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}}{}^{\top}\right)} \\[15pt] \boldsymbol{q}_i^{(s)} = \left[\boldsymbol{x}_i\boldsymbol{W}_{qc}^{(s)}\boldsymbol{W}_{kc}^{(s)}{}^{\top}, \boldsymbol{x}_i\boldsymbol{W}_{qr}^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\right]\in\mathbb{R}^{d_c + d_r}\\ \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}} = \left[\boldsymbol{c}_i, \boldsymbol{x}_i\boldsymbol{W}_{kr}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\right]\in\mathbb{R}^{d_c + d_r}\\ \boldsymbol{v}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}} = \boldsymbol{c}_i= \boldsymbol{x}_i \boldsymbol{W}_c\in\mathbb{R}^{d_c} \end{gathered} \\ \end{array} $$
Why use MLA at all? We've already discussed this question in two articles, Transformer Upgrade Path: 21. What's Good About MLA? (Part 1) and Transformer Upgrade Path: 21. What's Good About MLA? (Part 2), so we won't repeat it here. In short, we would like MLA to also have some means, analogous to QK-Norm, that can guarantee suppression of MaxLogit.
Getting Straight to the Target
In the meantime we also tried some indirect approaches, such as separately lowering the learning rate of $\boldsymbol{Q},\boldsymbol{K}$, or separately increasing their Weight Decay, but none of these worked. The closest we came to success was Partial QK-Norm: for MLA, $\boldsymbol{Q},\boldsymbol{K}$ splits into four parts — qr, qc, kr, kc — of which the first three can be materialized during decoding. So we added RMSNorm to these three parts, which did suppress MaxLogit, but the effect on length extrapolation was terrible.
After many failed attempts, we couldn't help but reflect: everything we'd tried so far was really just an "indirect approach" to suppressing MaxLogit. What would be a direct approach that could actually guarantee a solution to MaxLogit explosion? From the inequality $\eqref{eq:kexi}$, it's natural to think of applying singular value clipping to $\boldsymbol{W}_q,\boldsymbol{W}_k$, but this is fundamentally still an indirect approach, and the computational cost of singular value clipping is also not low.
But clearly, post-hoc rescaling of $\boldsymbol{W}_q,\boldsymbol{W}_k$ is in principle feasible; the question is when to rescale and by how much. Finally, one day it dawned on the author: MaxLogit itself is the most direct signal for triggering the rescaling! Specifically, when MaxLogit exceeds the expected threshold $\tau$, we directly multiply $\boldsymbol{Q}\boldsymbol{K}^{\top}$ by $\gamma = \tau / S_{\max}$, so that the new MaxLogit is guaranteed not to exceed $\tau$. The operation of multiplying by $\gamma$ can be absorbed separately into the weights $\boldsymbol{Q}\boldsymbol{K}$, giving us the initial version of QK-Clip:
$$\begin{aligned} &\boldsymbol{W}_t = \text{Optimizer}(\boldsymbol{W}_{t-1}, \boldsymbol{G}_t) \\ &\text{if }S_{\max}^{(l)} > \tau\text{ and }\boldsymbol{W} \in \{\boldsymbol{W}_q^{(l)}, \boldsymbol{W}_k^{(l)}\}: \\ &\qquad\boldsymbol{W}_t \leftarrow \boldsymbol{W}_t \times \sqrt{\tau / S_{\max}^{(l)}} \end{aligned}$$
where $S_{\max}^{(l)}$ is the MaxLogit of the $l$-th layer's Attention, and $\boldsymbol{W}_q^{(l)}, \boldsymbol{W}_k^{(l)}$ is the corresponding weight $\boldsymbol{Q},\boldsymbol{K}$. In other words, after the optimizer update, based on the magnitude of $S_{\max}^{(l)}$ we decide whether to clip the weight $\boldsymbol{Q},\boldsymbol{K}$; the amount of clipping is determined directly by the ratio between $S_{\max}^{(l)}$ and the threshold $\tau$, directly guaranteeing that the clipped matrix no longer causes MaxLogit explosion. And since this operates directly on the weights, it doesn't affect the inference mode, and is naturally compatible with MLA.
Fine-Tuning the Details
The initial version of QK-Clip did successfully suppress MaxLogit in MLA, but after carefully examining the model's "internals," we found that it suffered from an "over-clipping" issue. Fixing this issue gives us the final version of QK-Clip.
As we know, any variant of Attention has multiple Heads. Initially, we monitored only one MaxLogit metric per Attention layer, with the Logits of all Heads lumped together for taking the max — which meant QK-Clip also clipped all Heads together. However, once we started monitoring the MaxLogit of each Head separately, we found that in fact only a small number of Heads per layer exhibit MaxLogit explosion. If all Heads are clipped by the same ratio, then most Heads are being "punished for no reason" — that's what over-clipping means.
Simply put, the QK-Clip operation multiplies by a number less than 1, a number that just happens to exactly offset the growth trend for the Head that is experiencing MaxLogit explosion, but for other Heads it's simply a shrinkage (since they have no growth trend, or only a very weak one). Being multiplied by a number less than 1 for an extended period, for no good reason, easily drives values toward zero — this is what "over-clipping" looks like.
So, to avoid "collateral damage," we should monitor MaxLogit and apply QK-Clip on a per-Head basis. But there's another devilish detail hidden here: the initial version of QK-Clip spread the clipping factor evenly over $\boldsymbol{Q},\boldsymbol{K}$, but MLA's $\boldsymbol{Q},\boldsymbol{K}$ has four parts — qr, qc, kr, kc — of which kr is shared across all Heads. If we clip it too, we'll again run into the "collateral damage" problem. Therefore, for the pair (qr, kr), we should only clip qr.
After the above adjustments, the final version of QK-Clip is
$$\begin{aligned} &\boldsymbol{W}_t = \text{Optimizer}(\boldsymbol{W}_{t-1}, \boldsymbol{G}_t) \\ &\text{if }S_{\max}^{(l,h)} > \tau: \\ &\qquad\text{if }\boldsymbol{W} \in \{\boldsymbol{W}_{qc}^{(l,h)}, \boldsymbol{W}_{kc}^{(l,h)}\}: \\ &\qquad\qquad\boldsymbol{W}_t \leftarrow \boldsymbol{W}_t \times \sqrt{\tau / S_{\max}^{(l,h)}} \\ &\qquad\text{elif }\boldsymbol{W} \in \{\boldsymbol{W}_{qr}^{(l,h)}\}: \\ &\qquad\qquad\boldsymbol{W}_t \leftarrow \boldsymbol{W}_t \times \tau / S_{\max}^{(l,h)} \end{aligned}$$
where the superscript ${}^{(l,h)}$ denotes the $l$-th layer, $h$-th Head.
The Road to Scale-Up
At this point, we've fully covered the operational details of QK-Clip. It directly uses our expected MaxLogit as a signal, making as small a modification as possible to the weights of $\boldsymbol{Q},\boldsymbol{K}$, to achieve the effect of keeping the MaxLogit value within a specified threshold. And because this method directly modifies the weights, it has better compatibility than QK-Norm and can be used with MLA.
In the training of Kimi K2, we set the threshold $\tau$ to 100, with a total training run of about 220k steps. Starting from roughly step 7k, Heads whose MaxLogit exceeded $\tau$ began to appear. For a fairly long stretch after that, Muon Updates and QK-Clip engaged in a "tug of war" — Muon wanting to increase MaxLogit while QK-Clip wanted to decrease it — remaining in a delicate balance the whole time. Interestingly, after 70k steps, the MaxLogit of all Heads spontaneously dropped below 100, and QK-Clip stopped having any effect.
After nearly 70k steps of Muon-QK-Clip tug-of-war, MaxLogit spontaneously drops
This suggests that, under the influence of Weight Decay, as long as we can keep training stable, the model will very likely end up lowering MaxLogit on its own; the role of QK-Clip is precisely to help the model get through the early stage of training more smoothly. Some readers might worry that QK-Clip could hurt performance, but we ran comparative experiments on small models, and even when we used QK-Clip to squeeze MaxLogit down to a very small value (e.g., 30), we observed no substantive difference in performance. Combined with the observation that the model spontaneously lowers MaxLogit in the middle-to-late stages of training, we have good reason to believe that QK-Clip is lossless in terms of performance.
We also observed in our experiments that Muon is generally more prone to MaxLogit explosion than Adam, so in a sense, QK-Clip is an update rule tailored specifically to supplement Muon — it is one of the "secret techniques" that let Muon clear the level of ultra-large-scale training, which is also what the title of this article is getting at. With this in mind, we combined the modifications to Muon we proposed in Moonlight with QK-Clip, and gave it the name "MuonClip" ($\boldsymbol{W}\in\mathbb{R}^{n\times m}$):
$$\text{MuonClip}\quad\left\{\quad\begin{aligned} &\boldsymbol{M}_t = \mu \boldsymbol{M}_{t−1} + \boldsymbol{G}_t \\[8pt] &\boldsymbol{O}_t = \newcommand{msign}{\mathop{\text{msign}}}\msign(\boldsymbol{M}_t) \underbrace{\times \sqrt{\max(n,m)}\times 0.2}_{\text{Match Adam Update RMS}} \\[8pt] &\boldsymbol{W}_t = \boldsymbol{W}_{t−1} − \eta_t (\boldsymbol{O}_t + \lambda \boldsymbol{W}_{t-1}) \\[8pt] &\left.\begin{aligned} &\text{if }S_{\max}^{(l,h)} > \tau: \\ &\qquad\text{if }\boldsymbol{W} \in \{\boldsymbol{W}_{qc}^{(l,h)}, \boldsymbol{W}_{kc}^{(l,h)}\}: \\ &\qquad\qquad\boldsymbol{W}_t \leftarrow \boldsymbol{W}_t \times \sqrt{\tau / S_{\max}^{(l,h)}} \\ &\qquad\text{elif }\boldsymbol{W} \in \{\boldsymbol{W}_{qr}^{(l,h)}\}: \\ &\qquad\qquad\boldsymbol{W}_t \leftarrow \boldsymbol{W}_t \times \tau / S_{\max}^{(l,h)} \end{aligned}\quad\right\} \text{QK-Clip} \end{aligned}\right.$$
Note that "Muon is generally more prone to MaxLogit explosion than Adam" doesn't mean only Muon can experience MaxLogit explosion. We know that DeepSeek-V3 was trained with Adam, and we've also observed MaxLogit explosion in DeepSeek-V3's open-sourced model. There's also Gemma2, which used $\text{softcap}$ to prevent MaxLogit explosion and was also trained with Adam. So, although we've emphasized the value of QK-Clip for Muon, if readers insist on using Adam, it can also be combined with Adam to form AdamClip.
Reflections on the Cause
Why is Muon more prone to causing MaxLogit explosion? In this section the author attempts to offer a theoretical explanation, for readers' reference.
From the inequality $\eqref{eq:kexi}$, we can see that MaxLogit explosion often implies that the spectral norm of $\boldsymbol{W}_q$ or $\boldsymbol{W}_k$ is showing signs of exploding — indeed, the very definition of spectral norm involves taking a $\max$ operation, so the two are fundamentally connected. Thus, the question can be rephrased as "why is Muon more prone to causing spectral norm explosion?" We know that the spectral norm equals the largest singular value, so we can further narrow the question to "why does Muon tend to increase singular values more?"
What's the difference between Muon and Adam? The update given by Muon has gone through an $\msign$ operation, so all its singular values are equal — that is, its effective rank is full rank. A generic matrix, on the other hand, typically has singular values of varying magnitudes, dominated by a handful of leading singular values — from the point of view of effective rank, they are low-rank. We make the same assumption about Adam's update. This assumption isn't new — for example, high-order MuP likewise assumes low-rankness for Adam's updates.
In formula terms, let the SVD of parameter $\boldsymbol{W}_{t-1}$ be $\sum_i \sigma_i \boldsymbol{u}_i \boldsymbol{v}_i^{\top}$, the SVD of Muon's update be $\sum_j \bar{\sigma}\bar{\boldsymbol{u}}_j \bar{\boldsymbol{v}}_j^{\top}$, and the SVD of Adam's update be $\sum_j \tilde{\sigma}_j\tilde{\boldsymbol{u}}_j \tilde{\boldsymbol{v}}_j^{\top}$; then
\begin{gather} \boldsymbol{W}_t = \sum_i \sigma_i \boldsymbol{u}_i \boldsymbol{v}_i^{\top} + \sum_j \bar{\sigma}\bar{\boldsymbol{u}}_j \bar{\boldsymbol{v}}_j^{\top}\qquad (\text{Muon}) \\ \boldsymbol{W}_t = \sum_i \sigma_i \boldsymbol{u}_i \boldsymbol{v}_i^{\top} + \sum_j \tilde{\sigma}_j\tilde{\boldsymbol{u}}_j \tilde{\boldsymbol{v}}_j^{\top}\qquad (\text{Adam}) \\ \end{gather}
Clearly, if a singular vector pair $\boldsymbol{u}_i \boldsymbol{v}_i^{\top}$ happens to be close to some $\bar{\boldsymbol{u}}_j \bar{\boldsymbol{v}}_j^{\top}$ or $\tilde{\boldsymbol{u}}_j \tilde{\boldsymbol{v}}_j^{\top}$, they will add up directly, thereby increasing the singular value of $\boldsymbol{W}_t$. Since Muon's update is full-rank, its "chance of colliding" with $\boldsymbol{W}_{t-1}$ is far greater than Adam's, so Muon is more prone to increasing the singular values of the parameter.
Of course, the above analysis is general, and not limited to the weights of $\boldsymbol{Q},\boldsymbol{K}$. Indeed, in Moonlight we already verified that the singular-value entropy of weights trained with Muon is generally higher, which corroborates this hypothesis. What's special about Attention Logits is that they are a bilinear form, $\boldsymbol{q}_i\cdot \boldsymbol{k}_j = (\boldsymbol{x}_i \boldsymbol{W}_q)\cdot(\boldsymbol{x}_j \boldsymbol{W}_k)$, and the product $\boldsymbol{W}_q,\boldsymbol{W}_k$ makes explosion risk even greater, and can easily lead to a vicious cycle where "things that are bad get worse," ultimately resulting in MaxLogit explosion.
Comparison of singular-value entropy (equivalent to effective rank) of model weights trained with Muon vs. Adam
Lastly, "Muon's chance of collision is far greater than Adam's" is a relative statement — in practice, singular vectors actually colliding is still a low-probability event, which also explains why only a small fraction of Attention Heads exhibit MaxLogit explosion.
This perspective can also explain a phenomenon we noticed previously in Moonlight: when a model pretrained with Muon/Adam is fine-tuned with Adam/Muon instead, the result is usually suboptimal. This is because the weights trained by Muon have higher effective rank, while Adam's updates are low-rank — combining a high-rank weight with low-rank updates hurts fine-tuning efficiency. Conversely, Adam's trained weights have lower effective rank, but Muon's updates are full-rank, giving it a greater chance of intervening in small singular-value components, pushing the model away from its pretrained low-rank local optimum, and thereby also hurting fine-tuning efficiency.
Some Extensions
At this point, we've covered the important computational and experimental details of QK-Clip. We should also mention that, although the idea of QK-Clip is simple, implementing it in distributed training is somewhat tricky because it needs to be applied per-Head, and the parameter matrices at that point are often split into "fragmented pieces" (not too hard to modify on top of Muon, but somewhat more complicated on top of Adam).
For the author and the team, QK-Clip is not merely a specific method for solving the MaxLogit explosion problem — it also represents a moment of "sudden realization" after repeatedly trying, and failing, to solve the problem through indirect means: once you have a clear metric, you should look for an approach that can guarantee solving the problem directly, rather than wasting time on approaches — like lowering the LR, increasing Weight Decay, or partial QK-Norm — that might work, but aren't guaranteed to.
Methodologically speaking, the idea behind QK-Clip is not limited to solving MaxLogit explosion — it could be described as an "antibiotic" for many training-instability problems. What we mean by antibiotic is that it may not be the most elegant solution, but it is often one of the most direct and effective ways to solve a problem. QK-Clip has exactly this property, and can be generalized into the principle of "clip wherever things are unstable."
For example, in some cases a model may exhibit "MaxOutput explosion." In that case, we could consider clipping the weight $\boldsymbol{W}_o$ based on the value of MaxOutput. Analogous to the Per-Head operation in QK-Clip, here we'd also need to consider a Per-Dim operation, but the cost of Per-Dim clipping is clearly too high, so some compromise may be needed. In short, "clip wherever things are unstable" offers a unified approach to solving the problem, but the specific details are left for everyone to work out for themselves.
Finally, this style of manually devising update rules based on certain signals — as in QK-Clip — was to some extent inspired by DeepSeek's Loss-Free load balancing strategy. Here again, our respects to DeepSeek!
Summary
This article presents QK-Clip, a new approach to the MaxLogit explosion problem. Unlike QK-Norm, it is a post-hoc adjustment scheme applied to the Q and K weights, and it does not change the model's forward computation, giving it broader applicability. It is an important stabilization strategy for the "Muon + MLA" combination in ultra-large-scale training, and is also one of the key techniques behind our latest trillion-parameter model, Kimi K2.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.