---
title: 'MACRO: Msign-Aligned Riemannian Optimization'
url: https://www.emergentmind.com/topics/msign-aligned-constrained-riemannian-optimizer-macro
type: topic
---

# MACRO: Msign-Aligned Riemannian Optimization

The Msign-Aligned Constrained Riemannian Optimizer (MACRO) is a provably convergent, single-loop optimization framework designed for stochastic, nonconvex optimization problems with explicit manifold constraints, specifically motivated by large language model (LLM) pre-training. MACRO systematically unifies and subsumes heuristic stabilization techniques—such as explicit normalization layers (e.g., RMSNorm) and decoupled weight decay—by enforcing geometry-aware constraints on weight matrices. It achieves stability and competitive pre-training perplexity while rigorously guaranteeing exact Riemannian optimization, offering both theoretical and practical advancements over conventional methods [2605.04418].

## 1. Problem Formulation and Constraint Geometry

MACRO addresses the optimization problem:
\[
\text{minimize} \quad \mathcal{L}(W) = \mathbb{E}_{\xi \sim D}[\ell(W;\xi)] \quad \text{subject to} \quad W \in \mathcal{M}
\]
where $W \in \mathbb{R}^{n \times m}$ is a weight matrix, $\mathcal{L}$ is the expected loss, and $\mathcal{M}$ is a constraint manifold. Common choices for $\mathcal{M}$ include:

- **Frobenius sphere:** $\mathcal{M}_F(R)=\{W : \|W\|_F = R\}$
- **Spectral sphere:** $\mathcal{M}_S(R)=\{W : \|W\|_2 = R\}$
- **Oblique manifold:** per-row or per-column $\ell_2$-norm $=$ $R$

For any $W\in\mathcal{M}$, the tangent space $T_\mathcal{M}W$ is defined as $\{\Delta \in \mathbb{R}^{n \times m} : \langle \Delta, \nabla h(W) \rangle = 0\}$, with $h(W)$ reflecting the constraint (e.g., $h(W) = \|W\|_F^2$ for Frobenius). The projection onto the tangent space (Riemannian gradient) for the Frobenius sphere is
\[
\Pi_{T_{\mathcal{M}_F} W}(G) = G - \left(\frac{\langle G, W \rangle}{\|W\|_F^2}\right) W
\]
and for the spectral sphere, given $u_1, v_1$ the top singular vectors,
\[
\Pi_{T_{\mathcal{M}_S} W}(G) = G - (\langle G, u_1 v_1^\top \rangle) u_1 v_1^\top
\]
Retracting back to the manifold after an update,
\[
\mathcal{R}_{\mathcal{M}}(\hat{W}) = \arg\min_{W \in \mathcal{M}} \|W - \hat{W}\|_F
\]
which, in practice, is $\mathcal{R}_F(\hat{W}) = R \cdot \hat{W} / \|\hat{W}\|_F$ for the Frobenius sphere and $\mathcal{R}_S(\hat{W}) = R \cdot \hat{W} / \|\hat{W}\|_2$ (approximate) for the spectral sphere [2605.04418].

## 2. Msign-Aligned Update Rule

The core innovation of MACRO lies in its use of the matrix-sign (msign) operator, which is defined as $\text{msign}(G) = U V^\top$ for the SVD $G = U \Sigma V^\top$. The msign operator solves the linear minimization oracle
\[
\arg\max_{\|X\|_M \leq 1} \langle G, X \rangle
\]
The optimization step at iteration $t$ proceeds as follows:

1. **Project momentum onto tangent space:** $\Phi_t = \Pi_{T_{\mathcal{M}} W_t}(M_t)$
2. **Extract steepest direction via LMO:** $O_t = \text{msign}(\Phi_t)$
3. **Normalize and scale update direction:** $\widetilde{\nabla}_t = \frac{cR}{\|O_t\|_M + \epsilon} O_t$
4. **Descent plus retraction:** $W_{t+1} = \mathcal{R}_{\mathcal{M}}(W_t - \eta_t \cdot \widetilde{\nabla}_t)$

This procedure yields a Riemannian steepest-descent step aligned via the msign operator and combined with efficient, explicit retraction, thus performing a true single-loop geometric update [2605.04418].

## 3. Algorithmic Structure and Hyperparameters

Pseudocode for MACRO defines the following steps per iteration:

1. $G_t \leftarrow \nabla \ell(W_t; \xi_t)$
2. $M_t \leftarrow \beta M_{t-1} + (1-\beta) G_t$
3. $\Phi_t \leftarrow \Pi_{T_{\mathcal{M}} W_t}(M_t)$
4. $O_t \leftarrow \text{msign}(\Phi_t)$
5. $\widetilde{\nabla}_t \leftarrow (cR)/(\|O_t\|_M + \epsilon) \cdot O_t$
6. $W_{t+1} \leftarrow \mathcal{R}_{\mathcal{M}}(W_t - \eta_t \cdot \widetilde{\nabla}_t)$

Optimal theoretical hyperparameter settings include $\beta \approx 1 - 1/\sqrt{T}$, $c \in [0.5, 2]$, and $R$ selected via activation-control theory (detailed below). Batch sizes are typically 64–128. The single-loop structure provides practical efficiency, avoiding the need for double-loop exact projections as in some prior Riemannian solvers [2605.04418].

## 4. Theoretical Guarantees

Under the assumptions that:
- (A1) $\mathcal{M}$ is a compact $C^3$ manifold,
- (A2) $\mathcal{L}$ is lower-bounded, $L$-smooth, and the stochastic gradients are unbiased with variance $\sigma^2$,

it is established that, with suitable $\beta$ and learning-rate schedule $\eta_t \sim \Theta(\sqrt{\Delta / (L T^{3/2})})$, the convergence rate satisfies:
\[
\min_{0 \leq t < T} \mathbb{E}[\|\text{Grad}~\mathcal{L}(W_t)\|_{*, \mathcal{M}}] = O\left((\sqrt{L \Delta} \sigma) T^{-1/4}\right)
\]
This rate matches the optimal for stochastic, nonconvex, constrained problems. The proof leverages the retraction-smoothness lemma and momentum-variance control, with the msign operator ensuring that descent occurs along the direction of steepest Riemannian reduction [2605.04418].

## 5. Mechanistic Interpretations

**Activation Scale Control:** For a linear layer $Y = X W^\top$ and RMS$\|\text{vec}(X)\| \sim \Theta(1)$, manifold constraints guarantee bounded activations:
- Spectral sphere: choosing $R_{spec} = \Theta(\sqrt{d/T})$ ensures $\|\text{vec}(Y)\|_2 = O(1)$.
- Frobenius sphere: $R_{fro} = \Theta(\sqrt{d})$ ensures $\mathbb{E}_X[\|Y\|_F^2]$ is bounded, yielding RMS at $O(1)$. Hence, set $R_{spec} = r \sqrt{d/T}$, $R_{fro} = r \sqrt{d}$, $r \approx 1$ per activation-bound theory.

**Interplay with RMSNorm:** In transformer architectures, RMSNorm learns a per-layer scale $\gamma$. As $R$ increases, $\gamma$ shrinks to keep post-norm activation constant. If all learnable RMSNorms are removed, standard optimizers diverge, but MACRO remains stable at standard learning rates (3e-3–1e-2), demonstrating the sufficiency of explicit manifold constraints for scale control.

**Interaction with Weight Decay:** Conventional decoupled weight decay heuristically enforces:
- Relative learning rate $\eta_{rel,t} \approx \sqrt{2\lambda/\eta_t}$,
- Rotational equilibrium $\theta_t \to$ const.

MACRO enforces these exactly from initialization:
- Relative-LR: $\eta_{rel,t} = \|\Delta W_t\|_M / \|W_t\|_M = c \eta_t$
- Frobenius rotational equilibrium: rotational angle $\theta_t \approx c \eta_t$
- Spectral rotational equilibrium: $\theta_t \lesssim c \eta_t \cdot (R/(R-\sigma_2(W_{t+1})))$ where the spectral gap modulates rotation.

Thus, MACRO subsumes the functionality of both RMSNorm and weight decay via explicit geometry [2605.04418].

## 6. Empirical Evaluation

Evaluations on QWEN3-like architectures (RoPE+GQA+SwiGLU) at 120M, 330M, and 1B parameters (with OpenWebText as the dataset, and token budgets 3.7B/8.9B/50B) show that MACRO achieves validation perplexities competitive with or superior to baselines, including Muon, MuonH-fro/spec, SSO, and FSO. For example:

| Model | Muon | MuonH-fro | MuonH-spec | SSO | FSO | MACRO-fro | MACRO-spec |
|-------|------|-----------|------------|-----|-----|-----------|------------|
| 120M  |3.019 |3.007      |3.019       |3.011|3.001|3.005      |3.017       |
| 330M  |2.736 |2.717      |2.716       |2.712|2.726|2.718      |2.714       |
| 1B    |2.473 |2.468      |2.464       | --  | --  |2.467      |2.461       |

In normalization-free settings (removing learnable RMSNorm at 330M, LR=$10^{-2}$), Muon diverges, while MACRO remains stable with validation losses near 2.76 (fro) and 2.74 (spec).

Further, gradient norms under MACRO decay smoothly (approximately 30× during training) and there is no late-stage blowup seen in AdamW+weight-decay. Zero-shot μP transfer shows optimal learning rates remain stable under MACRO when scaling width. Tangent-space projection residuals remain low ($\sim 10^{-2}$–$10^{-3}$), and performance is comparable to double-loop methods despite its single-loop efficiency [2605.04418].

## 7. Implementation Insights and Recommendations

MACRO is best suited to scenarios requiring rigorous, geometry-aware constraints for stability in deep or large LLM pre-training, especially when minimal tuning of RMSNorm and weight decay is desired. Implementation guidelines include:
- Choose $R$ via activation-bound theory: $R_{fro}=r \sqrt{d}$, $R_{spec}=r \sqrt{d/T}$ with $r \in [0.5, 3]$.
- Set alignment $c \sim 1$ (explore $[0.5, 2]$).
- Use $\beta \approx 1 - 1/\sqrt{T}$ or 0.9 for momentum.
- Approximate spectral retraction by normalizing with $\|\cdot\|_2$.
- Monitor gradient norms and NaN incidence.
  
Possible pitfalls include setting $R$ too small (leading to under-capacity) or too large (leading to loss of scale control), and approximate spectral retractions can violate strict compactness in rare cases of nearly repeated singular values. Omitting the tangent space projection step negates Riemannian guarantees but is inexpensive when included.

In summary, MACRO delivers a unified, geometry-centric approach to LLM pre-training optimization, obviating the need for extensive heuristic tuning of normalization and weight regularization, and achieves strong stability and competitive perplexity in practice [2605.04418].

Source: https://www.emergentmind.com/topics/msign-aligned-constrained-riemannian-optimizer-macro