Demystifying Manifold Constraints in LLM Pre-training

Abstract: The empirical success of LLM pre-training relies heavily on heuristic stabilization techniques, such as explicit normalization layers and weight decay. While recent constrained optimization approaches that explicitly restrict weights may improve numerical stability and performance, the mechanism and motivation for adding constraints still remain elusive. This paper systematically demystifies the role of explicit manifold constraints in LLM pre-training. By introducing the Msign-Aligned Constrained Riemannian Optimizer (MACRO)-a provably convergent, single-loop optimization framework-our study disentangles weight regularization heuristics from interacting mechanisms like RMS normalization and decoupled weight decay. Theoretical analyses and comprehensive empirical evaluations reveal that manifold constraints independently bound forward activation scales and enforce stable rotational equilibrium, thereby subsuming the roles of these heuristic mechanisms. Evaluations on large-scale LLM architectures demonstrate that MACRO achieves highly competitive performance while rigorously preserving the theoretical guarantees of exact Riemannian optimization.

• The paper introduces MACRO, a provably convergent Riemannian optimizer that leverages manifold constraints to replace traditional normalization and weight decay.
• The paper demonstrates that explicit geometric constraints such as spectral and Frobenius spheres improve activation control and training stability in LLMs.
• The paper empirically validates that manifold constraints streamline LLM pre-training by reducing heuristic tuning while maintaining competitive performance.

Demystifying Manifold Constraints in LLM Pre-training: A Technical Overview

Context and Motivation

Modern LLM pre-training pipelines rely substantially on mechanisms that stabilize optimization and tune the geometry and scale of hidden representations. Canonically, these include normalization layers (e.g., RMSNorm, LayerNorm), heuristic regularization (e.g., weight decay), and various optimizer modifications. However, recent work has shown that explicitly constraining weights to reside on low-dimensional manifolds (e.g., spectral or Frobenius spheres) can improve numerical stability and generalization in LLMs. Despite empirical evidence, the mechanisms by which these manifold constraints act, and their interplay with existing heuristics, remain largely opaque.

This paper rigorously investigates the geometric and optimization-theoretic underpinnings of explicit manifold constraints in LLM training. By introducing the Msign-Aligned Constrained Riemannian Optimizer (MACRO), a provably convergent single-loop optimizer, the authors disentangle the direct effects of weight constraints from those of normalization and weight decay, providing new theoretical and empirical insights. The core claims are that manifold constraints subsume traditional normalization and regularization heuristics: they inherently bound forward activations and enforce stable updates, thereby offering a principled framework for robust LLM pre-training.

Manifold Constraints: Geometric Foundations

The paper investigates four key manifold constraints for weight matrices:

• Frobenius Sphere: $\|W\|_F = R$ — bounds the global energy of weights.
• Spectral Sphere: $\|W\|_2 = R$ — constrains the maximal singular value.
• Oblique Manifolds: row-wise (output) or column-wise (input) norm constraints, $\|W_{i,:}\|_2 = R_{\text{out}}$, $\|W_{:,j}\|_2 = R_{\text{in}}$; these act on individual feature dimensions.

Analyzing the effect of these constraints on activation propagation, the authors show that the spectral sphere directly controls the worst-case amplification of linear layers, while the Frobenius sphere regulates the average-case scale, grounded in input covariance structure and stable rank theory.

The practical guidelines for constraint radii are summarized as follows:

• Spectral Sphere: $r\sqrt{d_{\text{in}}/d_{\text{out}}}$
• Frobenius Sphere: $r\sqrt{d_{\text{in}}}$
• Output/Input Oblique: $r$ or $r\sqrt{d_{\text{in}}/d_{\text{out}}}$

MACRO: A Provably Convergent Riemannian Optimizer

To operationalize constrained LLM training, the authors propose MACRO, a single-loop Riemannian optimization algorithm that generalizes spectral steepest descent methods (notably, Muon). Key characteristics:

• Tangent Space Projection: Momentum terms are projected onto the tangent space of the constraint manifold at every step, ensuring updates are always feasible.
• Riemannian LMO Direction: Rather than using Euclidean gradients, updates are computed via a Linear Minimization Oracle (LMO) on the Riemannian gradient, optimally reflecting the manifold geometry.
• Update–Weight Ratio Alignment: Each update is normalized such that the update/parameter norm ratio is strictly tied to the learning rate, locking the effective learning rate.
• Projectional Retraction: After each update, weights are retracted precisely onto the manifold, structurally preventing drift.
• Theoretical Guarantees: Under standard smoothness and compactness assumptions, MACRO achieves $\mathcal{O}(T^{-1/4})$ convergence rates for general nonconvex stochastic objectives, matching the best-known results for this class of problems.

Effects on Training Dynamics

Activation Control: Spectral vs. Frobenius Constraint

Both constraints secure $\Theta(1)$-scaled activations, but their mechanism differs. The spectral constraint offers a non-adaptive upper bound (ideal for architectures lacking normalization layers), while Frobenius constraints depend on the empirical input covariance and are thus sensitive to rank and representation dynamics.

Figure 1: Rotational angle under the spectral constraint: The spectral rotation angles display high-frequency variance. Despite this, both the $\|W\|_2 = R$0 and $\|W\|_2 = R$1 trajectories remain well-balanced and stable. This dynamic behavior supports the Wedin bound analysis: the local spectral gap continuously modulates adaptive rotations.

Interplay with Normalization Layers

A striking empirical result is that, with manifold constraints alone, LLMs can be trained stably even when all learnable normalization layers (e.g., RMSNorm) are entirely removed. The optimizer automatically compensates, as increases in constraint radii $\|W\|_2 = R$2 induce adaptive shrinkage in parameters such as RMSNorm's $\|W\|_2 = R$3, suggesting redundancy. Crucially, MACRO maintains stability across regimes where standard optimizers diverge.

Figure 2: Training and validation loss for Muon, MACRO-fro, and MACRO-spec on the 330M QWEN3-like model with all learnable RMSNorm layers removed, at the chosen learning rate $\|W\|_2 = R$4.

Interaction with Weight Decay

Explicit manifold constraints naturally lock both the relative learning rate (update–weight ratio) and the rotational angle between parameter updates from the start of training. This geometric regularization bypasses the need for heuristic penalty tuning or gradual “rotational equilibrium” processes induced by weight decay. For example, under the Frobenius sphere, the rotation per step is statically determined by hyperparameters; under the spectral sphere, the rotation becomes anisotropic and adaptive, governed dynamically by the spectral gap.

Figure 3: 120M model, 3600-step schedule.

Figure 4: MACRO-spec.

Figure 5: Train and Validation Loss for 120M training.

Figure 6: The average of Per-step tangent space violation $\|W\|_2 = R$5, illustrating that MACRO's single-step projection yields small violation compared to double-loop projections, with negligible impact on final performance.

Empirical Validation and Ablation

Extensive experiments are performed on QWEN3-like architectures at multiple scales (120M to 1B parameters). MACRO achieves competitive or superior validation loss relative to both unconstrained (Muon) and other constrained (MuonH, SSO/FSO) baselines. Importantly, MACRO's single-loop design eliminates the computational overhead of double-loop projection algorithms without degrading constraint satisfaction or generalization. Notably, ablations reveal that oblique constraints provide no improvement over global Frobenius constraints, as they merely force unnecessary rigidity on feature directions.

Figure 7: Ablation study on geometric constraints.

Implications and Prospects

Practical Implications: The findings support replacing normalization layers and heuristic weight decay with principled manifold constraints, reducing reliance on sensitive architectural heuristics. This not only improves numerical stability and generalization, but also streamlines LLM architecture and training hyperparameter selection.

Theoretical Implications: The equivalence of manifold constraints and traditional normalization/regularization heuristics suggests a unifying geometric principle underpinning LLM training stability. By shifting the focus toward explicit geometric regularization, future work can develop scalable, theoretically grounded optimizers for even larger and more diverse models.

Future Directions:

• Layer-specific geometric constraints could close the minor performance gap observed when removing normalization, especially in deep or heterogeneous architectures.
• Extending this analysis to other norm constraints (e.g., nuclear norm, mixed norms) could expose further connections between optimization geometry and representation learning in deep networks.
• The design of new architectures, possibly omitting normalization entirely, leveraging manifold constraints as the primary stabilization technique.

Conclusion

This work provides a detailed account of how explicit manifold constraints fundamentally restructure training dynamics in LLM pre-training. By introducing the MACRO optimizer and delivering both theoretical and empirical justification, the authors establish that geometric constraints alone can replace normalization layers and weight decay, yielding principled, efficient, and robust training pipelines for large-scale LLMs. This principled approach to optimization on matrix manifolds marks a decisive advance for scalable and reliable LLM training.

Reference:

"Demystifying Manifold Constraints in LLM Pre-training" (2605.04418)