MALT：通过对角线预处理实现轻量级曲率感知 Muon

最近，Muon 通过使用 Newton-Schulz 迭代正交化动量矩阵，成为 AdamW 语言模型预训练的有前途的替代品。尽管 Muon 减轻了梯度各向异性，但它没有明确解释损失景观的曲率几何形状，因此可能仍然对曲率各向异性敏感。我们通过提出 MALT（轻量级两侧预处理增强 Muon）来弥补这一差距，它使用轻量级对角预处理器来降低 Muon 对曲率各向异性的敏感性。具体来说，MALT 使用具有低内存和计算开销的两侧对角预处理器来近似捕获损失景观的曲率几何形状。它使用 Newton-Schulz 迭代对预处理动量进行正交，并将结果映射回以定义更新方向，而范数嫁接则控制更新幅度。为了提高 MALT 对随机梯度噪声的鲁棒性，我们进一步提出了 MALTER (MALT with Adaptive stEpsize Rescaling)。在随机非凸设置中为 MALT 提供收敛保证。 GPT-2 小型、中型和大型预训练的实验表明，所提出的方法优于 Muon，同时保持几乎相同的内存占用和挂钟时间。