The Loss Does Not See the Basis, but Adam Does — part 3/12
Cleaned-text code points [31199, 46221). Bibliography omitted here; source.md includes it.
Fix ϵ > 0 \epsilon>0 and consider bias-corrected Adam using the update convention θ + = θ − η m ^ ⊘ ( v ^ + ϵ ) \theta^{+}=\theta-\eta\,\hat{m}\oslash(\sqrt{\hat{v}}+\epsilon) , with elementwise operations, shared stepsize η \eta , and initial moment estimates set to zero. Define the function D ϵ ( G ) ≔ G ⊘ ( | G | + ϵ ) , D_{\epsilon}(G)\coloneqq G\oslash\bigl(\lvert G\rvert+\epsilon\bigr), applied entrywise, with | ⋅ | \lvert\cdot\rvert the entrywise absolute value. Let G U = ∇ U L ( U 0 , V 0 ) G_{U}=\nabla_{U}L(U_{0},V_{0}) and G V = ∇ V L ( U 0 , V 0 ) G_{V}=\nabla_{V}L(U_{0},V_{0}) denote the gradients at initialization, and let ( U ~ 1 , V ~ 1 ) (\widetilde{U}_{1},\widetilde{V}_{1}) be the first iterate obtained from the gauge-rotated starting point ( U 0 Q , V 0 Q ) (U_{0}Q,V_{0}Q) . Define E Q ( G ) ≔ D ϵ ( G Q ) Q ⊤ − D ϵ ( G ) , D U ≔ D ϵ ( G U ) , D V ≔ D ϵ ( G V ) . E_{Q}(G)\coloneqq D_{\epsilon}(GQ)Q^{\top}-D_{\epsilon}(G),\qquad D_{U}\coloneqq D_{\epsilon}(G_{U}),\qquad D_{V}\coloneqq D_{\epsilon}(G_{V}). Then the difference in the resulting factorized products after one step satisfies the identity W ~ 1 − W 1 \displaystyle\widetilde{W}_{1}-W_{1} = − η [ E Q ( G U ) V 0 ⊤ + U 0 E Q ( G V ) ⊤ ] \displaystyle=-\eta\bigl[E_{Q}(G_{U})V_{0}^{\top}+U_{0}E_{Q}(G_{V})^{\top}\bigr] + η 2 [ E Q ( G U ) D V ⊤ + D U E Q ( G V ) ⊤ + E Q ( G U ) E Q ( G V ) ⊤ ] . \displaystyle+\eta^{2}\bigl[E_{Q}(G_{U})D_{V}^{\top}+D_{U}E_{Q}(G_{V})^{\top}+E_{Q}(G_{U})E_{Q}(G_{V})^{\top}\bigr]. Equivalently, the second-order term can be rewritten as η 2 [ D ϵ ( G U Q ) D ϵ ( G V Q ) ⊤ − D U D V ⊤ ] \eta^{2}\bigl[D_{\epsilon}(G_{U}Q)D_{\epsilon}(G_{V}Q)^{\top}-D_{U}D_{V}^{\top}\bigr] .
This shows that Adam’s failure of gauge equivariance manifests already at the level of the represented matrix in the first step—not merely as a mismatch in internal states. For example, with n = 1 n=1 , k = 2 k=2 , f ( w ) = 1 2 w 2 f(w)=\tfrac{1}{2}w^{2} , U 0 = V 0 = ( 1 , 0 ) U_{0}=V_{0}=(1,0) , and rotation matrix Q = 2 − 1 / 2 ( 1 1 − 1 1 ) Q=2^{-1/2}\begin{pmatrix}1&1\\ -1&1\end{pmatrix} , the two resulting products after one step are ( 1 − η / ( 1 + ϵ ) ) 2 (1-\eta/(1+\epsilon))^{2} and ( 1 − η / ( 2 − 1 / 2 + ϵ ) ) 2 (1-\eta/(2^{-1/2}+\epsilon))^{2} , which differ whenever 0 < η < 2 − 1 / 2 + ϵ 0<\eta<2^{-1/2}+\epsilon .
So far the classification is a list of examples. For memoryless update rules it can be made exact, and the exact answer is narrow: equivariance forces the update to be a left preconditioner whose only input is the Gram matrix.
###### Theorem 4.5 (Structure: equivariant memoryless rules are Gram-determined left preconditioners) .
Let Φ \Phi be a memoryless update rule defined on gradients G ∈ ℝ n × k G\in\mathbb{R}^{n\times k} with full column rank ( k ≤ n k\leq n ). Then Φ ( G Q ) = Φ ( G ) Q \Phi(GQ)=\Phi(G)Q holds for all Q ∈ O ( k ) Q\in\mathrm{O}(k) if and only if Φ ( G ) = H ( G G ⊤ ) G \Phi(G)\;=\;H\!\left(GG^{\top}\right)G for some matrix-valued function H H depending only on the Gram matrix G G ⊤ GG^{\top} —a quantity invariant under gauge transformations. The function H H can be canonically expressed as H ( G G ⊤ ) = Φ ( G ) G + H(GG^{\top})=\Phi(G)\,G^{+} , where G + G^{+} is the Moore–Penrose pseudoinverse. In the square invertible case ( k = n k=n ), this simplifies uniquely to H ( P ) = Φ ( P 1 / 2 ) P − 1 / 2 H(P)=\Phi(P^{1/2})\,P^{-1/2} . Full column rank is a generic property—an open, dense condition, holding with probability one under random initialization—so the rank-deficient set has measure zero and Φ \Phi is left unconstrained on it. The only point of that set the dynamics are forced through is G = 0 G=0 , which occurs at (i) perfect interpolation or (ii) a degenerate stationary point at which one factor has become zero. That case carries no freedom either, since the update vanishes identically there: equivariance forces Φ ( 0 ) = Φ ( 0 ) Q \Phi(0)=\Phi(0)Q for any Q ∈ O ( k ) Q\in\mathrm{O}(k) , so Φ ( 0 ) = 0 \Phi(0)=0 .
The algebra is elementary; the classification it buys is not. It demonstrates that the four cases in Proposition 4.2 are not ad hoc examples but the only shape available, up to the choice of representation. And it shows where the remaining flexibility sits. Right-gauge equivariance alone allows H H to be any matrix-valued function of G G ⊤ GG^{\top} . For instance, a fixed nonscalar left multiplier Φ ( G ) = D G \Phi(G)=DG preserves gauge symmetry without acting solely on individual singular values; even imposing left-orthogonal equivariance still permits reweighting in each singular direction to depend on the full spectrum (e.g. Φ ( G ) = ‖ G ‖ F G \Phi(G)=\left\lVert G\right\rVert_{F}\,G is bi-orthogonally equivariant).
When H H is further restricted to a fixed scalar function of the spectrum of G G ⊤ GG^{\top} , Proposition A.1 reduces it to a single univariate spectral transfer function h ( σ ) h(\sigma) : the update moves along the gradient’s own singular directions and rescales each one, with h ( σ ) = η σ h(\sigma)=\eta\sigma for GD and h ≡ η h\equiv\eta for the exact polar map. That one coordinate is the schedule axis of § 8 . Stateful methods instead require explicit handling of their evolving covariant state. Appendix A states the spectral proposition, the extensions to stateful rules and to deeper factorizations, and what balancedness does inside the class.
###### Theorem 4.6 (Transfer theorem) .
Let a ( t ) > 0 a(t)>0 be a measurable function such that 1 / a 1/a is locally integrable, and suppose a ( t ) a(t) is determined solely by gauge-invariant statistics of the trajectory up to time t t . Consider the scalar-preconditioned dynamics θ ˙ = − ∇ L ( θ ) / a ( t ) , θ = ( U , V ) , \dot{\theta}=-\nabla L(\theta)/a(t),\qquad\theta=(U,V), where θ \theta is locally absolutely continuous, and define the time transformation τ ( t ) = ∫ 0 t 𝑑 u / a ( u ) \tau(t)=\int_{0}^{t}du/a(u) . Then the reparameterized trajectory θ ~ ( τ ) ≔ θ ( t ( τ ) ) \widetilde{\theta}(\tau)\coloneqq\theta(t(\tau)) satisfies θ ~ ′ ( τ ) = − ∇ L ( θ ~ ( τ ) ) \widetilde{\theta}^{\prime}(\tau)=-\nabla L\bigl(\widetilde{\theta}(\tau)\bigr) for almost every τ \tau , meaning both trajectories follow the same path in parameter space. If in addition ∫ 0 ∞ 𝑑 u / a ( u ) = ∞ \int_{0}^{\infty}du/a(u)=\infty —which holds, for instance, when a a is bounded above—the transformed time reaches infinity, ensuring that if the gradient flow converges, so does the preconditioned version, to the same limit.
Consequently, any conclusion regarding a gradient flow which relies solely on the path or limit point is valid under the joint assumptions of the above theorem and the original result; this is true for the restricted-regime results ( Gunasekar et al., 2017 ; Arora et al., 2019 ; Li et al., 2021 ) . However, conclusions which are time-dependent, such as rates of convergence or hitting times, cannot be transferred because of the time rescaling. In the absence of the divergence condition, only the finite part of the path up to τ max = ∫ 0 ∞ 𝑑 u / a ( u ) \tau_{\max}=\int_{0}^{\infty}du/a(u) is shared.
###### Remark 4.7 (Scope of the transfer, and what equivariance does not buy) .
Theorem 4.6 applies to the memoryless, common-scalar continuous flow, so it reaches gradient descent and common-scalar-preconditioned GD and no further. Scalar-Adam carries a first-moment exponential moving average (EMA) and is therefore not strictly an instance of it; we treat its agreement with gradient flow as empirically observed rather than formally implied (Remark A.5 , § 5 ). Equivariance is what makes the transfer available, but on its own it is neither necessary nor sufficient for low-rank recovery. The second factor is the spectral schedule—the greedy, sequential growth of singular values from small initialization ( Li et al., 2021 ) . ScaledGD ( Tong et al., 2021 ) is the clean counterexample on one side: gauge-equivariant, yet it equalizes the convergence rate across all the singular values and gives up the greedy behavior we rely on ( erank ≈ 13.7 \mathrm{erank}\approx 13.7 on the zoo task, not tabulated). 1 1 1 Our implementation of ScaledGD uses a factor-aware preconditioner inspired by Tong et al. (2021) , but serves here as an equivariant control with equalized schedule, since we do not have a tuned reproducibility baseline for ScaledGD. Appendix D.10 supplies the counterexample on the other side, a non-equivariant rule (long-anneal signum) that recovers anyway. Within the class, how hard a method flattens the spectrum (empirically GD < < Shampoo < < Muon on our runs) is the axis that remains; for memoryless updates it is the h h of Proposition A.1 , and § 8 maps it.
## 5 The optimizer zoo
##### Design.
We evaluate nine optimization algorithms on the sensing task described in § 3 . Four of these are theoretically expected to be equivariant: gradient descent (GD), scalar-Adam (using the p = 0 p{=}0 setting from § 7 with a gauge-invariant root-mean-square (RMS) scalar), Muon ( Jordan et al., 2024 ) , and Shampoo ( Gupta et al., 2018 ) . The remaining five are predicted to break gauge symmetry: Adam ( Kingma & Ba, 2015 ) , RMSProp ( Tieleman & Hinton, 2012 ) , signSGD with momentum ( Bernstein et al., 2018 , “signum”;) , Lion ( Chen et al., 2023 ) , and Adafactor ( Shazeer & Stern, 2018 ) . Experiments use three paired seeds, with recovery averaged across seeds. For each such method, we pick the learning rate yielding the best recovery among those on that method’s own grid that reach interpolation. This is selection on the outcome, but it is applied identically to every method and it favors the coordinate-wise ones, so the class gap it produces is a lower bound. Appendix D.2 replaces the single selected number with the full learning-rate curves and the “can/cannot” reading of them. The four methods labeled (cosine) in Table 2 require learning-rate decay to achieve interpolation, as discussed in § 3 : Muon, signum, and Lion because of their constant update norm, and RMSProp because of its ϵ \epsilon -bounded denominator. All nine optimizers achieve the 10 − 7 10^{-7} threshold for interpolation (the loss is evaluated every 200 200 steps, so the faster methods are recorded several orders of magnitude below the threshold; for instance, Adam reaches a final loss of 1.2 × 10 − 11 1.2\times 10^{-11} ). To put these numbers in perspective, we also report the solution to the convex problem suggested by gradient-flow theory: the minimum of the nuclear norm min ‖ X ‖ ∗ \min\left\lVert X\right\rVert_{\ast} subject to ⟨ A i , X ⟩ = y i \left\langle A_{i},X\right\rangle=y_{i} , on the same three problem instances ( experiments/nuclear_norm_reference.py ). This forms the reference row of Table 2 ; the solve attains a feasibility residual of 7 × 10 − 16 7\times 10^{-16} and recovers X ∗ X^{\ast} on two of the three seeds, while on the third it finds a solution of strictly lower nuclear norm than X ∗ X^{\ast} ( 38.416 38.416 vs. 38.509 38.509 ), so X ∗ X^{\ast} is not the minimum-nuclear-norm solution at m = 2 dof m=2\,\mathrm{dof} for that seed’s measurements. Main results are displayed in Table 2 and Figure 1 ; further control experiments over the scheduling choices and selection criteria are deferred to Appendix D .
##### Result.
The two clusters of predicted methods are cleanly separated: all equivariant methods perform strictly better than any coordinate-wise method ( ≤ 0.286 \leq 0.286 vs. ≥ 0.42 \geq 0.42 ), leaving a 0.14 0.14 gap between the best performer in the latter group and the worst performer in the former at this computational budget (see Figure 1 ). Because the table quotes three-seed means with no dispersion, the separation claim itself is deferred to the 10 10 -seed ladder of Appendix D.9 , where it persists at every problem size up to n = 256 n{=}256 (GD, scalar-Adam, and Muon against the clean coordinate-wise methods, with two flagged exceptions). Since all nine methods manage to interpolate, the difference must be attributed to the choice of interpolant. The reference row lets us ground this in the absolute recovery error: compared to the min-nuclear-norm solution ( 0.0335 0.0335 ), Muon is better, while GD, scalar-Adam, and Shampoo are within an order of magnitude and the coordinate-wise approaches are all worse by more than an order of magnitude. Two points about that row are worth emphasizing. First, GD is worse than the convex benchmark and converges toward it as the step size decreases ( 0.131 → 0.113 0.131\to 0.113 , and 0.108 0.108 on the finer ladder grid) but never reaches it within this budget. This illustrates the distinction made in Remark A.5 between gradient flow at infinitesimal initialization and the discrete GD we actually run, at initialization scale 10 − 3 10^{-3} and on a finite budget: the nuclear-norm characterization is a restricted-regime statement about the former, while practitioners must contend with the latter. Second, while Muon outperforms the convex benchmark, this is not a contradiction, since on the single seed where nuclear-norm minimization did not return X ∗ X^{\ast} , Muon did recover it; hence these methods do more than just approximate the convex relaxation. Muon is behaving differently because of its equal-rate schedule ( Kang et al., 2026 ) , a bias not related to any norm—the reading § 8 develops.
We note three observations in particular. (i) Muon is capable of close to exact recovery (recovery error 6.8 × 10 − 6 6.8\times 10^{-6} , effective rank 2.95 2.95 instead of the true rank 3 3 ) on this task once cosine decay permits interpolation. This demonstrates a stronger inductive bias towards low-rank solutions than gradient descent shows at practical learning rates, as one would expect from its equal-rate spectral dynamics ( Kang et al., 2026 ) when the true matrix is exactly low-rank. (ii) Adafactor eliminates the low-rank bias as completely as Adam does; the per-coordinate anisotropy of the factored diagonal approximation makes this a matter of gauge dependence rather than of memory footprint. (iii) Finally, signum showcases why recovery, and not effective rank, is the metric to rely on: with the lowest effective rank of any coordinate-wise method ( 7.8 7.8 , against Shampoo’s 6.95 6.95 ) but a recovery error of 0.445 0.445 , signum still falls among the worst optimizers for this task. Its updates are concentrated along only a few directions, and those directions are rarely aligned with the true solution. (When doubling the length of the annealing phase, signum is in fact the only optimizer to return to the low-recovery regime, but through a different and more volatile mechanism than the flowing optimizers, rather than by preserving the flow-induced bias. This is discussed further in Appendix D.10 .)
##### Controls (Appendix D ).