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

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.

• The paper demonstrates that module-specific manifold constraints, using Stiefel for attention and DGram for MLP, improve validation loss over uniform or unconstrained strategies.
• The study employs Manifold Muon to enforce geometric regularization, revealing that DGram on attention layers induces instability due to unbounded singular value growth.
• Empirical evidence on GPT-2 shows that optimizing with heterogeneous constraints yields lower stable validation loss, emphasizing the need to tailor weight-space design to module function.

Module-Wise Manifold Geometry in Transformer Optimization

Introduction

This paper, "Different Layers, Different Manifolds: Module-Wise Weight-Space Geometry in Transformer Optimization" (2606.13276), systematically investigates the impact of module-specific manifold constraints on the training dynamics and stability of Transformer architectures, with a focus on GPT-2 pretraining. The central claim is that imposing uniform matrix manifold constraints (e.g., Stiefel or diagonal Gram [DGram]) across all weight matrices is suboptimal. Instead, assigning Stiefel geometry to attention layers and DGram to MLP/FFN layers results in improved stability and validation loss relative to both uniform constraint strategies and unconstrained optimization. The paper provides mechanistic analysis, empirical evidence, and architectural delineation to substantiate the necessity of module-wise geometry-aware optimization for large-scale LLMs.

Manifold Muon: The Geometric Optimization Framework

Muon and its manifold extension, Manifold Muon, enable structured matrix updates by leveraging matrix-normalized gradients. These updates can be restricted to specific smooth manifolds in weight space—most notably, the Stiefel manifold (enforcing $W^\top W = I$, imposing orthonormal columns and tightly bounding singular values) and DGram (requiring off-diagonal vanishing in $W^\top W$, thus columns are orthogonal but can grow unbounded in norm). These structures enforce not just local regularization but global geometric properties affecting the optimization landscape and spectral dynamics of learned weights.

Given the heterogeneous roles of Transformer modules—attention versus MLP/FFN—the authors posited that different modules likely require different geometric priors for optimal training. To test this, five constraint assignments were evaluated: fully unconstrained, all-Stiefel, all-DGram, heterogeneous (Stiefel for attention, DGram for MLP), and the inverse heterogeneous configuration.

Empirical Results: Layer-Wise Manifold Assignment

The principal experimental finding is a stark asymmetry in stability and performance between constraint strategies.

Figure 1: Validation loss curves for five manifold assignments. Hetero (Stiefel attention, DGram MLP) achieves the lowest stable validation loss; DGram-attention configurations exhibit sharp instability.

Under identical training conditions using GPT-2 small (124M parameters, OpenWebText):

• Hetero (Stiefel attention, DGram MLP) produces the lowest stable validation loss, outperforming unconstrained and all-Stiefel configurations.
• All-DGram and Hetero-Inv (DGram attention) become rapidly unstable, regardless of the choice for MLP constraints.
• All-Stiefel offers improvement over no constraint but underperforms heterogeneous manifold assignment.
• Unconstrained runs remain stable but yield the highest validation loss among successful configurations.

These results decisively reject the sufficiency of uniform manifold constraints and underscore the practical significance of module-wise geometry-aware design.

Mechanistic Analysis: Spectral Growth and Attention Instability

The mechanistic explanation centers on the spectral evolution of attention and MLP/FFN weights. The Stiefel constraint strictly bounds singular values, controlling the scale of projections and thus capping the magnitude of attention logits and softmax input. DGram, conversely, allows singular values to grow, retaining column orthogonality but permitting arbitrarily large norms.

Figure 2: Maximum singular value evolution for attention projections (left) and MLP/FFN projections (right) throughout training. DGram on attention induces explosive growth and instability; growth in MLP/FFN weights under DGram does not destabilize training if attention remains Stiefel-constrained.

When DGram is applied to attention projections, empirical traces confirm rapid, unbounded growth in the largest singular value, followed by unstable loss trajectories. This phenomenon is directly connected to the scale of the resulting query-key dot products fed into the softmax, inducing saturation effects where gradients through the softmax become vanishingly small—a well-understood degenerate regime for deep network optimization. The formal analysis demonstrates that, under DGram constraints, the weight-product-induced amplification can make the softmax arbitrarily sharp, yielding near-zero Jacobian norms, and thus degenerate gradients for attention probabilities.

In contrast, MLP/FFN weights can accommodate DGram's scale freedom due to their uncoupled, elementwise nonlinearity (e.g., GELU) and the mitigating influence of residual connections and LayerNorm, which collectively buffer against scale amplification in a way that attention softmax layers cannot.

Architectural Implications

Figure 3: Schematic of the HETERO module-wise manifold assignment. Within GPT-2 blocks, attention projections receive Stiefel Manifold Muon, MLP/FFN projections DGram Manifold Muon; other parameters use AdamW.

The heterogeneous manifold assignment—Stiefel for attention, DGram for MLP/FFN—is architecturally simple to implement and introduces negligible computational overhead relative to baseline Manifold Muon setups. This strategy addresses two competing desiderata: maintaining spectral control where gradient degeneracy is a threat (attention layers) and exploiting scale flexibility where it can enhance expressivity or training efficiency (MLP/FFN layers).

Numerical Results and Empirical Observations

The best validation loss achieved by the heterogeneous assignment is 3.3544, marginally better than all-Stiefel (3.3679) and unconstrained (3.3855), under identical optimizer and learning-rate schedules. Crucially, all configurations assigning DGram to the attention pathway become unstable, despite matching all other hyperparameters.

The study controls for confounds by using shared hyperparameters across all manifold strategies. While not exhaustive in terms of DGram-specific tuning (e.g., applying explicit scale regularization), the consistency of instability in DGram-attention runs, confirmed by robust singular value and loss diagnostics, provides strong evidence for the necessity of module-specific spectral control in attention projections.

Theoretical Implications and Future Directions

This research advances the understanding that transformer module optimization is not a geometry-invariant problem—weight-space symmetries and constraints must be tailored to the computational role of each module. The precise matching between spectral growth, softmax saturation, and loss spikes elucidates a key failure mode when attention is not tightly spectrally controlled.

Practically, this implies that geometry-aware optimization—potentially via automatic assignment of manifold constraints per module—is a viable avenue for improving transformer stability, convergence, and perhaps generalization. Theoretically, the results prompt renewed interest in the interplay between architectural nonlinearity, manifold regularization, and dynamic scale control.

Open questions include:

• The effect of manifold assignment at larger model and data scales,
• The efficacy of explicit scale regularization when using DGram attention,
• Methods for automatic geometry assignment tuned to both architecture and optimizer schedules.

Conclusion

The findings substantiate that module-wise manifold constraints, particularly Stiefel for attention and DGram for MLP/FFN, are preferable to uniform strategies in transformer pretraining. This reflects a deeper principle: the optimization geometry should reflect architectural function, not merely broad regularization. The implications are broad, affecting both the theoretical design of manifold-constrained optimizers and the practical implementation of scalable, stable training for LLMs.