QK-Clip：让Muon在Scaleup之路上更进一步 - 科学空间|Scientific Spaces

四个月前，我们发布了Moonlight，在16B的MoE模型上验证了Muon优化器的有效性。在Moonlight中，我们确认了给Muon添加Weight Decay的必要性，同时提出了通过Update RMS对齐来迁移Adam超参的技巧，这使得Muon可以快速应用于LLM的训练。然而，当我们尝试将Muon进一步拓展到千亿参数以上的模型时，遇到了新的“拦路虎”——MaxLogit爆炸。 为了解决这个问题，我们提出了一种简单但极其有效的新方法，我们称之为“QK-Clip”。该方法从一个非常本质的角度去看待和解决MaxLogit现象，并且无损模型效果，这成为我们最新发布的万亿参数模型“Kimi K2”的关键训练技术之一。 我们先来简单介绍一下MaxLogit爆炸现象。回顾Attention的定义 O=softmax(Q K ⊤ )V (1) 这里省略了缩放因子 1/ d − − √ ，因为它总可以吸收到 Q,K 的定义中。“MaxLogit爆炸”中的Logit，指的是Softmax前的Attention矩阵，即 Q K ⊤ ，而MaxLogit指的是全体Logit的最大值，我们将它记为 S max = max i,j q i ⋅ k j (2) 这里的 max 其实还要在batch_size维度上取，最终得到一个标量。而MaxLogit爆炸是指， S max 随着训练的推进一直往上涨，增长速度是线性甚至是超线性的，并且在相当长的时间内没有稳定的迹象。 MaxLogit爆炸现象 MaxLogit本质上是一个异常值指标，它的爆炸意味着异常值超出了可控范围。具体来说，我们有 | q i ⋅ k j |≤∥ q i ∥∥ k j ∥=∥ x i W q ∥∥ x j W k ∥≤∥ x i ∥∥ x j ∥∥ W q ∥∥ W k ∥ (3) 由于 x 通常会加RMSNorm，所以一般情况下 ∥ x i ∥∥ x j ∥ 是不会爆炸的，因此MaxLogit爆炸意味着谱范数 ∥ W q ∥,∥ W k ∥ 有往无穷大发展的风险，这显然不是一个好消息。 由于再大的数值经过Softmax后都变得小于1，所以比较幸运的情况下，这个现象不会带来太严重的后果，顶多是浪费了一个Attention Head，但比较糟糕的情况下，可能会引起Grad Spike甚至训练崩溃。因此，保险起

saved by

• Yudhister Joel Kumar
• Riya Tyagi