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

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

Xie et al. (2025) show that SGD’s optimization trajectory is rotationally invariant while Adam’s is basis-dependent, attributing Adam’s empirical efficacy to its alignment with coordinate ( ℓ ∞ \ell_{\infty} ) geometry. We examine the complementary aspect of this basis dependence, specifically isolating the selection cost it imposes. § 10.1 details the conditions under which these trade-offs favor one algorithm over the other. To the best of our knowledge, none of these works explore low-rank recovery or interpolant choice in factored models, the setting where Adam’s ℓ ∞ \ell_{\infty} -based coordinate structure interacts with Gunasekar-style implicit bias. That setting is the focus of this paper. DePavia et al. (2025) show that small rotations of data or feature space alter the richness bias of Adam in classification settings, and propose an orthogonal reparameterization to alleviate this. Our work is complementary to theirs: we study the internal gauge of parameter factorizations and its impact on solution selection, and we characterize the optimizer zoo—the nine deployed update rules of § 5 —in relation to this property.

##### Optimizer design, symmetry, and structured methods.

Table 1 summarizes the most closely related prior and concurrent studies, categorized by the symmetry examined and their relation to this paper. Appendix E covers this literature in full: transformation invariance in optimizer design, symmetry-driven accounts of implicit bias, and Muon and structured optimizers.

##### Our own prior work.

FlowAdam ( Singh & Sheel, 2026 ) enhanced Adam with clipped gradient-flow injection and demonstrated improvements on coupled problems using default settings. The current paper shifts focus from algorithm development to mechanistic explanation. In § 10 , we show that our symmetry criterion retrospectively clarifies FlowAdam’s design evolution: its original per-coordinate clipping introduced a gauge-breaking transformation that inadvertently constrained the very bias it aimed to introduce, whereas correcting this with global-norm clipping restores gradient descent’s bias within the hybrid method.

## 3 Setup: gauge, task, and protocol

##### Factored model and gauge transformation.

Let f : ℝ n × n → ℝ f:\mathbb{R}^{n\times n}\to\mathbb{R} be a differentiable function, and consider the factored objective L ⁡ ( U , V ) = f ⁡ ( U ​ V ⊤ ) L(U,V)=f(UV^{\top}) , where U , V ∈ ℝ n × k U,V\in\mathbb{R}^{n\times k} and the representation is overparameterized ( k k at least the rank of any solution of interest). For any orthogonal matrix Q ∈ O ⁡ ( k ) Q\in\mathrm{O}(k) , define the gauge transformation ρ Q ​ ( U , V ) = ( U ​ Q , V ​ Q ) . \rho_{Q}(U,V)=(UQ,\,VQ). Since ( U ​ Q ) ​ ( V ​ Q ) ⊤ = U ​ V ⊤ (UQ)(VQ)^{\top}=UV^{\top} , it follows that L ∘ ρ Q = L L\circ\rho_{Q}=L : this gauge transformation is a symmetry of all such factored objectives, regardless of f f . We consider the square case for notational simplicity; all subsequent claims hold for rectangular inputs f : ℝ d 1 × d 2 → ℝ f:\mathbb{R}^{d_{1}\times d_{2}}\to\mathbb{R} with L ⁡ ( U , V ) = f ⁡ ( U ​ V ⊤ ) L(U,V)=f(UV^{\top}) , U ∈ ℝ d 1 × k U\in\mathbb{R}^{d_{1}\times k} and V ∈ ℝ d 2 × k V\in\mathbb{R}^{d_{2}\times k} (as in the hyperspectral experiments of § 9 ), since the gauge only acts on the shared latent dimension k k .

###### Definition 3.1 (Gauge equivariance of an optimizer) .

Consider an optimizer with internal state s t s_{t} (e.g. momentum terms, preconditioning matrices) and update rule ( U t + 1 , V t + 1 , s t + 1 ) = 𝒜 ⁡ ( U t , V t , s t , ∇ L ) . (U_{t+1},V_{t+1},s_{t+1})=\mathcal{A}(U_{t},V_{t},s_{t};\,\nabla L). It is gauge-equivariant if for every Q ∈ O ⁡ ( k ) Q\in\mathrm{O}(k) there is a corresponding transformation σ Q \sigma_{Q} on the state such that initializing from ( ρ Q ​ ( U 0 , V 0 ) , σ Q ​ ( s 0 ) ) \bigl(\rho_{Q}(U_{0},V_{0}),\,\sigma_{Q}(s_{0})\bigr) leads to states ( ρ Q ​ ( U t , V t ) , σ Q ​ ( s t ) ) \bigl(\rho_{Q}(U_{t},V_{t}),\,\sigma_{Q}(s_{t})\bigr) at every step t t .

All methods we consider have zero-initialized states, and we impose σ Q ​ ( 0 ) = 0 \sigma_{Q}(0)=0 , ensuring that both trajectories start from the same initial conditions. That convention is not what excludes Adam: its entrywise second moment admits no state map σ Q \sigma_{Q} at all, whatever it does at zero, because squaring and right-rotation do not commute—which is exactly what the scalar version repairs (proof of Proposition 4.2 (2)). For an equivariant optimizer it follows that the product matrices W t = U t ​ V t ⊤ W_{t}=U_{t}V_{t}^{\top} are invariant under the gauge action; that invariance is exactly what the coordinate-wise rules lose. This defines a trajectory on the orthogonal orbit of the initial pair ( U 0 , V 0 ) (U_{0},V_{0}) , which is not uniquely defined by W 0 W_{0} alone (since the full symmetry group is GL ⁡ ( k ) \mathrm{GL}(k) , and even gradient descent depends on the balance term U 0 ⊤ ​ U 0 − V 0 ⊤ ​ V 0 U_{0}^{\top}U_{0}-V_{0}^{\top}V_{0} , which can vary within a GL ⁡ ( k ) \mathrm{GL}(k) orbit). We restrict to the orthogonal subgroup O ⁡ ( k ) \mathrm{O}(k) , which is the largest subgroup acting by isometries (Lemma 3.2 ), under which gradients transform covariantly and Frobenius-norm statistics are preserved (see the attention-related footnote in § 6 ).

###### Lemma 3.2 (Orthogonal gauge as maximal isometric symmetry) .

For any A ∈ GL ⁡ ( k ) A\in\mathrm{GL}(k) , the transformation ( U , V ) ↦ ( U ​ A , V ​ A − ⁣ ⊤ ) (U,V)\mapsto(UA,\,VA^{-\top}) preserves the product U ​ V ⊤ UV^{\top} . This transformation preserves the Euclidean product metric ‖ U ‖ F 2 + ‖ V ‖ F 2 \left\lVert U\right\rVert_{F}^{2}+\left\lVert V\right\rVert_{F}^{2} for all ( U , V ) (U,V) if and only if A ∈ O ⁡ ( k ) A\in\mathrm{O}(k) .

A critical property that makes this symmetry useful for verification is that gradients transform consistently with the factors: ∇ U L ​ ( U ​ Q , V ​ Q ) = ( ∇ U L ​ ( U , V ) ) ​ Q , \nabla_{U}L(UQ,VQ)=\bigl(\nabla_{U}L(U,V)\bigr)Q, and similarly for V V (Lemma 4.1 ). Optimization updates constructed from operations that respect right multiplication by Q Q will therefore exhibit equivariance, while updates that apply a nonlinear operation entrywise in the ( U , V ) (U,V) coordinates will generally not.

##### Task: recovery in the interpolation regime.

Our test-case task is matrix sensing under the implicit-bias setup of Gunasekar et al. (2017) and Recht et al. (2010) : sensing a planted matrix X ∗ = U ∗ ​ V ∗ ⁣ ⊤ X^{\ast}=U^{\ast}V^{\ast\top} ( n = 40 n{=}40 ) of rank r ∗ = 3 r^{\ast}{=}3 through m = 2 ​ dof m=2\,\mathrm{dof} Gaussian measurements y i = ⟨ A i , X ∗ ⟩ y_{i}=\left\langle A_{i},X^{\ast}\right\rangle with dof = r ∗ ​ ( 2 ​ n − r ∗ ) \mathrm{dof}=r^{\ast}(2n-r^{\ast}) . We consider full overparameterization ( k = n k{=}n ), small initialization ( 10 − 3 10^{-3} ), no weight decay, and optimize the squared loss on the measurements. The parameter space is underdetermined: infinitely many matrices (low- and full-rank) are compatible with any given measurements.

All methods are stopped at interpolation (loss below 10 − 7 10^{-7} , final loss reported), to focus the comparison on differences in the set of possible minima reached once the model has enough expressivity to fit the data. Accordingly, all methods achieve interpolation (see Appendix D ); but four of the nine considered—Muon, Lion, RMSProp, and signum (signSGD with momentum)—require learning-rate decay specifically shaped to their own dynamics (see below). We report recovery and effective rank, rec ( W ) = ‖ W − X ∗ ‖ F / ‖ X ∗ ‖ F , erank ( W ) = exp ( − ∑ i π i log π i ) , \mathrm{rec}(W)=\left\lVert W-X^{\ast}\right\rVert_{F}/\left\lVert X^{\ast}\right\rVert_{F},\qquad\mathrm{erank}(W)=\exp\Bigl(-\textstyle\sum_{i}\pi_{i}\log\pi_{i}\Bigr), with π i = σ i / ∑ j σ j \pi_{i}=\sigma_{i}/\sum_{j}\sigma_{j} and 0 ​ log ⁡ 0 ≔ 0 0\log 0\coloneqq 0 ( Roy & Vetterli, 2007 ) . Throughout, recovery means the error rec ⁡ ( W ) \mathrm{rec}(W) , so smaller is better and “best recovery” always means the smallest value; the spectral shares π i \pi_{i} are unrelated to the dial parameter p p of § 7 . Recovery is our main concern, as effective rank can be misleading: a method might reach a low-rank matrix in parameter space far from the true X ∗ X^{\ast} (see § 5 ).

##### Protocol.

Implicit-bias experiments are notoriously rife with pitfalls and subtle interactions; some of the most critical ones are addressed by the following controls, which were applied in all experiments reported in the main text and the appendices. More discussion is in Appendix D . • No weight decay in any implicit-bias experiment. Not only is implicit bias easily masked by explicit regularization, but interactions between the two (e.g. a regularizer placed in the loss versus decoupled from it) would require careful dissection to properly disentangle.

• Run to full interpolation, or report inability to do so. See Appendix D.7 for an example of a sign-descent method where early stopping (at loss 5 × 10 − 2 5\times 10^{-2} ) creates the illusion of good recovery.

• Symmetric scheduling. Interpolation either with a fixed learning rate (preferred) or with cosine decay; methods with constant update norm (Muon, signum, Lion) are compared against schedules that decay the step size to zero, to avoid an easy win for methods that cannot interpolate at a constant step size. RMSProp needs decay for a different reason, given in Appendix D.7 . This is a per-method criterion rather than a class-level one; it is informed by observations, rather than theory; and it is supplemented by all methods being compared with fully decayed schedules, so that differences in learning-rate decay cannot explain differences in outcomes (Appendix D.1 ).

• Sweeps of learning rates, not single values. Picking the best recovery value a posteriori is outcome bias; picking the biggest value compatible with convergence is speed bias. We therefore report full recovery-versus-learning-rate curves and frame discussion around asymmetries in achievable performance (Appendices D.2 and D.3 ), rather than points.

• Report recovery and rank jointly, not as separate metrics ; use the same random seeds for all methods; and use deterministic, full-batch updates (where relevant) to compare related methods (§ 6 ).

## 4 Theory: equivariance as a classifier of update rules

All proofs are provided in Appendix B . Once the orthogonal gauge is fixed, verifying covariance becomes straightforward. However, the implications go beyond algebra—they determine which update rules maintain the intrinsic geometry of the factorization, rather than depending on an arbitrary choice of basis.

###### Lemma 4.1 (Gradient covariance) .

Let L ⁡ ( U , V ) = f ⁡ ( U ​ V ⊤ ) L(U,V)=f(UV^{\top}) and let Q ∈ O ⁡ ( k ) Q\in\mathrm{O}(k) . Then ∇ U L ​ ( U ​ Q , V ​ Q ) = ( ∇ U L ​ ( U , V ) ) ​ Q , ∇ V L ​ ( U ​ Q , V ​ Q ) = ( ∇ V L ​ ( U , V ) ) ​ Q . \nabla_{U}L(UQ,VQ)=\bigl(\nabla_{U}L(U,V)\bigr)Q,\qquad\nabla_{V}L(UQ,VQ)=\bigl(\nabla_{V}L(U,V)\bigr)Q.

###### Proposition 4.2 (Equivariant class) .

The following update rules satisfy gauge equivariance as defined in Definition 3.1 . 1. Gradient descent and Polyak/Nesterov momentum : any linear combination of past gradients computed at appropriately transformed points.

2. Scalar-preconditioned Adam : a version of Adam where the per-coordinate second-moment estimate is replaced by a single scalar ν ^ t \hat{\nu}_{t} , obtained via bias-corrected exponential moving average (EMA) of gauge-invariant quantities (such as the average of all entries in g ⊙ 2 g^{\odot 2} across both factors). The update uses denominator a t = ν ^ t + ϵ a_{t}=\sqrt{\hat{\nu}_{t}}+\epsilon and step direction Δ ∝ m ^ t / a t \Delta\propto\hat{m}_{t}/a_{t} .

3. Muon : updates given by Δ = msign ⁡ ( M t ) \Delta=\mathrm{msign}(M_{t}) , where M t M_{t} is the momentum buffer and msign ⁡ ( M ) = A ​ B ⊤ \mathrm{msign}(M)=AB^{\top} comes from the compact SVD M = A ​ Σ ​ B ⊤ M=A\Sigma B^{\top} (with msign ⁡ ( 0 ) = 0 \mathrm{msign}(0)=0 ). Equivariance holds exactly under real arithmetic and remains valid for any finite-order Newton–Schulz approximation.

4. Shampoo : the step Δ = ( L t + λ I ) − 1 / 4 G t ( R t + λ I ) − 1 / 4 , L t = ∑ u ≤ t G u G u ⊤ , R t = ∑ u ≤ t G u ⊤ G u , \Delta=(L_{t}+\lambda I)^{-1/4}\,G_{t}\,(R_{t}+\lambda I)^{-1/4},\qquad L_{t}=\sum_{u\leq t}G_{u}G_{u}^{\top},\quad R_{t}=\sum_{u\leq t}G_{u}^{\top}G_{u}, with damping λ > 0 \lambda>0 , ensuring the matrix roots are well-defined even when matrices are rank-deficient early in training (our experiments use λ = 1 \lambda=1 ; for λ = 0 \lambda=0 , interpret the roots on their active subspace). The subscripts L t L_{t} and R t R_{t} always refer to Shampoo’s left and right accumulators, never to the objective L L . In particular, note that equivariance is achieved exactly in this setting even with damping.

In all cases, the resulting product of matrices W t = U t ​ V t ⊤ W_{t}=U_{t}V_{t}^{\top} is uniquely defined by the input data and the gauge orbit of the initial factors ( U 0 , V 0 ) (U_{0},V_{0}) , independently of the specific basis used for those factors. The results hold for the full-matrix cases described above, assuming real arithmetic. Implementations of practical interest that use variable splitting, mixed update rules, coordinate-wise clipping, or other non-equivariant operations on subsets of the parameters are a different matter: in these cases, equivariance is achieved only within the subsets of parameters updated by the stated rules.

###### Proposition 4.3 (Nonlinear coordinate-wise maps break the gauge) .

Assume k ≥ 2 k\geq 2 , and suppose that a memoryless update function Φ \Phi has the form Φ ​ ( G ) i ​ j = ϕ ⁡ ( G i ​ j ) \Phi(G)_{ij}=\phi(G_{ij}) , where ϕ : ℝ → ℝ \phi:\mathbb{R}\to\mathbb{R} is a fixed function. If the update function satisfies Φ ⁡ ( G ​ Q ) = Φ ⁡ ( G ) ​ Q for all matrices ​ G ​ and all orthogonal transformations ​ Q ∈ O ⁡ ( k ) , \Phi(GQ)=\Phi(G)\,Q\qquad\text{for all matrices }G\text{ and all orthogonal transformations }Q\in\mathrm{O}(k), then ϕ \phi must be linear, ϕ ⁡ ( x ) = c ​ x \phi(x)=cx for some real constant c c . In other words, any optimization algorithm that applies a fixed nonlinear function at the level of coordinates, at the zero optimizer state, cannot be gauge-equivariant.

The reason is that by Definition 3.1 , the zero state must be taken to itself by any gauge transformation. Both the original and the gauge-transformed run therefore start identically and apply the same update Φ \Phi , so equivariance requires Φ ⁡ ( G ​ Q ) = Φ ⁡ ( G ) ​ Q \Phi(GQ)=\Phi(G)Q for all G G , which is only possible if ϕ \phi is linear. Many standard optimization algorithms have such coordinate-wise nonlinearities in their first update—Adam, RMSProp, signSGD, Lion—while Adafactor carries them in its factored statistics and is analyzed separately in the proof. (Checking the equivariance condition for subsequent updates would require reasoning about how the optimizer state evolves under gauge transformations, but starting from the zero state avoids this complication.)

###### Proposition 4.4 (Adam has an exact first-step gauge defect) .