From Muon to Gradient Clipping: Some Thoughts on QK Stability
9 hours ago
- #Deep Learning
- #Optimization
- #Transformer
- Muon优化器在函数空间中工作，通过对权重更新的谱范数施加约束，旨在限制对任何输入的函数变化，这与Adam等参数空间优化器形成对比。
- 将标准Muon应用于Transformer的Q和K矩阵会导致训练不稳定，因为它分别约束W_Q和W_K的更新，未能控制它们在注意力分数中的耦合乘积，这可能导致谱范数爆炸。
- 将Muon修改用于QK的一种理论尝试是直接约束注意力分数矩阵S的变化，这引出了一个优化问题，其约束条件耦合了ΔW_Q和ΔW_K。
- 通过三角不等式解耦这一约束并解决子问题，得到了理论上简洁的更新公式，但它需要昂贵的伪逆计算，使其不切实际。
- 数学简化表明，梯度G_Q自然地包含W_K，这引出了一个涉及投影矩阵的几何解释，然而外部的伪逆仍然存在，未能解决计算瓶颈。
- 通过统一的双线性矩阵B = W_Q W_K^T 来看待注意力，将Muon更新简化为单矩阵约束，但这种方法由于参数爆炸（d^2对2dh参数）而不实用。
- 多头注意力可视为完整双线性注意力的低秩近似，这与LoRA的结构相联系，提示了开发用于参数高效微调的Muon-LoRA优化器的潜力。
- 由于计算障碍，实际解决方案转向经验近似，如梯度裁剪或Kimi的MuonClip，这些方法利用批信息来近似最坏情况行为，从而稳定训练。