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

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

This means the update progresses along the singular directions of the gradient itself, scaling each by a scalar factor h ⁡ ( σ i ) h(\sigma_{i}) , which we refer to as the spectral transfer function. (When an inverse or root of a singular positive-semidefinite matrix appears later, it is interpreted as acting on the support, with zero assigned on the nullspace.) For memoryless one-step methods, standard gradient descent (GD) corresponds to h ⁡ ( σ ) = η ​ σ h(\sigma)=\eta\sigma , while common-scalar-preconditioned GD gives h t ​ ( σ ) = η ​ σ / a t h_{t}(\sigma)=\eta\sigma/a_{t} , equivalent to a time-rescaled GD (Theorem 4.6 ). When applied to a single input matrix, both the exact polar map and support-restricted undamped one-step Shampoo fully flatten nonzero singular values, resulting in h ≡ η h\equiv\eta . In contrast, accumulated damped Shampoo, momentum-based methods, and Muon depend on internal state; thus, the above expressions serve as useful one-step approximations rather than full descriptions of their dynamics.

The fractional family introduced by Dong & Sawin (2026) , which replaces each σ i \sigma_{i} with σ i q \sigma_{i}^{\,q} , directly implements h ⁡ ( σ ) = η ​ σ q h(\sigma)=\eta\sigma^{q} . Their exponent q q is not our dial p p , and it runs the other way: q = 1 q{=}1 is GD and q = 0 q{=}0 is the polar map, whereas p = 1 p{=}1 is Adam and p = 0 p{=}0 is the shared scalar.

Two implications of Proposition A.1 are particularly relevant below.

(i) The ratio h ⁡ ( σ ) / h ⁡ ( σ max ) h(\sigma)/h(\sigma_{\max}) , measuring relative amplification of small versus large singular values, defines the “aggressiveness” axis explored in § 8 . However, this quantity alone does not fully rank practical optimizers: for instance, the one-step idealizations of the exact polar map and undamped Shampoo both yield h ≡ η h\equiv\eta , yet in actual trajectories, accumulation and damping cause Shampoo to flatten singular values less aggressively than Muon—an ordering (GD < < Shampoo < < Muon) based on empirical observation rather than the idealized form.

(ii) Near rank-deficient matrices, the Lipschitz constant of Φ \Phi between nearby matrices is at least | h ⁡ ( δ ) | / δ \lvert h(\delta)\rvert/\delta , as seen when comparing diagonal matrices diag ⁡ ( 1 , ± δ ) \mathrm{diag}(1,\pm\delta) for small δ \delta (distinct from Adam’s stabilizer ϵ \epsilon ). For the exact msign \mathrm{msign} operation, this quotient diverges as δ → 0 \delta\to 0 whenever h ⁡ ( 0 + ) > 0 h(0^{+})>0 , which includes the exact polar map and support-restricted one-step Shampoo. In contrast, GD and scalar-preconditioned GD maintain bounded ratios ( η \eta and η / a t \eta/a_{t} , respectively), provided a t a_{t} remains bounded away from zero; rules where the scaling vanishes with the gradient norm (e.g. a = ‖ G ‖ F a=\left\lVert G\right\rVert_{F} ) are excluded.

The practical finite Newton–Schulz iteration is equivariant under real arithmetic and Lipschitz continuous. Using the standard normalization X 0 = M / ( ‖ M ‖ F + ϵ NS ) X_{0}=M/(\left\lVert M\right\rVert_{F}+\epsilon_{\mathrm{NS}}) (with small stabilizer ϵ NS \epsilon_{\mathrm{NS}} , NS denoting Newton–Schulz) and first coefficient a = 3.4445 a=3.4445 as in Jordan et al., 2024 (see Appendix C ), after five iterations the slope at zero singular values in original units becomes a 5 / ( ‖ M ‖ F + ϵ NS ) a^{5}/(\left\lVert M\right\rVert_{F}+\epsilon_{\mathrm{NS}}) , conditioned on the global scale rather than a fixed univariate h h . As a result, the finite map can greatly enhance small singular values while not introducing the discontinuity of the exact counterpart. This behavior is fundamentally different from the effect of breaking the gauge symmetry, and the two can be distinguished experimentally using the twin protocol of § 6 . Remark A.2 generalizes this observation to state-dependent update rules.

###### Remark A.2 (Stateful rules and right preconditioners) .

With state-dependent update rules, the transformation properties of the various objects are still consistent: momenta transform covariantly ( M t ↦ M t ​ Q M_{t}\mapsto M_{t}Q ), Frobenius norms and left-side accumulators L t = ∑ u G u ​ G u ⊤ L_{t}=\sum_{u}G_{u}G_{u}^{\top} remain invariant, while right-side quantities—such as R t = ∑ u G u ⊤ ​ G u R_{t}=\sum_{u}G_{u}^{\top}G_{u} , or factor Gram matrices U ⊤ ​ U U^{\top}U and V ⊤ ​ V V^{\top}V —transform via conjugation as R t ↦ Q ⊤ ​ R t ​ Q R_{t}\mapsto Q^{\top}R_{t}Q . A sufficient construction preserving these properties is Δ = H ⁡ ( ℐ t ) ​ Z t ​ K ​ ( R t ) , \Delta=H(\mathcal{I}_{t})\,Z_{t}\,K(R_{t}), where Z t Z_{t} is a covariant quantity, like G t G_{t} or M t M_{t} , ℐ t \mathcal{I}_{t} is gauge-invariant, and K K is a spectral function. This captures the covariant structures present in momentum-based gradient descent ( Z t = M t Z_{t}=M_{t} ), scalar-Adam, Muon ( Z t = M t Z_{t}=M_{t} ), Shampoo, and ScaledGD. It does not capture all stateful update rules respecting equivariance, however, only those composed as a product of per-factor update rules and their gauge-invariant couplings.

###### Corollary A.3 (Deeper factorizations) .

Assume now a composite weight matrix W = U 1 U 2 ⋯ U D W=U_{1}U_{2}\cdots U_{D} of depth D D (written D D , not L L , which is the objective), so that the loss is invariant under the transformations at each factorization interface, ( U i , U i + 1 ) ↦ ( U i ​ Q i , Q i ⊤ ​ U i + 1 ) , Q i ∈ O ⁡ ( k i ) , (U_{i},U_{i+1})\mapsto(U_{i}Q_{i},\,Q_{i}^{\top}U_{i+1}),\qquad Q_{i}\in\mathrm{O}(k_{i}), with Q i Q_{i} orthogonal. Then gradients transform covariantly at each interface, with a change of basis on the side each factor acts. As a consequence, the earlier covariance and non-covariance conditions on the update rule carry over to each interface individually, with the caveat of transposition for the gradient terms. Any fixed nonlinear coordinate-wise update rule will fail to be equivariant at any interface where the internal dimension satisfies k i ≥ 2 k_{i}\geq 2 . On the other hand, if all factor gradients are scaled by the same positive scalar a ⁡ ( t ) a(t) , then the dynamics for all D D factors reduce to a single time-rescaled gradient flow, and the transfer theorem continues to hold under its original conditions. Equivariance at individual interfaces therefore does not by itself imply the transfer. The gauge freedom of the type discussed in § 6 is obtained when D = 2 D=2 .

###### Proposition A.4 (Balancedness: a special property of GD-type flows, not a general invariance) .

For standard gradient flow, and in particular in the setting of Theorem 4.6 , the quantity B t = U t ⊤ ​ U t − V t ⊤ ​ V t B_{t}=U_{t}^{\top}U_{t}-V_{t}^{\top}V_{t} is conserved during training. For diagonally preconditioned flows of the form U ˙ = − D U ⊙ ∇ U L \dot{U}=-D_{U}\odot\nabla_{U}L (and analogously for V V ), it is conserved when D U D_{U} and D V D_{V} are one and the same positive scalar, and drifts otherwise at a rate set by how far each departs from that common scalar. Agreeing up to a factor is not enough: D U ≡ 1 D_{U}\equiv 1 , D V ≡ 2 D_{V}\equiv 2 already breaks the cancellation. Nevertheless, conservation of B t B_{t} is not necessary for good low-rank recovery: neither Muon nor Shampoo conserves B t B_{t} , and both carry large final imbalances ‖ B T ‖ F \left\lVert B_{T}\right\rVert_{F} ( 1.47 1.47 – 1.65 1.65 ) next to gradient descent’s 0.06 0.06 (Table 2 ), yet both still recover well on the zoo task. The imbalance present at the 10 − 3 10^{-3} initialization is negligible at this precision, so those final values are accumulated drift. In particular, this suggests that balancedness is a property of flows with common-scalar preconditioning, not of equivariant optimization procedures in general.

###### Remark A.5 (Comments on discretization and step-size choice) .

Theorem 4.6 concerns continuous-time dynamics, and while the corresponding discrete-time statements hold under smoothness and stability assumptions, an explicit Euler discretization of the memoryless common-scalar flow is only guaranteed to track the trajectory up to an O ⁡ ( η ) O(\eta) error over a fixed time interval. In practice, this does not mean that decreasing the step size η \eta always improves performance, nor does it apply directly to scalar-Adam-style algorithms with fixed β \beta . The learning-rate trends of Appendix D.2 are specific to this task and empirical: both GD and scalar-Adam recover better at the smaller rates tested ( 0.131 → 0.113 0.131\to 0.113 for GD and 0.201 → 0.165 0.201\to 0.165 for scalar-Adam), while Adam’s recovery is almost insensitive to the choice of η \eta , staying within 0.56 0.56 – 0.58 0.58 across its grid.

## Appendix B Proofs

### B.1 Lemma 3.2 (maximal isometric gauge)

###### Proof.

For A ∈ GL ⁡ ( k ) A\in\mathrm{GL}(k) , ( U ​ A ) ​ ( V ​ A − ⁣ ⊤ ) ⊤ = U ​ A ​ A − 1 ​ V ⊤ = U ​ V ⊤ , (UA)(VA^{-\top})^{\top}=UAA^{-1}V^{\top}=UV^{\top}, so the full action preserves the represented matrix. If it is an isometry for every ( U , V ) (U,V) , take V = 0 V=0 and let U U have an arbitrary single nonzero row u ∈ ℝ 1 × k u\in\mathbb{R}^{1\times k} . Then ‖ u ​ A ‖ 2 = ‖ u ‖ 2 \left\lVert uA\right\rVert_{2}=\left\lVert u\right\rVert_{2} for every u u , hence A ​ A ⊤ = I AA^{\top}=I . Since A A is square, this is equivalent to A ∈ O ⁡ ( k ) A\in\mathrm{O}(k) . Conversely, if A ∈ O ⁡ ( k ) A\in\mathrm{O}(k) , both right multiplications preserve Frobenius norms, so the action is an isometry. ∎

### B.2 Lemma 4.1 (gradient covariance)

###### Proof.

Write L ⁡ ( U , V ) = f ⁡ ( U ​ V ⊤ ) L(U,V)=f(UV^{\top}) and let G = ∇ f ​ ( W ) ∈ ℝ n × n G=\nabla f(W)\in\mathbb{R}^{n\times n} at W = U ​ V ⊤ W=UV^{\top} . Then ∇ U L = G ​ V \nabla_{U}L=GV and ∇ V L = G ⊤ ​ U \nabla_{V}L=G^{\top}U . At the transformed point ( U ​ Q , V ​ Q ) (UQ,VQ) the represented matrix is unchanged, W ′ = U ​ Q ​ ( V ​ Q ) ⊤ = W W^{\prime}=UQ(VQ)^{\top}=W , so ∇ f ​ ( W ′ ) = G \nabla f(W^{\prime})=G and ∇ U L ​ ( U ​ Q , V ​ Q ) = G ⁡ ( V ​ Q ) = ( G ​ V ) ​ Q = ( ∇ U L ) ​ Q \nabla_{U}L(UQ,VQ)=G\,(VQ)=(GV)Q=(\nabla_{U}L)\,Q ; likewise ∇ V L ​ ( U ​ Q , V ​ Q ) = G ⊤ ​ ( U ​ Q ) = ( ∇ V L ) ​ Q \nabla_{V}L(UQ,VQ)=G^{\top}(UQ)=(\nabla_{V}L)\,Q . ∎

### B.3 Proposition 4.2 (equivariant class)

###### Proof.

Throughout, primed quantities denote the run initialized at ( U 0 ​ Q , V 0 ​ Q ) (U_{0}Q,V_{0}Q) , with the state initialized by σ Q \sigma_{Q} as constructed below. We show by induction that U t ′ = U t ​ Q , V t ′ = V t ​ Q , U^{\prime}_{t}=U_{t}Q,\qquad V^{\prime}_{t}=V_{t}Q, with the state transforming as stated. By Lemma 4.1 , if the iterates are gauge-related at time t t , then the factor gradients satisfy G t ′ = G t ​ Q G^{\prime}_{t}=G_{t}Q .

(1) GD and momentum. For GD, U t + 1 ′ = U t ′ − η ​ G t ′ = ( U t − η ​ G t ) ​ Q . U^{\prime}_{t+1}=U^{\prime}_{t}-\eta G^{\prime}_{t}=(U_{t}-\eta G_{t})Q. Momentum buffers are linear in the gradients, so by induction (with σ Q : M ↦ M ​ Q \sigma_{Q}:M\mapsto MQ ), M t ′ = β ​ M t − 1 ′ + G t ′ = M t ​ Q , M^{\prime}_{t}=\beta M^{\prime}_{t-1}+G^{\prime}_{t}=M_{t}Q, and the update again right-multiplies by Q Q . For Nesterov, the gradient is evaluated at the lookahead point U t + β ⁡ ( U t − U t − 1 ) U_{t}+\beta(U_{t}-U_{t-1}) . In the primed run this is the same point right-multiplied by Q Q , and Lemma 4.1 applies there verbatim. Hence the lookahead gradient is also covariant, and the induction goes through unchanged.

(2) Scalar-preconditioned Adam. The first moments transform as M t ′ = M t ​ Q M^{\prime}_{t}=M_{t}Q , since the EMA is linear. The second-moment state is the single scalar ν t = β 2 ​ ν t − 1 + ( 1 − β 2 ) ​ mean ​ ( G t ⊙ 2 ) , ν 0 = 0 , \nu_{t}=\beta_{2}\,\nu_{t-1}+(1-\beta_{2})\,\mathrm{mean}\bigl(G_{t}^{\odot 2}\bigr),\qquad\nu_{0}=0, where the mean is taken over all entries of both factors. Because the mean is linear, it commutes with the EMA, so this scalar equals mean ⁡ ( v t ) \mathrm{mean}(v_{t}) for the entrywise EMA v t v_{t} of g ⊙ 2 g^{\odot 2} . Storing the scalar therefore changes no number. It is also what makes the state map literal, since the entrywise v t v_{t} does not itself transform under a fixed map (squaring and right-rotation do not commute). Now, the entry mean of G ⊙ 2 G^{\odot 2} is ‖ G ‖ F 2 \left\lVert G\right\rVert_{F}^{2} divided by the entry count, so ν t \nu_{t} depends on the gradients only through their Frobenius norms. Since ‖ G t ​ Q ‖ F = ‖ G t ‖ F \left\lVert G_{t}Q\right\rVert_{F}=\left\lVert G_{t}\right\rVert_{F} , we get ν t ′ = ν t \nu^{\prime}_{t}=\nu_{t} . The bias-corrected read ν ^ t = ν t / ( 1 − β 2 t ) \hat{\nu}_{t}=\nu_{t}/(1-\beta_{2}^{t}) divides by a fixed scalar and is therefore invariant as well. The state map is σ Q : ( M U , M V , ν ) ↦ ( M U ​ Q , M V ​ Q , ν ) , \sigma_{Q}:(M_{U},M_{V},\nu)\mapsto(M_{U}Q,M_{V}Q,\nu), with σ Q ​ ( 0 ) = 0 \sigma_{Q}(0)=0 as Definition 3.1 requires. Finally, the update − η M ^ t / ( ν ^ t + ϵ ) -\eta\,\hat{M}_{t}/(\sqrt{\hat{\nu}_{t}}+\epsilon) is a scalar multiple of a covariant matrix, hence covariant.

(3) Muon. With momentum M t ′ = M t ​ Q M^{\prime}_{t}=M_{t}Q as in (1), it suffices to show that msign ⁡ ( M ​ Q ) = msign ⁡ ( M ) ​ Q \mathrm{msign}(MQ)=\mathrm{msign}(M)\,Q . Let M = A ​ Σ ​ B ⊤ M=A\Sigma B^{\top} be an SVD. Then M ​ Q = A ​ Σ ​ ( Q ⊤ ​ B ) ⊤ MQ=A\Sigma(Q^{\top}B)^{\top} is an SVD of M ​ Q MQ , so msign ⁡ ( M ​ Q ) = A ​ ( Q ⊤ ​ B ) ⊤ = A ​ B ⊤ ​ Q = msign ⁡ ( M ) ​ Q . \mathrm{msign}(MQ)=A(Q^{\top}B)^{\top}=AB^{\top}Q=\mathrm{msign}(M)\,Q. Next, consider the Newton–Schulz iteration X j + 1 = a ​ X j + b ⁡ ( X j ​ X j ⊤ ) ​ X j + c ​ ( X j ​ X j ⊤ ) 2 ​ X j X_{j+1}=aX_{j}+b\,(X_{j}X_{j}^{\top})X_{j}+c\,(X_{j}X_{j}^{\top})^{2}X_{j} with X 0 = M / ( ‖ M ‖ F + ϵ NS ) X_{0}=M/(\left\lVert M\right\rVert_{F}+\epsilon_{\mathrm{NS}}) . If X j ↦ X j ​ Q X_{j}\mapsto X_{j}Q , then X j ​ X j ⊤ X_{j}X_{j}^{\top} is invariant, so every term right-multiplies by Q Q , and the property is preserved at every iteration and at any truncation. Equivariance is therefore exact in real arithmetic for the finite practical algorithm, not only for the idealized msign \mathrm{msign} ; finite-precision sensitivity is a separate issue.

(4) Shampoo. The left accumulator L t = ∑ u ≤ t G u ​ G u ⊤ L_{t}=\sum_{u\leq t}G_{u}G_{u}^{\top} is invariant, since G u ​ Q ​ Q ⊤ ​ G u ⊤ = G u ​ G u ⊤ G_{u}QQ^{\top}G_{u}^{\top}=G_{u}G_{u}^{\top} , while R t = ∑ G u ⊤ ​ G u R_{t}=\sum G_{u}^{\top}G_{u} transforms as R t ′ = Q ⊤ ​ R t ​ Q R^{\prime}_{t}=Q^{\top}R_{t}Q . For any spectral function h h we have h ⁡ ( Q ⊤ ​ R ​ Q ) = Q ⊤ ​ h ​ ( R ) ​ Q h(Q^{\top}RQ)=Q^{\top}h(R)\,Q . Taking h ( R ) = ( R + λ I ) − 1 / 4 h(R)=(R+\lambda I)^{-1/4} , with the damping included since Q ⊤ ​ R ​ Q + λ ​ I = Q ⊤ ​ ( R + λ ​ I ) ​ Q Q^{\top}RQ+\lambda I=Q^{\top}(R+\lambda I)Q , we obtain ( L ′ t + λ I ) − 1 / 4 G ′ t ( R ′ t + λ I ) − 1 / 4 \displaystyle(L^{\prime}_{t}+\lambda I)^{-1/4}\,G^{\prime}_{t}\,(R^{\prime}_{t}+\lambda I)^{-1/4} = ( L t + λ I ) − 1 / 4 G t Q Q ⊤ ( R t + λ I ) − 1 / 4 Q \displaystyle=(L_{t}+\lambda I)^{-1/4}\,G_{t}Q\,Q^{\top}(R_{t}+\lambda I)^{-1/4}Q = [ ( L t + λ I ) − 1 / 4 G t ( R t + λ I ) − 1 / 4 ] Q . \displaystyle=\big[(L_{t}+\lambda I)^{-1/4}G_{t}(R_{t}+\lambda I)^{-1/4}\big]Q. ∎

### B.4 Proposition 4.3 (coordinate-wise rules break the gauge)

###### Proof.