不同模块应采用不同几何约束，提升Transformer训练稳定性与性能
Different Layers, Different Manifolds: Module-Wise Weight-Space Geometry in Transformer Optimization
- 按模块分配不同几何约束：注意力层用Stiefel，MLP层用DGram
- 该组合在GPT-2预训练中表现最佳，其他配置易失稳
- 发现DGram约束会引发注意力权重奇异值增长，导致软最大饱和
权重空间几何在神经网络优化中起关键作用，但通常对所有权重矩阵施加统一的流形约束。本文探讨Transformer不同模块是否偏好不同流形几何。研究基于GPT-2预训练，对比了注意力和MLP模块在Stiefel与DGram约束下的分层配置。结果表明存在明显不对称性：将注意力层约束于Stiefel几何、MLP层采用DGram几何的配置表现最优；而反向配置及全DGram配置在相同超参数设置下趋于不稳定。我们追溯其失败原因在于DGram约束导致注意力权重奇异值持续增长，进而放大注意力分数并引发softmax饱和。研究提示，Transformer的优化应具备模块特异性，而非统一的对称与几何感知策略。
原文摘要 · Abstract (English)
Weight-space geometry plays a central role in neural network optimization, yet manifold constraints are often applied uniformly across all weight matrices. In this work, we ask whether different transformer modules prefer different manifold geometries. We study Manifold Muon for GPT-2 pretraining and compare layer-wise assignments of Stiefel and DGram constraints across attention and MLP blocks. Our results show a clear asymmetry: constraining attention layers with Stiefel geometry while assigning DGram geometry to MLP layers gives the best performance among the tested configurations, whereas the inverted assignment and all-DGram configuration become unstable under the shared hyperparameter setting. We trace this failure to singular value growth in DGram-constrained attention weights, which can amplify attention logits and induce softmax saturation. These findings suggest that symmetry-aware and geometry-aware optimization for transformers should be module-specific rather than uniform.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。