From Muon to Gradient Clipping: Some Thoughts on QK Stability - Hanchi Sun’s Personal Website You are using an outdated browser. Please upgrade your browser to improve your experience.
Hanchi Sun<br>Large Language Model
Follow Bethlehem, PA, USA<br>Lehigh University<br>Email<br>Google Scholar<br>Github<br>LinkedIn<br>X (formerly Twitter)<br>Zhihu
Prelude: an elegant theory and a stubborn problem<br>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.<br>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.<br>The elegance of the theory runs into a wall in practice. That naturally raises two questions:<br>What is the root cause behind this phenomenon?<br>Can we follow Muon’s first principles and propose a modified update idea suitable for QK?<br>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.<br>1. Theory review: why standard Muon does not fit QK<br>To understand this mismatch, we first need to understand Muon’s original intention and how it differs fundamentally from optimizers such as Adam.<br>The derivation and essence of Muon<br>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:<br>\[\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:<br>\[\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.<br>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:<br>\[\|\Delta W\|_F \le \eta.\]Its geometric meaning is intuitive: the Euclidean distance moved in parameter space should not be too large.<br>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$.<br>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:<br>\[\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.<br>Now Muon’s optimization problem is clear:<br>\[\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,<br>\[\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.<br>The geometric conflict in QK updates<br>With Muon’s mechanism in mind, we can revisit its failure on QK updates. As Su Jianlin analyzes in his blog post...