Technical Summary: Different Layers, Different Manifolds

Problem Statement

Weight-space geometry is a critical factor in neural network optimization, yet current manifold-constrained training methods typically apply a single, uniform geometric constraint (e.g., orthogonality or Stiefel constraints) across all weight matrices in a Transformer. This work investigates whether different Transformer modules—specifically attention layers and Multi-Layer Perceptron (MLP) blocks—prefer distinct weight-space geometries. The authors posit that the computational roles and nonlinear mechanisms of these modules differ significantly, suggesting that a uniform manifold assignment may be suboptimal or even destabilizing.

Methodology

The study utilizes Manifold Muon, an optimizer that applies matrix-normalized updates to hidden-layer weights while constraining them to specific matrix manifolds. The authors compare two primary constraint families:

1. Stiefel Constraint: Enforces column orthonormality ($W^\top W = I$), strictly bounding singular values and providing strong spectral control.
2. DGram Constraint: A weaker Gram-space constraint requiring the off-diagonal entries of the Gram matrix to be zero ($\text{Off}(W^\top W) = 0$). This preserves column orthogonality but allows column norms (and thus singular values) to vary, offering additional scale degrees of freedom.

The authors evaluate five distinct assignment strategies on a GPT-2 small model (124M parameters) pre-trained on OpenWebText:

• UNCONSTRAINED: No manifold constraints.
• ALL-STIEFEL: Stiefel constraints on both Attention and MLP weights.
• ALL-DGRAM: DGram constraints on both Attention and MLP weights.
• HETERO: Stiefel constraints on Attention, DGram on MLP.
• HETERO-INV: DGram constraints on Attention, Stiefel on MLP.

All configurations share identical model architectures, data, training schedules, and optimizer hyperparameters to isolate the effect of manifold assignment.

Key Results

The experiments reveal a clear asymmetry in how different modules respond to manifold constraints:

• Optimal Configuration: The HETERO assignment (Stiefel for Attention, DGram for MLP) achieved the lowest stable validation loss (3.3544), outperforming both the unconstrained baseline and the uniform ALL-STIEFEL configuration.
• Instability in DGram-Attention: Configurations assigning DGram constraints to attention layers (ALL-DGRAM and HETERO-INV) became unstable under the shared hyperparameter setting. These runs exhibited explosive growth in validation loss and failed to recover.
• Stability in DGram-MLP: Conversely, assigning DGram constraints to MLP layers (when paired with Stiefel attention) remained stable and improved performance compared to uniform Stiefel constraints.

Mechanistic Analysis

The paper provides a mechanistic explanation for the instability observed in DGram-constrained attention:

1. Singular Value Growth: Unlike Stiefel constraints, DGram permits the singular values of weight matrices to grow arbitrarily large. In attention layers, the query ($W_Q$) and key ($W_K$) projections multiply to form attention logits ($Z = X W_Q W_K^\top X^\top / \sqrt{d}$).
2. Logit Amplification: Uncontrolled growth in the spectral norms of $W_Q$ and $W_K$ directly amplifies the magnitude of attention logits.
3. Softmax Saturation: Large logit magnitudes push the softmax function into a saturated regime, where the output approaches a one-hot vector. This causes the Jacobian of the softmax to vanish exponentially, leading to gradient degeneracy and training instability.
4. MLP Tolerance: MLP layers tolerate DGram constraints because their activation functions (e.g., GELU) operate element-wise and do not couple token scores through a normalized probability simplex. Consequently, scale amplification in MLP weights does not force the entire module into a saturated routing regime.

Contributions

The authors identify three primary contributions:

1. Empirical Evaluation: They evaluate module-wise manifold assignments for Manifold Muon in GPT-2 pretraining, demonstrating that uniform constraints are insufficient.
2. Asymmetry Discovery: They show that attention and MLP layers prefer different weight-space geometries, with attention requiring strong spectral control (Stiefel) and MLP benefiting from weaker scale-preserving freedom (DGram).
3. Instability Mechanism: They identify singular value growth and subsequent softmax saturation as a plausible mechanism for the instability of DGram-constrained attention.

Significance and Limitations

The paper argues that symmetry-aware and geometry-aware optimization for Transformers should be module-specific rather than uniform. This suggests that weight-space geometry is a practical design axis for improving training stability and dynamics in modern networks.

The authors maintain a modest stance regarding their claims:

• They do not claim that DGram attention is inherently unstable under all conditions, but rather that it exposes an unstable pathway under the tested shared hyperparameter setting without additional scale control.
• The experiments are limited to GPT-2 small and a specific set of seeds; the validation loss gap between the best heterogeneous and uniform configurations is modest.
• The mechanistic analysis serves as an explanation for observed instability rather than a complete proof that DGram attention cannot be stabilized (e.g., via retuned learning rates or explicit scale regularization).

Future work is suggested to develop automatic methods for assigning manifold constraints across different modules, architectures, and scales.