The Loss Does Not See the Basis, but Adam Does — part 8/12
Cleaned-text code points [107230, 123172). Bibliography omitted here; source.md includes it.
Consider a memoryless entrywise map Φ ( G ) i j = ϕ ( G i j ) \Phi(G)_{ij}=\phi(G_{ij}) . Equivariance for all Q ∈ O ( k ) Q\in\mathrm{O}(k) requires Φ ( G Q ) = Φ ( G ) Q \Phi(GQ)=\Phi(G)\,Q for all G G . Take k ≥ 2 k\geq 2 and let Q θ Q_{\theta} be a rotation by θ \theta in the first two coordinates. Apply the requirement to matrices G G supported on a single row ( x , y , 0 , … ) (x,y,0,\dots) : it reads ϕ ( x cos θ − y sin θ ) = ϕ ( x ) cos θ − ϕ ( y ) sin θ , ϕ ( x sin θ + y cos θ ) = ϕ ( x ) sin θ + ϕ ( y ) cos θ \phi(x\cos\theta-y\sin\theta)=\phi(x)\cos\theta-\phi(y)\sin\theta,\qquad\phi(x\sin\theta+y\cos\theta)=\phi(x)\sin\theta+\phi(y)\cos\theta for all x , y , θ x,y,\theta . Taking x = y = 0 x=y=0 in the first identity gives ϕ ( 0 ) = ϕ ( 0 ) ( cos θ − sin θ ) \phi(0)=\phi(0)(\cos\theta-\sin\theta) for every θ \theta , hence ϕ ( 0 ) = 0 \phi(0)=0 (take θ = π / 2 \theta=\pi/2 ). Setting y = 0 y=0 in both identities then gives ϕ ( x cos θ ) = ϕ ( x ) cos θ \phi(x\cos\theta)=\phi(x)\cos\theta and ϕ ( x sin θ ) = ϕ ( x ) sin θ \phi(x\sin\theta)=\phi(x)\sin\theta ; as θ \theta ranges over [ 0 , 2 π ) [0,2\pi) its cosine and sine sweep all of [ − 1 , 1 ] [-1,1] , so ϕ ( c x ) = c ϕ ( x ) \phi(cx)=c\,\phi(x) for every c ∈ [ − 1 , 1 ] c\in[-1,1] and every x ∈ ℝ x\in\mathbb{R} . For any x x and any u u with 0 < | u | ≤ | x | 0<|u|\leq|x| , choosing c = u / x ∈ [ − 1 , 1 ] c=u/x\in[-1,1] gives ϕ ( u ) = ( u / x ) ϕ ( x ) \phi(u)=(u/x)\,\phi(x) , i.e. ϕ ( u ) / u = ϕ ( x ) / x \phi(u)/u=\phi(x)/x . Since any two nonzero reals are both dominated in absolute value by a common third real, the ratio ϕ ( x ) / x \phi(x)/x takes a single value c c on ℝ ∖ { 0 } \mathbb{R}\setminus\{0\} , whence ϕ ( x ) = c x \phi(x)=c\,x for all x x : linearity, with no appeal to additivity. Thus any nonlinear fixed entrywise map fails equivariance already for a 2 × 2 2\times 2 rotation. The standard zero-state first steps of Adam and RMSProp are nonlinear maps of this form, as are signSGD and Lion’s outer-sign update. An explicit counterexample with θ = π / 4 \theta=\pi/4 is immediate (with the usual convention sign ( 0 ) = 0 \mathrm{sign}(0)=0 ). Take the single row G = ( 1 , 1 ) G=(1,1) , so G Q π / 4 = ( 0 , 2 ) GQ_{\pi/4}=(0,\sqrt{2}) . Applying sign after the rotation gives sign ( G Q π / 4 ) = ( 0 , 1 ) \mathrm{sign}(GQ_{\pi/4})=(0,1) , whereas rotating the sign image gives sign ( G ) Q π / 4 = ( 1 , 1 ) Q π / 4 = ( 0 , 2 ) \mathrm{sign}(G)\,Q_{\pi/4}=(1,1)\,Q_{\pi/4}=(0,\sqrt{2}) ; the two disagree, so sign \mathrm{sign} is not gauge-equivariant.
Adafactor. Its second-moment surrogate is not entrywise but rank-one factored, so we check it directly. Adafactor replaces the per-coordinate second moment by the reconstruction v ^ i j = r i c j / ∑ l r l \hat{v}_{ij}=r_{i}c_{j}/\sum_{l}r_{l} from the row and column sums r i = ∑ j v i j r_{i}=\sum_{j}v_{ij} , c j = ∑ i v i j c_{j}=\sum_{i}v_{ij} of the accumulator v = EMA ( g 2 ) v=\mathrm{EMA}(g^{2}) , with update g i j / v ^ i j g_{ij}/\sqrt{\hat{v}_{ij}} . Taking a single step (so that v ^ = g ⊙ 2 \hat{v}=g^{\odot 2} after bias correction; any common positive factor here rescales both sides of the test alike), the row sums are gauge-invariant, r i = ( G G ⊤ ) i i r_{i}=(GG^{\top})_{ii} , but the column sums are the basis-dependent diagonal c j = ( Q ⊤ G ⊤ G Q ) j j c_{j}=(Q^{\top}G^{\top}G\,Q)_{jj} of the conjugated Gram: already the θ = π / 4 \theta=\pi/4 rotation of G = diag ( 1 , δ ) G=\mathrm{diag}(1,\delta) with 0 < δ ≠ 1 0<\delta\neq 1 leaves r = ( 1 , δ 2 ) r=(1,\delta^{2}) fixed while sending c = ( 1 , δ 2 ) ↦ 1 2 ( 1 + δ 2 ) ( 1 , 1 ) c=(1,\delta^{2})\mapsto\tfrac{1}{2}(1+\delta^{2})(1,1) . Writing Q = 2 − 1 / 2 ( 1 1 − 1 1 ) Q=2^{-1/2}\begin{pmatrix}1&1\\ -1&1\end{pmatrix} , the resulting zero-state updates are Φ ( G ) = 1 + δ 2 diag ( 1 , δ − 1 ) , Φ ( G Q ) = ( 1 1 − 1 1 ) , \Phi(G)=\sqrt{1+\delta^{2}}\,\mathrm{diag}(1,\delta^{-1}),\qquad\Phi(GQ)=\begin{pmatrix}1&1\\ -1&1\end{pmatrix}, whereas Φ ( G ) Q = 1 + δ 2 2 ( 1 1 − δ − 1 δ − 1 ) ≠ Φ ( G Q ) . \Phi(G)Q=\frac{\sqrt{1+\delta^{2}}}{\sqrt{2}}\begin{pmatrix}1&1\\ -\delta^{-1}&\delta^{-1}\end{pmatrix}\neq\Phi(GQ). The two sides differ at order 1 1 for fixed δ \delta , while a small scalar stabilizer added to v ^ \hat{v} perturbs both sides continuously, so a sufficiently small stabilizer does not reconcile them. Adafactor therefore breaks equivariance through a factored, coordinate-dependent second moment. ∎
### B.5 Proposition 4.4 (exact first-step Adam defect)
###### Proof.
For either factor and its gradient G G , zero initial state and bias correction give m 1 = ( 1 − β 1 ) G , v 1 = ( 1 − β 2 ) G ⊙ 2 , m ^ 1 = G , v ^ 1 = G ⊙ 2 . m_{1}=(1-\beta_{1})G,\qquad v_{1}=(1-\beta_{2})G^{\odot 2},\qquad\hat{m}_{1}=G,\qquad\hat{v}_{1}=G^{\odot 2}. Hence the first Adam update is exactly − η D ϵ ( G ) -\eta D_{\epsilon}(G) ; in particular, it is independent of β 1 , β 2 \beta_{1},\beta_{2} under the stated convention. By Lemma 4.1 , the gradients in the rotated run are G ~ U = G U Q \widetilde{G}_{U}=G_{U}Q and G ~ V = G V Q \widetilde{G}_{V}=G_{V}Q . Gauge-align its first iterate back to the reference coordinates: U ~ 1 Q ⊤ = U 0 − η D ϵ ( G U Q ) Q ⊤ = U 0 − η [ D U + E Q ( G U ) ] , \widetilde{U}_{1}Q^{\top}=U_{0}-\eta D_{\epsilon}(G_{U}Q)Q^{\top}=U_{0}-\eta\bigl[D_{U}+E_{Q}(G_{U})\bigr], and likewise V ~ 1 Q ⊤ = V 0 − η [ D V + E Q ( G V ) ] . \widetilde{V}_{1}Q^{\top}=V_{0}-\eta\bigl[D_{V}+E_{Q}(G_{V})\bigr]. The gauge-aligned and reference products are W ~ 1 \displaystyle\widetilde{W}_{1} = ( U ~ 1 Q ⊤ ) ( V ~ 1 Q ⊤ ) ⊤ , \displaystyle=(\widetilde{U}_{1}Q^{\top})(\widetilde{V}_{1}Q^{\top})^{\top}, W 1 \displaystyle W_{1} = ( U 0 − η D U ) ( V 0 − η D V ) ⊤ . \displaystyle=(U_{0}-\eta D_{U})(V_{0}-\eta D_{V})^{\top}. Subtracting and expanding gives the stated identity. The alternative quadratic expression follows from D ϵ ( G Q ) Q ⊤ = D ϵ ( G ) + E Q ( G ) D_{\epsilon}(GQ)Q^{\top}=D_{\epsilon}(G)+E_{Q}(G) for each factor. For the displayed witness, G U = G V = ( 1 , 0 ) G_{U}=G_{V}=(1,0) while G U Q = G V Q = ( 2 − 1 / 2 , 2 − 1 / 2 ) G_{U}Q=G_{V}Q=(2^{-1/2},2^{-1/2}) ; direct substitution gives the two scalar products in the statement. ∎
### B.6 Theorem 4.5 (structure)
###### Proof.
( ⇐ \Leftarrow ) ( G Q ) ( G Q ) ⊤ = G G ⊤ (GQ)(GQ)^{\top}=GG^{\top} , so Φ ( G Q ) = H ( G G ⊤ ) G Q = Φ ( G ) Q \Phi(GQ)=H(GG^{\top})\,GQ=\Phi(G)\,Q . ( ⇒ \Rightarrow ) For full-column-rank G G let G + = ( G ⊤ G ) − 1 G ⊤ G^{+}=(G^{\top}G)^{-1}G^{\top} (the Moore–Penrose pseudoinverse for full column rank) and define X ( G ) ≔ Φ ( G ) G + ∈ ℝ n × n X(G)\coloneqq\Phi(G)\,G^{+}\in\mathbb{R}^{n\times n} . Then X ( G ) G = Φ ( G ) ( G + G ) = Φ ( G ) X(G)\,G=\Phi(G)(G^{+}G)=\Phi(G) , and X X is gauge-invariant: ( G Q ) + = Q ⊤ G + (GQ)^{+}=Q^{\top}G^{+} , so X ( G Q ) = Φ ( G ) Q Q ⊤ G + = X ( G ) X(GQ)=\Phi(G)\,Q\,Q^{\top}G^{+}=X(G) . It remains to see that X X depends on G G only through G G ⊤ GG^{\top} , i.e., that full-column-rank matrices with equal Grams lie on one gauge orbit: if G 1 G 1 ⊤ = G 2 G 2 ⊤ G_{1}G_{1}^{\top}=G_{2}G_{2}^{\top} , set Q ≔ G 1 + G 2 Q\coloneqq G_{1}^{+}G_{2} ; then Q ⊤ Q = G 2 ⊤ ( G 1 G 1 ⊤ ) + G 2 = G 2 ⊤ ( G 2 G 2 ⊤ ) + G 2 = I k Q^{\top}Q=G_{2}^{\top}(G_{1}G_{1}^{\top})^{+}G_{2}=G_{2}^{\top}(G_{2}G_{2}^{\top})^{+}G_{2}=I_{k} (full column rank), and G 1 Q = G 1 G 1 + G 2 = P range ( G 1 ) G 2 = G 2 G_{1}Q=G_{1}G_{1}^{+}G_{2}=P_{\mathrm{range}(G_{1})}\,G_{2}=G_{2} since range ( G i ) \mathrm{range}(G_{i}) and range ( G i G i ⊤ ) \mathrm{range}(G_{i}G_{i}^{\top}) coincide. Hence X ( G 1 ) = X ( G 2 ) X(G_{1})=X(G_{2}) , and H ( G G ⊤ ) ≔ X ( G ) H(GG^{\top})\coloneqq X(G) is well defined, giving Φ ( G ) = H ( G G ⊤ ) G \Phi(G)=H(GG^{\top})\,G . In the square invertible case, the polar decomposition G = ( G G ⊤ ) 1 / 2 O G=(GG^{\top})^{1/2}O with O ∈ O ( n ) O\in\mathrm{O}(n) gives directly Φ ( G ) = Φ ( ( G G ⊤ ) 1 / 2 ) O = H ( G G ⊤ ) G \Phi(G)=\Phi\big((GG^{\top})^{1/2}\big)\,O=H(GG^{\top})\,G with H ( P ) = Φ ( P 1 / 2 ) P − 1 / 2 H(P)=\Phi(P^{1/2})\,P^{-1/2} ; uniqueness follows by evaluating at G = P 1 / 2 G=P^{1/2} . ∎
### B.7 Proposition A.1 (spectral transfer)
###### Proof.
With the compact SVD G = ∑ i σ i u i v i ⊤ G=\sum_{i}\sigma_{i}u_{i}v_{i}^{\top} we have G G ⊤ = ∑ i σ i 2 u i u i ⊤ GG^{\top}=\sum_{i}\sigma_{i}^{2}\,u_{i}u_{i}^{\top} , so ψ ( G G ⊤ ) = ∑ i ψ ( σ i 2 ) u i u i ⊤ + ψ ( 0 ) P 0 \psi(GG^{\top})=\sum_{i}\psi(\sigma_{i}^{2})\,u_{i}u_{i}^{\top}+\psi(0)\,P_{0} , where P 0 = I − ∑ i u i u i ⊤ P_{0}=I-\sum_{i}u_{i}u_{i}^{\top} projects onto ker ( G G ⊤ ) \ker(GG^{\top}) (nontrivial whenever G G is rank-deficient, in particular when k < n k<n ). Since P 0 G = 0 P_{0}G=0 that term drops out, and ψ ( G G ⊤ ) G = ∑ i ψ ( σ i 2 ) σ i u i v i ⊤ \psi(GG^{\top})\,G=\sum_{i}\psi(\sigma_{i}^{2})\,\sigma_{i}\,u_{i}v_{i}^{\top} . The instances follow by inspection; for support-restricted one-step undamped Shampoo, ( G G ⊤ ) supp − 1 / 4 G ( G ⊤ G ) − 1 / 4 = ∑ i σ i − 1 / 2 σ i σ i − 1 / 2 u i v i ⊤ = ∑ i u i v i ⊤ = msign ( G ) . (GG^{\top})_{\mathrm{supp}}^{-1/4}\,G\,(G^{\top}G)^{-1/4}=\sum_{i}\sigma_{i}^{-1/2}\sigma_{i}\sigma_{i}^{-1/2}u_{i}v_{i}^{\top}=\sum_{i}u_{i}v_{i}^{\top}=\mathrm{msign}(G). For the Lipschitz claim (consequence (ii) stated after the proposition), fix 0 < δ < 1 0<\delta<1 (a small singular value, distinct from Adam’s stabilizer ϵ \epsilon ) and evaluate at G ± = diag ( 1 , ± δ ) G_{\pm}=\mathrm{diag}(1,\pm\delta) (embedded in the top-left 2 × 2 2\times 2 block): the second singular pair of G ± G_{\pm} is ( σ , u , v ) = ( δ , e 2 , ± e 2 ) (\sigma,u,v)=(\delta,e_{2},\pm e_{2}) , so Φ ( G + ) − Φ ( G − ) = 2 h ( δ ) e 2 e 2 ⊤ \Phi(G_{+})-\Phi(G_{-})=2h(\delta)\,e_{2}e_{2}^{\top} while G + − G − = 2 δ e 2 e 2 ⊤ G_{+}-G_{-}=2\delta\,e_{2}e_{2}^{\top} : the corresponding Lipschitz quotient is at least | h ( δ ) | / δ \lvert h(\delta)\rvert/\delta , which diverges as δ → 0 \delta\to 0 whenever h ( 0 + ) > 0 h(0^{+})>0 , and equals η \eta for h ( σ ) = η σ h(\sigma)=\eta\sigma . ∎
### B.8 Corollary A.3 (deeper factorizations)
###### Proof.
Write L ( U 1 , … , U D ) = f ( U 1 ⋯ U D ) L(U_{1},\dots,U_{D})=f(U_{1}\cdots U_{D}) , fix an interface i i , and set A = U 1 ⋯ U i − 1 A=U_{1}\cdots U_{i-1} , B = U i + 2 ⋯ U D B=U_{i+2}\cdots U_{D} , G = ∇ f ( W ) G=\nabla f(W) . Then ∇ U i L = A ⊤ G ( U i + 1 B ) ⊤ \nabla_{U_{i}}L=A^{\top}G\,(U_{i+1}B)^{\top} and ∇ U i + 1 L = ( A U i ) ⊤ G B ⊤ \nabla_{U_{i+1}}L=(AU_{i})^{\top}G\,B^{\top} . Under ( U i Q , Q ⊤ U i + 1 ) (U_{i}Q,\,Q^{\top}U_{i+1}) the represented W W , hence G G , is unchanged, so ∇ U i L ↦ ( ∇ U i L ) Q \nabla_{U_{i}}L\mapsto(\nabla_{U_{i}}L)\,Q and ∇ U i + 1 L ↦ Q ⊤ ( ∇ U i + 1 L ) \nabla_{U_{i+1}}L\mapsto Q^{\top}\,(\nabla_{U_{i+1}}L) : the gradient transforms covariantly on the acted side of each factor, which is the only property the cited proofs use (transposed for the left action). ∎
### B.9 Theorem 4.6 (transfer theorem)
###### Proof.
Let θ ( t ) \theta(t) be a locally absolutely continuous solution of θ ˙ = − ∇ L ( θ ) / a ( t ) \dot{\theta}=-\nabla L(\theta)/a(t) with a ( t ) > 0 a(t)>0 determined by gauge-invariant trajectory statistics. Define τ ( t ) = ∫ 0 t 𝑑 u / a ( u ) \tau(t)=\int_{0}^{t}du/a(u) , a strictly increasing bijection onto [ 0 , τ max ) [0,\tau_{\max}) with τ max = ∫ 0 ∞ 𝑑 u / a ( u ) \tau_{\max}=\int_{0}^{\infty}du/a(u) , and θ ~ ( τ ) = θ ( t ( τ ) ) \tilde{\theta}(\tau)=\theta(t(\tau)) . Then, for a.e. τ \tau , d θ ~ d τ = θ ˙ ⋅ d t d τ = ( − ∇ L ( θ ) / a ) ⋅ a = − ∇ L ( θ ~ ) \frac{d\tilde{\theta}}{d\tau}=\dot{\theta}\cdot\frac{dt}{d\tau}=\big(-\nabla L(\theta)/a\big)\cdot a=-\nabla L(\tilde{\theta}) : θ ~ \tilde{\theta} is a Carathéodory gradient-flow solution with the same initial condition, and the two trajectories traverse the same set of points under a monotone reparameterization: the whole gradient-flow path when τ max = ∞ \tau_{\max}=\infty , and otherwise its prefix θ ~ ( [ 0 , τ max ) ) \tilde{\theta}([0,\tau_{\max})) . So any property of the gradient-flow path on the traversed portion holds for the scalar-preconditioned flow. When τ max = ∞ \tau_{\max}=\infty , which holds in particular whenever a a is bounded above (since then ∫ ∞ 𝑑 u / a ( u ) \int^{\infty}du/a(u) diverges), the reparameterization also reaches the gradient-flow limit, so the limit-point characterizations of Gunasekar et al. (2017) ; Arora et al. (2019) ; Li et al. (2021) under their hypotheses transfer as well; if τ max < ∞ \tau_{\max}<\infty only the finite-time path is shared. Since a a is computed from gauge-invariant statistics, the reparameterization is itself gauge-independent, so the statement holds uniformly over the gauge orbit. ∎
### B.10 Proposition A.4 (balancedness)
###### Proof.
Under gradient flow, U ˙ = − G V \dot{U}=-GV , V ˙ = − G ⊤ U \dot{V}=-G^{\top}U with G = ∇ f ( W ) G=\nabla f(W) : d d t ( U ⊤ U ) = − V ⊤ G ⊤ U − U ⊤ G V , d d t ( V ⊤ V ) = − U ⊤ G V − V ⊤ G ⊤ U , \tfrac{d}{dt}(U^{\top}U)=-V^{\top}G^{\top}U-U^{\top}GV,\qquad\tfrac{d}{dt}(V^{\top}V)=-U^{\top}GV-V^{\top}G^{\top}U, which are equal, so B ˙ = 0 \dot{B}=0 . (Time reparameterization by 1 / a ( t ) 1/a(t) multiplies both by the same scalar; conservation persists for the whole class of Theorem 4.6 .) Under a diagonal preconditioned flow U ˙ = − D U ⊙ ( G V ) \dot{U}=-D_{U}\odot(GV) , V ˙ = − D V ⊙ ( G ⊤ U ) \dot{V}=-D_{V}\odot(G^{\top}U) , B ˙ = − [ ( D U ⊙ G V ) ⊤ U + U ⊤ ( D U ⊙ G V ) ] + [ ( D V ⊙ G ⊤ U ) ⊤ V + V ⊤ ( D V ⊙ G ⊤ U ) ] . \dot{B}=-\big[(D_{U}\odot GV)^{\top}U+U^{\top}(D_{U}\odot GV)\big]+\big[(D_{V}\odot G^{\top}U)^{\top}V+V^{\top}(D_{V}\odot G^{\top}U)\big]. Write D U = c 11 ⊤ + E U D_{U}=c\,\mathbf{1}\mathbf{1}^{\top}+E_{U} and D V = c 11 ⊤ + E V D_{V}=c\,\mathbf{1}\mathbf{1}^{\top}+E_{V} for a common scalar c c , where 𝟏𝟏 ⊤ \mathbf{1}\mathbf{1}^{\top} is the all-ones matrix of the factor’s shape. The common-scalar terms cancel, while the remainder obeys ‖ B ˙ ‖ F ≤ 2 ‖ E U ‖ ∞ ‖ G V ‖ F ‖ U ‖ F + 2 ‖ E V ‖ ∞ ‖ G ⊤ U ‖ F ‖ V ‖ F , \left\lVert\dot{B}\right\rVert_{F}\leq 2\left\lVert E_{U}\right\rVert_{\infty}\left\lVert GV\right\rVert_{F}\left\lVert U\right\rVert_{F}+2\left\lVert E_{V}\right\rVert_{\infty}\left\lVert G^{\top}U\right\rVert_{F}\left\lVert V\right\rVert_{F}, by ‖ E ⊙ Z ‖ F ≤ ‖ E ‖ ∞ ‖ Z ‖ F \left\lVert E\odot Z\right\rVert_{F}\leq\left\lVert E\right\rVert_{\infty}\left\lVert Z\right\rVert_{F} (with ‖ ⋅ ‖ ∞ \left\lVert\cdot\right\rVert_{\infty} the entrywise maximum) and the Frobenius submultiplicativity bound. Thus, at bounded factor and gradient norms, the drift vanishes as both factor preconditioners approach the same scalar. ∎
### B.11 Proposition 8.1 (greedy vs. equal-rate phase boundary)
###### Proof.
Each decoupled mode is a balanced two-factor scalar factorization w = u v w=uv with u = v = w u=v=\sqrt{w} and per-mode loss 1 2 ( u v − s ) 2 \tfrac{1}{2}(uv-s)^{2} , so ∂ u L = v ( w − s ) \partial_{u}L=v(w-s) and ∂ v L = u ( w − s ) \partial_{v}L=u(w-s) . The balanced set u = v u=v is invariant under both flows below, since u ˙ = v ˙ \dot{u}=\dot{v} there, so the reduction to the single variable w w is consistent. (The tied model w = u 2 w=u^{2} gives identical dynamics at twice the rate in (i), a uniform time-rescaling.)