0. Towards Optimizing Compositions

Today, many of the most capable language models are trained with the Muon optimizer [2]. From Kimi K2 [20] to DeepSeek V4 [21], Muon has demonstrated strong performance at scale over the previous mainstream optimizer AdamW [22][8].

A key insight behind Muon is to treat matrices in neural networks as linear operators, instead of collections of independent scalars. In particular, while Adam treats every weight as its own optimization problem, updating each entry proportionally to its gradient and second moment, Muon's optimization problem directly utilizes the matrix structure of a parameter it is optimizing.

Muon’s update rule follows from formulating steepest descent for a single linear map, $y = Wx$. But many natural model operations are compositional: their behavior depends on how several pieces interact. In those cases, the object of interest is the composition itself, so we should formulate the optimization problem at the level of the composed operation rather than only at the level of its individual factors.

In this manner, we introduce Compositional Muon (CM) which extends Muon’s core principle to composed circuits. Instead of controlling the operator norm and trust region of each weight matrix separately, where applicable it controls the optimization of the composition. For the QK and OV products in attention, this leads to dynamic partner-whitened update rules that adapt step sizes for each matrix based on the geometry of its partner.

By designing around composed circuits, CM improves the spectral behavior of these updates and suggests a rich design space for compositional optimizers.

1. The Compositional Stability Problem

1.1 Steepest Descent for a Single Linear Layer

Consider a single linear layer $y = Wx$. The parameter $W$ and the input-output function are the same object: any perturbation $\Delta W$ directly changes what the layer computes. The resulting optimization question is then: what update $\Delta W$ decreases the loss most rapidly, subject to the layer's function not changing too much?

The operator norm is the right constraint since it bounds the worst-case output perturbation over all unit inputs. The solution is $\Delta W^* = -\epsilon \, UV^\top$, where $G_W = U\Sigma V^\top$ is the (thin) SVD of the gradient. This solution replaces all singular values with 1s, yielding a unit-operator-norm update maximally aligned with the gradient [3]. The resulting optimizer is Muon [2] (in practice, it (semi-)orthogonalizes the momentum buffer rather than the raw gradient).

One of the key properties that makes this work is that for a single linear layer, the parameter perturbation is the functional perturbation. For compositions of multiple operations, this is no longer necessarily true.

1.2 Composed Functional Perturbations

In a neural network, no layer operates in isolation. The function the loss evaluates is a composition of maps, and perturbations to one layer's parameters affect the behavior of all subsequent layers. Beyond controlling the operator norm of each parameter's update separately, one possible solution would be to control the perturbation of the entire composition itself.¹

Take the QK product in one attention head of a Transformer [1] (ignoring RoPE and QK-norm for now; see Appendix F for discussion on these design choices). The loss sees $W_Q$ and $W_K$ only through their product $M = W_Q W_K^\top$. A perturbation $(\Delta W_Q, \Delta W_K)$ induces (ignoring the higher-order term) a first-order functional change

Controlling $\|\Delta W_Q\|_{\text{op}}$ and $\|\Delta W_K\|_{\text{op}}$ separately only bounds $\|\Delta M\|_{\text{op}}$ through the partner norms:

In particular, the same factor-level update size can induce very different functional perturbations depending on the current spectral geometry of the partner matrix.

The “compositional steepest descent problem” asks for the update that controls the functional perturbation directly:

This is harder than the single-layer case, but it admits tractable formulations. The same composition principle applies to the OV product $W_O W_V$, and analogues exist for diagonal RMSNorm scales [9] (Appendix G) and activation-gated MLPs (Appendix H).

2. Composed Perturbations

2.1 The Product Structure and Rank Obstruction

Recall that a perturbation $(\Delta W_Q, \Delta W_K)$ induces the first-order functional change (2)

and the steepest descent problem in terms of the product gradient $G_M = \nabla_M \mathcal{L}$ is

where $T_M := \{\Delta W_Q W_K^\top + W_Q \Delta W_K^\top\}$ is the set of achievable first-order product perturbations.

If we could optimize over $\Delta M$ freely, the solution would be

where $\operatorname{msign}(A)=UV^\top$ for the thin SVD $A=U\Sigma V^\top$. The natural next step would be to decompose $\Delta M^*$ into factor updates but this fails for rank deficiency reasons: each of $\Delta W_Q W_K^\top$ and $W_Q \Delta W_K^\top$ has rank at most $d_h$. In particular, the achievable perturbations have rank at most $2d_h \ll d$, while $\operatorname{msign}(G_M)$ is generically full-rank.²

Instead, we optimize directly over $(\Delta W_Q, \Delta W_K)$ from the start. The difficulty is that $\|\Delta M\|_{\text{op}} \leq \epsilon$ constrains the operator norm of a sum, $\Delta W_Q W_K^\top + W_Q \Delta W_K^\top$, so the feasible updates are jointly determined rather than separable factor-wise; moreover, each piece carries the partner matrix, so how severely a given $\Delta W_Q$ or $\Delta W_K$ perturbs the product depends on the spectral geometry of $W_K$ and $W_Q$.

2.2 The Achievable Set as a Tangent Space

Before continuing, it is worth seeing what $T_M$ actually is. It is the tangent space at $M$ to the manifold of rank-$d_h$ matrices,

Differentiating a factorization $X = AB^\top$ produces $\{\dot A B^\top + A \dot B^\top\}$, and at $A = W_Q$, $B = W_K$ this is exactly $T_M$. The rank obstruction from §2.1 is then just the statement that tangent vectors to $\mathcal{M}_{d_h}$ have rank at most $2d_h$, so the full-rank $\operatorname{msign}(G_M)$ is off the tangent space and cannot be a feasible step.

This viewpoint also suggests a Manifold-Muon [4]-inspired update is possible, whereby one mandates that the update lie in the tangent space to the current point on the weight manifold.

3. Half-Split Composed Muon

In this section, we develop the half-split and hybrid half-split composed Muon update rules. We begin with a discussion of polar decompositions for the QK product, and then turn to the OV pathway, where the head structure forces a "hybrid" update rule. See Table 1 for the concrete update rules.

3.1 Polar Decompositions and Half-Split Composed Muon

Observe that the product-gradient objective decomposes as

This decomposition suggests a simple surrogate: split the update into two separately budgeted steps, one for $Q$ and one for $K$. This is generally not the optimal solution to the original problem, since the original constraint is on the sum of the two product perturbations, not on each term separately.

Let us treat the first term on the right-hand side first. If we allocate equal perturbation budgets to the two terms, the corresponding optimization problem is

Using polar decompositions, we can formulate an equivalent problem that is closer in structure to the original Muon problem. Write

Here $U_K \in \mathbb{R}^{d \times d_h}$ and $C_K \in \mathbb{R}^{d_h \times d_h}$ is symmetric positive semi-definite. We have

The objective becomes

where $G_Q = G_M W_K = G_M U_K C_K$.

Since $U_K^\top \in \mathbb{R}^{d_h \times d}$ has orthonormal rows, right-multiplication by it preserves operator norms, so

Thus, with $Z_Q = \Delta W_Q C_K$, the original problem may be reformulated as

The problem is now in the standard Muon form, with the solution

$Z_Q^* = -\frac{\epsilon}{2} \, \operatorname{msign}(G_Q C_K^{-1})$. Undoing the change of variables ($\Delta W_Q = Z_Q C_K^{-1}$):

A similar analysis can be applied to the second term on the right-hand side, yielding

This analysis yields the first composition-derived update rule, which we call “half-split”.³.

The main limitation of this approach is that the 50/50 budget split is arbitrary. If the gradient is more informative on one side, half the budget is wasted on the other. One can drop the per-side budgets in favor of a single jointly constrained spectral step on the concatenated whitened factors, sidestepping the allocation question at the cost of a slightly looser worst-case product bound; we work this variant out in Appendix B. For grouped-query attention, the half-split idea suggests allocating budgets by query group; see Appendix D.

3.2 The Hybrid Case: $V$ Per-Head, $O$ Per-Matrix

For QK, everything above was deduced per-head: each head has its own query-key product, produces its own attention matrix, and therefore defines its own composed map

Thus the natural way to perform the optimization is per-head.

The OV pathway is different. The value matrices still act head-by-head, but the output projection is the operation that merges all heads back into the residual stream. This creates an asymmetry: the $V$-side perturbation is naturally head-local, while the $O$-side perturbation is naturally matrix-global.

For token $i$, write the head-$h$ context and value feature as

The OV output is

This equation suggests a hybrid approach to the OV pathway. Before $W_O$, the heads are still separate; at $W_O$, they are not. The model sees one stacked value vector and one output projection into the residual stream.

First consider a $V$-side perturbation. Since $W_{V,h}$ is applied before head $h$'s value feature is merged with the others, the first-order change is

Each term has its own input $\tilde{x}_{h,i}$, its own value matrix $W_{V,h}$, and its own output block $W_{O,h}$. The heads may attend to different tokens, so their effective context vectors $\tilde{x}_{h,i}$ should be treated as separate directions. This gives the headwise bound

The natural $V$-side question is therefore local to each head. Applying the half-split rule to each head gives

where

Here $\epsilon$ is the optimizer scale for the head-local rule. The formula controls the shape of each head's value update; the overall scale of the summed OV perturbation is set by the learning rate and CM multiplier.

For an $O$-side perturbation, the value features have already been formed. The first-order change is

Now the heads enter as blocks of one stacked vector. The output projection is one linear map from the concatenated head space back to the residual stream. Thus the natural $O$-side question is global.

For each head, define

Whiten the $O$ gradients per head, concatenate them, and take one matrix sign:

Writing

the corresponding blockwise update is

The difference from the $V$ rule is where the matrix sign is taken. On the $V$ side, each head is orthogonalized separately because the value map acts before head aggregation. On the $O$ side, the whitened gradients are concatenated and orthogonalized once because $W_O$ is the map that performs the aggregation.

3.3 Isotropic Approximation

The full half-split rules use inverse square roots of partner Gram matrices. However, we may also attempt to approximate the partner Gram matrix by a scaled isometry. For a tall factor

if the columns of $W$ are approximately orthogonal and have comparable norms, then

Thus

and under this assumption, we may approximate the QK partners via

and the $V$-side OV partner via

On the $O$-side of the OV update, the partner is

and the relevant Gram matrix is therefore the row Gram

The isotropic approximation in this case is

Under this approximation, every partner inverse root becomes a scalar multiplier:

and because $\operatorname{msign}$ is scale-invariant, the scalar inside the spectral sign drops out. What remains is ordinary Muon, but with a partner-dependent effective learning rate.

For QK, the full rule

becomes

and similarly

For OV, the $V$-side update becomes

The $O$-side remains global across heads, but the per-head whitening factors become scalars:

Empirically, the isotropic approximation performs similarly to the full CM update.

To check whether the isotropic assumption holds in practice, we plot the relevant partner Gram for a representative layer/head below.

The scalar factors also give a useful effective-learning-rate interpretation. For example, the $V$-side effective learning rate is

where $\eta$ is the raw WSD learning rate, $1/2$ is the half-split quota, $mp$ is the CM multiplier, and $c_{O,h}$ is the scalar approximation to the $O$-partner Gram root.

4. Gauge Motion and Momentum

So far we have treated the factors $W_Q$ and $W_K$ as if they were fixed coordinates for the product update. But the loss does not see these coordinates directly. It sees only $M = W_QW_K^\top$, and this representation is not unique: for any invertible matrix $R\in GL(d_h)$,

leaves the product unchanged:

Differentiating this action gives the corresponding pure-coordinate, or vertical, directions

These directions are nonzero in parameter space, but invisible to the product to first order:

More specifically, if we let $\mathcal P = \mathbb R^{d\times d_h}\times\mathbb R^{d\times d_h}$ be the space of factor pairs, and $\pi(W_Q,W_K)=W_QW_K^\top$ be the product map, then the natural object is not $\mathcal P$, but rather the quotient $\mathcal P / GL(d_h)$.

However, in practice we utilize objects that at implementation time are stored in factor coordinates, e.g. momentum buffers. In particular, even though one can show (see Appendix E) that the undamped half-split update is not coordinate-dependent, the resulting optimizer with momentum, damping, and finite Newton-Schulz iterations is coordinate dependent.

Hence, we need to provide a prescription for how to perform quotient-level motion in redundant coordinates. We may do so via a connection. Namely, given a point $(W_Q, W_K)$ on $\mathcal P$, one defines the vertical directions

A connection then smoothly assigns a complementary horizontal subspace $H_{(W_Q,W_K)}$, so that every factor update decomposes as

where $\Delta^V$ is pure gauge motion and $\Delta^H$ is the chosen representative of the product-level motion.

One simple choice of connection is to declare horizontal directions to be those Frobenius-orthogonal to the gauge directions.

For tangent vectors $(\delta_Q,\delta_K)$ and $(\delta_Q',\delta_K')$, we have the flat Frobenius inner product

An update $(\Delta W_Q,\Delta W_K)$ is then horizontal, by definition, if

Expanding the traces gives the horizontality condition

For ordinary gradient descent, the factor gradients are

Therefore

So the raw negative gradient is horizontal for the Frobenius connection.

However, in the case of momentum, even if it is horizontal at one step, it is not generally horizontal at the next step, because the horizontal subspace itself depends on the current factors. If

at the current point, then after an update

the same buffer would need to satisfy

Using the old horizontality condition, this reduces to

which need not hold. So momentum can accumulate vertical components even when each fresh raw gradient is horizontal.

This creates two practical issues. First, the buffer can drift into vertical directions because the horizontal space moves with the factors. Second, a whitened update is tied to the partner geometry at the step where it was computed. A term whitened using $C_K^{(t-1)}$ is not naturally comparable to one whitened using $C_K^{(t)}$.

We handle the second issue by accumulating raw gradients rather than already-whitened updates:

The half-split update is then formed using the current partner geometry:

Vertical drift can be handled separately by projecting the buffer or update back onto the horizontal space. Concretely, given $(m_Q,m_K)$, choose $X$ so that

satisfies

Writing $S=X^\top$, this gives

The correction is vertical, so it leaves the first-order product perturbation unchanged. It only changes the representative of that perturbation in factor coordinates.

5. Experiments

5.1 Architecture ablations

We ran experiments to ablate two common design choices in the QK pathway, QK-norm [7] and the positional method (ALiBi vs. RoPE [6][5]), while also comparing whether CM is applied to OV only or to both OV+QK. We evaluated three architecture settings:

• No QK-norm + ALiBi
• No QK-norm + RoPE
• QK-norm + RoPE

In each setting, we report the best validation loss and the corresponding hyperparameters.

Across all three settings, CM on both OV+QK beats vanilla Muon by a clear margin. In the cleanest setup for the theory (no QK-norm + ALiBi, where the learned QK logits remain a plain matrix product), OV-only CM already outperforms Muon, and adding QK CM further improves. With RoPE (no QK-norm + RoPE), OV-only CM underperforms Muon, but OV+QK CM still wins; this suggests a non-trivial interaction between CM and RoPE that we do not yet fully understand (TODO: root cause). In the most realistic setup (QK-norm + RoPE), OV-only CM improves over Muon, while adding QK CM does not provide additional gains, consistent with the possibility that QK-norm reduces or removes the benefit of applying CM on the QK circuit.

To understand the effect of CM on the composed update $\Delta M_h$, we track its singular values and Frobenius norm throughout training. This is the quantity the loss actually sees through the circuit, so stable composed-update spectra are evidence that CM is controlling the functional step rather than merely the individual factor steps. Under CM, both the spectrum and the norm stay stable, and the coefficient of variation of the per-head update norm across heads within a layer remains small:

5.2 1B pretraining

To check that Compositional Muon’s gains survive scale, we repeat the comparison at 1B parameters (QK-norm + RoPE, 70B tokens). We test two CM parameterizations against a tuned Muon baseline: (i) full-whitening OV+QK, and (ii) isotropic OV gauge, which uses the isotropic approximation plus a scalar gauge rebalance on the V↔O pair. Both beat Muon throughout training: the curves separate early and remain separated through the end.

We also tested at 1B parameters (no QK-norm + ALiBi, 70B tokens), and observed even larger gains from CM on both OV+QK—suggesting an additional benefit from applying CM to the QK circuit in this setting.

We summarize the final loss of the two settings as follows.

5.3 Hyperparameter sweeps

The following shows a CM multiplier sweep comparing CM and isotropic CM; the results agree well across different mp values.

The following shows a CM multiplier sweep comparing isotropic CM (no gauge) and isotropic CM (gauge); the results agree well across different mp values.

5.4 NanoGPT optimization track

We entered Compositional Muon (CM) in the modded-nanogpt speedrun track 3 optimization benchmark. Starting from the #309 baseline, our submission (PR #311) adds a single coupled update for the attention $W_V$↔$W_O$ pair. Over 20 seeds, it reaches the leaderboard threshold at step 2875 (mean loss 3.27901) vs. step 2890 for the baseline (~15 steps later).

5.5 Effective learning rate on AdamW

To isolate which part of isotropic CM matters, we port the same partner-dependent rescaling idea to AdamW [8] (“partner-rescaled AdamW”). Concretely, we keep AdamW’s update rule fixed and only modify the OV leg’s per-head scale using the partner norm (the same scalar that appears in the isotropic approximation). This lets us test whether the gains are specific to Muon’s spectral sign, or more generally due to a better effective learning-rate allocation across heads/layers. Figure 13 shows that partner-rescaled AdamW improves over AdamW, though the gains are smaller than in the Muon setting; this may be due to suboptimal hyperparameter choices.

5.6 Downstream evals

We evaluate the best 340M checkpoints on four standard benchmarks (ARC-Challenge 25-shot, HellaSwag 10-shot, MMLU 5-shot, Winogrande 5-shot).

On these pretraining benchmarks CM and vanilla Muon are largely comparable, without a clear winner. For MMLU at this scale of model, the scores roughly amount to random guessing. We include these results for completeness and to avoid overstating pretraining gains from CM.

6. Related Work

Certainly, this work owes quite a bit of intellectual debt to Muon [2][15]. Muon Split in GLM-5 [19] was, to our knowledge, the first work to apply Muon at scale with per-head Q/K/V updates and matrix-level $W_O$ updates.

There are also several recent works studying how optimizer behavior depends on parameterization. Riemannion [11] makes this point for LoRA, where many low-rank factorizations represent the same adapter. Ji-Ha Kim, independently of us, discusses the analogous gauge issue for Transformer heads, showing that equivalent head bases can train differently under coordinate-based optimizers [23]. Modular duality and modular norm optimization [17][18] argue that updates should respect module structure, while symmetry-compatible optimizer design [12] derives optimizer rules from parameter-block equivariances.

7. Conclusion

The main perspective offered by Compositional Muon is that, when possible, one should directly control the functional perturbation of a composed circuit, rather than the individual pieces which constitute the circuit. For attention, this means treating $W_QW_K^\top$ and $W_OW_V$ as the natural objects of optimization, and then deriving factor-space update rules that approximate spectral steepest descent on those composed maps. The half-split and hybrid rules developed here are simple surrogates for this more general problem: they do not exhaust the possible choices, but they show that even a relatively direct approximation can produce better spectral behavior and potential pretraining gains over Muon. In particular, we suspect there are likely better surrogate problems than the ones discussed in this post.

More broadly, compositional optimization suggests a(nother) path toward optimizer-architecture co-design. Some circuits may be easier to optimize not only because of their expressivity, but because their composition-level perturbation geometry admits tractable update rules. This perspective could be combined with other optimizer ideas, such like Shampoo and SOAP, by asking not only how to precondition an individual parameter, but how to precondition the composed map in which that parameter participates.

8. Appendices

Appendix A. Regularizing the Overdetermined Sylvester System

The decomposition

cannot generically exist for

because the right-hand side is typically full-rank while the left-hand side is tangent to the rank-$d_h$ manifold. Still, it is useful to understand the associated linear system.

Let

The adjoint is

A minimum-norm solution, when it exists, lies in the image of this adjoint. Thus one can write

for some $H\in\mathbb R^{d\times d}$. Substituting gives

This Sylvester equation is singular because both Gram matrices have rank at most $d_h<d$ and hence have zero eigenvalues.

A regularized version is

Then the induced factor updates are

The achieved product perturbation is

Thus the residual $2\lambda H_\lambda$ is the price of regularization.

Appendix B. Joint Whitened Compositional Muon

The half-split rule solves two whitened subproblems separately, each with budget $\epsilon/2$. An alternative is to concatenate the whitened variables and apply one spectral sign.

Recall

with whitened gradients

Concatenate

Then

The joint whitened problem is

Its solution is

Splitting $Z^*=[Z_Q^*\mid Z_K^*]$ and unwhitening gives

The joint SVD orthogonalizes shared directions across the Q and K whitened gradients together. If the two whitened gradients have orthogonal column spaces, the joint rule reduces to separate half-split steps at the same scale. If their column spaces overlap, the joint rule treats shared directions once rather than independently on each side.

The tradeoff is that the worst-case product bound becomes looser:

Appendix C. Hybrid OV Derivation

The OV pathway differs from QK because the attention block returns a single residual-stream contribution, not $H$ independent outputs. For target token $i$, the output is the sum over heads of the output projection applied to each head's attention-weighted value vectors:

where $x_j$ is the input vector of token $j$, and $\alpha_{h,ij}$ is the scalar attention weight. Let $\tilde{x}_{h,i} = \sum_{j=1}^N \alpha_{h,ij} x_j \in \mathbb{R}^d$ represent the attention-weighted context vector entering head $h$. Because the attention weights are constrained by a softmax, $\tilde{x}_{h,i}$ is strictly a convex combination. Assuming bounded token inputs $\|x_j\|_2 \leq 1$, the effective head inputs are bounded within the unit ball:

A first-order perturbation of the factor matrices $(\Delta W_O, \Delta W_V)$ induces a functional change in the output:

If we allocate equal perturbation budgets to the $O$ and $V$ pathways, the two sides require different geometric treatments.

C.1 The Global O-Side

Examine the $O$-side perturbation: $\Delta y_{i, O} = \sum_{h=1}^H \Delta W_{O,h} (W_{V,h} \tilde{x}_{h,i})$. The full output projection $W_O = [\,W_{O,1}\mid \cdots \mid W_{O,H}\,] \in \mathbb{R}^{d \times H d_h}$ is a single matrix. Because all head contributions sum directly within the residual stream and act on a shared spatial embedding, we can stack the active value features $v_{h,i} = W_{V,h} \tilde{x}_{h,i}$ vertically into a single vector $v_i$.

This allows the total $O$-side change to be bounded by a single linear map. Using polar decompositions $W_{V,h} = C_{V,h} U_{V,h}^\top$ (where $C_{V,h} = (W_{V,h} W_{V,h}^\top)^{1/2}$), define the horizontally concatenated, partner-whitened block

The aggregate $O$-side change can be written as $\Delta y_{i,O}=Z_Ow_i$, where the stacked blocks are

Therefore

Thus $\|Z_O\|_{\mathrm{op}}$ gives a global bound relative to the stacked head-input norm. If one wants an $\epsilon/2$ worst-case output bound under only per-head unit bounds, the rule should use $\|Z_O\|_{\mathrm{op}}\leq \epsilon/(2\sqrt H)$; in practice this constant is folded into the learning-rate scale. Consequently, a per-matrix msign on the full concatenated $O$-gradient can coordinate updates across heads.

The update rule is: whiten per head, then orthogonalize per matrix. For each head, define $\widetilde G_{O,h} = G_{O,h} C_{V,h}^{-1}$. Concatenate these horizontally to form $\widetilde G_O \in \mathbb{R}^{d \times H d_h}$, apply one spectral sign to the full matrix $S = \operatorname{msign}(\widetilde G_O)$, and split $S$ back into head blocks $S_h$. The update is:

C.2 The Head-Local V-Side

Conversely, the $V$-side perturbation acts on the raw tokens before the multi-head aggregation occurs:

Crucially, because the attention weights $\alpha_{h,ij}$ are computed independently across heads via separate QK projections, the effective inputs $\tilde{x}_{h,i}$ can vary independently across heads. There is no single shared input vector before the value step on which a global $W_V$ operator norm acts.

This bound is tight when the headwise worst-case output directions align, but it is not generally an equality for fixed weights. The important point is that the bound decomposes by head: each $V$ block is naturally controlled through its own partner $W_{O,h}$, rather than through one global matrix-level orthogonalization across all value blocks.

For each head, define $C_{O,h} = (W_{O,h}^\top W_{O,h})^{1/2}$. The $V$-side update remains strictly local:

This hybrid rule gives a global stacked-input guarantee on the $O$-side and a head-local guarantee on the $V$-side. It should not be read as an $\epsilon/2$ bound on the summed $V$-side residual perturbation unless the per-head budgets are divided or otherwise allocated across heads; the worst-case aggregate bound scales like $H$, while the $O$-side stacked bound scales like $\sqrt H$.

Appendix D. Grouped-Query Attention

The QK derivation in the main text treats each attention head independently. In grouped-query attention (GQA), multiple query heads share the same key and value head. This changes the natural budget allocation.

For QK, suppose a key head $K_g$ is shared by a group $\mathcal H_g$ of query heads. The learned products are

A perturbation of the shared key affects all products in the group:

Let $G_g=|\mathcal H_g|$ be the number of query heads sharing this key head. The $Q$-side remains head-local, so each query head can use the usual half-split budget for its own term $\Delta W_{Q,h}W_{K,g}^\top$. The shared $K$-side is different: the same update $\Delta W_{K,g}$ appears in all $G_g$ products.

If these per-product bounds are enforced, then the summed shared-key contribution is bounded by $\epsilon/2$ by the triangle inequality. This is a conservative budget allocation, not a complete closed-form solution for the shared $\Delta W_{K,g}$; the actual shared update still involves the collection of partner matrices $\{W_{Q,h}\}_{h\in\mathcal H_g}$. Less conservative variants could allocate the same total group budget by gradient norm, or solve a grouped joint-whitened subproblem analogous to Appendix B.

The same principle applies to the OV pathway when values are shared by a group of query heads. Head-local parameters keep head-local budgets; shared parameters get a group budget, divided by the number of query heads that use them.

Appendix E. Equivariance of Exact Half-Split

The undamped full-whitening half-split rule is equivariant under the full $GL(d_h)$ gauge action. We show this for the Q-side; the K-side is symmetric.

Write the polar decomposition

The Q-side half-split update is