From Muon to Gradient Clipping: Some Thoughts on QK Stability

Prelude: an elegant theory and a stubborn problem

In the toolbox of deep learning, the optimizer is the engine. While most of us have become used to tuning Adam and its variants, another optimizer built on a rather different philosophy has gradually come to researchers’ attention: the Muon optimizer. Keller Jordan and Jeremy Bernstein first proposed it in their work [1]. Its depth lies in the fact that it no longer stays within the familiar parameter space. Instead, it starts from the perspective of function space and builds its update rule from there.

This idea has recently been advanced and successfully put into practice in work such as Kimi K2 [7], showing substantial promise. Yet the public material around K2 also points out a difficult issue: if the original Muon optimizer is applied directly to the Query (Q) and Key (K) weight matrices of a Transformer, training can become extremely unstable, and can even collapse.

The elegance of the theory runs into a wall in practice. That naturally raises two questions:

What is the root cause behind this phenomenon?

Can we follow Muon’s first principles and propose a modified update idea suitable for QK?

This post records my attempt to think through that path. It is a winding journey: from strict theoretical derivation, to engineering bottlenecks, and finally to feasible heuristic schemes. It is less a polished answer than a record of thinking across theory, engineering, and intuition.

1. Theory review: why standard Muon does not fit QK

To understand this mismatch, we first need to understand Muon’s original intention and how it differs fundamentally from optimizers such as Adam.

The derivation and essence of Muon

Everything starts from the most basic objective in optimization. For a weight matrix $W$, suppose we have computed its gradient $G$ with respect to loss $L$. We want to find an update $\Delta W$ that reduces the loss as much as possible. Mathematically, the first-order Taylor expansion says:

To make $L$ decrease as fast as possible, we want $\Delta W$ to point as much as possible in the opposite direction of $G$. This can be formalized as a constrained optimization problem:

Clearly, without any constraint, the entries of $\Delta W$ could go to infinity, which is meaningless. Every optimizer must impose some constraint that limits the “size” of the update step. The difference lies in how this size is defined.

The traditional constraint, in parameter space, is what we know from SGD, Adam, and similar optimizers. It is imposed directly on $\Delta W$ itself. The most common choice is the Frobenius norm:

Its geometric meaning is intuitive: the Euclidean distance moved in parameter space should not be too large.

Muon’s constraint, in function space, is entirely different. Muon says that we should not care primarily about how far the parameter $W$ moves. We should care about how much this movement changes the model’s function. For a linear layer, the function maps an input $x$ to an output $Wx$. Therefore, the functional effect of a parameter update is the output change $\Delta W x$.

The core idea of Muon is that this functional change should be bounded, and that the bound should hold for every possible input. To remove the influence of the magnitude of $x$, we consider only unit-norm inputs. Thus Muon’s constraint becomes:

This says: for any input vector of length 1, after transformation by the update matrix $\Delta W$, the length of the output vector should not exceed $\eta$. Anyone who has studied matrix theory will recognize that this supremum is precisely the definition of the spectral norm $|\Delta W|_2$, which equals the largest singular value of the matrix.

Now Muon’s optimization problem is clear:

This optimization problem under a spectral norm constraint has a beautiful analytic solution: the characteristic matrix sign function,

The $\operatorname{msign}$ function preserves the singular vectors of $G$, but changes all of its singular values to 1. This is an extreme form of directional control. It distributes the update “energy” evenly across all singular directions.

The geometric conflict in QK updates

With Muon’s mechanism in mind, we can revisit its failure on QK updates. As Su Jianlin analyzes in his blog post [2], the problem lies in the bilinear nature of attention.

A first, relatively shallow intuition is functional coupling. The functions of Q and K appear as a pair. They only produce meaningful attention scores through the product $QK^T$. Independently constraining the functional changes of $W_Q$ and $W_K$ does not guarantee that the functional change of their coupled product is controlled.

The deeper geometric explanation, following Su Jianlin’s conjecture [2], is related to Muon’s full-rank update behavior. The observed “MaxLogit explosion” in practice often means that the spectral norm of $W_Q$ or $W_K$ is showing signs of explosion. So the question becomes: why is Muon more likely to cause spectral norm explosion? Since the spectral norm is the largest singular value, this becomes: why does Muon tend to increase singular values?

What is the difference between Muon and Adam? Muon’s update is processed by $\operatorname{msign}$, so all singular values of the update are equal; in this sense its effective rank is full. A generic matrix, by contrast, usually has singular values of different sizes and is dominated by the first few. From the perspective of effective rank, it is low-rank. We may make a similar low-rank assumption about Adam updates. This assumption is not new; higher-order muP makes a similar assumption about the low-rank nature of Adam updates.

In formula form, let the SVD of the parameter $W$ be $U \Sigma V^T$, the SVD of the Muon update be $U_m \Sigma_m V_m^T$, and the SVD of the Adam update be $U_a \Sigma_a V_a^T$. The updated weight can be written schematically as:

If a singular-vector pair of $W$ is close to a singular-vector pair of the update, then the corresponding singular values tend to add directly, increasing the singular value of $W$ in that direction.

Since Muon’s update is full-rank, its energy is spread evenly over all singular directions. Therefore, it has a much higher chance of “colliding” with the singular vectors of $W$ than a low-rank Adam update. Muon is naturally more likely to increase the singular values of the parameters.

This effect is amplified in attention. The core of attention scores is the bilinear form $QK^T$. If Muon is simultaneously increasing the spectral norms of $W_Q$ and $W_K$, the increase is multiplied when the final dot product is computed. The risk of explosion rises sharply. A “worse gets worse” feedback loop forms and can eventually collapse training.

At this point, we have a theoretical explanation for why such an elegantly designed optimizer becomes unsuitable for Transformer QK weights. The root issue is that we must constrain the change of the product term, not the two independent changes. This insight is the starting point for the next step.

2. A principled attempt: customizing Muon for QK

The first section gave us the key insight: simply constraining updates to $W_Q$ and $W_K$ separately does not work. In attention, the two are functionally coupled. What ultimately matters is their product. Independent constraints may allow the norms of the two vectors to change little, while their directions become aligned, causing their dot product, that is, the attention score, to explode.

Therefore, a principled modification that truly follows Muon’s philosophy must start from the root problem: directly constrain the change in the final functional output, namely the attention score matrix. This is the starting point of the theoretical attempt.

2.1 A new optimization objective: from constraining parameters to constraining function

Our goal is no longer to constrain the weight matrices $W_Q$ and $W_K$ themselves. Instead, we constrain the change in the attention score matrix $S$ induced by changes in them.

Let us define the objects involved:

• Input matrix $X$: an $n \times d$ matrix, where $n$ is the sequence length and $d$ is the embedding dimension. Each row is the vector representation of one token.
• Weight matrices $W_Q$ and $W_K$: two $d \times h$ matrices, where $h$ is the head dimension. They map the $d$-dimensional input into query and key spaces.
• Query and key matrices: $Q = XW_Q$ and $K = XW_K$, both $n \times h$.
• Attention score matrix: $S = QK^T = XW_QW_K^TX^T$. This is an $n \times n$ matrix. Its $(i,j)$ entry is the similarity between the $i$-th query and the $j$-th key.

After updating the weights, we want the new attention score matrix $S’$ not to differ too much from the old $S$. The change $\Delta S$ can be written as:

Expanding the terms gives:

Because $\Delta W_Q$ and $\Delta W_K$ are small updates, the product $\Delta W_Q \Delta W_K^T$ is a second-order infinitesimal term and can be ignored. This linearization is the basis of almost all optimization theory. We obtain the first-order approximation:

Now we constrain the size of this change. In Muon, the most natural norm is the spectral norm, denoted $|\cdot|_2$. It measures the maximum factor by which a matrix can stretch a vector. Applying a spectral norm constraint means that we wish to control the worst-case change in attention scores.

The new constraint is:

where $\epsilon$ is a small positive number representing the budget of functional change we allow.

To simplify the problem and make it independent of a specific input $X$, Muon makes a core assumption: the constraint should hold for any possible input. We therefore consider the worst case, where $X$ has the largest possible amplification capacity. Using the submultiplicativity of matrix norms,

Assuming the input is normalized, $|X|_2 \le 1$, the constraint reduces to a direct constraint on the intermediate weight expression:

This is the Muon-style constraint tailored to QK. It captures the coupling effect between $W_Q$ and $W_K$.

2.2 The complete optimization problem

The final goal is to make the loss $L$ decrease as fast as possible while satisfying the above constraint. By the principle of gradient descent, the first-order change in the loss is approximately:

where $G_Q$ and $G_K$ are the gradients of the loss with respect to $W_Q$ and $W_K$. The trace is the matrix-inner-product generalization. To make $\Delta L$ as negative as possible, we minimize this expression.

Putting everything together, the optimization problem becomes:

subject to

2.3 The rough road to a solution: decoupling and the pseudoinverse

This optimization problem looks elegant, but it is difficult to solve. The main difficulty is that the constraint tightly couples $\Delta W_Q$ and $\Delta W_K$, so we cannot solve for them independently as in standard gradient descent.

The first step is decoupling. We use a common and effective strategy: relax the constraint with the triangle inequality. We impose a stronger condition than the original constraint; if the stronger one holds, the original one must hold:

The benefit is that we can allocate the update budget $\epsilon$ between the two parts, decomposing a complicated coupled problem into two independent subproblems. A natural and fair allocation is to give each half:

Subproblem 1, for $\Delta W_Q$:

Subproblem 2, for $\Delta W_K$:

The two subproblems have the same structure, so we focus on the first.

The second step is variable substitution. The constraint $|\Delta W_Q W_K^T|_2$ is still troublesome, because $\Delta W_Q$ is multiplied by $W_K^T$. The key trick is to define a new variable:

Then the constraint becomes clean:

Now we need to express the old variable $\Delta W_Q$ in terms of the new variable $Z$. Solving from $Z = \Delta W_Q W_K^T$ requires the Moore-Penrose pseudoinverse. For a matrix $A$, its pseudoinverse is denoted $A^+$. It generalizes the ordinary inverse and exists even when $A$ is non-square or non-invertible.

Using the pseudoinverse, we obtain schematically:

Substituting this into the objective of subproblem 1 gives:

Using the cyclic property of trace, $\operatorname{Tr}(ABC)=\operatorname{Tr}(BCA)$, we can put the problem into the standard Muon form:

This is exactly the standard form of Muon. The solution is known. When minimizing a trace inner product under a spectral norm constraint, the optimal direction is the negative matrix sign of the gradient-like term.

Applying this result, we obtain:

The third step is restoring the solution. We have solved for the intermediate variable $Z$ and now substitute it back to obtain the desired $\Delta W_Q$:

Ignoring the constant factor for compactness, the proportional relation is:

Through a step-by-step derivation, we obtain a custom Muon update rule. In theory, this formula is beautiful. It answers the question we began with.

But when we leave the mathematical idealization and return to engineering reality, an enormous obstacle stands in the way: the pseudoinverse.

The formula contains two pseudoinverse operations. Computing a pseudoinverse usually requires singular value decomposition (SVD), which is expensive. For a matrix of size roughly $d \times h$, the complexity is costly enough that performing it at every training step is unacceptable.

This is the first major setback. We have derived a theoretically impeccable tool, but it is almost unusable in practice because of its computational complexity. The result is frustrating, but it clearly marks the gap between theory and practice. It also sets up the next question: can we bypass or simplify this expensive pseudoinverse computation?

Returning to the original road: separating $W_Q$ and $W_K$

During the derivation, we used the triangle inequality to decouple the coupled constraint. Even after that, the subproblem constraint $|\Delta W_Q W_K^T|_2$ still required the pseudoinverse. A natural idea is to take one further approximation and completely eliminate the pseudoinverse.

Using submultiplicativity again, we have:

Thus subproblem 1 can be relaxed into:

This form is familiar. It has degraded all the way back into the standard Muon optimizer, except that the learning rate, or upper bound of the spectral norm constraint, is now determined by the dynamic term $|W_K^T|_2$. Its solution is simply $\operatorname{msign}(G_Q)$ with a dynamic scale.

But the price of this apparently clever simplification is huge. It essentially gives up modeling the relationship between $W_Q$ and $W_K$. The reason we worked hard to derive the pseudoinverse formula was precisely to capture their geometry. The subtlety of that solution is that it can sense the degree of alignment between the update direction $\Delta W_Q$ and the current weight $W_K$. If the direction of $\Delta W_Q$ is roughly orthogonal to the row space of $W_K$, then its product has little effect on attention scores and the optimizer can allow a larger update. Conversely, if their directions align, then a tiny parameter movement can induce a large attention-score change, so the optimizer should be more conservative. This delicate geometric regulation is carried by the pseudoinverse term.

The new approximation discards all of that geometric information. It only cares about the size of $W_K$ itself and ignores the direction of the update. My guess is that this approximation gap may be too large. It simplifies computation, but it also throws away the essence of the problem, so it is unlikely to work well.

Another fork: bypassing SVD with iteration?

When the cost of SVD seems hopeless, a natural question appears: must we use SVD to compute the pseudoinverse? The answer is no. Linear algebra has iterative methods based purely on matrix multiplication and addition that can also converge to a pseudoinverse. Such methods are friendly to parallel devices such as GPUs.

One typical example is the Newton-Schulz iteration. To compute the pseudoinverse $A^+$ of a matrix $A$, we can start from a suitable initial value $X_0$ and iterate:

A good initial guess is often a scaled transpose, for example $X_0 = \alpha A^T$, where $\alpha$ is a scaling constant that may be approximated using the reciprocal of a spectral norm.

This method looks attractive. It replaces the seemingly uncomputable SVD with a sequence of matrix multiplications, which deep learning frameworks do very well. It appears to give us a shortcut around the computational bottleneck.

But this shortcut leads to a new problem: numerical stability. Such iterative methods are sensitive to numerical precision. In deep learning training we often use low-precision floating-point formats such as bfloat16. At low precision, the iteration may have difficulty converging to a sufficiently accurate solution, or may not converge at all.

Does this mean the path is completely blocked? Perhaps not. Some engineering tricks might help.

One idea is caching. During training, $W_K$ does not change violently at every step. Therefore its pseudoinverse may also be relatively stable. We may not need to recompute it at every step. Instead, we could compute it every, say, 10 steps and reuse the cached result, amortizing the expensive computation.

Another subtler idea is warm start. We could use the pseudoinverse computed at the previous step as the initial guess for the current step’s iteration. Since $W_K$ changes little between adjacent steps, the previous solution should be an excellent approximation to the current solution, significantly reducing the number of iterations needed.

Interestingly, this warm-start strategy is usually ineffective for computing $\operatorname{msign}$. The result of $\operatorname{msign}$ depends only on the singular vectors of the input matrix, that is, its directions, and not on the sizes of its singular values. Under a small perturbation, singular vectors can change sharply, especially when singular values are close. The previous step’s direction may therefore be very different from the current step’s direction. The pseudoinverse, however, depends on both singular vectors and singular values. My guess is that this dependence on singular values may make the pseudoinverse function itself smoother, so that the previous step’s solution remains a good initial point for the current step.

These engineering techniques, especially warm-start iteration, provide concrete and potentially feasible directions for reducing the computational bottleneck. But that is a further investigation. For now, we stop here. We seem to have reached the end of this path of principled correction.

It may be time to give up this obsession and look at the problem from a different angle.

In the second section, after many twists, we derived a theoretically clean but almost impractical update formula. It is like a sharp sword locked behind glass: we can see its edge, but cannot use it.

Just before giving up, we made one last attempt. Is there really no way to simplify the complicated formula? Perhaps some hidden algebraic structure lies inside it.

This brings us to the third act: an unexpected turn, followed by an even deeper disappointment.

3. A turn for the better? An unfinished mathematical simplification

We return to the core part of the intimidating update formula, namely the argument inside $\operatorname{msign}$:

At first glance, this argument looks like a pile of unrelated matrices tangled together. But then we notice that the gradient $G_Q$ is not free-floating. Its own structure is determined by the chain rule of backpropagation, and it is deeply related to $W_K$.

3.1 The unusual structure of $G_Q$

Consider how $G_Q$ is computed. The gradient propagates backward through the computation graph:

By the chain rule, the gradient can be decomposed. Let $E = \partial L / \partial S$, the core gradient arriving from the attention-score layer. Matrix differential rules give:

and therefore:

This reveals the internal structure of $G_Q$:

For simplicity, I state the conclusion directly here. The derivation involves matrix-trace differentials and the chain rule, but the core point is that the gradient of $W_Q$ must pass through the $K$ branch, so the final expression necessarily contains $W_K$.

This is the first sign of hope. The gradient $G_Q$ is not an independent term. It naturally contains the very $W_K$ that we were trying to escape. Could it interact with the pseudoinverse term in the formula and produce some kind of cancellation?

3.2 Internal cancellation and the birth of a projection matrix

With some excitement, substitute the expression for $G_Q$ into the argument of $\operatorname{msign}$. Define a few components:

• Core gradient term: $C = X^T E X$. This $d \times d$ matrix can be viewed as the core gradient information propagated back from the attention-score layer.
• Gradient expression: $G_Q = CW_K$.
• Transposed pseudoinverse: we need $((W_K^T)^+)^T$. If $W_K$ is full column rank, then a standard expression for the pseudoinverse gives a form involving $(W_K^T W_K)^{-1}$.

Putting these pieces together, the argument can be reorganized. The original exploration contained a false start that tried to simplify into a form involving $(W_K W_K^T)^{-1}$, but that matrix is singular when $d > h$. The corrected derivation points to the standard projection matrix built from the invertible $h \times h$ matrix $W_K^T W_K$.

The expression can be arranged into the elegant geometric form:

The part on the right,

is a famous object in linear algebra: a projection matrix. It projects a vector onto the subspace spanned by the columns of $W_K$.

Thus the once-terrifying argument of $\operatorname{msign}$ becomes:

This is a very satisfying simplification. It says that the update direction for $W_Q$ is essentially determined by first projecting the core gradient $C$ into the subspace related to $W_K$, and then taking its matrix sign. This matches our intuition: the update to $W_Q$ should proceed only along directions that can effectively interact with $W_K$.

3.3 Hope, but not quite: the fatal external pseudoinverse

It feels as though we are close to victory. The formula has become concise and geometric. We cannot help asking: have we eliminated the pseudoinverse?

The answer is cruel: no.

First, the projection matrix $P_K$ itself contains the inverse $(W_K^T W_K)^{-1}$. We have not eliminated the inverse; we have hidden it inside a geometrically meaningful object.

But the more fatal blow comes afterward. Return to the final update formula:

The simplification above only concerns the argument inside $\operatorname{msign}$. Outside it, the lone pseudoinverse term $(W_K^T)^+$ still sits there, waiting to be multiplied by the result of $\operatorname{msign}$.

This means the full computation is still:

1. Compute the core gradient $C$.
2. Compute the inverse $(W_K^T W_K)^{-1}$.
3. Use that inverse to construct the projection matrix $P_K$.
4. Compute $\operatorname{msign}(C P_K)$, which may itself require an iterative computation.
5. Use the pseudoinverse $(W_K^T)^+$ again.
6. Multiply the results to obtain the final update.

We have not solved the problem. We have merely made it more elaborate.

Interim conclusion

This attempted simplification was an interesting intellectual exercise. It greatly deepened the understanding of the geometric nature of QK updates: the update happens in a projected subspace.

But from the perspective of engineering practice, it offers no real help. We have not removed the dependence on expensive inverse or pseudoinverse operations. We may even need to compute them twice. This theoretical side path, which looked promising, ultimately led to a thicker wall. Once again, we seem to have reached the end of the principled-correction route.

4. Bilinear attention

After the difficult derivation in section 2 and the “false” simplification in section 3, the exploration seems to have entered a dead end. The theoretically perfect formula is not practical to implement. This dilemma forces us to take a step back and revisit the essence of the problem. A simple but fundamental question appears: why must we view $W_Q$ and $W_K$ separately?

4.1 A thought experiment: one unified bilinear matrix

Let us run a thought experiment. The core computation of attention scores is $qk^T$, where $q = xW_Q$ and $k = yW_K$. Substituting these in, the interaction between two input vectors $x$ and $y$ can be written as:

Notice the term in the middle, $W_QW_K^T$. This is a large $d \times d$ matrix. Define it as a unified attention interaction matrix:

With this definition, the whole attention-score computation becomes a classical bilinear form:

If we replaced the Q and K projection layers in the model architecture with such a unified matrix $B$, then for an entire input sequence $X$, the attention matrix would become:

Now apply Muon’s philosophy to this new architecture. What needs to be constrained is the functional change. Here the function is the bilinear transformation defined by $B$. The functional update induced by $\Delta B$ affects the output as $X\Delta B X^T$. By Muon’s first principle, we constrain the worst-case functional change:

This simplifies to a spectral norm constraint on the weight update itself:

Then the optimization problem becomes transparent:

where $G_B$ is the gradient of the loss with respect to the unified matrix. The solution is exactly the standard Muon update:

This path is theoretically perfect. It solves the QK coupling problem at the root, reducing a complex two-factor constraint into a single-matrix constraint and producing a simple analytic solution. All the complications from sections 2 and 3, such as pseudoinverses and decoupling, disappear.

4.2 The real cost: exploding parameters and computation

Theoretical perfection usually comes with a real cost. Why does the standard Transformer not use this design? The answer is efficiency.

The unified interaction matrix $B$ has shape $d \times d$, where $d$ is the embedding dimension. In a standard attention head, $W_Q$ and $W_K$ have shape $d \times h$, where $h$ is the head dimension.

The parameter counts are:

• Standard attention head: $W_Q$ plus $W_K$, or $2dh$ parameters.
• Unified bilinear attention: $B$, or $d^2$ parameters.

In a typical LLM, $d$ is much larger than $h$. For example, if $d = 4096$ and $h = 128$, the standard head has:

parameters, while the unified matrix has:

parameters.

That is a 16-fold difference for one attention head alone. For a multi-head, multi-layer Transformer, this parameter explosion is unacceptable. In addition, computing $XBX^T$ is much more expensive than computing $QK^T$.

Therefore, the theoretically perfect Full Bilinear Attention scheme is impractical because of its enormous parameter and compute cost.

4.3 A new view: MHA as low-rank approximation

This apparently failed thought experiment gives us a deeper way to understand multi-head attention.

Look again at the definition $B = W_Q W_K^T$. This is a $d \times d$ matrix obtained by multiplying a $d \times h$ matrix and an $h \times d$ matrix. By linear algebra, the rank of $B$ cannot exceed $h$.

Since in practice the head dimension $h$ is always much smaller than the embedding dimension $d$, the interaction matrix $B$ inside standard attention is naturally low-rank.

This gives us a key conclusion: standard multi-head attention can be viewed as a low-rank approximation of a full-size, full-rank bilinear attention mechanism.

The model architecture enforces a structural low-rank constraint by decomposing $B$ into $W_QW_K^T$. This greatly reduces parameter count and computational complexity. The difficulty in all the earlier derivations comes from the fact that we tried to optimize the two factors, $W_Q$ and $W_K$, while preserving this low-rank factorization structure.

4.4 A corollary: a bridge toward Muon-LoRA

This low-rank approximation view also gives an unexpected connection to another popular area: parameter-efficient fine-tuning (PEFT).

Consider the core idea of LoRA, Low-Rank Adaptation. When fine-tuning a large pretrained weight matrix $W_0$, LoRA does not directly update $W_0$. Instead, it learns a low-rank update $\Delta W$ to adapt to a new task:

The low-rank update is also decomposed into a product of two small matrices:

where $A$ is roughly $d \times r$, $B$ is roughly $r \times k$, and the rank $r$ is usually small.

Compare the two structures:

• Attention head: the interaction matrix is decomposed as $W_QW_K^T$.
• LoRA: the weight update is decomposed as $AB$.

Their mathematical essence is the same: use the product of two thin matrices to represent a low-rank matrix.

This leads to a natural corollary. If we can derive a principled and computationally efficient Muon update rule for a low-rank factorization such as $W_QW_K^T$, then the same methodology can be transferred directly to design a new LoRA fine-tuning optimizer based on Muon’s function-space idea. We might call it Muon-LoRA.

This discovery is an unexpected byproduct. It raises the value of solving the QK update problem. It is no longer merely a stability trick for training large models. It may open a new function-space optimization path for parameter-efficient fine-tuning.

4.5 Revisiting the analytic solution

The low-rank approximation view gives us an excellent prism through which to revisit the complicated analytic solution derived in sections 2 and 3, and to understand why it is so heavy.

Our original target was to find an optimal solution satisfying:

Now we can see more clearly that the analytic solution paid a huge price to solve the problem while preserving the $W_QW_K^T$ factorized form. Its essence can be summarized as a cumbersome multistep process.

First, it operates in a high-dimensional space. The core step of the analytic solution, the matrix sign function, is effectively performed in a conceptual $d \times d$ matrix space. This corresponds to the expensive Full Bilinear Attention matrix discussed above. It attempts to apply a Muon update to this large matrix conceptually.

Second, it uses a heuristic budget allocation. To decouple the constraint, we used the triangle inequality and split the functional-update budget $\epsilon$ evenly between $W_Q$ and $W_K$. This is a heuristic rule rather than an optimal decoupling.

Third, it requires complicated projection operations. Most importantly, to split the update to that conceptual $d \times d$ matrix back into updates for the two low-rank factors, the formula introduces complex projection machinery.

It first uses the pseudoinverse to construct a projection matrix, restricting gradient information to the low-rank subspace associated with the other factor. Then, after computing the update direction, it uses the pseudoinverse again to project the update back into the original $W_Q$ or $W_K$ space according to its effect on the other factor.

In short, the analytic solution is complicated because it uses general mathematical tools, such as pseudoinverses and projections, to force a solution for an optimization problem constrained by low-rank structure. It not only contains expensive pseudoinverses computationally; at a deeper conceptual level, it takes a long detour. It lifts the problem into a high-dimensional $d \times d$ space, performs a Muon update there, and then pushes it back into the low-rank space through complex and costly projections. This is neither direct nor efficient.

This suggests that a more pragmatic direction may be to search for a simpler approximate solution directly in the low-rank space.

Interlude

I cannot guarantee that the first four sections are completely correct, but I think the overall direction is not wrong. I do not think it is embarrassing to work these ideas out over a weekend, and I do not think it is embarrassing to write them down.

In the remaining part, I try to propose some of my own ideas. They are of course immature and may not be entirely correct, but they at least explore a class of approaches that seem more feasible from an engineering perspective. I do not think failed attempts are worthless.

5. Returning to the starting point: more direct solutions

The theoretical exploration above is mathematically principled, but it ultimately leads to an analytic solution involving complex operations such as pseudoinverses. The computational cost is high, making it nearly infeasible in practice. This forces us to ask: must we pursue a theoretically perfect and universal solution?

From theoretical guarantee to empirical approximation

The core of Muon’s idea is to seek a worst-case theoretical guarantee that holds for all possible inputs. This leads us to constrain the spectral norm of the weight update. But this universal guarantee comes at a tremendous computational price.

Deep learning practice suggests another path. During training, the direct cause of instability is often that the current data batch produces extreme values in the attention scores, which then create large gradients.

This suggests a change in perspective: perhaps the “worst case” within the current data batch can be used as an empirical approximation to the theoretical worst case. This compromise gives up theoretical completeness, but may gain engineering feasibility. We no longer seek a universal but complicated update rule for the weight matrices. Instead, we intervene directly at key points in the optimization process.

The intervention design space

Once we abandon the complicated analytic solution, we can intervene in several key quantities in the optimization process from a more direct engineering perspective. These interventions share the same goal: smoothing QK weight updates without adding too much computational overhead.

One scheme I would like to try is to intervene directly in the gradient of the attention scores.

This opens a broad design space:

• Backward-pass interventions. These methods modify gradients after they are computed but before the optimizer updates weights.
• Clipping-based methods. The most direct example is to impose fixed upper and lower bounds on the entries of the attention-score gradient.
• Gradient normalization. Use the gradient spectral norm estimated from the max logit gradient to rescale the gradient and control its overall magnitude.
• Parameter-correction interventions. These methods do not directly modify gradients. Instead, after the optimizer updates weights, they apply a post-hoc correction to the parameters based on the forward-pass results.

Within this design space, we can study backward interventions such as gradient clipping, and also examine parameter-correction ideas.

A preliminary attempt: a gradient-clipping heuristic

As a first step, consider a gradient clipping scheme among backward-pass interventions. This is not a mature conclusion. It is more a starting example showing the general shape of this class of heuristics.

The rough procedure is:

1. Run the forward and backward passes normally and compute the original attention-score gradient $\partial L / \partial S$.
2. Clip the gradient matrix elementwise to obtain a clipped gradient, for example $G_S^{\text{clip}} = \operatorname{clip}(G_S, -c, c)$, where $c$ is a tunable hyperparameter.
3. Use the standard optimizer, such as Adam, with these modified gradients to update $W_Q$ and $W_K$.

Gradient clipping vs. MuonClip

In the technical report for Kimi K2, the Kimi AI team also proposed a very successful solution called MuonClip. It is an excellent example of a parameter-correction intervention. Comparing the two side by side reveals clever choices under different engineering philosophies.

Kimi’s MuonClip can be summarized from the public technical details as follows.

Empirical probing: after each optimization step, use the current batch to compute the maximum attention logit produced by the updated $Q$ and $K$. Denote it by $M$.

Dynamic rescaling: set a predetermined logit threshold $\tau$. If $M$ exceeds $\tau$, compute a scaling factor, for example $s = \tau/M$ when $M > \tau$.

Parameter correction: apply this scaling factor back to the weight matrices as a post-hoc correction. Schematically,

where $\alpha$ is a balancing hyperparameter, with $\alpha = 1/2$ as a natural special case.

This correction has two effects.

First, it normalizes the weights, or implicitly adjusts the learning rate. Scaling $W_Q$ and $W_K$ normalizes the parameter norm and also implicitly affects the effective learning rate in subsequent steps. In that sense, it is a kind of normalization inside the intervention process.

Second, it adjusts the temperature of the forward pass. Unlike the gradient-clipping idea, MuonClip also intervenes in the forward computation. The attention scores computed in the next forward pass are globally scaled by $s$. This is equivalent to changing the softmax temperature from 1 to $1/s$. When the attention logit is too large, the temperature rises, the output distribution becomes smoother, and extreme values are suppressed.

A collision of philosophies

The common point is an empirical turn. Both approaches notice that trying to find a universal update rule for all theoretical inputs is too expensive. Both therefore turn to current-batch information as a proxy for the real world, approximating the theoretical worst case.

The core difference lies in scope and mechanism.

In scope, gradient clipping is purely a backward-path intervention. It occurs before gradients are passed to the optimizer. It affects only the current update and does not change the forward computation. MuonClip intervenes in both backward and forward dynamics in effect: it modifies weights after the optimizer update, with the central goal of changing the next forward computation, namely adjusting temperature.

In mechanism, gradient clipping is clipping. It cuts off extreme gradient values rather bluntly. MuonClip is scaling. It stabilizes training by adjusting attention-logit temperature and weight norms.

There is also a process-versus-global distinction. Gradient clipping limits only a local change quantity. MuonClip directly limits a global observable.

In short, gradient clipping and MuonClip share a family resemblance, but they represent two different design choices within a broad exploration space.

How accurate is max logit as a spectral-norm estimate?

Using the current batch’s max logit to approximate the theoretical spectral norm naturally raises a concern: how accurate is this empirical estimate? After all, we are using finite samples to estimate a maximum over an infinite input space. In theory, this could encounter the curse of dimensionality. In high-dimensional space, randomly sampled points may struggle to cover the worst direction that realizes the norm.

In this setting, however, several strong factors make the estimate fairly reliable.

First, the dimension is not that high. In typical Transformer models, the attention head dimension is often 128. This is not low, but it is not a number that immediately triggers catastrophic dimensionality issues.

Second, the sample size is huge. The sample size used for estimation far exceeds the usual batch size. A batch containing millions of tokens produces billions of token pairs, which are the samples used to compute logits. Such a huge sample count gives a high probability of seeing near-worst-case inputs inside the batch.

Third, the effective sample size can be further enlarged. If we do not rely only on the current batch, but track the maximum logit using a strategy similar to an exponential moving maximum, the effective sample size spans multiple training steps and becomes much larger and more stable, further reducing estimation error.

We can draw an analogy with contrastive learning. In contrastive learning, embedding dimensions are often much higher, for example 2048, while the in-batch sample count used to find the hardest negatives is much smaller, for example 32K or 64K. Even in that more challenging setting of higher dimension and smaller sample count, models can still train very well. This supports the practical robustness and effectiveness of empirical methods that search for and suppress hard cases through in-batch sampling.

Therefore, despite theoretical concerns, using max logit to drive correction is a reasonably accurate and efficient proxy for the spectral norm in practice.

Potential problems and challenges

These empirical heuristics may be more feasible in engineering terms, but they introduce new problems and uncertainties.

First, there is the gap between empirical behavior and theory. How robust is it to approximate a global stability guarantee using batch-local information? Elementwise clipping, for example, may over-clip in some gradient directions and under-clip in others, introducing new bias.

Second, there is hyperparameter sensitivity. These methods introduce new hyperparameters, such as the clipping threshold $c$ or the MuonClip threshold $\tau$. Final model performance may be very sensitive to them, and systematically finding the optimal values is itself difficult.

Third, there is interaction with the optimizer. How do these operations interact with the internal mechanisms of stateful optimizers such as Adam, including second moments? Modifying gradients before they enter the optimizer, or correcting parameters after an update, may have hidden conflicts and could even harm optimization.

Answering these questions requires careful experimental validation.

These are the thoughts from one weekend of speculation. I hope they are useful to others thinking about the same problem.

References

[1] Jordan, K., Bernstein, J., et al. (2024). “Muon: An optimizer for hidden layers in neural networks.” Blog post. Available at: https://kellerjordan.github.io/posts/muon/

[2] Su Jianlin. (2025). “QK-Clip: Letting Muon Go Further on the Road to Scaleup.” kexue.fm. Available at: https://spaces.ac.cn/archives/11126

[3] Su Jianlin. (2025). “Muon Sequel: Why Did We Choose to Try Muon?” kexue.fm. Available at: https://spaces.ac.cn/archives/10739

[4] Liu, J., Su, J., et al. (2025). “Muon is Scalable for LLM Training.” arXiv preprint arXiv:2502.16982.

[5] Jordan, K., et al. (2024). “modded-nanogpt: Speedrunning the NanoGPT baseline.” GitHub repository. Available at: https://github.com/KellerJordan/modded-nanogpt

[6] Jordan, K., et al. (2024). “Muon Optimizer.” GitHub repository. Available at: https://github.com/KellerJordan/Muon

[7] Moonshot AI. (2025). “Kimi K2: Open Agentic Intelligence.” Project page. Available at: https://moonshotai.github.io/Kimi-K2/