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:
[
\Delta L \approx \langle G, \Delta W \rangle = \operatorname{Tr}(G^T \Delta W).
]
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:
[
\min_{\Delta W} \operatorname{Tr}(G^T \Delta W) \quad \text{subject to a size constraint on } \Delta W.
]
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:
[
|\Delta W|_F \le \eta.
]
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:
[
\sup_{|x|_2 = 1} |\Delta W x|_2 \le \eta.
]
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:
[
\min_{\Delta W} \operatorname{Tr}(G^T \Delta W) \quad \text{subject to} \quad |\Delta W|_2 \le \eta.
]
This optimization problem under a spectral norm constraint has a beautiful analytic solution: the characteristic matrix sign function,
[
\Delta W^* = -\eta , \operatorname{msign}(G).
]
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:
[
W_{\text{new}} = W + \Delta W.
]
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_Q W_K^T X^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:
[
\Delta S = X(W_Q + \Delta W_Q)(W_K + \Delta W_K)^T X^T - XW_Q W_K^T X^T.
]
Expanding the terms gives:
[
\Delta S = X(\Delta W_Q W_K^T + W_Q \Delta W_K^T + \Delta W_Q \Delta W_K^T) X^T.
]
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:
[
\Delta S \approx X(\Delta W_Q W_K^T + W_Q \Delta W_K^T) X^T.
]

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:
[
|\Delta S|_2 \le \epsilon,
]
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,
[
|X A X^T|_2 \le |X|_2^2 |A|_2.
]
Assuming the input is normalized, $|X|_2 \le 1$, the constraint reduces to a direct constraint on the intermediate weight expression:
[
|\Delta W_Q W_K^T + W_Q \Delta W_K^T|_2 \le \epsilon.
]
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:
[
\Delta L \approx \operatorname{Tr}(G_Q^T \Delta W_Q) + \operatorname{Tr}(G_K^T \Delta W_K),
]
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:
[
\min_{\Delta W_Q, \Delta W_K} \operatorname{Tr}(G_Q^T \Delta W_Q) + \operatorname{Tr}(G_K^T \Delta W_K)
]
subject to
[
|\Delta W_Q W_K^T + W_Q \Delta W_K^T|_2 \le \epsilon.
]

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:
[
|\Delta W_Q W_K^T|_2 + |W_Q \Delta W_K^T|_2 \le \epsilon.
]
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$:
  [
  \min_{\Delta W_Q} \operatorname{Tr}(G_Q^T \Delta W_Q) \quad \text{subject to} \quad |\Delta W_Q W_K^T|_2 \le \epsilon/2.
  ]
• Subproblem 2, for $\Delta W_K$:
  [
  \min_{\Delta W_K} \operatorname{Tr}(G_K^T \Delta W_K) \quad \text{subject to} \quad |W_Q \Delta W_K^T|_2 \le \epsilon/2.
  ]

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:
[
Z = \Delta W_Q W_K^T.
]
Then the constraint becomes clean:
[
|Z|_2 \le \epsilon/2.
]
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:
[
\Delta W_Q = Z (W_K^T)^+.
]
Substituting this into the objective of subproblem 1 gives:
[
\operatorname