From Muon to Gradient Clipping: Some Thoughts on QK Stability

17 days ago

Muon optimizer works in function space by constraining the spectral norm of weight updates, aiming to limit the functional change for any input, contrasting with parameter-space optimizers like Adam.
Applying standard Muon to Transformer's Q and K matrices leads to training instability because it separately constrains updates to W_Q and W_K, failing to control their coupled product in attention scores, which can cause spectral norm explosion.
A principled attempt to modify Muon for QK involves directly constraining the change in attention score matrix S, leading to an optimization problem with a constraint coupling ΔW_Q and ΔW_K.
Decoupling this constraint via the triangle inequality and solving subproblems results in a theoretically clean update formula, but it requires expensive pseudoinverse computations, making it impractical.
A mathematical simplification reveals that the gradient G_Q naturally contains W_K, leading to a geometric interpretation involving a projection matrix, yet the external pseudoinverse remains, not solving the computational bottleneck.
Viewing attention through a unified bilinear matrix B = W_Q W_K^T simplifies the Muon update to a single-matrix constraint, but this approach is impractical due to parameter explosion (d^2 vs. 2dh parameters).
Multi-head attention can be seen as a low-rank approximation of full bilinear attention, connecting to LoRA's structure, suggesting potential for a Muon-LoRA optimizer for parameter-efficient fine-tuning.
Due to computational hurdles, practical solutions shift to empirical approximations like gradient clipping or Kimi's MuonClip, which use batch information to approximate worst-case behavior and stabilize training.

Hasty Briefsbeta

transformer

TokenTown: A visual way to understand how LLMs work

From Muon to Gradient Clipping: Some Thoughts on QK Stability