MUD: Faster Transformer Training via Triangular Momentum Whitening

Paper: 2603.17970 Authors: Ben S. Southworth, Stephen Thomas Categories: cs.LG, math.NA, math.OC

The Gap

AdamW is the default optimizer for transformer training, but it treats each parameter independently — it has no idea that weight matrices have rows and columns that interact. Muon (2024) fixed this by orthogonalizing the momentum matrix before applying the update, which acts like a cheap second-order method. The catch: Muon’s orthogonalization uses a polar decomposition approximation (Newton-Schulz iterations), which requires several large matrix multiplications per step. On modern hardware, those multiplications are not free — they can eat30–60% of your step time, and the overhead scales badly with model size and batch configuration.

So the gap is: we want the statistical benefits of momentum whitening (faster convergence per step) without paying the compute cost of polar decomposition. Nobody had asked whether a cheaper triangular factorization could do the same job.

Problem: AdamW ignores matrix structure -> slow convergence
    |
    v
Prior fix (Muon): polar decomposition whitening
    |
    * Bottleneck: multiple large matmuls per step
    * Overhead: 30-60% of step time, hardware-dependent
    |
    v
Gap: Can we whiten momentum cheaply without polar decomp?
    |
    v
Assumption: Triangular (Cholesky-like) factors are enough
    |
    v
Method (MUD): Gauss-Seidel iteration on Gram matrix
    |
    v
Evidence: 1.3-3x tokens/s improvement over Muon
    |
    v
Conclusion: Triangular whitening matches polar quality at lower cost

The Increment

One sentence: Before this paper, getting Muon-quality training meant paying Muon’s compute tax; after this paper, you can get most of the benefit with a triangular update that’s 1.3–3x cheaper on wall-clock time.

Core Mechanism

MUD’s core idea is to whiten the momentum matrix using a triangular factor instead of a full orthogonal one. Here’s what that means concretely. You have a momentum matrix M (same shape as your weight matrix). You want to “decorelate” its rows so that the update doesn’t waste steps on correlated directions. Muon does this by computing the polar factor $U = M(M^T M)^{-1/2}$, which requires iterating a matrix polynomial. MUD instead finds a lower-triangular matrix L such that L M has approximately orthonormal rows — meaning (LM)(LM)^T ≈ I.

The key operation is a Gauss-Seidel sweep over the Gram matrix G = MM^T. You update L row by row: each row of L is chosen to decorelate the corresponding row of M against all previously processed rows. This is exactly the structure of Gram-Schmidt orthogonalization, but done in a way that’s equivalent to one sweep of symmetric Gauss-Seidel preconditioning on G. The authors prove this converges quadratically near the fixed point (row-orthonormal matrices), which means once you’re close, you need very few inner iterations.

In practice, MUD runs 1–3 inner Gauss-Seidel sweps per optimizer step. Each sweep is O(n^2 d) for an n×d matrix, same asymptotic cost as Muon’s Newton-Schulz, but the constant is much smaller because triangular solves are memory-bandwidthfriendly and don’t require the same synchronization overhead as repeated full matrix multiplications.

Momentum M (n x d)
    |
    v
Gram matrix G = M^T   (n x n, cheap if n < d)
    |
    v
Gauss-Seidel sweep on G:
  for i = 1..n:
    L[i,i] = 1 / sqrt(G[i,i] - sum(L[i,j]^2 for j<i))
    L[i,j] = -(G[i,j] + sum(L[i,k]*L[k,j])) * L[i,i]
    |
    v
Whitened update W = L M   (rows approx orthonormal)
    |
    v
Apply W as gradient update (like Muon's polar factor)

Think of it like tuning a choir. Muon’s approach is to record all the singers simultaneously, compute a full SVD of the sound, and rotate everyone into perfect harmony — mathematically clean, but expensive. MUD’s approach is the choir director going down the row one singer at a time: “You, sing your note. You, adjust so you don’t clash with her. You, adjust so you don’t clash with either of them.” Each singer only needs to hear the people already seated. That sequential, triangular correction is Gram-Schmidt. It’s not as globally optimal as the SVD rotation, but it gets the choir close enough to in-tune, much faster, because you never have to hold the whole sound in memory at once.

The method components map directly: the choir = weight matrix rows, the director’s sequential correction = Gauss-Seidel sweep, “don’t clash” = decorelating the Gram matrix off-diagonals, “close enough to in-tune” = quadratic convergence near the orthonormal fixed point.

Key Concepts

• Momentum whitening: When you accumulate gradient momentum across steps, the momentum vector tends to develop correlations — some directions get reinforced together, so your update wastes capacity moving along corelated axes. Whitening transforms the momentum so its directions are uncorrelated and equally scaled, like rotating an ellipse into a circle. This is why second-order methods converge faster: they implicitly whiten the gradient. MUD does this cheaply for matrix-shaped parameters.

• Gram matrix and Gauss-Seidel: The Gram matrix G = MM^T captures all pairwise inner products between rows of M. If G = I, the rows are orthonormal — perfectly decorelated. Gauss-Seidel is a classical iterative solver that fixes one variable at a time while holding others fixed. Applied here, it means: zero out one off-diagonal entry of G at a time by applying a triangular correction. It’s the numerical linear algebra equivalent of fixing one equation at a time in a system of equations.

• Polar decomposition vs triangular factorization: Every matrix M can be written as M = UP where U has orthonormal rows and P is symmetric positive definite (polar decomposition). Muon extracts U. MUD instead finds L (lower triangular) such that LM has orthonormal rows — this is the QR decomposition from the left. Both achieve whitening; the difference is purely computational. Triangular factors are cheaper to compute iteratively because each step only touches a triangular subproblem.

Framework Shift

Before (Muon / polar decomposition):    After (MUD / triangular whitening):

  M --> [Newton-Schulz iterations]         M --> [Gram matrix G = MM^T]
         |                                        |
         | 5-10 full matmuls                | 1-3 Gauss-Seidel sweps
         |                                        | (triangular, sequential)
         v                                        v
  Polar factor U                L M  (approx row-orthonormal)
  (globally optimal orthogonalization)     (locally convergent, cheaper)
         |                        |
         v                                v
  High quality, high overhead              Similar quality, lower overhead
  ~30-60% of step time                     ~10-20% of step time

From polar decomposition to triangular factorization, the core shift is: you don’t need global optimality in the whitening step — local, sequential correction is good enough and much cheaper.

Expert Assessment

Problem choice: This is a real gap. Muon’s overhead is a genuine pain point that practitioners have noticed, and the question “can we whiten more cheaply?” is natural and well-posed. It sits at the intersection of numerical linear algebra and deep learning optimization, which is underexplored. Not a manufactured gap.

Method maturity: The insight is clever — connecting Gram-Schmidt to Gauss-Seidel preconditioning is a nice observation, and the quadratic convergence proof gives it theoretical grounding. That said, the method is essentially “do QR instead of polar decomp,” which is a known idea in numerical linear algebra. The novelty is in recognizing this works well enough for optimizer whitening and proving it. It’s not brute force, but it’s also not a paradigm shift — it’s a well-executed engineering insight with solid math backing.

Experimental integrity: The baselines look fair — they tune AdamW and Muon carefully and report wall-clock time rather than just step count, which is the right metric for this claim. The 1.3–3x improvement range is wide, which is honest (hardware and model size matter a lot). The ESM-2 protein LM experiment is a nice out-of-distribution test. One mild concern: the paper doesn’t deeply ablate the number of inner Gauss-Seidel sweps vs quality tradeoff across all settings, so it’s not fully clear when1 sweep is enough vs3.

Writing quality: The theoretical section (fixed points, convergence proof) is solid but dense. The weakest part is the experimental section’s narrative — the authors present numbers but don’t always explain *why MUD wins more on some architectures than others. A paragraph connecting model width/depth to the overhead ratio would make the results much more interpretable and actionable.

Verdict: weak accept — solid engineering contribution with honest experiments, but the novelty ceiling is “smart application of classical numerical methods” rather than a new idea.

Takeaways

• If you’re using Muon, try MUD as a drop-in replacement. The overhead reduction is real and the convergence quality is comparable. Worth a weekend experiment.
• The broader lesson: optimizer overhead is a first-class concern, not an afterthought. When evaluating optimizers, always measure wall-clock time-to-perplexity, not just steps-to-perplexity.
• The Gauss-Seidel framing of Gram-Schmidt is a useful mental model for anyone building custom preconditioners — sequential triangular corrections are often cheaper than global factorizations and converge quadratically near fixed points.
• For protein LMs and other non-standard transformer architectures, this paper gives evidence that momentum whitening generalizes beyond language modeling, which is useful if you’re working in those domains.