BMuon：面向结构化可压缩性的 Muon 谱几何预算分配

arXiv:2610.23571v1 发布类型：新 摘要：Muon 已成为神经网络训练中的一种强大优化器，并且可以解释为在谱范数几何下进行最速下降。然而，我们发现，在训练后对 Transformer 的 MLP 隐藏通道进行剪枝时，使用 Muon 训练的模型性能可能会急剧下降。为协调 Muon 的优化性能与结构化可压缩性，我们提出 Budget Muon (BMuon)，通过修改 Muon 的谱几何，使不同模块中的隐藏通道围绕共享的分配预算展开竞争。所得预言机恰好是在一族范数下归一化的最速下降，这族范数在 Muon 的几何与结构化稀疏几何之间进行插值。我们主要的稀疏性控制量是一个可解释的全局通道分配预算，而不是依赖具体设置的稀疏性惩罚系数。对于线性分解模型，我们证明 Muon 在隐藏层宽度存在冗余时可以得到成本最低的精确表示；而在 BMuon 预言机几何和预算足够稀缺的情况下，每个成本最低的表示都具有最小宽度。在实际训练中，我们为精确分配问题推导了一个下界最大化代理目标。在使用 ViT-S/16 的 ImageNet-1K 实验中，BMuon 在渐进式结构化 MLP 剪枝下优于 Muon，并且其学习到的分配方案比权重幅值提供了更有效的剪枝标准。在 FineWeb 上训练的一个 1B 参数 LLaMA 风格模型中，BMuon 在训练后 MLP 通道剪枝下优于 Muon。