语言模型头的第一性原理更新几何

Muon 激励围绕每个参数块的函数设计优化器几何结构，并使用隐藏线性层的谱范数。对于语言模型头来说，谱范数并不是功能变化的忠实度量。 Softmax 删除了共享的 logits 移位，而谱范数可以为不改变输出概率的更新分配任意大的大小。因此，我们将 LM head 和 softmax 视为一个模块，并为其组合导出更新几何结构。希尔伯特的投影距离尊重这种不变性，因为它测量成对对数赔率的最大变化。对于标记行 $s_i^\top$ 的更新 $S$，我们表明 $\left\lVert h\right\rVert_2\leq H$ 上的最大希尔伯特距离恰好是 $H D(S)$，其中 $D(S)=\max_{i<j}\left\lVert s_i - s_j\right\rVert_2$ 是欧几里德行直径。该直径取代了由此产生的 Muon 式最速下降问题中的谱范数。精确的解决方案是可能的，但其直接公式为每个标记对包含一个 $d$ 维向量变量。对于大约 $50$k 的词汇量，这意味着超过 10 亿个标记对，使得每个训练步骤的计算都不切实际。相反，我们施加更强的共球约束，并将投影 RowNorm 导出为 $O(Vd)$ 解决方案。对于精确的 RowNorm 预言机，我们证明其一阶递减至少是精确直径约束最优值的 $1/\sqrt{2}$。以 Muon 为骨干，在 190M、380M 和 640M 参数下对三个种子进行的实验表明，RowNorm 将平均最终步长直径和经验 Hilbert RMS 扰​​动分别降低了 $45$--$60$ 和 $12$--$15$，而平均最终验证损失仅增加了 $0.0057$-$0.0153$。