Compositional Muon — partner-whitened optimizer that controls composed functional updates in attention

TL;DR

Compositional Muon (CM) is a Muon-family optimizer from Tilde Research that controls the update to composed functional objects — specifically the $M_{QK} = W_Q W_K^\top$ and $M_{OV} = W_V W_O$ products inside attention — rather than to each weight matrix independently. Solving the steepest-descent problem under an operator-norm constraint on the composed product gives partner-whitened update rules: each factor’s step is dynamically rescaled by the inverse Gram root of its partner, so heads whose partner has large singular values get smaller steps and vice versa. CM beats Muon at 340M and 1B/70B-token scale, and a single coupled OV update sets a new modded-nanoGPT speedrun record (2875 vs 2890 baseline).

Key claims

Modern optimizers control trust regions on individual matrices (Adam: per-entry; Muon: per-matrix spectral norm), but the loss interacts with $W_Q$ and $W_K$ only through the product $M_{QK} = W_Q W_K^\top$ — this “compositional gap” means the trust region is mis-located relative to where the functional change actually lives [thread §1–2].
Bounding factor updates separately bounds $\|\Delta M\|_{\text{op}}$ only through partner norms ( $\|\Delta M\|_{\text{op}} \le \|\Delta W_Q\|_{\text{op}}\|W_K\|_{\text{op}} + \|W_Q\|_{\text{op}}\|\Delta W_K\|_{\text{op}}$ ), so a unit-norm factor update on a head with a large-spectrum partner causes a much larger functional change than on a head with a small-spectrum partner — and vanilla Muon has no way to know this [thread §3].
The compositional steepest-descent problem doesn’t decompose cleanly (the constraint couples the two factors). CM solves it via a “half-split” surrogate that allocates the budget equally between the two sides and solves each subproblem in closed form using the polar decomposition of the partner [thread §4–5].
The resulting rule is partner-whitened: $\Delta W_Q = -(\epsilon/2)\,\mathrm{msign}(G_Q C_K^{-1})\,C_K^{-1}$ and symmetrically for $W_K$ , where $C_K = (W_K^\top W_K)^{1/2}$ . When the partner has large singular values, the update is suppressed; when small, amplified [thread §6].
The OV branch needs a hybrid rule because V-side perturbations hit raw tokens per-head (worst case can fire independently for each head, so V is treated per-head through its partner $W_{O,h}$ ) while O-side perturbations sum into the shared residual stream (worst case is a single global operator, so O is orthogonalized globally after per-head scalar whitening) [thread §7].
Full partner-whitening requires matrix square roots per head per step — expensive. Under an isotropic Gram approximation (singular values tightly clustered, which holds empirically under Muon + weight decay), $\mathrm{msign}$ ‘s scale-invariance makes the whitening factor inside collapse to a scalar, and the update reduces to standard Muon scaled by $c_{O,h}^{-1} = (\|W_{O,h}\|_F^2 / d_h)^{-1/2}$ — i.e. Muon with a dynamic per-head learning rate set by the partner’s Frobenius norm [thread §8–9].
The isotropic approximation captures nearly all gains of full CM empirically [thread §9].
Product parameterizations have a gauge symmetry ( $W_Q \mapsto W_Q R$ , $W_K \mapsto W_K R^{-\top}$ leaves $M$ invariant), which is a genuinely new failure mode versus single-matrix Muon: along gauge orbits the optimizer can enter a feedback loop where one factor shrinks and the other grows, making isotropic CM’s partner-Frobenius scaling imbalanced. A scalar gauge correction restores balance, though empirically the gains are small [thread §10–12].
At 340M small-scale pretraining, CM-on-both-QK-and-OV gives the largest gains in the cleanest configuration (no RoPE, no QK-Norm). With RoPE alone, OV-only CM still wins and QK-CM helps further. With QK-Norm + RoPE, OV-CM still wins but QK-CM gains shrink — consistent with QK-Norm partially absorbing the same control [thread §13, Image 3].
CM directly regulates per-head composed-update spectra: tracked $\|\Delta M_{OV,h}\|$ and $\|\Delta M_{QK,h}\|$ stay tightly controlled across heads, while under vanilla Muon the cross-head coefficient of variation is large — direct evidence the optimizer is regulating the functional step rather than just the factor step [thread §14].
At 1B parameters / 70B tokens, both full-whitening CM and isotropic CM separate from Muon early in training and stay separated through the end; the gap is larger without QK-Norm [thread §15, Image 4].
modded-nanoGPT speedrun: adding a single coupled OV update to the baseline reduces steps-to-threshold from 2890 to 2875 — new record on track 3 [thread §16, Image 5].
The compositional rescaling idea transfers to AdamW (partner-rescaled AdamW improves over AdamW), but gains are smaller than the Muon version — suggesting part of the benefit is better per-head LR allocation and part is interaction with Muon’s spectral-sign step specifically [thread §17].
CM is positioned as a new design axis (compositional viewpoint) orthogonal to the base optimizer — “compositional Shampoo” is suggested as a natural next direction. Also generalizes to RMSNorm scales, gated MLPs, and any structured composition in the network [thread §18].

Method

The thread frames optimizer design as choosing where to place the trust region. Adam puts it on individual entries; Muon on individual matrices (under the spectral norm); CM on composed functional objects (matrix products that the loss actually sees). For the QK product, the compositional steepest-descent problem is

$\min_{\Delta W_Q, \Delta W_K} \langle G_Q, \Delta W_Q \rangle + \langle G_K, \Delta W_K \rangle \quad \text{s.t.} \quad \|\Delta M\|_{\text{op}} \le \epsilon$

The constraint couples the two factors through $\Delta M = \Delta W_Q W_K^\top + W_Q \Delta W_K^\top$ . CM solves a “half-split” surrogate that splits the budget equally between the two sides and solves each independently. Using the polar decomposition $W_K = U_K C_K$ and the fact that right-multiplying by the orthogonal factor $U_K$ doesn’t change operator norms, each subproblem reduces to a partner-whitened spectral steepest descent that admits a closed form: $\Delta W_Q = -(\epsilon/2)\,\mathrm{msign}(G_Q C_K^{-1})\,C_K^{-1}$ .

For the OV branch, the V-side and O-side need different treatments because they hit different downstream structures (per-head raw tokens vs shared residual stream). V is treated per-head through partner $W_{O,h}$ ; O is orthogonalized globally across heads after per-head scalar whitening.

The full rule requires a matrix square root per partner per step (expensive). The practical version uses an isotropic Gram approximation: when partner singular values are tightly clustered (empirically the regime modern Muon + WD keeps you in), $C_K^{-1}$ collapses to a scalar inside $\mathrm{msign}$ (scale-invariant), and the update becomes standard Muon with a per-head learning rate set by the inverse partner Frobenius norm. A gauge correction ensures the per-head LRs don’t drift along gauge orbits.

Results

340M small-scale pretraining: CM beats Muon across architecture settings. Largest gain in the theory-cleanest configuration (no RoPE, no QK-Norm). With RoPE only, OV-CM still wins and QK-CM on top helps further. With QK-Norm + RoPE, OV-CM still improves over Muon; QK-CM gains shrink — consistent with QK-Norm acting on the same composition [thread §13, Image 3].
Composed-update stability (340M): tracked $\|\Delta M\|$ per head per layer is tightly controlled under CM, while vanilla Muon shows substantial cross-head variance in update norms within a layer [thread §14].
1B / 70B-token pretraining: both full-whitening CM and isotropic CM separate from Muon early and stay separated; gap larger without QK-Norm [thread §15, Image 4].
modded-nanoGPT speedrun (track 3): single coupled OV update → 2875 steps to threshold vs 2890 baseline — new record [thread §16, Image 5].
Partner-rescaled AdamW: improves over AdamW but by less than CM improves over Muon — partial portability of the compositional idea to non-Muon bases [thread §17].

Why it’s interesting

Compositional Muon is the third optimizer-side entry in the last two months from Tilde Research alone, and it occupies a distinct slot from Aurora: A Leverage-Aware Optimizer for Rectangular Matrices: Aurora locates Muon’s defect in tall MLP matrices (row-leverage anisotropy of the polar factor), while CM locates a defect in attention specifically (factor-level trust regions don’t see the QK and OV compositions). Both are pathology-first derivations — name a specific functional object the optimizer is mis-controlling, derive the correction — but they target non-overlapping pathologies, which raises a sharp open question for Training stability at scale: do CM (compositions) + Aurora (tall matrices) compose, or does each already saturate the same headroom?

The framing also sharpens the running disagreement between How to Set the Learning Rate for Large-Scale Pre-training? (QK-Norm makes the µP motivation stale) and Controlled LLM Training on Spectral Sphere (the optimizer is what’s missing). CM’s empirical result that QK-Norm partially absorbs QK-side CM gains but does not absorb OV-side CM gains is direct evidence that QK-Norm is locus-specific — it controls the same thing CM-QK controls but not the same thing CM-OV controls, so “normalization vs optimizer” is not a binary choice but a per-locus tradeoff.

The single coupled OV update beating the modded-nanoGPT record by ~15 steps is small in absolute terms, but it’s done with a one-line change (replace OV’s Muon step with the isotropic partner-rescaled variant). That’s the cheapest CM ablation to run on any existing training stack, and the bar to try it is low — particularly because the isotropic approximation is just “scale Muon’s per-head OV update by inverse partner Frobenius norm.” The paper and code links in the thread are empty placeholders at filing time, so this filing is the thread itself; a paper/repo follow-up should append pointers here when it lands.