The Loss Does Not See the Basis, but Adam Does — part 4/12

Cleaned-text code points [46221, 60903). Bibliography omitted here; source.md includes it.

The observed splitting is not due to learning-rate scheduling: under a single shared cosine decay we still get 9 / 9 9/9 on classification, and giving GD and Adam the same benefit of the doubt as Muon—the full 2 × 10 4 2\times 10^{4} -step cosine decay horizon, no early stopping—does not improve their results (Table 6 ). The same is true for initialization scales of 10 − 3 10^{-3} and 3 × 10 − 3 3\times 10^{-3} , though it is less pronounced at 10 − 2 10^{-2} , where every method but Muon has lost the small-initialization edge (Appendix D.4 ). We also observe the same phenomenon in a larger setting ( 60 × 60 60{\times}60 , rank 5 5 ), not tabulated: 0.000 0.000 (Muon), 0.077 0.077 (GD), 0.465 0.465 (Adam); see experiments/zoo_size_check.py . The only optimizer for which scheduling makes a difference is RMSProp: because of the ϵ \epsilon -bounded denominator in its update, it cannot reach the interpolation threshold at a fixed step size, and it recovers 0.527 0.527 either way (Appendix D.7 ).

## 6 The same gauge lives in attention heads

##### Attention has the same symmetry.

In an attention head with queries q = W Q ​ x q=W_{Q}x and keys κ = W K ​ x \kappa=W_{K}x ( Vaswani et al., 2017 ) (written κ \kappa here, since k k is the gauge dimension throughout), the logits are invariant to the choice of basis in the head dimension: they only depend on the weights through q ⊤ ​ κ = x ⊤ ​ W Q ⊤ ​ W K ​ x q^{\top}\kappa=x^{\top}W_{Q}^{\top}W_{K}x . For any orthogonal matrix A h A_{h} specific to the head, replacing ( W Q , W K ) (W_{Q},W_{K}) with ( A h ​ W Q , A h ​ W K ) (A_{h}W_{Q},A_{h}W_{K}) leaves all logits unchanged, hence all outputs of the model. In other words, picking a basis in the head dimension is a gauge freedom, and W Q ⊤ ​ W K W_{Q}^{\top}W_{K} plays the role of U ​ V ⊤ UV^{\top} (under the convention U = W Q ⊤ U=W_{Q}^{\top} , V = W K ⊤ V=W_{K}^{\top} ). 2 2 2 The full function-preserving group of a head’s logits is larger (any invertible A A acts by ( W Q , W K ) ↦ ( A ​ W Q , A − ⁣ ⊤ ​ W K ) (W_{Q},W_{K})\mapsto(AW_{Q},A^{-\top}W_{K}) ); as in Definition 3.1 we work with its orthogonal subgroup, the maximal part acting by isometries of parameter space (Lemma 3.2 ). Hence all statements from § 4 carry over to individual attention heads.

This symmetry has been previously noted and deliberately broken: in particular, Silverstein et al. (2026) induce a preferred direction by adding an unlearned bias to q q and v v , resampled at random each batch, in the spirit of a Hamiltonian/Noether framework for memory-efficient optimizers. We instead retain the gauge symmetry and study the interplay between standard optimizers and it.

For a given head, define U = W Q ⊤ U=W_{Q}^{\top} , V = W K ⊤ V=W_{K}^{\top} , and Q = A h ⊤ Q=A_{h}^{\top} (this Q Q is the gauge element of § 3 ; the Q Q in W Q W_{Q} is the query label, and the two are unrelated). Then, by Proposition 4.4 , the exact first-step deviation in the invariant W Q ⊤ ​ W K W_{Q}^{\top}W_{K} under Adam at zero optimizer state is given by the expression in that proposition. The subsequent drift in logits is then an empirical consequence of this deviation at the product level, which the twin protocol isolates.

##### Twin protocol.

We train two 2 2 -layer 4 4 -head transformers with embedding dimension d = 64 d{=}64 on modular addition modulo 47 47 ( Power et al., 2022 , the arithmetic task of) , using full-batch training on deterministic CPU, with no bias terms or qk-norm (to retain the gauge symmetry). In the second model, we replace each head’s ( W Q , W K ) (W_{Q},W_{K}) by ( A h ​ W Q , A h ​ W K ) (A_{h}W_{Q},\,A_{h}W_{K}) at initialization, with A h A_{h} a random orthogonal matrix drawn per head. As both models realize the same function at initialization, any deviation during training is due to the optimizer being sensitive to the internal basis. We also have a twin with A = I A{=}I (the identity matrix), which should be bit-identical and serve as a sanity check for reproducibility; a noise twin with the same basis but with W Q W_{Q} and W K W_{K} perturbed by a Gaussian of magnitude 10 − 7 10^{-7} , probing the chaos amplification due to small functional differences; and runs with equivariant optimizers. Of the latter, SGD is applied with heavy-ball momentum 0.9 0.9 , and Muon follows the standard hybrid procedure: it updates the matrix parameters—where W Q W_{Q} , W K W_{K} , and hence the head gauge live—while Adam optimizes the embeddings and the output head. The gauge acts only on the parameters Muon updates, and the remaining gradients are identical across twins, so the hybrid is equivariant for the symmetry under test (see Appendix C ).

##### Results (Table 3 , Figure 2 ).

The Adam twins diverge right away: after one step, their logit difference is already 3.6 × 10 − 3 3.6\times 10^{-3} , as opposed to 2.9 × 10 − 7 2.9\times 10^{-7} for the noise twin at the same step (a difference of four orders of magnitude). The relative logit distance then saturates at 0.77 0.77 . Meanwhile, the W Q ⊤ ​ W K W_{Q}^{\top}W_{K} invariants have a relative Frobenius distance of 56 % 56\% (QK drift: mean relative Frobenius distance of per-head invariants, see Appendix C ; not shown in Table 3 , which reports logit distances), which manifests as a substantial shift in the gauge-invariant structure of the model, beyond mere rotations of individual heads. That last number is the one with a downstream consequence, which we take up at the end of this section.

The two perturbations are not directly comparable, as they act on different scales: the gauge transformation leaves the function of the model intact, but changes the parameters by O ⁡ ( 1 ) O(1) . The logit differences at step 0 0 for the two twins are 1.8 × 10 − 7 1.8\times 10^{-7} vs. 2.6 × 10 − 7 2.6\times 10^{-7} (Table 3 ; the gauge twin’s step- 0 0 value is rounding, one function evaluated in two bases, which is why the same-basis A = I A{=}I twin is exactly 0 0 ), so the noise twin serves as a probe of how sensitive the optimizer is to a function difference of that magnitude when no basis has changed. In short, what Adam learns in an attention head depends on an arbitrary choice of coordinates that has no effect on the model’s input–output behavior.

SGD and scalar-Adam twins sit at the harness’s numerical-noise level at step 1 1 and stay four to five orders of magnitude below Adam’s structural split thereafter, so replacing Adam’s per-coordinate denominator by a shared scalar—the p = 0 p{=}0 end of the dial we build next (§ 7 )—kills the basis dependence here as well as on the sensing task. Muon’s gauge twin also diverges, but its noise twin tracks the same trajectory up to a certain point (Appendix D.5 ): the agreement is within 0.05 0.05 decades averaged over the trajectory past step 100 100 (at step 100 100 itself the two still differ by 0.24 0.24 decades, Table 3 ), and both twins sit at the noise floor for about 50 50 steps, before numerical chaos takes over due to the instability of msign \mathrm{msign} , not basis dependence.

Proposition A.1 shows that any spectrum-flattening schedule massively amplifies small singular values: unboundedly for the exact msign \mathrm{msign} , and in practice by a large but finite Newton–Schulz factor around zero, which is where msign \mathrm{msign} acts in early training with near-zero momentum. To disentangle these effects, note that the noise twin separates chaotic amplification by the schedule from genuine basis dependence during training. We thus propose using the triplet (gauge twin, A = I A{=}I twin, noise twin) as a diagnostic for claims of equivariance during training.

The phenomena are reliably observed across different random seeds and initialization draws: six (init × \times gauge-draw) pairs for Adam and two for each equivariant method, with the noise twin run out to a 100 × 100\times stronger perturbation, past which Adam’s gauge split still stands clear of it by over 600 × 600\times (Appendix D.5 ).

##### Scale, and real stochastic training.

These phenomena are not specific to toy models or CPU precision. They occur in both full-batch and stochastic settings, and are not specific to model depth (see Appendix D.5 for the same analysis in a GPU float64 setting). When rerun in float64 on GPU (deterministic, with A = I A{=}I twins exactly zero), for modular addition over the integers mod 97 97 , Adam’s gauge twins diverge on step 1 1 , by 8.2 8.2 – 11 × 10 − 3 11\times 10^{-3} ( 4 4 layer/ d = 128 d{=}128 ) and 6.5 6.5 – 7.2 × 10 − 3 7.2\times 10^{-3} ( 6 6 layer/ d = 256 d{=}256 ), for six init × \times draw pairs each, saturating at 0.61 0.61 – 0.75 0.75 gauge drift. By contrast, all the equivariant methods are confined to 2.8 × 10 − 16 2.8\times 10^{-16} – 1.1 × 10 − 15 1.1\times 10^{-15} at step 1 1 : about machine epsilon in float64 , 12 12 – 13 13 orders of magnitude below Adam.

Similar findings obtain for the stochastic setting: for a 6 6 -layer, d = 256 d{=}256 character-level language model (tiny-Shakespeare, Karpathy, 2015 ), with all methods receiving identical minibatch streams, Adam’s gauge twins split at step 1 1 ( 2.2 2.2 – 2.4 × 10 − 3 2.4\times 10^{-3} ), while the equivariant twins stay at 3 3 – 8 × 10 − 16 8\times 10^{-16} . Critically, while the two Adam runs have very similar validation losses ( 1.579 1.579 vs. 1.585 1.585 ), they have very different functions; their logit distributions end up 0.36 0.36 – 0.37 0.37 apart in relative distance, depending on the gauge pair (Table 9 ).

At larger depths, the final gauge drift stops being an informative summary, as most methods’ twins separate by the end; SGD’s are the exception and stay at float level throughout, though at 6 6 layer/ d = 256 d{=}256 SGD never leaves chance accuracy, so its zero drift reports an untrained model rather than a preserved gauge (Appendix D.5 ). These observations are not specific to stochastic training: a similar effect is seen in the deterministic full-batch mod- 97 97 experiments at 6 6 layer/ d = 256 d{=}256 , where even the equivariant twins start to separate, and it again manifests in the character-level LM. Wherever a method’s twins do separate, each one lands where its own noise twin lands (Appendix D.5 tabulates both), suggesting that late-time gauge divergence occurs due to Lyapunov instability amplified onto the gauge symmetry via finite precision arithmetic; the equivariance failure modes are subdominant to the amplified floating-point error. For Muon, the instability is the Newton–Schulz iteration’s sensitivity to small singular values, while for scalar-Adam it is limited precision arithmetic (the depth and minibatching acting as additional amplifiers). Only Adam moved at step 1 1 (a change 39 × 39\times larger than even the 100 × 100\times stronger noise twin), while the equivariant methods were confined to float precision. In this picture, early basis-dependence is the discriminating signal, and this trend appears to generally hold (Appendix D.5 , Table 9 ).

##### A consequence for model merging.

Aligning two networks’ internal bases before averaging them (permutation re-basin, Ainsworth et al., 2023 , and its rotational analogue here) presupposes that alignment makes the networks agree. The twin experiments make this premise optimizer-dependent inside every attention head. Alignment can at most reconcile the gauge, so two runs are head-mergeable only if the gauge-invariant product W Q ⊤ ​ W K W_{Q}^{\top}W_{K} agrees.

In the shallow deterministic regime, where the dynamics stay non-chaotic, our equivariant twins agree in the invariant to floating-point error and are mergeable by one Procrustes step; at 6 6 layers they separate as well, but only as far as their own noise twins do (Appendix D.5 ). Adam’s twins instead disagree in the invariant itself by 56 % 56\% , so no per-head rotation can reconcile them, even though the runs began as the same function. Basis-dependent training does not simply pick among equivalent solutions; it removes the agreement that alignment-based merging needs. We have not tested merging end to end, so this is a consequence of the measured invariant rather than a merging experiment.

## 7 Isolating preconditioner anisotropy with a dial

The zoo sorts the optimizers, and the attention twins show the same sorting inside a transformer. Neither says which ingredient of Adam does the damage. Here, we employ a one-parameter family to isolate the specific factor driving these differences. We define Adam- p p by modifying Adam’s denominator v ^ + ϵ \sqrt{\hat{v}}+\epsilon to ( v ^ + ϵ ) p ​ ( s ¯ + ϵ ) 1 − p (\sqrt{\hat{v}}+\epsilon)^{\,p}\,(\bar{s}+\epsilon)^{\,1-p} , where s ¯ = mean ⁡ ( v ^ ) \bar{s}=\sqrt{\operatorname{mean}(\hat{v})} represents a shared root-mean-square scalar, and ϵ = 10 − 8 \epsilon=10^{-8} is the stabilizing term used in Adam. The choice of s ¯ \bar{s} is significant. Computing the arithmetic mean of v ^ \hat{v} across the entries in both factors gives a gauge-invariant quantity, since the Frobenius-norm statistic to which it is equivalent is preserved under the right-orthogonal transformation of Lemma 4.1 . As such, when p = 0 p{=}0 the denominator reduces to the gauge-invariant scalar of Proposition 4.2 (2). Adam- 0 0 then coincides with the scalar-Adam procedure of § 5 up to an additive O ⁡ ( ϵ ) O(\epsilon) term in the denominator: it is the same algorithm, and the results are identical ( 0.201 0.201 , erank 5.43 5.43 ). In contrast, the geometric mean of v ^ \sqrt{\hat{v}} would fail to be gauge-invariant, rendering the p = 0 p{=}0 case itself dependent on an arbitrary choice of factorization basis. We adhere to the RMS convention throughout, giving the geometric-mean alternative only as a consistency check. 3 3 3 The geometric-mean dial traces the same monotone curve and agrees with the RMS dial at the endpoint to within 0.03 0.03 (recovery 0.229 0.229 , erank 5.8 5.8 at p = 0 p{=}0 ). It is an approximate rather than an exact dial, because its p = 0 p{=}0 limit is not equivariant.

When p = 1 p{=}1 , we recover standard Adam, which is coordinate-wise and gauge-breaking. For intermediate values of p ∈ ( 0 , 1 ) p\in(0,1) , the preconditioner is less anisotropic while retaining the same adaptation, momentum, and noise properties as standard Adam. The overall size of the denominator is controlled by p p , so for the purposes of the experimental sweeps it makes sense to test either a shared set of learning rates for all p p or a separately optimized learning rate for each p p (see Algorithm 1 for the exact procedure, including the two places in which ϵ \epsilon is used).