# Different Layers, Different Manifolds: Module-Wise Weight-Space Geometry in Transformer Optimization

arXiv: https://arxiv.org/abs/2606.13276 · HTML: https://mlsanity.com/paper/2606.13276/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 9 of its 16 references that are in the map.
- First posted on arXiv 11 Jun 2026.

## Builds on (most cited first, 9 of 9)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Git Re-Basin: Merging Models modulo Permutation Symmetries](https://mlsanity.com/paper/2209.04836/) (2022)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [The Role of Permutation Invariance in Linear Mode Connectivity of Neural Networks](https://mlsanity.com/paper/2110.06296/) (2021)
- [Orthogonal Weight Normalization: Solution to Optimization over Multiple Dependent Stiefel Manifolds in Deep Neural Networks](https://mlsanity.com/paper/1709.06079/) (2017)
- [Muon Optimizes Under Spectral Norm Constraints](https://mlsanity.com/paper/2506.15054/) (2025)
- [Mano: Restriking Manifold Optimization for LLM Training](https://mlsanity.com/paper/2601.23000/) (2026)
- [Manifold constrained steepest descent for smooth and closed-set optimization](https://mlsanity.com/paper/2601.21487/) (2026)
- [Feedback Gradient Descent: Efficient and Stable Optimization with Orthogonality for DNNs](https://mlsanity.com/paper/2205.08385/) (2022)

## Built on by (most cited first, 1 of 1)

- [Modular Norm RandOpt: Population-Efficient Ensembling through Architecture-Aware Perturbations](https://mlsanity.com/paper/2609.25745/) (2026)

## Abstract (arXiv)

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.