From Muon to Gradient Clipping: Some Thoughts on QK Stability
9 hours ago
- #Deep Learning
- #Optimization
- #Transformer
- 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.