aiqbench analysis
This research demonstrates that applying uniform manifold constraints across all transformer layers is suboptimal, as attention blocks benefit from Stiefel geometry while MLP layers perform better under DGram constraints. The study reveals that mismatched geometries in attention heads can cause singular value growth and softmax saturation, leading to training instability. These findings advocate for module-specific optimization strategies when employing advanced weight-space regularization techniques like Manifold Muon.
Different Layers, Different Manifolds: Module-Wise Weight-Space Geometry in Transformer Optimization
arXiv:2606.13276v1 Announce Type: cross Abstract: 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.
Abstract: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.
Machine Learning (cs.LG); Artificial Intelligence (cs.AI)
arXiv:2606.13276 [cs.LG]
(or arXiv:2606.13276v1 [cs.LG] for this version)
https://doi.org/10.48550/arXiv.2606.13276 Focus to learn more arXiv-issued DOI via DataCite (pending registration)
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Full text under the publisher's open licence · 230 words
Related coverage
J++ Lens: Jacobian Filtering Enables More Faithful Workspace Lenses
LessWrong · Oct 8, 2026
Enabling DeepSeek-V4-Flash Training on AMD Instinct MI355X GPUs with Primus
AMD · Sep 3, 2026
Game Design Creativity evals
LessWrong · Oct 12, 2026
PIRAMID: Progress and Plans
LessWrong · Jul 27, 2026