MACRO: Demystifying Manifold Constraints as the Multi-Tool of LLM Pre-training

TL;DR

Is the complexity of modern LLMs—with their layers of RMSNorm and finicky weight decay—actually necessary? This paper introduces MACRO, a Riemannian optimizer that proves manifold constraints can independently handle stability and regularization. By constraining weights to Spectral or Frobenius spheres, authors successfully trained stable 330M models without any learnable normalization layers, while matching SOTA performance.

The Hidden Interplay of LLM Stability

Standard LLM training is a balancing act of heuristics:

1. RMSNorm: Prevents activation "explosion" or "vanishing."
2. Weight Decay: Prevents weights from drifting into unoptimal regions.

The authors observed that these are essentially "patches" for a lack of geometric control. If we explicitly constrain weight matrices to a fixed manifold (like a sphere), we gain direct control over the rotation and scale of the representational space.

Methodology: MACRO — Convergent and Efficient

The core contribution is the Msign-Aligned Constrained Riemannian Optimizer (MACRO). Unlike previous constrained methods that required expensive "inner loops" (bisection searches) to satisfy constraints, MACRO uses a single-loop update that is both provably convergent and computationally light.

The Update Logic:

1. Tangent Projection: Projects the gradient momentum onto the local tangent space of the manifold.
2. Steepest Descent: Applies a matrix sign function (msign) to find the most efficient direction on the manifold.
3. Manifold Retraction: Maps the updated weights back to the manifold to prevent "constraint drift."

(Note for User: Refer to Section 3, Algorithm 1 in the paper for the specific logic flow)

Key Insights: Why Constraints Work

1. Abolishing Normalization Layers

One of the most striking results is the Interplay with Normalization. The authors found that as the constraint radius $R$ increases, the learnable parameter $\gamma$ in RMSNorm naturally shrinks to compensate. By removing learnable $\gamma$ entirely, MACRO remained stable while standard Muon diverged (NaN).

• Spectral Sphere: Controls the worst-case activation scale.
• Frobenius Sphere: Controls the average-case activation scale.

2. Replacing Weight Decay

Weight decay is traditionally used to reach a "rotational equilibrium." MACRO achieves this from iteration zero. By locking the relative learning rate ($\Vert \Delta W\Vert /\Vert W\Vert$), MACRO ensures that the optimization trajectory is a pure rotation, preventing the gradient amplification often seen at the end of training.

Experimental Proof: SOTA and Scaling

The authors tested MACRO on QWEN3-like architectures from 120M to 1B parameters.

• Performance: MACRO-spec (Spectral) consistently outperformed the unconstrained Muon baseline.
• $\mu$P Transfer: Because the relative learning rate is "locked" by the manifold geometry, hyperparameters (like learning rate) tuned on a small model could be transferred to larger models with zero modification.

(Note for User: Refer to Table 4 for the quantitative comparison against Muon and SSO)

Professional Analysis & Conclusion

This work shifts the perspective on manifold optimization from a "theoretical niche" to a "practical powerhouse" for large-scale pre-training.

Takeaways:

• Geometric over Heuristic: Manifold constraints provide a principled way to manage the internal dynamics of deep networks.
• Efficiency: The single-loop approximation in MACRO proves that we don't need to sacrifice speed for mathematical rigor.
• Future Scope: The success of the "Normalization-Free" experiment suggests we might be able to simplify Transformer blocks significantly, potentially reducing parameter counts without losing expressive power.

Limitations: While MACRO uses a global radius, the authors admit that layer-specific radii (co-designed with the specific module role, e.g., Attention vs. FFN) could further close the performance gap with standard non-constrained models.