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.
翻译:权重空间几何在神经网络优化中扮演着核心角色,然而流形约束通常被均匀应用于所有权重矩阵。本研究探讨了不同Transformer模块是否偏好不同的流形几何。我们针对GPT-2预训练研究了Manifold Muon,并比较了注意力层与MLP模块中Stiefel与DGram约束的逐层分配方案。结果显示显著不对称性:对注意力层施加Stiefel几何约束,同时对MLP层分配DGram几何约束,在测试配置中取得最佳性能;而反向分配方案及全DGram配置在共享超参数设置下变得不稳定。我们将此失效归因于DGram约束下注意力权重的奇异值增长——这会放大注意力对数几率并导致softmax饱和。这些发现表明,Transformer的对称感知与几何感知优化应具有模块特异性,而非统一处理。