Demystifying Manifold Constraints in LLM Pre-training

In 2026, the open-weight model Kimi K3 has about 2.8 trillion total parameters.¹ Pretraining is still scaling. But data and compute are not unlimited. The training recipe deserves as much attention as model size. Can we get more from pretraining by changing how weights learn, rather than only adding data or compute?

One clue comes from the weights themselves. Kaiyue Wen reported an unusual pattern in Muse Glimmer: the matrices he examined had nearly the same root-mean-square (RMS) value, about 0.006, across layers and matrix types. RMS measures the typical size of their entries.²

Screenshot of Kaiyue Wen's post reporting near-constant per-matrix weight RMS in Muse Glimmer. — Kaiyue Wen's post about weight scales in Muse Glimmer.

The comparison below shows the weight-RMS patterns in Qwen3-8B and Muse Glimmer.

Comparison between Qwen3-8B and Muse Glimmer

In these plots, weight RMS varies across Qwen3-8B layers and matrix types. Muse Glimmer shows a much more uniform pattern near 0.006.

Qwen3-8B weight RMS varies across layers and matrix types; its vertical-axis range is larger than the Muse plot's. — Qwen3-8B: weight RMS by layer and matrix type. From the same thread.

Qwen3-8B weight RMS: enlarged view of the same original chart. — Qwen3-8B: weight RMS by layer and matrix type. From the same thread.

Muse-Glimmer-30B weight RMS: the plotted matrix types stay near 0.006 across layers. — Muse-Glimmer-30B: weight RMS by layer and matrix type. From Kaiyue Wen's post.

Muse-Glimmer-30B weight RMS: enlarged view of the same original chart. — Muse-Glimmer-30B: weight RMS by layer and matrix type. From Kaiyue Wen's post.

The plots use different vertical-axis ranges. Compare the patterns and labeled values, not bar heights alone.

Why might keeping weight scale fixed help a model learn? That's the puzzle we want to explore.

Normalization itself is not new. What matters here is what we normalize. LayerNorm and RMSNorm act on activations—the signals passed through a model. Other methods act on the updates or the weights themselves.³⁴

What are we normalizing?

Activations — signals inside the model: LayerNorm and RMSNorm normalize activations. nGPT also normalizes hidden states.³⁴⁵
Updates — proposed changes to weights: Muon approximately orthogonalizes momentum matrices. Mano uses row- and column-wise geometry to project and normalize momentum.⁶⁷
Weights — parameters learned during training: nGPT normalizes parameter vectors. MuonH fixes matrix Frobenius norms, while SSO fixes spectral norms.⁵⁸⁹

Some methods act on more than one object. Hyperball and SSO, for example, also control update norms.⁸⁹

But adding a weight constraint is more than a rescaling trick. The entries can still change, but not in every possible way. Take the Stiefel constraint discussed in Thinking Machines Lab's Modular Manifolds. It keeps a matrix's columns orthonormal: perpendicular to each other and unit length. Such a matrix cannot become rank-deficient and collapse an input direction to zero. This restricts the transformations the layer can learn, not just the size of its weights.¹¹

Different constraints leave different kinds of freedom. When the allowed weights form a smooth manifold, training becomes a manifold-constrained optimization problem. MCSD, for example, provides algorithms for such problems, including Stiefel constraints.¹⁰

From normalization to constrained optimization

Let $W$ be a weight matrix, and let $L (W)$ be the training loss. A constraint limits the allowed weights to a set $M$ :

min_{W} L (W) ⟶ min_{W \in M} L (W) .

For a matrix $W \in R^{m \times n}$ with $m \geq n$ , the Stiefel constraint is

W^{⊤} W = I_{n} .

Its columns are orthonormal, so all $n$ singular values are one. The matrix has full column rank and preserves input-vector lengths. A fixed Frobenius norm does not impose these same restrictions.¹¹

Many manifold algorithms choose a direction in the tangent space. This space describes changes that preserve the constraint to first order. For the Stiefel manifold, a tangent direction $A$ satisfies

W^{⊤} A + A^{⊤} W = 0 .

A finite step can still leave the manifold. A retraction maps the candidate weights back to it. Section 0 introduces the update procedure used in our study.¹¹¹²

One question is how to design an optimizer for these constraints. Another question is how these constraints change training.¹² In this blog, we answer the following question to explain the advantages of weight normalization.

How do manifold constraints affect the training dynamics of modern LLMs?

We study activation scales and changes in weight direction, in addition to the final loss.¹² First, we define the constrained problem and introduce the optimizer that we use to study it.

Section 0

Preparation: Candidate Constraints and the Optimizer

Before we explore what weight constraints change, let's define the training problem and choose an optimizer.

Let $W \in R^{p \times n}$ be a weight matrix. Its input dimension is $n$ , and its output dimension is $p$ . We minimize the training loss $L (W)$ , while keeping $W$ in a constraint set $M$ :

min_{W \in M} L (W) .

We start with two ways to fix matrix scale. The Frobenius norm is the square root of the sum of squared entries. The spectral norm is the largest singular value. It measures the largest factor by which the matrix can stretch an input vector.

The two constraint sets are

M_{F} (R) = {W : {‖ W ‖}_{F} = R}, M_{S} (R) = {W : {‖ W ‖}_{2} = R} .

The radius $R > 0$ sets the required norm. Both constraints leave the individual entries free to change.¹

Both constraints act on the whole matrix. Later, we will ask whether each row or column needs its own norm constraint.

To train under these constraints, we propose the Msign-Aligned Constrained Riemannian Optimizer (MACRO). Its two variants are MACRO-fro and MACRO-spec.

Why not simply use MuonH or SSO? We want a practical testbed with a convergence guarantee. MuonH normalizes weights and updates, without an explicit tangent-space projection. SSO solves a constrained direction problem through an inner iterative search.⁴⁵

MACRO builds on the MCSD approach to balance manifold optimization with training cost. It uses a single-loop update, with a convergence guarantee under the paper's assumptions. Here is how the pieces fit together.⁶¹

How does MACRO update the weights?

At step $t$ , MACRO forms momentum $M_{t}$ from the current gradient and previous momentum. It then projects $M_{t}$ onto the tangent space at $W_{t}$ , giving $Φ_{t}$ . This explicit projection distinguishes MACRO from MuonH and is central to our convergence analysis.¹⁴

MACRO then follows the spectral steepest-descent framework to compute the update direction:

O_{t} = msign (Φ_{t}) .

For a full-rank matrix, this operation keeps the singular vectors and sets the singular values to one.

Next, we scale this direction with the constraint radius:

{\tilde{\nabla}}_{t} = \frac{R O_{t}}{{‖ O_{t} ‖}_{M} + ε} .

Here, ${‖ \cdot ‖}_{M}$ is the Frobenius norm for MACRO-fro and the spectral norm for MACRO-spec.

The constraint fixes the weight scale. This rescaling controls the proposed update size relative to it. We return to this update-to-weight ratio in Section 3, where small updates can disappear during rounding.¹

We then take a step with learning rate $η_{t}$ . A retraction returns the candidate weights to the constraint set:

W_{t + 1} = P_{M} (W_{t} - η_{t} {\tilde{\nabla}}_{t}) .

Explore the five stages below to see what each part contributes.¹

Figure 0.1

Select a stage to see its role.

M_{t} = β M_{t - 1} + (1 - β) G_{t} .

Compute the exponential moving average of the gradients.

Momentum

M_{t} = β M_{t - 1} + (1 - β) G_{t} .

Compute the exponential moving average of the gradients.

Tangent projection

Φ_{t} = Π_{T_{W_{t}} M} (M_{t}) .

Project the momentum onto the tangent space at $W_{t}$ .

This projection is central to our convergence analysis.

Orthogonalize

Φ_{t} = U Σ V^{⊤}, O_{t} = msign (Φ_{t}) = U V^{⊤} .

Form $O_{t} = U V^{⊤}$ from the singular vectors of $Φ_{t}$ .

This gives the direction for a spectral steepest-descent step.

Rescale

{\tilde{\nabla}}_{t} = \frac{R O_{t}}{{‖ O_{t} ‖}_{M} + ε} .

Scale the update direction using the radius $R$ .

This controls the proposed step size relative to the weight scale.

Retract

Z_{t} = W_{t} - η_{t} {\tilde{\nabla}}_{t}, W_{t + 1} = P_{M} (Z_{t}) .

Map the updated weights back to the constraint set.

The next step starts from weights that satisfy the constraint.

Convergence result

Theorem 1 gives a convergence bound under the stated assumptions. With the specified learning rate and momentum, the bound is $O (T^{- 1 / 4})$ after $T$ steps. This matches the rate for stochastic optimization of general smooth, nonconvex functions.¹

How do we choose the radius?

We now have a constraint and an update rule. But how large should the constraint be?

The same radius can mean different things for different matrices. The Frobenius norm combines squared values from all entries. The spectral norm measures the largest stretch along an input direction. So we choose the radius from the layer's effect on its inputs, not from the norm value alone.

Our aim is to control activation scale as layer widths change. This shares a goal with width-aware parameterizations such as μP.⁷

Consider the linear layer $Y = X W_{0}^{⊤}$ . Let $r_{X}$ and $r_{Y}$ be the RMS values of its inputs and outputs. We choose dimension-dependent radii and leave a separate multiplier $r > 0$ for tuning. For normalized Gaussian initialization, the rules are:²

Frobenius

R_{fro} = r \sqrt{p}

Expected squared output RMS:

E_{G} [r_{Y}^{2} ∣ X] = r^{2} r_{X}^{2}

Spectral

R_{spec} = r \sqrt{p / n}

Output RMS bound:

r_{Y} = r r_{X} with high probability

For the assumptions and probability bounds, open Proposition 1 below. The formulas account for the input dimension $n$ and output dimension $p$ . The multiplier $r$ is independent of these dimensions. Does this extra hyperparameter make constraints harder to use? In Section 1, we show that training is robust to changes in this multiplier when RMSNorm and weight constraints work together.

Proposition 1: Radius selection at initialization

Let $T, n, p \geq 1$ , $r > 0$ , and $X \in R^{T \times n}$ . Define

r_{X} = \frac{{‖ X ‖}_{F}}{\sqrt{T n}}, r_{Y} = \frac{{‖ X W_{0}^{⊤} ‖}_{F}}{\sqrt{T p}} .

Let $G \in R^{p \times n}$ have independent standard normal entries. The matrix $G$ is independent of $X$ .

Frobenius constraint. Set

W_{0} = r \sqrt{p} \frac{G}{{‖ G ‖}_{F}} .

Then

E_{G} [r_{Y}^{2} ∣ X] = r^{2} r_{X}^{2} .

Spectral constraint. Set

W_{0} = r \sqrt{\frac{p}{n}} \frac{G}{{‖ G ‖}_{2}} .

The output satisfies the deterministic bound

r_{Y} \leq r r_{X} .

For any $ε \in (0, 1)$ and $a > 0$ ,

r_{Y} \geq \frac{r \sqrt{1 - ε}}{\sqrt{n / p} + 1 + a} r_{X}

with conditional probability at least

1 - \exp (- p ε^{2} / 4) - \exp (- a^{2} p / 2) .

The probability is over $G$ , for a fixed $X$ .

Theorem 2 in Appendix C.2 gives concentration bounds for both initializations. With $r, r_{X} = Θ (1)$ and $n / p = O (1)$ , these bounds give $r_{Y} = Θ (1)$ with high probability as the widths grow. The aspect-ratio condition is stated in Remark 1 of that appendix.²

What if we constrain each row or column?

Earlier, we fixed one norm for the whole matrix. Would finer control help?

The Frobenius sphere fixes the total squared norm. It does not fix how that total is divided among rows or columns. Oblique manifolds add constraints at this finer level. An output-oblique constraint gives every row the same fixed norm. An input-oblique constraint does the same for columns.³

We match each oblique variant to the same total Frobenius norm as MACRO-fro. This keeps the overall scale fixed while changing how it is shared.

In these tests, neither oblique variant improves validation loss over the matched Frobenius sphere. We therefore focus on the Frobenius and spectral spheres in the following sections.³

Compare Frobenius and oblique constraints

The matched radii are

Constraint	Required norm	Radius
Frobenius sphere	One norm for the matrix	$R_{fro} = r \sqrt{p}$
Output-oblique	One fixed norm for every row	$R_{out} = r$
Input-oblique	One fixed norm for every column	$R_{in} = r \sqrt{p / n}$

Each choice gives ${‖ W ‖}_{F}^{2} = r^{2} p$ . The first panel compares validation loss at each tested learning rate. Positive differences mean that the oblique variant has higher loss. The second panel compares maximum row norms during training.

(a)

Validation-loss difference from MACRO-fro

MACRO-oobl − MACRO-froMACRO-iobl − MACRO-fro

Validation-loss difference

Positive values mean that the oblique method has higher validation loss.

Select a point to see the loss values.

(b)

Maximum row norm during training

MACRO-froMACRO-iobl

Maximum row norm

0.0010.0030.0050.010.03

Learning rate

Select a curve to highlight it.

Figure 0.2. (a) Validation-loss differences from MACRO-fro. (b) Maximum row norms for MACRO-fro and MACRO-iobl during training.

Neither oblique variant reduces validation loss in this comparison. MACRO-fro and MACRO-iobl also show similar maximum row norms during training.³

With the constraints and radii in place, we can return to the question from the introduction. RMSNorm already rescales activations. What changes when we also fix weight scale?

Section 1

Manifold Constraints and RMSNorm

The radius analysis in Section 0 links weight scale to the output scale of a linear layer. Weight constraints control weight scale through the radius. RMSNorm can adjust activation scale through its learned parameter $γ$ . Both mechanisms can therefore affect activation scale, but through different parts of the model.¹

For an input vector $x$ , RMSNorm applies a learned scale $γ$ :

RMSNorm (x) = γ ⊙ \frac{x}{RMS (x)} .

The symbol $⊙$ denotes element-wise multiplication. RMSNorm divides the input by its RMS value, then scales each output coordinate with $γ$ .

If we increase the constraint radius, does RMSNorm adjust in response? Figure 1.1 tracks ${‖ γ_{t} ‖}_{2}$ across different constraint radii, where $t$ is the training step.

Figure 1.1 · RMSNorm scale

Radius setting

RMSNorm scale norm ‖γₜ‖₂

Compare the curves at the same training step.

Figure 1.1. The norm of the learned RMSNorm scale during training, for different constraint radii.

Larger radii correspond to smaller learned values of ${‖ γ_{t} ‖}_{2}$ for both constraints. In other words, RMSNorm adjusts its learned scale in the opposite direction. This compensation can reduce the effect of weight-scale changes on activations.¹

The two mechanisms therefore have overlapping roles. We next test the constraints after we remove the standard learnable RMSNorm layers.

Do constraints help without the standard learnable RMSNorm layers?

We use a modified Qwen3-like 330M model for this test. We remove all learnable RMSNorm layers. However, we keep RMS normalization after some nonlinear blocks, including attention blocks. We also keep QK-norm. Weight constraints control the activation scale only indirectly.

We replace SwiGLU with Norm-Gated SwiGLU. It normalizes the Swish branch before the element-wise product. This keeps the output scale of the same order as the output scale of ReLU. The following table shows the ablation architecture.²

Component	Standard model	Modified model
Standard learnable RMSNorm layers	Present	Removed
Feed-forward activation	SwiGLU	Norm-Gated SwiGLU
Parameter-free RMSNorm after attention	Absent	Present

Norm-Gated SwiGLU

Let $X$ be the input, and let $W_{1}, W_{2}$ be the two projection matrices. Define the two branches:

A = Swish (X W_{1}^{⊤}), B = X W_{2}^{⊤} .

Norm-Gated SwiGLU computes

Y = \frac{A}{{‖ A ‖}_{RMS}} ⊙ B .

The division controls the RMS scale of the Swish branch before multiplication.²

Figure 1.2 compares final validation loss across learning rates from $0.005$ . Lower loss is better.

Figure 1.2 · Modified 330M model

MuonMACRO-froMACRO-spec

Validation loss

Lower loss is better.

Figure 1.2. Final validation loss for the modified 330M model. Muon diverges at a learning rate of

0.03

Muon diverges at $η = 0.03$ . Both MACRO variants remain finite at all five learning rates. The lowest losses are 2.813 for Muon, 2.758 for MACRO-fro, and 2.739 for MACRO-spec.³

These results show that weight constraints improve training stability and validation loss in the ablation architecture.

How sensitive is training to the radius?

Weight constraints help stabilize training, but our modified model still uses parameter-free normalization. The two mechanisms have overlapping roles, rather than one simply replacing the other.²

That overlap may offer another benefit when both mechanisms are present. Recall Figure 1.1: larger radii corresponded to smaller learned RMSNorm scales. This compensation suggests that training may be less sensitive to the radius.¹

To test this, we return to standard Qwen3 models with learnable RMSNorm. We vary the radius multiplier $r$ in the 240M and 600M models. Each method keeps the same learning rate across radius settings.

Figure 1.3 compares final validation perplexity (PPL) at $r = 0.5, 1, 2, 5$ . Lower PPL is better.⁴

Figure 1.3 · Standard architecture

MACRO-fro · Learning rate 0.016MACRO-spec · Learning rate 0.008

Validation perplexity (PPL)

Lower PPL is better.

Figure 1.3. Final validation PPL across radius multipliers in the standard architecture with learnable RMSNorm.

At 600M, PPL changes little over $r \in [1, 5]$ . MACRO-fro ranges from 15.296 to 15.329. MACRO-spec ranges from 15.277 to 15.312. At 240M, MACRO-spec changes more than MACRO-fro across the same range.⁴

This result agrees with the compensation in Figure 1.1. Learned RMSNorm scales can offset changes in weight scale. This helps explain why a range of radii gives similar performance.

Takeaway

Manifold constraints and RMSNorm have aligned roles: constraints improve stability in the ablation architecture, while performance with standard RMSNorm is relatively insensitive to the radius multiplier.

Section 2

Manifold Constraints and Weight Decay

RMSNorm is not the only part of the training recipe that controls scale. Weight decay acts directly on the weights. If we now fix their norm, what is left for weight decay to do?

How does weight decay affect rotation?

To see the connection, consider one weight matrix as a vector $w_{t}$ . Let $s_{t}$ be its optimizer increment, including the learning rate but not weight decay. The update is

w_{t + 1} = (1 - η λ) w_{t} + s_{t},

where $η$ is the learning rate and $λ$ is the decay coefficient.

The increment can change both the size and direction of the weights. We split it into a radial component $s_{t}^{∥}$ , parallel to $w_{t}$ , and a tangential component $s_{t}^{⊥}$ , perpendicular to it. Weight decay combines with the radial component. Figure 2.1 shows the same step before and after this split.

Figure 2.1 · Radial and tangential updates

(a) Optimizer update and decay

Add the optimizer update and the weight-decay update.

(b) Radial and tangential components

Combine the radial component with weight decay. Keep the tangential component unchanged.

Figure 2.1. Two views of the same step. Weight decay combines with the radial component. The arc marks the initial norm.

The tangential component turns the weight vector. A finite tangential step can also increase its norm. For a small total step, the rotation angle in radians is approximately

θ_{t} \approx \frac{{‖ s_{t}^{⊥} ‖}_{2}}{{‖ w_{t} ‖}_{2}} .

The same tangential update gives a larger angle when the weight norm is smaller. So weight scale also affects how far the weights turn.

What happens to that scale over many steps? In fixed-learning-rate experiments, Kosson et al. observed stable weight norms and average rotation after an initial transient.²

A simple random-update model illustrates how this balance can form. Under its assumptions, the expected squared weight norm approaches a fixed level. This is a balance that develops during training, not a norm enforced after every step.

A simple model of norm equilibrium

We use a random-update model to explain the balance. The model uses a fixed learning rate and a fixed decay coefficient.

Let $s_{t}$ be the update before weight decay. The weight update is

w_{t + 1} = b w_{t} + s_{t}, b = 1 - η λ, 0 < b < 1 .

The model uses two conditions at every step:

E ⟨ w_{t}, s_{t} ⟩ = 0, E {‖ s_{t} ‖}_{2}^{2} = S^{2} < \infty .

The expectation is over different sequences of random updates. The first condition sets the expected inner product to zero. Individual updates need not be perpendicular to the weight.

The second condition keeps the expected value of ${‖ s_{t} ‖}_{2}^{2}$ constant across steps. Individual update sizes can still change.

Let $Q_{t} = E {‖ w_{t} ‖}_{2}^{2}$ , with $Q_{0} < \infty$ . The update equation gives

\begin{matrix} Q_{t + 1} & = b^{2} Q_{t} + 2 b E ⟨ w_{t}, s_{t} ⟩ + E {‖ s_{t} ‖}_{2}^{2} \\ = b^{2} Q_{t} + S^{2} . \end{matrix}

Thus,

Q_{t + 1} - Q_{t} = \underset{increase from updates}{\underset{⏟}{S^{2}}} - \underset{decrease from decay}{\underset{⏟}{(1 - b^{2}) Q_{t}}} .

The balance point is

Q_{*} = \frac{S^{2}}{1 - b^{2}} .

When $Q_{t} < Q_{*}$ , the increase is more than the decrease. When $Q_{t} > Q_{*}$ , the decrease is more than the increase.

The difference from the balance point satisfies

Q_{t} - Q_{*} = b^{2 t} (Q_{0} - Q_{*}) .

Because $0 < b < 1$ , the absolute difference decreases to zero. This result concerns the expected squared norm. The norm can still change at each step.

MACRO takes a different route: it fixes the required norm after each update. Recall from Section 0 that it also scales the proposed update with this norm. Together, these choices give us a more direct way to control rotation.

How does the Frobenius constraint control rotation?

Start with MACRO-fro. It keeps ${‖ W_{t} ‖}_{F} = R$ and scales the proposed step to approximately $η_{t} R$ in the same norm. Retraction restores the radius without changing the candidate's direction.

We measure the angle between consecutive weight matrices, treating each matrix as a vector. For a small step that is nearly perpendicular to $W_{t}$ ,

θ_{t}^{fro} \approx \frac{η_{t} R}{R} = η_{t} .

The radius cancels. Under these conditions, a fixed step size gives a nearly fixed angle.³ Figure 2.2 checks this prediction across different fixed step sizes.

Figure 2.2 · Frobenius constraint

Whole-matrix rotation

Learning rate

Rotation angle (rad)

Select a curve to inspect its rotation angle.

Figure 2.2. Whole-matrix rotation under the Frobenius constraint, at different fixed step sizes.

The reported angles remain nearly constant at each setting. Larger step sizes give larger angles. The paper calls this static global rotation: the whole-matrix angle stays nearly steady in these runs.³

The exact Frobenius angle

Let ${‖ W_{t} ‖}_{F} = R$ . The whole-matrix angle is

θ_{t}^{fro} = \arccos (\frac{{⟨ W_{t + 1}, W_{t} ⟩}_{F}}{{‖ W_{t + 1} ‖}_{F} {‖ W_{t} ‖}_{F}}) .

Let $O_{t}$ be the matrix-sign direction from Section 0. Define its alignment with the weights:

α_{t} = \frac{{⟨ W_{t}, O_{t} ⟩}_{F}}{R {‖ O_{t} ‖}_{F}} .

For the normalized update without the small numerical stabilizer, Appendix C.3 gives

\cos θ_{t}^{fro} = \frac{1 - η_{t} α_{t}}{\sqrt{1 - 2 η_{t} α_{t} + η_{t}^{2}}} .

For small $η_{t}$ ,

θ_{t}^{fro} \approx η_{t} \sqrt{1 - α_{t}^{2}} .

When $| α_{t} |$ is also small, this reduces to $θ_{t}^{fro} \approx η_{t}$ . The approximation therefore uses both a small step and small alignment.³

How does the spectral constraint change rotation?

Does fixing the largest singular value give the same kind of control? The strongest stretch is fixed, but the input and output directions associated with it can still change.

Here, we track the leading right singular vector $v_{t}$ , the input direction that the matrix stretches most. The leading left singular vector $u_{t}$ gives the corresponding output direction. We measure how these directions change between steps, rather than the whole-matrix angle used above.

The spectral gap separates the largest and second singular values. A smaller gap permits larger changes in the leading directions for the same step size. Their rotation can therefore depend on the changing spectrum, not only on the learning rate.⁴

Figure 2.3 follows these input and output directions during training.

Figure 2.3 · Spectral constraint

Leading singular-direction rotation

Learning rate

Left singular direction u

Principal angle (rad)

Right singular direction v

Principal angle (rad)

Compare the left and right angles at the same setting.

Figure 2.3. Rotation of the leading left and right singular directions under the spectral constraint.

The measured angles vary during training, with similar trends on the left and right. The paper calls this adaptive subspace rotation.⁴ Fixing a norm does not give both constraints the same rotation behavior.

How we measure spectral rotation

Let $u_{t}$ and $v_{t}$ be the leading left and right singular vectors of $W_{t}$ . The spectral constraint gives

W_{t} v_{t} = R u_{t} .

We measure their changes by

θ_{u, t} = \arccos (| u_{t + 1}^{⊤} u_{t} |), θ_{v, t} = \arccos (| v_{t + 1}^{⊤} v_{t} |) .

The absolute values treat opposite signs as the same singular direction. These angles are different from the whole-matrix angle in Figure 2.2.

Set $θ_{t}^{spec} = max {θ_{u, t}, θ_{v, t}}$ . For small angles, the paper gives the estimate

θ_{t}^{spec} ≲ η_{t} \frac{R}{R - σ_{2} (W_{t + 1})} .

Here, $σ_{2}$ is the second singular value. The denominator is the positive spectral gap used in this estimate. A smaller gap permits a larger rotation for the same step size. This is an estimate, not an equality.⁴

Does weight decay still improve training?

We can now return to the question at the start. Once the constraint controls weight scale, does weight decay still add a benefit?

We compare $λ = 0$ and $λ = 0.1$ on a 330M Qwen3-like model. Each pair uses the same learning rate and all other settings. The table reports validation loss.

Table 2.1. Each cell shows loss without decay / loss with decay.

Optimizer	$η = 0.007$	$η = 0.010$	$η = 0.030$
MACRO-fro	2.742 / 2.742	2.730 / 2.729	2.718 / 2.718
MACRO-spec	2.728 / 2.727	2.715 / 2.715	2.783 / 2.783

The paired losses differ by at most $0.001$ at the reported precision. Additional weight decay gives little benefit in these tests.⁵

Other constrained optimizers

The paper also tests MuonH-fro and MuonH-spec with the same decay settings.

Optimizer	$η = 0.007$	$η = 0.010$	$η = 0.030$
MuonH-fro	2.740 / 2.740	2.722 / 2.722	2.717 / 2.717
MuonH-spec	2.728 / 2.729	2.715 / 2.715	2.766 / 2.766

These paired losses also differ by at most $0.001$ .⁵

Weight-decay tuning for Muon

A separate 240M experiment varies both the learning rate and weight decay for Muon. The table reports validation perplexity (PPL).

Learning rate	$λ = 0.01$	$λ = 0.05$	$λ = 0.10$	$λ = 0.50$
0.008	21.348	21.020	20.845	21.168
0.016	20.972	20.680	20.676	21.874
0.032	20.691	20.583	20.848	23.345

The best tested decay coefficient changes with the learning rate. The lowest Muon PPL is 20.583. MACRO-fro reaches 20.579 without decay on the constrained matrices. MACRO-spec reaches 20.502.⁶

Takeaway

MACRO fixes weight scale after each update. Frobenius constraints give nearly steady whole-matrix angles in the fixed-step runs. Spectral constraints allow changing leading-direction angles. Additional weight decay gives little benefit in these tests.

So far, we have asked how far the weights turn. Next, we ask whether a small proposed update changes the stored weight at all.

Section 3

Manifold Constraints and Low-Precision Training

We have looked at how far an update turns the weights. But that assumes the update changes the stored values. With limited precision, rounding can cancel a small update before it takes effect.

A common mixed-precision setup uses low precision for much of the forward and backward passes. It still keeps FP32 master weights—the parameter values that receive optimizer updates. This can speed up computation and reduce activation memory.¹

We go a step further: we store the master weights themselves in BF16, with no extra FP32 copy. This saves additional weight-storage memory, but each updated weight must be rounded back to BF16. Here, we use round-to-nearest.²

Start with one stored weight $w$ and a proposed update $δ$ . The optimizer proposes $w + δ$ , but BF16 stores only certain values. Rounding selects the nearest one. If the proposed value stays inside the original rounding interval, the stored weight does not change. We call this update cancellation.

At $w = 1.5$ , adding $0.002$ rounds back to $1.5$ . The update is nonzero; the stored change is zero.²

Figure 3.1 · BF16 rounding

Current weight wProposed value w + δStored weight w⁺

Proposed update δ +0.002

Proposed value 1.502Stored weight 1.5Actual change 0

Rounding cancels this update. Local ULP: 0.0078125.

Figure 3.1. A small update can leave the stored weight unchanged. Dashed lines mark the rounding boundaries.

Now keep the update at $0.004$ and change the weight. At $w = 1.5$ , the stored value changes to $1.5078125$ . At $w = 3$ , the same update causes no change because the gap between nearby BF16 values is larger.

The update is the same. Only the weight scale changed. So update size alone is not enough: we need to compare it with weight size.

Rounding thresholds

Let $Q_{BF16}$ denote rounding to the nearest BF16 value. The stored result is

w^{+} = Q_{BF16} (w + δ) .

The local gap between neighboring values is a unit in the last place (ULP). If the update is smaller than half the gap in its direction, rounding keeps the original weight.

For the values in Figure 3.1, this condition is

\frac{| δ |}{| w |} < \frac{ULP (w)}{2 | w |} .

At $w = 1.5$ , the local gap is $0.0078125$ . Half that gap is $0.00390625$ , so the update $0.002$ is canceled. At an exact midpoint, the example uses round-to-nearest, ties-to-even.²

The relative threshold for normal BF16 values is about $0.2 %$ to $0.4 %$ . The example compares neighboring values directly, including the direction of the update. This is a coordinate-level rule, not a fixed cancellation threshold for a whole matrix.

How do constraints control the relative update size?

For one coordinate, the relative size is $| δ | / | w |$ . Recall Step 4: Rescale in Section 0. It controls the proposed step size relative to the weight scale. For a matrix, we measure this relationship with the update-to-weight ratio:

ρ_{t} = \frac{{‖ {Δ W}_{t} ‖}_{F}}{{‖ W_{t} ‖}_{F}} .

Here, $Δ W_{t}$ is the proposed update before retraction and rounding. For Muon, it excludes weight decay. The ratio also equals the ratio of their RMS norms, so it compares typical entry sizes.²

MACRO-fro fixes the weight norm at $R$ and scales the proposed step with the same radius. For the idealized normalized update,

ρ_{t} = \frac{η_{t} R}{R} = η_{t} .

The radius cancels. The relative update follows the learning rate rather than weight growth.

Muon does not tie the update size to the weight norm in this way. If its weights grow without a matching increase in update size, $ρ_{t}$ falls. MACRO-spec also controls relative scale, but its ratio depends on the weight spectrum.²

Ratios for the three optimizers

For a matrix $W_{t} \in R^{p \times n}$ , let $d = \min (p, n)$ . The Frobenius and RMS ratios are equal:

ρ_{t} = \frac{{‖ {Δ W}_{t} ‖}_{F}}{{‖ W_{t} ‖}_{F}} = \frac{{‖ {Δ W}_{t} ‖}_{RMS}}{{‖ W_{t} ‖}_{RMS}} .

For ideal full-rank matrix-sign directions, the three ratios are

ρ_{t} = {\begin{cases} \frac{{\tilde{η}}_{t} \sqrt{d}}{{‖ W_{t} ‖}_{F}} & Muon, \\ η_{t} & MACRO-fro, \\ η_{t} \sqrt{κ (W_{t})} & MACRO-spec. \end{cases}

The effective Muon coefficient ${\tilde{η}}_{t}$ includes its update scaling. The Muon update excludes weight decay. Define the spectral factor

κ (W_{t}) = \frac{d {‖ W_{t} ‖}_{2}^{2}}{{‖ W_{t} ‖}_{F}^{2}} .

This factor describes the weight spectrum. The formulas use exact norm constraints and omit the small numerical stabilizer.²

For Figure 3.2, Muon ratios use recorded BF16 weight norms and an idealized update estimate. They are estimates of proposed updates, not measurements of the final stored change. The weight norm is measured; the proposed-update ratio is estimated.

This gives us two changes to watch: weights can grow, or updates can shrink. Either can make a proposed update small relative to its weight.

How does weight growth affect Muon?

Does this scale effect appear in training? Figure 3.2 compares weight norms, estimated relative updates, and the loss gap over the same training period.

What about weight decay from Section 2? Its shrinkage can also be too small to change a stored BF16 weight.³

The upper plot shows a selected matrix. The lower plot shows the training-loss gap, defined as BF16 loss minus FP32 loss. A positive gap means that BF16 training has higher loss. Compare Weight norm and Update-to-weight ratio with the gap below.

Figure 3.2 · Muon during training

Show

Muon matrix · layer 0

MuonMACRO-froMACRO-spec

Solid: BF16. Dashed: FP32.

Weight norm ‖Wₜ‖F

Measured

Training-loss gap · BF16 − FP32

Selected step 9,200

Step 9,200 · 9.647B tokensMuon · Loss gap +1.1307Weight norm: BF16 138.2 · FP32 11.42

Move across either plot to compare the same training step. Click to lock the cursor. Use the slider on a touch screen.

Chart settings and values

Choose the first and last steps. Select the scale option to fit the loss axis to the selected range.

First stepLast step

Fit the gap axis to the visible time range

Hold Shift and drag to zoom.

Figure 3.2. Muon weight scale and the BF16 training-loss gap. The ratio view uses an estimate of the proposed update norm.

In the key-projection example, the BF16 weight norm rises early and becomes much larger than the FP32 norm. Muon also develops a large training-loss gap early. The smaller relative update provides a scale-based explanation for this gap.³

But weight growth is only one way to reduce the ratio. Even with a fixed norm, the learning rate can still fall.

Why does the MACRO loss gap increase late in training?

Fixing the weight norm does not fix the relative update throughout training. Recall the idealized MACRO-fro relation $ρ_{t} = η_{t}$ . Cosine decay lowers the learning rate, so it also lowers $ρ_{t}$ .

How much update magnitude might be lost as this ratio falls? We use a simple Gaussian model to connect relative scale with cancellation. The model uses the same ratio $ρ_{t}$ and estimates the fraction of absolute-update magnitude that remains.²

This fraction is update retention. It measures how much proposed magnitude survives the model's cancellation rule, not how many coordinates survive. As $ρ_{t}$ falls, the model retains less of the proposed magnitude.

Inside the cancellation model

We model a weight coordinate and its proposed update as independent Gaussian variables:

w \sim N (0, σ_{w}^{2}), δ \sim N (0, σ_{δ}^{2}) .

Their means are zero. We set their standard deviations to the corresponding matrix RMS norms. This gives the same relative scale as the training analysis:

\frac{σ_{δ}}{σ_{w}} = \frac{{‖ {Δ W}_{t} ‖}_{RMS}}{{‖ W_{t} ‖}_{RMS}} = ρ_{t} .

The model cancels an update when $| δ | \leq τ | w |$ . It leaves the other updates unchanged:

δ_{τ} = δ 1 {| δ | > τ | w |} .

We use $τ = 1 / 510 \approx 0.002$ as a BF16 reference. This threshold rule models cancellation, not every effect of BF16 rounding.²

The retained absolute-update fraction is

R (ρ) = \frac{E | δ_{τ} |}{E | δ |} = \frac{1}{\sqrt{1 + {(τ / ρ)}^{2}}} .

For example, $R = 0.8$ means that 80% of the proposed absolute-update magnitude remains on average. This is not the fraction of coordinates that survive. The fraction of canceled coordinates is

p_{drop} (ρ) = \frac{2}{π} \arctan (\frac{τ}{ρ}) .

At $ρ = τ$ , the model cancels half the coordinates but retains about 70.7% of the absolute-update magnitude. The canceled coordinates contribute less magnitude on average.

For the idealized MACRO-fro relation $ρ_{t} = η_{t}$ ,

R_{t}^{fro} = {[1 + {(\frac{τ}{η_{t}})}^{2}]}^{- 1 / 2} .

When $η_{t}$ is much larger than $τ$ , the model retains most of the update magnitude. As $η_{t}$ falls, retention decreases. Thus, fixed weight scale does not remove the effect of learning-rate decay.

Figure 3.3 uses this relation for MACRO-fro. It compares the estimate with measured loss gaps; it does not predict a numerical loss gap. The figure has no MACRO-spec retention curve because matched spectral-factor records are unavailable.

Figure 3.3 applies the model to the MACRO-fro learning-rate schedule. The upper plot shows estimated retention. The lower plot shows measured training-loss gaps for both MACRO variants. Select Late training to compare the two patterns.

Figure 3.3 · Cancellation and the MACRO loss gap

MACRO-froMACRO-spec

Retention measures absolute-update magnitude. Cancellation counts coordinates.

MACRO-fro · retained absolute-update fraction

Model estimate

Loss-gap scale

Symlog shows zero and negative gaps, with a linear region from −0.001 to +0.001.

MACRO training-loss gap · BF16 − FP32

Selected step 9,000

Step 9,000MACRO-fro · Loss gap +0.0165MACRO-spec · Loss gap +0.0310Model for MACRO-fro: retained magnitude 9.29% · canceled coordinates 94.07%

Chart settings and values

First stepLast step

Fit loss-gap scale to the visible MACRO range

Figure 3.3. Estimated update retention for MACRO-fro and measured training-loss gaps for both MACRO variants. The two plots show the same training period.

After warmup, the MACRO loss gaps stay small through most of training, then increase near the end. For MACRO-fro, this period also has lower estimated retention. The model helps explain why a small gap can remain even with controlled weight scale.²

The constraint does not change BF16 precision. As the learning rate falls, some updates can still become too small.

How much does this matter in the final model? The table compares final validation loss between each method's BF16 and FP32 runs:

{Δ L}_{val} = L_{val}^{BF16} - L_{val}^{FP32} .

Optimizer	240M	600M
Muon	+1.1470	+0.9394
MACRO-fro	+0.0172	+0.0239
MACRO-spec	+0.0317	+0.0476

Both MACRO variants have a much smaller BF16 penalty in these comparisons.⁴

Takeaway

Weight scale affects whether small updates survive rounding. MACRO controls the update-to-weight ratio and shows much less BF16 degradation than Muon in these tests. Learning-rate decay can still leave small updates at risk late in training.

Section 4

Standard Pre-training Results

A smaller BF16 penalty is useful. But does MACRO also improve standard pretraining, and what does it cost?

We started this blog by asking whether a different training recipe could get more from pretraining. We now return to that practical question.

How does MACRO perform across model sizes?

We train Qwen3 models with 240M, 600M, and 1.2B parameters on C4. Each configuration uses 40 training tokens per parameter and five random seeds. Figure 4.1 compares final validation perplexity (PPL); lower is better.¹

Figure 4.1 · Standard pre-training

Final validation perplexity

Lower is better. Select a model size.

Figure 4.1. Final validation PPL for standard Qwen3 pre-training. The figure preserves the reported ± values.

Both MACRO variants have lower PPL than Muon at all three tested sizes. The best constrained method varies with model size.¹

Benchmark details

The standard benchmark uses Qwen3 models on C4, with TorchTitan as the training framework. Each configuration uses 40 tokens per parameter and five random seeds.¹

Figure 4.1 retains the reported means and ± values. It compares Muon, MuonH-fro, MuonH-spec, FSO, SSO, MACRO-fro, and MACRO-spec where results are available.

MACRO-spec has the lowest mean PPL at 240M. MuonH-fro has the lowest mean at 600M. MACRO-fro has the lowest mean at 1.2B.

FSO and SSO are not reported at 1.2B. The source attributes this omission to their high per-step costs. Missing results remain labeled Not reported.

The 240M efficiency test uses plain implementations without specialized kernels. Metrics are averaged over training steps and five seeds.²

End-to-end step time measures seconds per training step. TPS/GPU reports token throughput per GPU. MFU means model FLOPs utilization. Lower step time is better; higher throughput and MFU are better.

What is the per-step cost?

In Section 0, we chose a single-loop optimizer to keep the testbed practical. What does that choice cost per training step?

Figure 4.2 compares the 240M model using plain implementations without specialized kernels. Start with end-to-end step time. The other views show token throughput and model FLOPs utilization.²

Figure 4.2 · 240M efficiency

Single-loop and double-loop cost

Lower is better.

Plain implementations · no specialized kernels · average over steps and five seeds

Figure 4.2. Efficiency for the 240M model. Lower E2E time is better. Higher throughput and MFU are better.

MACRO-fro takes $0.6411$ seconds per step, compared with $0.6311$ for Muon. That is about $10$ ms more, or $1.6 %$ . MACRO-spec costs more because its tangent projection and retraction use power iteration. Both MACRO variants remain much faster than the double-loop solvers in this test.²

Does MACRO require learning-rate warmup?

We have already revisited weight decay. One more part of the training recipe is worth checking: how much warmup do these methods need?

Figure 4.3 compares four warmup fractions in the 240M and 600M models. It reports final validation PPL.³

Figure 4.3 · Additional training result

Warmup sensitivity

Lower PPL is better. Lines connect only the tested settings.

MuonMACRO-froMACRO-spec○ Best tested

Final validation PPL

Figure 4.3. Final validation PPL for different warmup fractions. The 600M values at 10% warmup come from author slide 44. Paper Table 12 leaves those entries unreported.

Both MACRO variants give their lowest PPL at $0 %$ warmup in these sweeps. Muon does best at $5 %$ for 240M. At 600M, its results at $0 %$ and $5 %$ are nearly equal. MACRO can therefore train without warmup in the tested configurations.³

Takeaway

A fixed weight norm does more than limit weight magnitude. It interacts with activation normalization, regulates rotation, and controls the relative scale of low-precision updates. These interactions help explain the benefits of adding constraints, beyond a lower final validation loss.