Broken scale symmetries in undercomplete linear autoencoders

Extended Abstract Track 1–35, 2026

Symmetry and Geometry in Neural Representations

Broken scale symmetries in undercomplete linear autoencoders Farhad Pashakhanloo

fpasha@fas.harvard.edu

Center for Brain Science Harvard University

arXiv:2610.03640v1 [cs.LG] 2 Oct 2026

Jacob A. Zavatone-Veth

jzavatoneveth@fas.harvard.edu

Center for Brain Science and Society of Fellows Harvard University

Editors: List of editors’ names

Abstract Neural network loss landscapes have many symmetries, which are preserved by gradient flow but broken by finite-stepsize stochastic gradient descent (SGD). A canonical example of such a symmetry is scale in homogeneous networks: one can scale up the parameters in one layer and down in the next without changing the network output. Previous work has documented cases in which SGD breaks this symmetry in favor of balancing gradient noise or minimizing fluctuations. Here, we show that the solution geometry of undercomplete linear autoencoders instead selects a preferred sign for scale drift: on the PCA solution manifold, SGD favors large decoder weights. This directed scale drift occurs on a slow timescale, and its dynamics admit an analytically-tractable effective description. However, it cannot continue indefinitely: increasing scale eventually drives the dynamics towards a finite-stepsize stability boundary. The resulting solutions are sharper than a balanced baseline in the sense of the maximum eigenvalue of the loss Hessian, but different sharpness measures can move in opposing directions. Thus, undercomplete autoencoders give a concrete illustration of how loss geometry can convert residual gradient noise into directed motion along a manifold of functionally-equivalent solutions. Keywords: stochastic gradient descent, representation learning, parameter symmetries

1. Introduction How does SGD navigate symmetry-induced flat directions in neural network loss landscapes? Previous works have proposed a range of perspectives on this inductive bias, from effective regularizers (Li et al., 2022; Ziyin et al., 2025b) and equilibrium balance conditions (Ziyin et al., 2025a,b), to manifold diffusion (Pashakhanloo and Koulakov, 2023) and sparsifying collapse (Chen et al., 2024). However, it remains unclear what determines the direction and endpoint of SGD motion along continuous symmetry orbits. Here, we study a solvable model—undercomplete linear autoencoders—in which we can give a precise answer to this question in the specific case of scale symmetries. We show that in this setting solution geometry selects a preferred direction for drift along scale-symmetry directions, and solve for its slow dynamics. © 2026 F. Pashakhanloo & J.A. Zavatone-Veth.

Pashakhanloo Zavatone-Veth

Figure 1: (A) Continual minibatch training drives scale imbalance in linear autoencoders. Frobenius norms of decoder and encoder weights, ∥W ∥F and ∥U ∥F , for contractive and expansive autoencoders, respectively (see App. F and Fig. S1 for additional details). (B) Scale growth in a contractive linear autoencoder for different learning rates (η) and minibatch size (B). s1 = (W ⊤ W )11 is the dominant scale. Theory predictions are based on Eq. 6. (C) Anisotropic scale growth in a contractive autoencoder. Different scale modes, si = (W ⊤ W )ii , grow at different rates. Theory predictions are based on Eq. 6. (D-F) Edge of Stability phenomenon driven by imbalancing dynamics. (D) Scale growth and saturation for different modes. (E) Sequential saturation of eigenvalues. (F) Saturation as a function of the inverse of the learning rate and comparison to theory.

2. Undercomplete solution geometry results in directed scale drift We consider a tractable toy setting: linear autoencoders f (x) = W U x, (1) k×d d×k with encoder weights U ∈ R and decoder weights W ∈ R . We train these models to minimize an L2 reconstruction loss L = 12 Ex [∥f (x) − x∥2 ], assuming zero-mean data x ∈ Rd with covariance Σ = E[xx⊤ ]. Our focus is on the undercomplete regime k < d, where exact reconstruction is impossible if Σ is of full rank. These networks have a GL(k) symmetry group that preserves the function output: mapping (W, U ) 7→ (W A, A−1 U ) for A ∈ GL(k) leaves f (x) invariant. Under gradient flow, it is well-known that the Noether charge Q = W ⊤ W − U U ⊤ ∈ Rk×k (2) corresponding to this symmetry is conserved, i.e., ∂t Q = 0 (App. A; Tanaka and Kunin, 2021; Kunin et al., 2021). Q measures imbalances in the incoming and outgoing weights for each latent neuron. In contrast, under SGD with learning rate η and loss L̂t = 1 2 2 Ex∼Bt [∥f (x) − x∥ ] for a minibatch Bt of size B and covariance Σ̂t , Q evolves as Qt+1 = Qt + η 2 [(∇W L̂t )⊤ (∇W L̂t ) − (∇U L̂t )(∇U L̂t )⊤ ] (3) 2

Broken scale symmetries in undercomplete linear autoencoders

Therefore, SGD will generically break the symmetry, as has been noted in general terms in prior works (Tanaka and Kunin, 2021; Kunin et al., 2021; Ziyin et al., 2025a). We now specialize to undercomplete autoencoders, for which (3) simplifies substantially. Let Σ have eigenvalues λ1 ≥ λ2 ≥ · · · ≥ λd ≥ 0 with a strict spectral gap λk > λk+1 separating the k-principal subspace from the remaining directions, and let P ∈ Rk×d be an orthogonal projection to the principal subspace, with P P ⊤ = Ik . Then, the global loss minima form the manifold (W, U ) = (P ⊤ A, A−1 P ) for A ∈ GL(k); in this sense linear autoencoders learn to implement PCA by projecting their inputs to the principal subspace (Baldi and Hornik, 1989; Kunin et al., 2019; Pashakhanloo, 2025). We next come to the key geometric fact. On the PCA solution manifold, the residual r = (W U − Id )x for an arbitrary datapoint x lies entirely in the complement of the principal subspace, as Id − W U = Id − P P ⊤ ≡ P⊥ projects to the complement. This forces ∇U L̂t = 0 to be identically zero, while ∇W L̂t = −P⊥ Σ̂t P ⊤ A−⊤ depends on the product of in- and t out-of-subspace fluctuations, and does not vanish (App. A). Thus, Qt+1 − Qt = η 2 (∇W L̂t )⊤ (∇W L̂t ) ⪰ 0 (4) on the PCA solution manifold; Q evolves in a preferred direction within the cone of PSD matrices. This is a directional bias towards decoder-heavy weights, as can be seen by considering the scalar charge q = tr(Q) = ∥W ∥2F − ∥U ∥2F associated with global rescalings (W, U ) 7→ (aW, a−1 U ) for a > 0 (Figure 1A). No noise-balanced equilibrium, where the two terms in (3) cancel on average (Ziyin et al., 2025a), is possible. In contrast, for an expansive autoencoder perfect reconstruction is possible, meaning that there is not a fundamental geometric reason for directed drift (Figure 1A). We now analyze the scale dynamics.

3. Effective scale dynamics We can derive a simple set of effective dynamics that allow us to predict the anisotropic growth in decoder scale. Under SGD with minibatch size B, changes to the balancing matrix Q accumulate over Tslow ∼ B/η 2 steps, which is slow relative to the Tfast ∼ 1/η steps over which fluctuations away from the solution manifold relax back towards it, as Tfast /Tslow ∼ η/B ≪ 1 when η is small (see App. C for details). By leveraging this separation of timescales, we show in App. C that the dynamics of the matrix S = A⊤ A summarizing the scale degrees of freedom are given in the small-learning-rate limit η ↓ 0 by the flow ∂τ S + S −1 ∂τ SS −1 = S −1/2 CS −1/2 , (5) 2 ⊤ ⊤ where C = E[∥x⊥ ∥ x∥ x∥ ] is defined in terms of x∥ = P x and x⊥ = (Id −P P )x. If S(τ = 0) and C commute, then this matrix differential equation can be solved analytically. Working in a basis where S = diag(s1 ,p . . . , sk ) and C = diag(c1 , . . . , ck ), we have the solution sa (τ ) = W [sa (0)2 exp(sa (0)2 + 2ca τ )], a ∈ [k], (6) where W (·) is the principal branch of the Lambert W -function (DLMF). Therefore, we expect anisotropic of growth along the different scale directions. In the Gaussian case, Prates d we have ca = λa j=k+1 λj , i.e., the rate of growth along the a-th direction is determined by the population variance along that direction multiplied by the total variance outside of the principal subspace. The fact that scale drift is driven by the coupling between within-subspace and out-of-subspace directions mirrors Pashakhanloo (2025)’s finding that this same coupling governs the rate of drift along symmetry directions corresponding to 3

Pashakhanloo Zavatone-Veth

rotations rather than re-scalings. This effective model produces excellent agreement with simulations at small but finite stepsize (Figure 1C).

4. Sharpness and the edge of stability Naı̈vely, these effective dynamics suggest that the scales should grow indefinitely. However, this cannot be the case: indefinite scale growth would destabilize the underlying discrete SGD dynamics. In particular, the assumption that the dynamics relax stably back to the PCA manifold will fail if the eigenvalues µ of the population Hessian become too large— that is, if ηµ exceeds some threshold κ > 0, which for full-batch GD would be exactly equal to 2 (Cohen et al., 2021). For S = diag(s1 , . . . , sk ) with ordered entries s1 > s2 > · · · > sk , we show in Appendix C that the Hessian of the population loss has dangerous modes ∥ ∥ ∥ ∥ µ11 > µ12 > · · · > µ1k given by µ1a = λ1 (s−1 1 + sa ) for a ∈ [k]. Therefore, as the scales increase anisotropically, the curvature along each of these modes will increase until it hits the ∥ threshold ηµ1a = κ at which the underlying discrete dynamics become unstable. Thus, we predict that all scales will eventually saturate at a common value given by sa = κ/ηλ1 +O(η). We show that this prediction is consistent with simulations, with κ ≃ 2, in Figure 1D-F. Consistent with the parameterization-dependence of the Euclidean Hessian (Dinh et al., 2017), not all notions of sharpness increase with scale. Most simply, considering P global re-scalings with S = sIk , the Hessian on the PCA manifold has trace tr(H) = ds kj=1 λj + P P P ks dj=1 λj , which is minimized at s2⋆ = ( k1 kj=1 λj )/( d1 dj=1 λj ) ≥ 1 (App. B). More strikingly, Andreyev and Beneventano (2025) proposed a measure of “batch sharpness” that they suggest increases up to a threshold over SGD training; we show in Appendix D that this measure is decreasing in global scale on the PCA manifold. Moreover, we can prove that, for Gaussian data and a balanced initialization, their simplified “modified batch sharpness” measure is strictly decreasing under the effective dynamics, even as the population Hessian sharpens.

5. Generalization to homogeneous nonlinear autoencoders Introducing nonlinear activation functions breaks the GL(k) symmetry of linear autoencoders. For homogeneous activation functions like ReLU, a neuron-wise scaling symmetry remains, with A now being a diagonal matrix with positive diagonal entries rather than a general element of GL(k) (App. E). Our preliminary results suggest that undercomplete ReLU autoencoders show a similar directed drift towards decoder-heavy weights (Fig. S2). We leave an in-depth investigation of the geometric basis for this bias to future work.

6. Discussion We have shown that undercomplete linear autoencoders trained with SGD display directed scale drift towards decoder-heavy weights. The endpoint of this drift is selected not by an equilibrium statistical mechanism like balancing of gradient noise (Ziyin et al., 2025a), but instead dynamically by the fact that scale drift increases curvature up to the point where the solution would become unstable. 4

Broken scale symmetries in undercomplete linear autoencoders

From a broader optimization perspective, this resembles the general phenomenon of progressive sharpening in full-batch gradient descent, which posits that the top eigenvalues of the Hessian increase sequentially over training up to the stability threshold, after which they remain near the edge of stability (Cohen et al., 2021, 2025; Damian et al., 2023). How to precisely define the edge of stability in the stochastic setting remains an active area of research (Andreyev and Beneventano, 2025; Emmanouilidis et al., 2026). What is notable in our setting is the fact that the edge of stability is reached purely through a slow effective drift along a manifold of functionally-equivalent solutions, after the PCA solution manifold is reached. Therefore, progressive sharpening need not reflect continued learning of the function. Whether and how this phenomenology transfers to other optimization settings is an interesting direction for future work.

Acknowledgments F.P. was supported by the Harvard Center for Brain Science (CBS)-NTT Fellowship Program on the Physics of Intelligence. J.A.Z.-V. and this research were supported by the Office of the Director of the NIH under Award Number DP5OD037354, and additionally by a Junior Fellowship from the Harvard Society of Fellows.

Declaration of LLM Usage The development of this research did not involve LLMs in its core aspects. However, we used LLMs (OpenAI’s Codex) to generate codes for the simulations results presented here. We subsequently checked the generated code and plots.

5

Pashakhanloo Zavatone-Veth

References Arseniy Andreyev and Pierfrancesco Beneventano. Edge of stochastic stability: Revisiting the edge of stability for SGD. arXiv, 2025. URL https://arxiv.org/abs/2412.20553. Pierre Baldi and Kurt Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Neural Networks, 2(1):53–58, 1989. ISSN 0893-6080. doi: https://doi.org/10.1016/0893-6080(89)90014-2. URL https://www. sciencedirect.com/science/article/pii/0893608089900142. Feng Chen, Daniel Kunin, Atsushi Yamamura, and Surya Ganguli. Stochastic collapse: how gradient noise attracts SGD dynamics towards simpler subnetworks. Journal of Statistical Mechanics: Theory and Experiment, 2024(10):104023, oct 2024. doi: 10.1088/1742-5468/ ad6429. URL https://doi.org/10.1088/1742-5468/ad6429. Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum? id=jh-rTtvkGeM. Jeremy Cohen, Alex Damian, Ameet Talwalkar, J Zico Kolter, and Jason D. Lee. Understanding optimization in deep learning with central flows. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/forum? id=sIE2rI3ZPs. Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization: The implicit bias of gradient descent at the edge of stability. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=nhKHA59gXz. Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1019–1028. PMLR, 06–11 Aug 2017. URL https: //proceedings.mlr.press/v70/dinh17b.html. DLMF. NIST Digital Library of Mathematical Functions. http://dlmf.nist.gov/, Release 1.1.1 of 2021-03-15, 2021. URL http://dlmf.nist.gov/. F. W. J. Olver, A. B. Olde Daalhuis, D. W. Lozier, B. I. Schneider, R. F. Boisvert, C. W. Clark, B. R. Miller, B. V. Saunders, H. S. Cohl, and M. A. McClain, eds. Simon Du, Wei Hu, and Jason D Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https://proceedings.neurips.cc/paper_files/paper/2018/file/ fe131d7f5a6b38b23cc967316c13dae2-Paper.pdf. Konstantinos Emmanouilidis, Lachlan MacDonald, Salma Tarmoun, and Rene Vidal. SGD at the edge of stability: Stochastic stabilization with large learning rates. arXiv, 2026. URL https://arxiv.org/abs/2606.30930. 6

Broken scale symmetries in undercomplete linear autoencoders

Daniel Kunin, Jonathan Bloom, Aleksandrina Goeva, and Cotton Seed. Loss landscapes of regularized linear autoencoders. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 3560–3569. PMLR, 09–15 Jun 2019. URL https://proceedings.mlr.press/v97/kunin19a.html. Daniel Kunin, Javier Sagastuy-Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. In International Conference on Learning Representations, 2021. URL https: //openreview.net/forum?id=q8qLAbQBupm. Harold J. Kushner and G. George Yin. Stochastic Approximation and Recursive Algorithms and Applications. Springer New York, New York, NY, 2003. ISBN 978-0-387-21769-7. doi: 10.1007/0-387-21769. URL https://doi.org/10.1007/0-387-21769. Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? –a mathematical framework. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=siCt4xZn5Ve. Farhad Pashakhanloo. Contribution of task-irrelevant stimuli to drift of neural representations. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. URL https://openreview.net/forum?id=jAoqtT58G4. Farhad Pashakhanloo and Alexei Koulakov. Stochastic gradient descent-induced drift of representation in a two-layer neural network. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 27401–27419. PMLR, 2023. URL https://proceedings.mlr.press/v202/ pashakhanloo23a.html. Hidenori Tanaka and Daniel Kunin. Noether’s learning dynamics: Role of symmetry breaking in neural networks. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 25646–25660. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper_files/paper/2021/file/ d76d8deea9c19cc9aaf2237d2bf2f785-Paper.pdf. Liu Ziyin, Hongchao Li, and Masahito Ueda. Noise balance and stationary distribution of stochastic gradient descent. Phys. Rev. E, 111:065303, Jun 2025a. doi: 10.1103/PhysRevE.111.065303. URL https://link.aps.org/doi/10.1103/PhysRevE. 111.065303. Liu Ziyin, Yizhou Xu, and Isaac Chuang. Neural thermodynamics: Entropic forces in deep and universal representation learning. In D. Belgrave, C. Zhang, H. Lin, R. Pascanu, P. Koniusz, M. Ghassemi, and N. Chen, editors, Advances in Neural Information Processing Systems, volume 38, Main Conference, pages 56057–56089. Curran Associates, Inc., 2025b. doi: 10.52202/ 085713-1877. URL https://proceedings.neurips.cc/paper_files/paper/2025/ file/51173cf34c5faac9796a47dc2fdd3a71-Paper-Conference.pdf. 7

Pashakhanloo Zavatone-Veth

Appendix A. Linear autoencoder basics Here, we give a detailed discussion of the basic setup for linear autoencoders. Recall that the basic model architecture is f (x) = W U x, (A.1) k×d d×k where U ∈ R and W ∈ R . From the definition of the model, we can see that there is a GL(k) symmetry: the function output is invariant under transformations of the weights of the form (W, U ) 7→ (W A, A−1 U ), A ∈ GL(k). (A.2) Associated with this symmetry is the Noether charge (Kunin et al., 2021; Tanaka and Kunin, 2021) Q = W ⊤W − U U ⊤. (A.3)

A.1. Population loss and gradient flow For a set of data with population covariance E[xx⊤ ] = Σ, the population reconstruction loss is 1 1 L = Ex [∥f (x) − x∥2 ] = Ex tr[(W U − Id )xx⊤ (W U − Id )⊤ ] (A.4) 2 2 1 = tr[(W U − Id )Σ(W U − Id )⊤ ]. (A.5) 2 Under population gradient flow, the parameters then evolve as ∂t W = −∇W L = −(W U − Id )ΣU ⊤ (A.6) ∂t U = −∇U L = −W ⊤ (W U − Id )Σ. The charge Q = W ⊤ W − U U ⊤ is conserved under these dynamics: ∂t Q = (∂t W )⊤ W + W ⊤ ∂t W − (∂t U )U ⊤ − U (∂t U )⊤ = 0, as we can see that W ⊤ ∂t W = (∂t U )U ⊤ .

(A.7) (A.8)

A.2. Minibatch stochastic gradient descent Now consider stochastic gradient descent with minibatch size B. Let the matrix of data in the minibatch at time t be Xt ∈ Rd×B , and 1 (A.9) Σ̂t = Xt Xt⊤ B be the empirical covariance of the minibatch. Then, the minibatch loss is 1 L̂t = tr[(W U − Id )Σ̂t (W U − Id )⊤ ], (A.10) 2 so the minibatch gradients are ∇W L̂t = (Wt Ut − Id )Σ̂t Ut⊤ (A.11) ∇U L̂t = Wt⊤ (Wt Ut − Id )Σ̂t . and the update to the parameters is Wt+1 = Wt − η∇W L̂t = Wt − η(Wt Ut − Id )Σ̂t Ut⊤

(A.13)

Ut+1 = Ut − η∇U L̂t = Ut − ηWt⊤ (Wt Ut − Id )Σ̂t .

(A.14)

8

(A.12)

Broken scale symmetries in undercomplete linear autoencoders

Using the fact that Wt⊤ ∇W L̂t = (∇U L̂t )U ⊤ , we find that the corresponding update to the balancing matrix Qt = Wt⊤ Wt −  Ut Ut⊤ is  Qt+1 − Qt = η 2 (∇W L̂t )⊤ (∇W L̂t ) − (∇U L̂t )(∇U L̂t )⊤ .

(A.15)

Substituting in the definition of ∇W L̂t and ∇U L̂t in terms of Σ̂t , and defining the residual Rt = (Wt Ut − Id )Xt , (A.16) we can re-write this as η2 Qt+1 = Qt + 2 (Ut Xt Rt⊤ Rt Xt⊤ Ut⊤ − Wt⊤ Rt Xt Xt⊤ Rt⊤ Wt ). (A.17) B A.3. PCA solutions Our interest is in PCA-type solutions (Baldi and Hornik, 1989; Kunin et al., 2019; Pashakhanloo, 2025) (W, U ) = (P ⊤ A, A−1 P ), (A.18) k×d ⊤ where P ∈ R is an orthogonal projection P P = Ik and A ∈ GL(k) includes the symmetry degrees of freedom. For such a solution, Id − W U = Id − P ⊤ P = P⊥ (A.19) ⊤ is a projection onto the orthogonal complement of the span of P . As a result, W ⊤ (W U − Id ) = A⊤ P (P ⊤ P − Id ) = A⊤ (P − P ) = 0 (A.20) ⊤ as P P = Ik , meaning that ∇U L̂t = 0 (A.21) on the PCA solution manifold. Importantly, this is a statement that the gradient is identically zero on the solution manifold. In contrast, ∇W L̂t = (Wt Ut − Id )Σ̂t Ut⊤ = (P ⊤ P − Id )Σ̂t P ⊤ A−⊤ (A.22) t is not identically zero. Indeed, if we introduce the decomposition Xt = P ⊤ Xt,∥ + Xt,⊥ (A.23) such that P Xt = Xt,∥ and P⊥ Xt = Xt,⊥ , (A.24) we find that 1 ⊤ −⊤ ∇W L̂t = − Xt,⊥ Xt,∥ At . (A.25) B Together, this leads to η2 ⊤ ⊤ −⊤ Qt+1 − Qt = η 2 (∇W L̂t )⊤ (∇W L̂t ) = 2 A−1 Xt,∥ Xt,⊥ Xt,⊥ Xt,∥ At ⪰ 0 (A.26) B t on the PCA solution manifold.

Appendix B. Spectrum of the linear autoencoder Hessian We now derive the Hessian of the linear autoencoder reconstruction loss. This computation largely parallels previous Hessian computations by Baldi and Hornik (1989); Pashakhanloo and Koulakov (2023); Pashakhanloo (2025). Our general assumption on the data is that E[x] = 0 and E[xx⊤ ] = Σ, (B.1) 9

Pashakhanloo Zavatone-Veth

where Σ has an orthogonal eigendecomposition d X Σ= λi vi vi⊤ ,

(B.2)

i=1

with λ1 ≥ λ2 ≥ · · · ≥ λd ≥ 0 and vi⊤ vj = δij . Without loss of generality, we may assume that the eigenvectors vi are the standard Euclidean basis vectors; we will impose this assumption when convenient. We will also often assume that the eigenvalues satisfy spectral gap conditions of the form λk > λk+1 for some fixed k; this is required for the k-principal subspace to be well-defined. B.1. Warm-up: Computing the top PC As a warm-up, consider the case k = 1, in which we want to compute the top principal component. Then, the loss is 1 1 L = Ex ∥wu⊤ x − x∥2 = tr[(wu⊤ − Id )Σ(uw⊤ − Id )] (B.3) 2 2 1 1 = ∥w∥2 (u⊤ Σu) − u⊤ Σw + tr(Σ), (B.4) 2 2 from which we can easily see the Hessian is  (u⊤ Σu)Id 2wu⊤ Σ − Σ . (B.5) H= 2Σuw⊤ − Σ ∥w∥2 Σ Consider a general PCA-type solution for some fixed i ∈ [d] and a scale factor s > 0: w = s1/2 vi , u = s−1/2 vi , (B.6) which projects the data onto the i-th principal direction. Then, the Hessian becomes P   λi s−1 Id λi vi vi⊤ − j̸=i λj vj vj⊤ P P . (B.7) H= s j λj vj vj⊤ λi vi vi⊤ − j̸=i λj vj vj⊤ As a result, if we expand a general perturbation as ! Pd   α v δw j j = Pj=1 , (B.8) d δu j=1 βj vj we have the expansion    ⊤   X d  ⊤ αj δw δw αj = Hj , (B.9) H βj δu δu βj j=1

where Hi = λi

 −1  s 1 1 s

(B.10)

and, for j ̸= i,  −1  λi s −λj Hj = . −λj sλj This means that we may equivalently write     α1  β1  H1       H2   α2   ..  .  β2  .. . 10

(B.11)

(B.12)

Broken scale symmetries in undercomplete linear autoencoders

in these coordinates. This means that we can diagonalize the full Hessian by diagonalizing the 2 × 2 blocks. The block Hi , which corresponds to fluctuations in the principal subspace, has eigenvalue µi,− = 0, corresponding to the scale symmetry and the eigenvector  1/2  1 s √ (B.13) −1 −s−1/2 s+s and µi,+ = λi (s + s−1 ), (B.14) corresponding to the eigenvector  1/2  1 s √ (B.15) −1/2 . −1 s s+s For j ̸= i, we have that det Hj = λj (λi − λj ), (B.16) and eigenvalues p λi s−1 + λj s ± (λi s−1 + λj s)2 − 4λj (λi − λj ) (B.17) µj,± = 2 corresponding to eigenvectors   1 −λj q (B.18) −1 . λ2j + (µj,± − λi s−1 )2 µj,± − λi s This allows us to see several things. First, the Hessian is positive-semidefinite only if λi ≥ λj for all j ̸= i, and positive-definite on the orthogonal complement of the direction corresponding to the scaling symmetry only if λi > λj > 0. This is what we expect from the PCA solution. We now consider the effect of rescaling. First, it is clear that µi,+ = λi (s+s−1 ) is always increased by breaking the scale balance, but it behaves identically irrespective of whether the solution is decoder-heavy or encoder-heavy. Now, for j ̸= i, we have " # r 2 λi −1 −1 µj,± = s + λ̃j s − 4λ̃j (1 − λ̃j ) s + λ̃j s ± 2 after factoring out λi > 0 and defining λj λ̃j = . λi Assuming λ̃j > 0, the function 1 t(s) = tr(Hj ) = s−1 + λ̃j s λi has a unique minimum over s > 0 at −1/2 s⋆ = λ̃j ,

(B.19)

(B.20)

(B.21) (B.22)

1/2 where it takes the value t(s⋆ ) = 2λ̃j . From this, we find that µj,+ decreases in s on −1/2 −1/2 (0, λ̃j ), has a global minimum at s⋆ = λ̃j where it takes value

µj,+

1/2

−1/2 s⋆ =λ̃j

= λi (λ̃j

11

+ λ̃j ),

(B.23)

Pashakhanloo Zavatone-Veth

−1/2

and increases on (λ̃j

−1/2

, ∞). Analogously, µj,− increases on (0, λ̃j

), has a global maxi-

−1/2 mum at s⋆ = λ̃j where it takes value

µj,−

1/2

−1/2 s⋆ =λ̃j

= λi (λ̃j

− λ̃j ),

(B.24)

−1/2

and decreases on (λ̃j , ∞). This shows that changing the scale has, in general, a rather non-obvious effect on the spectrum of the Hessian, even in this simple case. We now prove several statements about the ordering of the eigenvalues. To do so, we first derive several preliminary results. Let p 1 µ± (λ) = [s−1 + λs ± (s−1 + λs)2 − 4λ(1 − λ)] (B.25) 2 be functions defining the normalized eigenvalues, with s > 0 fixed and λ ∈ [0, 1]. By inspection, we have µ± (λ) ≥ 0 for all λ ∈ [0, 1]. First consider µ+ (λ). At the boundaries, we have µ+ (λ = 0) = s−1 and µ+ (λ = 1) = s + s−1 . (B.26) Moreover, differentiating in λ, we have ! 1 (s2 + 4)λ − 1 ′ µ+ (λ) = , (B.27) s+ p 2 (s−1 + λs)2 − 4λ(1 − λ) which is positive on the interval λ ∈ (0, 1) for any s. To show this, consider two cases. First, if the numerator of the fraction inside the square brackets is non-negative, than the result is immediate. If the numerator is negative, then we must have 0 < s2 [(s−1 + λs)2 − 4λ(1 − λ)] − [(s2 + 4)λ − 1]2 = 4λ[2 − (s2 + 4)λ]. (B.28) 2 But, for the numerator to be negative, we must have λ < 1/(s + 4), meaning that 2 − (s2 + 4)λ > 1, and thus that the desired inequality is in fact satisfied. Now consider µ− (λ). At the boundaries, we see that µ− (λ = 0) = µ− (λ = 1) = 0. (B.29) By differentiating, we find that µ− (λ) has a unique maximum in λ at 2 λ⋆ = 2 , (B.30) s +4 where it takes value s µ− (λ⋆ ) = 2 . (B.31) s +4 Note that for any s > 0, s < s−1 , (B.32) 2 s +4 as upon multiplying through by s(s2 + 4) this is equivalent to the inequality s2 < s2 + 4. Therefore, assuming that i = 1 is the unique maximal eigenvalue, we conclude three things about the functions µ± (λ): µ+ (λ) is monotonically increasing, µ− (λ) is non-monotonic, and we have the strict uniform-bound separation s 0 ≤ µ− (λ) ≤ 2 < s−1 ≤ µ+ (λ′ ) ≤ s + s−1 (B.33) s +4 ′ for any pair λ, λ ∈ [0, 1]. With this lemma in hand, we find that for any pair (j, k) ∈ [d]2 such that λj > λk , the corresponding upper-branch eigenvalues of the Hessian satisfy µj,+ > µk,+ . (B.34) 12

Broken scale symmetries in undercomplete linear autoencoders

Moreover, we have a strict separation between the two branches: for any pair (j, k) we have that 1 s 1 µj,− ≤ 2 < s−1 ≤ µk,+ . (B.35) λ1 s +4 λ1 However, the lower-branch eigenvalues µj,− are not ordered simply by the corresponding eigenvalues λj of Σ. We remark that this result guarantees strong separation of bands when s < 1, but allows the gap to in principle become small when s ≫ 1. As a result, assuming that i = 1 is the index of the maximum eigenvalue, the leading eigenvalue of the Hessian is always µi,+ = λi (s + s−1 ). Thus, the leading eigenvalue is minimized by the balanced solution, and in this sense balancing minimizes sharpness. However, it does not favor imbalance in a particular direction, as it is invariant under s 7→ 1/s. Asymmetric effects on the Hessian therefore must come from lower eigenvalues, which are divided into two non-interlacing bands. One way to capture the asymmetry is through the trace of the Hessian: X tr(H) = tr(Hi ) + tr(Hj )

(B.36)

j̸=i

= λi (s + s−1 ) +

X (λj s + λi s−1 )

(B.37)

j̸=i −1

= tr(Σ)s + dλi s . (B.38) We can then see that the trace is invariant under s 7→ 1/s only if λi = λj for all j, i.e., if all the eigenvalues of Σ are identical. Otherwise, we have the strict inequality dλi > tr(Σ), (B.39) which means that the trace of the Hessian increases more as one moves towards a encoderheavy solution (s < 1) as compared to a decoder-heavy solution (s > 1). We note that tr(H) is minimized by taking   tr(Σ) −1/2 s⋆ = ≥ 1; (B.40) dλi it is decreasing on (0, s⋆ ) and increasing on (s⋆ , ∞). This means that in the sense of tr(H), decoder-heavy solutions can be flatter than encoder-heavy solutions or even balanced solutions.

B.2. Hessian spectrum for a general undercomplete linear autoencoder We now consider a general undercomplete linear autoencoder with k > 1. We will find that the Hessian has a similar structure to the k = 1 case—being divided into sectors corresponding to within-principal-subspace and out-of-subspace fluctuations. Recall that the loss, after averaging over some set of examples with E[x] = 0 and E[xx⊤ ] = Σ, is 1 1 L(W, U ) = Ex ∥W U x − x∥2 = tr[(W U − Id )Σ(U ⊤ W ⊤ − Id )]. (B.41) 2 2 13

Pashakhanloo Zavatone-Veth

Consider a generic point (W, U ), and perturb this point to (W + δW, U + δU ). Then, L(W + δW, U + δU ) 1 = tr[(W U + W δU + δW U + δW δU − Id )Σ 2 × (U ⊤ W ⊤ + δU ⊤ W ⊤ + U ⊤ δW ⊤ + δU ⊤ δW ⊤ − Id )] (B.42) = L(W, U ) + tr[(W δU + δW U )Σ(W U − Id )⊤ ] 1 + tr[(W δU + δW U )Σ(δU ⊤ W ⊤ + U ⊤ δW ⊤ )] + tr[(δW δU )Σ(W U − Id )⊤ ] 2 + cubic and quartic terms. (B.43) As a result, viewing the Hessian as a quadratic form acting on the perturbation (δW, δU ), we have H(δW, δU ) = tr[(W δU + δW U )Σ(W δU + δW U )⊤ ] + 2 tr[(δW δU )Σ(W U − Id )⊤ ]. (B.44)

Now, we specialize to a rank-k undercomplete PCA-type solution, with W = P ⊤ A,

(B.45)

−1

U = A P, (B.46) k×d where A is a general invertible matrix and P ∈ R is a projection matrix satisfying ⊤ P P = Ik . Specifically, we assume that that P projects onto the eigenspace spanned by the top k eigenvectors, assuming that it is unique (i.e., that λk > λk+1 ).

Choose coordinates for the input space such that Σ is diagonal, and then for the latent space such that
P = Ik 0 . (B.47) In these coordinates, decompose as  the perturbations 
δW∥ δW = and δU = δU∥ δU⊥ , (B.48) δW⊥ where δW∥ , δU∥ ∈ Rk×k and δW⊥⊤ , δU⊥ ∈ Rk×(d−k) . (B.49) With this choice of coordinates,    

δW∥ A δU∥ δU⊥ + A−1 0 W δU + δW U = (B.50) 0 δW⊥   δW∥ A−1 + AδU∥ AδU⊥ = . (B.51) δW⊥ A−1 0 14

Broken scale symmetries in undercomplete linear autoencoders

Then, letting Λ+ = diag(λ1 , . . . , λk ) and Λ− = diag(λk+1 , . . . , λd ), we have tr[(W δU + δW U )Σ(W δU + δW U )⊤ ] (B.52)     −1 ⊤ −1 ⊤ −1 δW∥ A + AδU∥ AδU⊥ (δW∥ A + AδU∥ ) (δW⊥ A ) Λ+ 0 = tr 0 Λ− δW⊥ A−1 0 (AδU⊥ )⊤ 0 (B.53)    −1 −1 ⊤ −1 ⊤ δW∥ A + AδU∥ AδU⊥ Λ+ (δW∥ A + AδU∥ ) Λ+ (δW⊥ A ) = tr (B.54) −1 ⊤ δW⊥ A 0 Λ− (AδU⊥ ) 0 = tr[(δW∥ A−1 + AδU∥ )Λ+ (δW∥ A−1 + AδU∥ )⊤ ] + tr[AδU⊥ Λ− (AδU⊥ )⊤ ] + tr[δW⊥ A−1 Λ+ A−⊤ δW⊥⊤ ] and, as

(B.55)

  0 0 Σ(W U − Id ) = , 0 −Λ−

(B.56)

tr[(δW δU )Σ(W U − Id )⊤ ] = − tr[δW⊥ δU⊥ Λ− ].

(B.57)

⊤

Therefore, the Hessian decomposes into two independent quadratic forms H(δW, δU ) = H∥ (δW∥ , δU∥ ) + H⊥ (δW⊥ , δU⊥ ), (B.58) with H∥ (δW∥ , δU∥ ) = tr[(δW∥ A−1 + AδU∥ )Λ+ (δW∥ A−1 + AδU∥ )⊤ ] (B.59) determined by within-subspace fluctuations, and H⊥ (δW⊥ , δU⊥ ) = tr[AδU⊥ Λ− δU⊥⊤ A⊤ ] + tr[δW⊥ A−1 Λ+ A−⊤ δW⊥⊤ ] − 2 tr[δW⊥ δU⊥ Λ− ] (B.60) determined by out-of-subspace fluctuations. We can therefore diagonalize each separately.

First, consider the out-of-subspace fluctuations. Using the vectorization convention tr(A⊤ B) = vec(A)⊤ vec(B), vec(ABC) = (C ⊤ ⊗ A) vec(B), (B.61) we can write the three relevant terms as tr[AδU⊥ Λ− δU⊥⊤ A⊤ ] = vec(δU⊥ )⊤ (Λ− ⊗ A⊤ A) vec(δU⊥ ) (B.62) tr[δW⊥ A−1 Λ+ A−⊤ δW⊥⊤ ] = vec(δW⊥⊤ )⊤ (Id−k ⊗ A−1 Λ+ A−⊤ ) vec(δW⊥⊤ )

(B.63)

tr[δW⊥ δU⊥ Λ− ] = vec(δW⊥⊤ )⊤ (Λ− ⊗ Ik ) vec(δU⊥ ),

(B.64)

hence

 ⊤    Id−k ⊗ A−1 Λ+ A−⊤ −Λ− ⊗ Ik vec(δW⊥⊤ ) vec(δW⊥⊤ ) H⊥ (δW⊥ , δU⊥ ) = . vec(δU⊥ ) vec(δU⊥ ) −Λ− ⊗ Ik Λ− ⊗ A ⊤ A (B.65) We can then see that there is a further splitting into eigenproblems for the 2k × 2k block matrices  −1  A Λ+ A−⊤ −λj Ik , (B.66) −λj Ik λ j A⊤ A for each j ∈ {k + 1, . . . , d}. This generalizes our result from the k = 1 case before. 15

Pashakhanloo Zavatone-Veth

We now at last come to H∥ . Expanding out all of the terms, H∥ = tr[(δW∥ A−1 + AδU∥ )Λ+ (δW∥ A−1 + AδU∥ )⊤ ] −1

−⊤

(B.67)

δW∥⊤ + δU∥⊤ A⊤ )]

(B.68) = tr[(δW∥ A + AδU∥ )Λ+ (A   ⊤ ⊤ −⊤ ⊤ −1 ⊤ ⊤ −1 −⊤ ⊤ = tr δW∥ A Λ+ A δW∥ + δW∥ A Λ+ δU∥ A + AδU∥ Λ+ A δW∥ + AδU∥ Λ+ δU∥ A (B.69) ⊤

−1

= vec(δW∥ ) (A ⊤

Λ+ A

−⊤

⊤

−1

⊗ Ik ) vec(δW∥ ) + 2 vec(δW∥ ) (A

Λ+ ⊗ A) vec(δU∥ )

⊤

+ vec(δU∥ ) (Λ+ ⊗ A A) vec(δU∥ ), (B.70) which means that  ⊤  −1   vec(δW∥ ) vec(δW∥ ) A Λ+ A−⊤ ⊗ Ik A−1 Λ+ ⊗ A H∥ (δW∥ , δU∥ ) = , (B.71) vec(δU∥ ) vec(δU∥ ) Λ+ A−⊤ ⊗ A⊤ Λ+ ⊗ A ⊤ A and thus the remaining eigenvalues determined by the spectrum of the 2k 2 × 2k 2 matrix   −1 are−⊤ −1 A Λ+ A ⊗ I k A Λ+ ⊗ A . (B.72) Λ+ A−⊤ ⊗ A⊤ Λ+ ⊗ A ⊤ A ⊤ We note that we can factorize  this matrix as J J, where  2

2

1/2 k ×2k −⊤ ⊗ I J = Λ1/2 . k Λ+ ⊗ A ∈ R + A

(B.73)

This exposes the fact that H∥ has k 2 zero modes, which come from solutions to δW∥ A−1 + AδU∥ = 0, (B.74) and thus result from the internal symmetry. Noting that 1/2 1/2 JJ ⊤ = Λ+ (AA⊤ )−1 Λ+ ⊗ Ik + Λ+ ⊗ AA⊤ (B.75) is positive-definite assuming A is invertible, the non-zero eigenvalues in this sector are determined by the eigenvalues of this k 2 × k 2 matrix. Therefore, in the end, we find that the eigenvalues of H split into sectors: one k 2 dimensional sector determined by the eigenvalues of 1/2 1/2 Λ+ (AA⊤ )−1 Λ+ ⊗ Ik + Λ+ ⊗ AA⊤ , (B.76) one 2k-dimensional sector determined by the eigenvalues of   −1 A Λ+ A−⊤ −λj Ik (B.77) −λj Ik λ j A⊤ A for each j ∈ {k + 1, . . . , d} (giving d − k sectors in all), and a remaining sector of k 2 zero modes. This gives, in all k 2 + k 2 + 2k(d − k) = 2kd eigenvalues, as we expect. We note that, as its lower two blocks we have  −1commute, A Λ+ A−⊤ −λj Ik det = det(λj A−1 Λ+ A−⊤ A⊤ A − λ2j Ik ) (B.78) −λj Ik λ j A⊤ A = det(λj A−1 Λ+ A − λ2j Ik ) = λkj

k Y

(λi − λj )

(B.79) (B.80)

i=1

upon applying the Weinstein–Aronszajn identity in passing to the third line. This shows that the Hessian on each of these sectors is positive-definite if λi > λj i ∈ {1, . . . , k}, j ∈ {k + 1, . . . , d}. (B.81) 16

Broken scale symmetries in undercomplete linear autoencoders

We can then compute that, in general, the Hessian has trace  −1  d X A Λ+ A−⊤ −λj Ik 1/2 ⊤ −1 1/2 ⊤ tr(H) = tr[Λ+ (AA ) Λ+ ⊗ Ik + Λ+ ⊗ AA ] + tr −λj Ik λ j A⊤ A j=k+1

(B.82) ⊤ −1

⊤

= d tr[Λ+ (AA ) ] + tr(Σ) tr(AA ). With the gauge-fixing A = S 1/2 , tr(H) = d tr(Λ+ S −1 ) + tr(Σ) tr(S).

(B.83) (B.84)

Though the eigenvalues in general depend on A in full detail, there is an important case in which they simplify. Suppose that we fix the positive gauge A = S 1/2 , S = diag(s1 , . . . , sk ), (B.85) with s1 ≥ s2 ≥ · · · ≥ sk ≥ 1, λ1 ≥ λ2 ≥ · · · ≥ λk > 0. (B.86) matching the ordering of the eigenvalues. These assumptions will be satisfied in the effective model we derive later. Then, the non-zero eigenvalues in the within-subspace sector are   1 ∥ , a, b ∈ [k], (B.87) µab = λa sb + sa together with the k 2 zero modes coming from the internal GL(k) symmetry. For each fixed a, ∥ ∥ ∥ µa1 ≥ µa2 ≥ · · · ≥ µak , (B.88) since ∥ ∥ µab − µad = λa (sb − sd ). (B.89) Moreover, the diagonal is ordered: ∥ ∥ ∥ µ11 ≥ µ22 ≥ · · · ≥ µkk , (B.90) −1 because s 7→ s + s is increasing on [1, ∞). For the out-of-subspace sector, each pair  a ≤ k < jgives a 2 × 2 block λa /sa −λj ⊥ Haj = . −λj λj sa Following our analysis in the its eigenvalues are  k = 1 case,s   2 1 λa λa µaj,± =  + λj sa ± + λj sa − 4λj (λa − λj ) . 2 sa sa

(B.91)

(B.92)

⊥ , and since λ ≤ λ , we have Since µaj,+ ≤ tr Haj j a   λa λa 1 ∥ µaj,+ ≤ + λ j sa ≤ + λa sa = λa sa + = µ∥aa ≤ µa1 . (B.93) sa sa sa Thus no out-of-subspace eigenvalue exceeds the largest within-subspace eigenvalue in the corresponding retained row. Consequently,   1 ∥ . (B.94) λmax (H) = max µa1 = max λa s1 + a≤k a≤k sa In particular,   1 λ1 s1 + ≤ λmax (H) ≤ λ1 (s1 + 1), (B.95) s1

17

Pashakhanloo Zavatone-Veth

and hence, in the strongly decoder-heavy regime s1 ≫ 1, λmax (H) = λ1 s1 + O(λ1 ). (B.96) The assumptions λ1 ≥ · · · ≥ λk , s1 ≥ · · · ≥ sk ≥ 1 (B.97) do not totally order the within-subspace eigenvalues. Indeed, for  fixed b, λa λc ∥ ∥ − , (B.98) µab − µcb = sb (λa − λc ) + sa sc and the second term need not have the same sign as the first. A product partial order follows under the additional ratio-ordering condition λ1 λ2 λk ≥ ≥ ··· ≥ . (B.99) s1 s2 sk Under this extra assumption, ∥ ∥ a ≤ c, b ≤ d =⇒ µab ≥ µcd . (B.100) B.3. Hessian spectrum for a general overcomplete linear autoencoder Now consider an overcomplete PCA-type solution with k > d, for which we have W U = Id . Our starting point is again H(δW, δU ) = tr[(W δU + δW U )Σ(W δU + δW U )⊤ ] + 2 tr[(δW δU )Σ(W U − Id )⊤ ] (B.101) = tr[(W δU + δW U )Σ(W δU + δW U )⊤ ], (B.102) where we note that the second term vanishes thanks to the assumption that the reconstruction is exact. We now vectorize the first term. As vec(W δU + δW U ) = (Id ⊗ W ) vec(δU ) + (U ⊤ ⊗ Id ) vec(δW ), (B.103) we can write the Hessian as  ⊤   vec(δW ) vec(δW ) ⊤ . (B.104) J J H(δW, δU ) = vec(δU ) vec(δU ) for
J = Σ1/2 U ⊤ ⊗ Id Σ1/2 ⊗ W . (B.105) Now, as W U = Id , both W and U must have rank d, hence J must have at most d2 , a bound which is attained assuming that Σ is full-rank. Then, noting that JJ ⊤ = Σ1/2 U ⊤ U Σ1/2 ⊗ Id + Σ ⊗ W W ⊤ (B.106) 2 2 is then a full-rank d × d matrix, we find that the Hessian has 2kd − d2 zero modes, plus d2 non-zero eigenvalues given by the eigenvalues of JJ ⊤ . If Σ has zero modes, then they will generate zero modes of JJ ⊤ . If we again write the parameters at the optimum as W = P ⊤A (B.107) U = A−1 P, where P now satisfies P ⊤ P = Id , we find that JJ ⊤ = Σ1/2 P ⊤ (AA⊤ )−1 P Σ1/2 ⊗ Id + Σ ⊗ P ⊤ AA⊤ P. In this case, the Hessian has the generally simple expression tr(H) = d tr(ΣU ⊤ U ) + tr(Σ) tr(W W ⊤ ). If we consider a gauge in which W W ⊤ = S and U ⊤ U = S −1 , then this becomes tr(H) = d tr(ΣS −1 ) + tr(Σ) tr(S). 18

(B.108) (B.109) (B.110) (B.111)

Broken scale symmetries in undercomplete linear autoencoders

Further supposing that S and Σ commute, this gives  d  X λj + λ̄sj , tr(H) = d sj

(B.112)

j=1

where λ̄ = d1 tr(Σ). √ For global re-scalings of the balanced solution, with A = sIk , we have the reduction JJ ⊤ = (s + s−1 )(Σ ⊗ Id ), (B.113) showing that in this case all non-zero eigenvalues are invariant under s 7→ 1/s.

Appendix C. Detailed derivation of the effective scale dynamics for PCA Here, we derive the effective model for scale dynamics that we presented in the main text. The accuracy of the predictions of these effective dynamics rely on two facts, which we prove below: First, that fluctuations normal to the solution manifold contract back to the manifold at a fast timescale tη. Second, that Qt evolves at a slow timescale tη 2 /B, and on that timescale that we can neglect stochasticity in its evolution. This both justifies the effective model we present in the text, and shows that the PCA solution is locally stable. The fundamental mechanism for this is that, as we showed in Appendix B, the normal component of the Hessian is strictly positive-definite if the solution is a projection to the k-principal subspace, while otherwise it has modes with negative eigenvalue. In this derivation, we will assume that the batches are i.i.d., and that each is composed of B i.i.d. datapoints. We assume that the data has bounded moments of all orders, i.e., E[∥x∥2p ] < ∞ for all p > 0. Alternatively, we may assume that the data distribution is subGaussian. In practice, we will use Gaussian data in the simulations against which we will compare the model’s predictions. We will also assume that the data covariance E[xx⊤ ] = Σ is diagonal. Consider a generic point on the solution manifold, with (W, U ) = (P ⊤ A, A−1 P ). Using the polar decomposition of A, we gauge-fix such that A = S 1/2 for S a symmetric PSD matrix. Then, (W, U ) = (P ⊤ S 1/2 , S −1/2 P ). Fix a compact subset K of the cone of symmetric positive definite matrices, defined for positive constants C > c > 0 by K = {0 ≺ cI ⪯ S ⪯ CI}, (C.1) and let MK = {(P ⊤ A, A−1 P ) : A⊤ A ∈ K} (C.2) be the corresponding submanifold of the solution manifold. The assumption that S is bounded will technically be required at several points in the proof. C.1. Fast relaxation of normal fluctuations We first want to control fluctuations in the normal space. We will do so via a rather abstract and generic derivation, which parallels previous analyses of how SGD behaves near solution manifolds, e.g. by Li et al. (2022); Pashakhanloo and Koulakov (2023); Pashakhanloo (2025). In a small neighborhood of MK , we can decompose the full parameter vector θt as θt = mt + ξt , (C.3) 19

Pashakhanloo Zavatone-Veth

where mt ∈ MK and ξt ∈ Nmn MK is in the normal space to MK at mn . Let Π be the local projection onto MK , and Π⊥ (θ) = θ−Π(θ) the corresponding normal-coordinate map. Thus, ξt = Π⊥ (θ). Write the SGD iteration as θt+1 = θt + δt , (C.4) where δt = −η(∇L(θt ) + ζt (θt )) (C.5) where ζt (θt ) is the SGD noise. Thus, applying Π⊥ , we have ξt+1 = Π⊥ (θt + δt ) (C.6) 2 (C.7) = ξt + DΠ⊥ θt [δt ] + O(∥δt ∥ ), where in the second line we have made a Taylor expansion. Now, at a given point m on the manifold, the derivative of the normal-coordinate map is the projection onto the normal space: DΠ⊥ (C.8) m = PNm , ⊥ so, assuming that Π is smooth in a neighborhood around the manifold, ⊥ (C.9) DΠ⊥ θt = DΠmt +ξt = PNmt + O(∥ξt ∥). This gives ξt+1 = ξt + PNmt δt + O(∥ξt ∥ ∥δt ∥) + O(∥δt ∥2 ). (C.10)

By definition, ∇L(m) = 0 for all m ∈ MK . Then, if we Taylor expand the population loss in normal coordinates, we have ⊥ PNm ∇L(m + ξ) = Hm ξ + O(∥ξ∥2 ), (C.11) ⊥ where Hm is the projection of the Hessian onto the normal space at m. Moreover, we have that ∇L(mt + ξt ) = O(∥ξt ∥). Therefore, letting ζt⊥ = PNmt ζt (C.12) be the projection of the SGD noise onto the normal space, we have that ⊥ PNmt δt = −ηHm ξ − ηζt⊥ + O(η∥ξt ∥2 ) (C.13) t t and ∥δt ∥ = O(η(∥ξt ∥ + ∥ζt ∥)). (C.14) Thus, we at last find that ⊥ ξt+1 = (I − ηHm )ξt − ηζt⊥ + ϵt , t

(C.15) where the remainder term ϵt = O(η∥ξt ∥2 ) + O(η∥ξt ∥ ∥ζt ∥) + O(η 2 (∥ξt ∥ + ∥ζt ∥)2 ) (C.16) has three contributions with distinct origins: the first comes from nonlinear contributions to the population gradient normal to the solution manifold, the second from the variation of the projection PNm away from the manifold, and the third from the expansion of the normal-coordinate map. We now use this recurrence to argue that fluctuations away from the solution manifold relax rapidly. We do so supposing that the normal Hessian satisfies the uniform bound ⊥ Hm ⪰ µI (C.17) t for some µ > 0, which in the case of interest is true provided that the data has a spectral gap separating the k-principal subspace from the remaining directions (Appendix B). Under this assumption, the stochastic dynamics for ξt are zero-reverting, and rapidly equilibrate. 20

Broken scale symmetries in undercomplete linear autoencoders

Given that the noise has covariance of order 1/B, we can see that we should have r  η (C.18) ∥ξt ∥ = O B in probability. If we instead considered a solution not corresponding to the k-principal subspace, the Hessian would not be positive-definite, and the process would not be zeroreverting. This is the basic reason why the PCA solution manifold is locally stable under SGD, while non-principal subspace solutions are locally unstable (Baldi and Hornik, 1989). To show this carefully, we fix a radius ρ > 0, and suppose that ξt lies within that tube, so that the remainder terms in ϵt are controlled. To be precise, let τρ = inf{t : ∥ξt ∥ ≥ ρ} (C.19) be the time of first exit from the tube. Then, using the assumption that the Hessian is uniformly positive-definite and supposing that η is small, there exist constants c, C > 0 independent of η and B such that η2 E[∥ξt+1 ∥2 1t+1<τρ ] ≤ (1 − cη)E[∥ξt ∥2 1t+1<τρ ] + C , (C.20) B where the effect of the remainder terms can be absorbed into the constants provided that the tube radius ρ and learning rate η are small enough. This assumes that ∥ζt⊥ ∥4 is uniformly bounded in expectation, which is satisfied for the datasets of interest given that ξt lies within the tube. Iterating this recurrence, we find that η E[∥ξt+1 ∥2 1t+1<τρ ] ≤ e−cηt ∥ξ0 ∥2 + C , (C.21) B which means that the mean-squared normal displacement is of order η/B after a time on p 1 1 the order of η log η . This implies that ∥ξt ∥ = O( η/B). We now want to argue that the probability that the exit time τρ will be reached over the order-B/η 2 time interval of interest is small. Under the assumption that the data has finite moments of all orders, we can see that we can use the same argument as above to show that  η p E[∥ξt+1 ∥2p 1t+1<τρ ] ≤ C (C.22) B for any p ≥ 1, after the fast transient of order η1 log η1 steps. Let T1 be large enough such that the transient is evaded, and T2 be of order B/η 2 , i.e., let T2 = ⌊τ B/η 2 ⌋ for some order-one τ . Then, by a union bound and Markov’s inequality,   X T2  η p P max ∥ξt ∥ > ρ ≤ ≤ Cη p−2 B 1−p , (C.23) P(∥ξt ∥ > ρ0 ) ≤ C(T2 − T1 ) T1 ≤t≤T2 B t=T1

where the constant changes between the last two inequalities, and depends on τ . This goes to zero with η for any p > 2, hence the probability of exiting the tube tends to zero. Therefore, we can use the result above for the stopped process to show that the meansquared displacement of ξt over the interval of interest is o(1). The important point here is that this relaxation is fast relative to the slow timescale of the evolution of Q. On the manifold, the mean update to Q is of order η 2 /B, so for an order-one change to accumulate we must wait B/η 2 steps. In contrast, we need to wait only 1/η steps for the fluctuations away from the manifold to relax back to the manifold, meaning that the two timescales have a relative separation η/B, which tends to zero as η tends to zero. This shows that the contribution of these fluctuations to the evolution of Q is o(1). 21

Pashakhanloo Zavatone-Veth

C.2. Stochastic approximation for flow along the manifold We now leverage the separation of timescales proved above to derive an effective flow describing the slow scale growth. This is essentially a standard exercise in stochastic approximation (Kushner and Yin, 2003). As there is a separation of timescales, we can evaluate the evolution of Qt on the solution manifold and incur only an o(1) error on the relevant timescale. On the manifold, we have Wt⊤(Wt Ut − Id ) = 0, which gives the immediate simplification  Qt+1 − Qt = η 2 Ut Σ̂t (Wt Ut − Id )⊤ (Wt Ut − Id )Σ̂t Ut⊤

+ o(1).

(C.24)

We now want to evaluate the mean and variance of the remaining update with respect to the batch. To do so, we decompose each datapoint as x = P ⊤ x∥ + x ⊥ , (C.25) where x⊥ lies outside the principal subspace, and satisfies P x⊥ = 0. This leads to the decomposition of the full data matrix as Xt = P ⊤ X∥,t + X⊥,t (C.26) −1/2

and thus, using the fact that Wt Ut − Id = P ⊤ P − Id and Ut = St P on the solution manifold given the gauge-fixing assumption, 1 −1/2 ⊤ ⊤ (Wt Ut − Id )Σ̂t Ut⊤ = (P ⊤ P − Id )(P ⊤ X∥,t + X⊥,t )(X∥,t P + X⊥,t )P ⊤ St (C.27) B 1 ⊤ −1/2 = X⊥,t X∥,t St (C.28) B 1 X −1/2 = x⊥,i x⊤ . (C.29) ∥,i St B i∈Bt

where in the last line we introduce indexing of the datapoints, with the minibatch defined by an index set Bt . Then, we have 1 −1/2 X ⊤ −1/2 , (C.30) (x⊥,i x⊥,j )x∥,i x⊤ Ut Σ̂t (Wt Ut − Id )⊤ (Wt Ut − Id )Σ̂t Ut⊤ = 2 St ∥,j St B i,j∈Bt so η 2 −1/2 X ⊤ −1/2 Qt+1 − Qt = 2 St + o(1). (C.31) (x⊥,i x⊥,j )x∥,i x⊤ ∥,j St B i,j∈Bt

Now, assuming that the datapoints are independent with mean zero and diagonal covariance, the mean update is   2 X η η2 ⊤  E 2 C, (C.32) (x⊤ ⊥,i x⊥,j )x∥,i x∥,j = B B i,j∈Bt

where we define C = E[∥x⊥ ∥2 x∥ x⊤ ∥ ]. We can then write

η 2 −1/2 η 2 −1/2 −1/2 −1/2 St CSt + St Zt St + o(1), B B 1 X ⊤ (x⊥,i x⊥,j )x∥,i x⊤ Zt = ∥,j − C B

Qt+1 − Qt = where

(C.33) (C.34) (C.35)

i,j∈Bt

is a noise term that is independent across timesteps. Assuming that the data has bounded eighth moments, Zt has bounded second moment, where the constant may depend on the 22

Broken scale symmetries in undercomplete linear autoencoders

batch size B. Thus, assuming that St is bounded (which holds if the solution stays within the submanifold MK ), over T ∼ B/η 2 steps, we have !   2  T X η η 2 −1/2 η4 −1/2 var , (C.36) S Zt St =O T 2 =O B t B B t=1 which tends to zero as η ↓ 0 at fixed batch size. Combining these results, we find that if we take η ↓ 0 while keeping τ = tη 2 /B to be of order one, the discrete-time dynamics tend in probability to the solution of the corresponding limiting ODE by standard Grönwall-type arguments (Kushner and Yin, 2003). This relies on the fact that Q(S) = S − S −1 is a Lipschitz function of S provided that the solution remains within MK , i.e., that S is uniformly bounded from below and above in Löwner order. Assuming that this condition holds, the limiting dynamics are d (S − S −1 ) = S −1/2 CS −1/2 , (C.37) dτ or, equivalently, dS dS + S −1 S −1 = S −1/2 CS −1/2 . (C.38) dτ dτ C.3. Solution assuming co-diagonalizability Now, we show that the effective dynamics for S can be solved analytically provided that S commutes with C at some reference time, which we define as τ = 0 without loss of generality. We claim that if S and C commute at τ = 0 then they do throughout time under the effective dynamics. This follows from the fact that if S commutes with C, then so does dS/dτ . To see this, we write the effective   dynamics as dS LS = S −1/2 CS −1/2 , (C.39) dτ where we define the linear map LS (X) = X + S −1 XS −1 . (C.40) −1 We can see that LS is well-defined for any PSD S, as in vectorized form we can write vec(LS (X)) = (Ik ⊗ Ik + S −1 ⊗ S −1 ) vec(X), (C.41) so we have the explicit formula −1 vec(L−1 ⊗ S −1 )−1 vec(Y ). (C.42) S (Y )) = (Ik ⊗ Ik + S This allows us to write dS −1/2 = L−1 CS −1/2 ). (C.43) S (S dτ −1 We can now show that LS (Y ) commutes with C if both S and Y do. If LS (X) = Y for [Y, C] = 0, we have 0 = [Y, C] = [LS (X), C] = [X, C] + [S −1 XS −1 , C] = [X, C] + S −1 [X, C]S −1 = LS ([X, C]) (C.44) as by assumption S commutes with C. But, as LX is invertible, this implies that [X, C] = 0. Therefore, we conclude the desired claim. As a result, if we assume that [S(0), C] = 0, then we can choose a basis where S(τ ) = diag(s1 (τ ), . . . , sk (τ )) and C = diag(σ1 , · · · , σk ) (C.45) 23

Pashakhanloo Zavatone-Veth

for all times τ ≥ 0. Projecting into this basis, we have the decoupled equations dsi 1 dsi σi (C.46) + 2 = , i ∈ [k]. dτ si si dτ Multiplying by si and recognizing  that    d 1 2 1 dsi s + log si = si + , (C.47) dτ 2 i si dτ we arrive at the equation   d 1 2 (C.48) s + log si = σi . dτ 2 i Integrating, we have 1 1 si (τ )2 + log si (τ ) = si (0)2 + log si (0) + σi τ. (C.49) 2 2 We now want to solve this for si (τ ). Multiplying by 2 and exponentiating, 2 2 si (τ )2 esi (τ ) = si (0)2 esi (0) +2σi τ . (C.50) Now, recall that the principal branch of the Lambert W function, which we will write simply as W , satisfies (DLMF) W (z)eW (z) = z. (C.51) Thus, we have the solution p si (τ ) = W [si (0)2 exp(si (0)2 + 2σi τ )], (C.52) as claimed in the main text. We note that W (x) = log x + log log x + o(1) as x → ∞ (DLMF), so at long times this gives √ si (τ ) ∼ 2σi τ . (C.53) For Gaussian data with covariance eigenvalues λ , . . . , λ , we have 1  k  d X λj  λi , σi =  (C.54) j=k+1

which predicts that the growth rates along different PCs are ordered.

We can in fact prove a stronger ordering statement, which we will use later. Suppose we start from a balanced initial condition with si (τ = 0) = 1. Then, fixing a time τ , let s(λ) solve s(λ)2 + 2 log s(λ) = 1 + aλ, (C.55) where d X a = 2τ λj , (C.56) j=k+1

such that si (τ ) = s(λi ). We can show that λ s(λ)

(C.57)

as(λ) . 2(s(λ)2 + 1)

(C.58)

λ 7→ is increasing in λ. By implicit differentiation, s′ (λ) =

24

Broken scale symmetries in undercomplete linear autoencoders

Then,   d λ 1 λs′ (λ) = 1− (C.59) dλ s(λ) s(λ) s(λ)   1 aλ = 1− (C.60) s(λ) 2(s(λ)2 + 1)   s(λ)2 + 2 log s(λ) − 1 1 1− (C.61) = s(λ) 2(s(λ)2 + 1) where in the last equation we have used the equation that defines s(λ). Therefore, λ/s(λ) is increasing in λ if s2 + 2 log s − 1 < 2(s2 + 1), (C.62) or if 2 log s < s2 + 3. (C.63) But, as log s < s − 1 for all s > 0, and as 2(s − 1) < s2 + 3 for all s > 0 as it is equivalent to 2s < s2 + 5, we conclude that this inequality is satisfied. Therefore, for any time τ , we have the ordering λ1 λ2 λk ≥ ≥ ··· ≥ . (C.64) s1 (τ ) s2 (τ ) sk (τ ) We note also that the specialization of these results to global scale growth is easy to derive. Then, the object of interest is −1 [tr(St+1 ) − tr(St+1 )] − [tr(St ) − tr(St−1 )] = ∥x⊥ ∥2 tr[St−1 x∥ x⊤ (C.65) ∥] Making the global scale Ansatz S = sIk ,this leads to  d 1 s− = σ̄, (C.66) dτ s where we let 1 σ̄ = tr(C). (C.67) k Then, it is easy to see that the same p analysis as above leads to s(τ ) = W [s(0)2 exp(s(0)2 + 2σ̄τ )]. (C.68)

C.4. Why is this a good effective model? We now return to the question of why this might be a good effective model for scale growth, as it is not the only way to obtain a set of effective dynamics. In particular, suppose we tried to derive an effective model directly from the dynamics of W ⊤ W . Expanding out the update, we have ⊤ ⊤ ⊤ 2 2 ⊤ ⊤ Wt+1 Wt+1 = Wt⊤ Wt − η(Wt⊤ rt x⊤ (C.69) t Ut + Ut xt rt Wt ) + η ∥rt ∥ Ut xt xt Ut . ⊤ Now, on the solution manifold, we have Wt rt = 0, and moreover that rt = −xt,⊥ and Ut xt = A−1 t xt,∥ . This leads to the simplified dynamics ⊤ −⊤ Wt+1 Wt+1 − Wt⊤ Wt = η 2 A−1 [∥xt,⊥ ∥2 xt,∥ x⊤ (C.70) t,∥ ]A Under the same stochastic approximation approach as before, assuming we can neglect fluctuations away from the solution manifold, this would lead to the dynamics dS = S −1/2 CS −1/2 (C.71) dτ 25

Pashakhanloo Zavatone-Veth

for S = A⊤ A. This again preserves the condition that [S(τ ), C] = 0, meaning that we can diagonalize to obtain dsi si = σi (C.72) dτ for each i ∈ [k]. As   dsi d 1 2 si = s , (C.73) dτ dτ 2 i this leads to p si (τ ) = si (0)2 + 2σi τ . (C.74) √ Thus, this approach gives the same 2σi τ growth at long times as our previous approach, but its early-time dynamics are different. Why is this the case, and which result should we prefer? Well, in deriving effective dynamics from the order-η 2 update from the dynamics of Q, we used the assumption that we are on the solution manifold only to simplify the update itself. When working with the W ⊤ W dynamics directly, we used this assumption to argue that the order-η term vanishes and also to simplify the order-η 2 term. Thus, we expect the latter effective model to incur errors as a result of fluctuations away from the solution manifold that yield non-zero values of the O(η) update; the two models are affected in similar ways by how those fluctuations change the O(η 2 ) term. This means that we should expect the stochastic approximation from the dynamics of Q to give a better effective model. This is consistent with the fact that we see that the Lambert-W solution matches experiment very closely. C.5. The edge of stability for the effective model The effective model as derived above predicts indefinite anisotropic scale growth si (τ ) ∼ √ 2σi τ . However, this cannot be a true property of the underlying discrete-time dynamics, as the maximum eigenvalue of the Hessian of the population loss is scale-dependent, so indefinite growth would destabilize the discrete-time dynamics.1 For an eigenvalue µi (H) of the Hessian, we recall that the general edge of stability criterion is that ηµi (H) ≈ 2, with the threshold being likely somewhat less than 2 for SGD (Cohen et al., 2021; Damian et al., 2023; Andreyev and Beneventano, 2025; Emmanouilidis et al., 2026). Recall from our analysis in Appendix B that the within-subspace sector Hessian eigenvalues are   1 ∥ µab = λa sb + , a, b ∈ [k], (C.75) sa and that we expect these to be the dangerous modes. We have shown above that the scales satisfy the ordering λj λi ≥ (C.76) si sj if i < j, if we start from a balanced initial condition. Thus, given the ordering results from Appendix B, we conclude that for any a ≤ c and b ≤ d, we have ∥ ∥ µab ≥ µcd . (C.77) 1. This would also violate the assumption of boundedness of S that we used in the derivation of the effective model, but that technical issue can be resolved by proving the result over fixed time intervals with a bound that grows in time. The destabilization of the SGD dynamics is a more substantial violation of the assumptions.

26

Broken scale symmetries in undercomplete linear autoencoders

∥

∥

∥

For a given b, we then have that µab ≥ µcb whenever a ≤ c, which shows that µ1b is the dangerous mode for each b. Given the ordering of the scales si , we expect them to sequentially hit the stability threshold. Write the threshold condition as ηλi (H) = κ for some 0 < κ ≤ 2. First, s1 will hit a threshold at s1 = s∗1 determined by   1 ∗ ηλ1 s1 + ∗ = κ. (C.78) s1 Subsequently, the lower modes will hit thresholds determined by  1 ηλ1 + s∗j = κ, j > 1, (C.79) s∗1 This leads to a common ceiling for all scale directions. To show this very explicitly, let κ̃ = κ/ηλ1 be the normalized stability threshold, which we assume to be large as we are working at small learning rate. Then, the above conditions state that the scales satisfy 1 1 (C.80) s∗1 = κ̃ − ∗ and s∗j = κ̃ − ∗ , j > 1. s1 s1 From this, we can see immediately that we must have all s∗j = s∗1 at the threshold, where by solving the equation for s∗1 we have √ κ̃ + κ̃2 − 4 ∗ ∗ sj = s1 = = κ̃ + O(κ̃−1 ), (C.81) 2 where the asymptotic expansion holds for κ̃ ≫ 1.

Appendix D. Batch sharpness decreases with scale Andreyev and Beneventano (2025) proposed that networks trained with minibatch SGD operate near the Edge of Stochastic Stability, defined by when the so-called “batch sharpness” saturates near 2/η. The batch sharpness is" defined by # BS = EB

∇L̂⊤ B ĤB ∇L̂B ∥∇L̂B ∥2

(D.1)

where L̂B is the loss on minibatch B, and ĤB is its Hessian. In the present setting, consider a minibatch XB ∈ Rd×B with empirical covariance Σ̂B = B1 XB XB⊤ . Then, we have 1 L̂B = tr[(W U − Id )Σ̂B (W U − Id )⊤ ], (D.2) 2 so the minibatch gradients are ∇W L̂B = (W U − Id )Σ̂B U ⊤ (D.3) ∇U L̂B = W ⊤ (W U − Id )Σ̂B . (D.4) Using the general result from Appendix B, the minibatch Hessian is compactly expressed as a bilinear form acting on a perturbation (δW, δU ): ĤB (δW, δU ) = tr[(W δU + δW U )Σ̂B (W δU + δW U )⊤ ] (D.5) + 2 tr[(δW δU )Σ̂B (W U − Id )⊤ ]. We now compute the batch sharpness at a PCA-type solution. Recall that such a solution satisfies (W, U ) = (P ⊤ A, A−1 P ) (D.6) 27

Pashakhanloo Zavatone-Veth

where P ∈ Rk×d is an orthogonal projection (P P ⊤ = Ik ) and A ∈ GL(k). Then, W U − Id = P ⊤ P − Id = −P⊥

(D.7)

is a projection, and moreover one that is annihilated by P . This shows that ∇U L̂B = 0, while ∇W L̂B = −P⊥ Σ̂B P ⊤ A−⊤ . As before, introduce coordinates for the data matrix XB = P ⊤ X∥,B + X⊥,B , (D.8) such that P XB = X∥,B and P⊥ XB = X⊥,B . (D.9) Then, 1 1 ∥∇L̂B ∥2 = ∥∇W L̂B ∥2 = 2 ∥P⊥ XB XB⊤ P ⊤ A−⊤ ∥2 = 2 tr(KB HB ). (D.10) B B where we have defined the B × B matrices ⊤ ⊤ (AA⊤ )−1 X∥,B and HB = X⊥,B X⊥,B . (D.11) KB = X∥,B Moreover, ∇W L̂B U = −P⊥ Σ̂B P ⊤ (AA⊤ )−1 P, (D.12) hence 1 ĤB (∇W L̂B , ∇U L̂B ) = tr[(∇W L̂B U )Σ̂B (∇W L̂B U )⊤ ] = 3 tr[KB2 HB ]. (D.13) B Therefore, we can write the batch sharpness  on the2 PCA  solution manifold as 1 tr[KB HB ] BS = EB . (D.14) B tr(KB HB ) Under a global re-scaling AA⊤ 7→ sAA⊤ , we have KB 7→ s−1 KB , and thus 1 (D.15) BS 7→ BS. s Therefore, on the PCA solution manifold the batch sharpness is decreasing in global scale; it is in fact precisely homogeneous of degree −1. This property is shared by what Andreyev and Beneventano (2025) define as the modified batch sharpness, which amounts to separately averaging the numerator and denominator: EB [∇L̂⊤ 1 EB [tr(KB2 HB )] B ĤB ∇L̂B ] MBS = . (D.16) = B EB [tr(KB HB )] EB [∥∇L̂B ∥2 ] The advantage of considering this quantity is that it can be computed analytically for the Gaussian data model on which we focus. As in Appendix C, work in a basis where the data covariance Σ and the scale matrix S are jointly diagonal. Then, X∥,B and X⊥,B are independent, and thus HB is independent of KB . Then, using the fact that EB [HB ] = P ( dj=k+1 λj )IB , we find that 1 EB [tr(KB2 )] MBS = , (D.17) B EB [tr(KB )] as the overall constant of proportionality cancels. Using Isserlis’ theorem for Gaussian moments, we find that EB [tr(KB )] = Ba1 , (D.18) EB [tr(KB2 )] = Ba21 + (B + 1)a2

(D.19)

where we let a1 =

k X λi i=1

si

and

28

a2 =

k X λ2 i

s2 i=1 i

.

(D.20)

Broken scale symmetries in undercomplete linear autoencoders

This gives 1 MBS = B



a2 a1 + (B + 1) a1

 .

(D.21)

Before the edge of stability is reached, we can use the solution of the effective model to study the time evolution of the modified batch sharpness. In fact, we can show that, if we start from a balanced initial condition, d MBS < 0 (D.22) dτ along the flow. To show this, we recall from Appendix C that, assuming Gaussian data, the dynamics under the effective model are dsi si , si (0) = 1 (D.23) = σ⊥ λi 2 dτ si + 1 P assuming a balanced initial condition, where σ⊥ = dj=k+1 λj is the variance outside the principal subspace. Moreover, we recall that the explicit solution is given implicitly by s2i + 2 log si − 1 = 2σ⊥ λi τ. (D.24) Then, we see immediately that d log si > 0, (D.25) dτ so k k X X d d λi λi d a1 = =− log si < 0. (D.26) dτ dτ si si dτ i=1

i=1

What requires more care is to show that a2 /a1 is decreasing. Expanding out the derivatives, we can see that we need to show that k k X X d a2 λ2i d λi d = 2a1 log s − a log si (D.27) −a21 i 2 2 dτ a1 si dτ si dτ i=1

i=1

k X

X λi λ2j λ3i d = log si + Fij s3 dτ s s2 i=1 i i<j i j

(D.28)

is positive, where we have  defined    λi sj d λi sj d Fij = 2 −1 log si + 2 − log sj . (D.29) si λj dτ si λj dτ As suggested by our definition, our remaining task is to show that Fij > 0, as the first sum is clearly positive. Focus on some entry with i < j. We know from Appendix C that if i ≤ j then λj λi ≥ . (D.30) si sj Therefore, λi sj α= ≥ 1, (D.31) si λj in terms of which d d Fij = (2α − 1) log si + (2 − α) log sj . (D.32) dτ dτ 29

Pashakhanloo Zavatone-Veth

If 1 ≤ α ≤ 2, then both coefficients 2α − 1 and 2 − α are non-negative, and at least one is strictly positive, so clearly Fij > 0. If α > 2, then more work is required. Let d ri = log si . (D.33) dτ We claim that if i < j, then ri > 12 rj ; we will prove this below. Given this claim, for any α ≥ 1 we have 1 3 Fij = (2α − 1)ri + (2 − α)rj > (2α − 1) rj + (2 − α)rj = rj > 0. (D.34) 2 2 To prove the claim that if i < j, then ri > 12 rj , we start by letting yi = s2i , and recall from before that if i < j then yi ≥ yj ≥ 1. The dynamics of si imply that 1 ri = σ⊥ λi . (D.35) yi + 1 At the initial condition τ = 0, all of the scales yi = 1, so we have ri (τ = 0) = 12 σ⊥ λi , meaning that ri ≥ rj ≥ 0 for any i < j given the eigenvalues are ordered, hence the claim follows trivially. For τ > 0, all yj > 1, and we can use the explicit solution s2i + log s2i − 1 = 2σ⊥ λi τ to write σ⊥ λi in terms of the solution yi , which leads to 1 y + log y − 1 ri = ϕ(yi ) for ϕ(y) = . (D.36) 2τ y+1 y−1 , We are then left with a self-contained problem about the scalar function ϕ(y) = y+log y+1 1 defined for y ≥ 1: we must prove that if x ≥ y > 1, then ϕ(x) > 2 ϕ(y). By inspection, ϕ(y = 1) = 0, and moreover ϕ(y) ≥ 0 for all y ∈ [1, ∞). Differentiating, we see that 3 + y −1 − log(y) ϕ′ (y) = (D.37) (1 + y)2 has a unique zero at y⋆ = 1/W (e−3 ), the unique solution to 3 + y −1 − log(y) = 0. From this, we find that ϕ is strictly increasing on [1, y⋆ ] and strictly decreasing on [y⋆ , ∞). An −3 elementary lower bound on y⋆ is that y⋆ > e3 , which holds because ϕ′ (e3 ) = (e3e+1)2 > 0. At the unique global maximum y⋆ , we have ϕ(y⋆ ) = 1 + W (e−3 ) ≤ 1 + W (1) < 2, and from this we see that ϕ(y) < 2 for all y. This upper bound also can be directly proved by noting that it is equivalent to the statement that y + log y − 1 < 2y + 2, which in turn is equivalent to the obvious inequality log y < y + 3.

We can now divide the proof of the desired claim into two cases. If ϕ(x) ≥ ϕ(y), then we can immediately conclude that ϕ(x) > 21 ϕ(y) because we know that ϕ(y) ≥ 0 for any y ≥ 1. If instead ϕ(x) < ϕ(y) for x ≥ y > 1, then we must have x > y⋆ . Now, as noted before, we know that ϕ(y) < 2, so the claim follows if ϕ(x) ≥ 1. We can see that ϕ(x) ≥ 1 is equivalent to the inequality x + log x − 1 ≥ x + 1, which is in turn equivalent to log x ≥ 2, or x ≥ e2 . But, we know that x > y⋆ , and from before that y⋆ > e3 , which is greater than e2 . As a result, we conclude that x > e2 and thus that ϕ(x) ≥ 1, and thus the desired result. This concludes the proof.

Appendix E. Nonlinear autoencoders As we focus on linear autoencoders, a natural question is whether any of the phenomenology generalizes to nonlinear models. 30

Broken scale symmetries in undercomplete linear autoencoders

E.1. Extending the PCA manifold to homogeneous nonlinear models Consider a general nonlinear autoencoder fθ,φ (x) = gθ (hφ (x)), (E.1) d k k d where hφ : R → R and gθ : R → R are the encoder and decoder, with parameter vectors φ and θ, respectively. We will usually suppress the subscripts, and write simply f , g, and h. Suppose that the model has a scale symmetry, in the sense that its output fθ,φ (x) is invariant for any example x under transformations of the form (θ, φ) 7→ (aθ, a−1 φ), a > 0. (E.2) 1 2 2 If we train the model to minimize an L reconstruction loss L = 2 Ex [∥f (x) − x∥ ] via gradient flow ∂t θ = −∇θ L, ∂t φ = −∇φ L, then it is well-known that the scale symmetry (E.2) means that the corresponding Noether charge q = ∥θ∥2 − ∥φ∥2 (E.3) is conserved, i.e., ∂t q = 0 (Du et al., 2018; Kunin et al., 2021; Tanaka and Kunin, 2021). To see this, note that our assumption that f (x) = g(h(x)) is invariant under transformations (θ, φ) 7→ (aθ, a−1 φ) for a > 0 implies that the per-example loss ℓ(x) = 21 ∥f (x) − x∥2 is invariant. This implies that d 0= ℓ(x) = ⟨∇θ ℓ(x), θ⟩ − ⟨∇φ ℓ(x), φ⟩, (E.4) da a=1 which in turn implies that ∂t q = 2(⟨θ, ∂t θ⟩ − ⟨φ, ∂t φ⟩) (E.5) = −2E[⟨θ, ∇θ ℓ(x)⟩ − ⟨φ, ∇φ ℓ(x)⟩]

(E.6)

= 0.

(E.7)

Under SGD updates θt+1 = θt − η∇θ L̂t , φt+1 = φt − η∇φ L̂t for L̂t = Ex∼Bt [ℓ(x)], we have (E.8) qt+1 = qt − 2ηEx∼Bt [⟨θ, ∇θ ℓ(x)⟩ − ⟨φ, ∇φ ℓ(x)] + η 2 (∥∇θ L̂t ∥2 − ∥∇φ L̂t ∥2 ) = qt + η 2 (∥∇θ L̂t ∥2 − ∥∇φ L̂t ∥2 ), (E.9) where the order-η term vanishes identically by (E.4). We now observe that we can partially extend the PCA solution geometry to these more general nonlinear models. Consider a particular example x. Let its latent encoding be z = h(x), and the residual be r = g(z) − x, such that the reconstruction error is ℓ(x) = 12 ∥r∥2 . By the chain rule, the parameter gradients appearing in the update (E.9) for qt are ∇θ ℓ = (∂g/∂θ)⊤ r and ∇φ ℓ = (∂h/∂φ)⊤ (∂g/∂z)⊤ r. For an overcomplete autoencoder, perfect reconstruction is possible, leading to r = 0 and qt+1 = qt . In contrast, for an undercomplete autoencoder with k < d, it is not generally possible to make the residual vanish for an arbitrary example. Geometrically, we can then decompose the residual r into components normal and tangential to the decoder manifold D = {g(z) : z ∈ im h}, which we write as r⊥ and r∥ , respectively. By definition, (∂g/∂z)⊤ r⊥ = 0. Thus, the encoder gradient ∇φ ℓ is sensitive only to r∥ , while the decoder gradient ∇θ ℓ is sensitive to both r∥ and r⊥ . Assuming that k ≪ d, it is possible for the undercomplete autoencoder to make the k-dimensional tangential component r∥ quite small. As a result, as the decoder gradient has an additional contribution from r⊥ relative to the encoder gradient, we expect undercomplete autoen31

Pashakhanloo Zavatone-Veth

coders to generically favor decoder-heavy solutions so long as ∥r∥ ∥ ≪ ∥r⊥ ∥. In the extreme, we can consider what we might refer to as projection-type solutions that have r∥ = 0; these are in some sense an analog of the PCA solutions in the linear case. On the manifold of such solutions, the imbalance dynamics reduce to qt+1 = qt + η 2 ∥∇θ L̂t ∥2 , (E.10) that is, there is a systematic drift towards larger decoder scales. E.2. Per-neuron scale symmetries in ReLU autoencoders In the section above we considered the case in which the GL(k) symmetry of a linear autoencoder is reduced all the way down to a single R>0 . However, for a standard ReLU autoencoder f (x) = W σ(U x), (E.11) where σ(t) = max(0, t) acts elementwise, has a larger symmetry group defined by re-scaling each hidden neuron, in addition to permutation symmetries. The continuous component of the symmetry group is simply the neuron-wise rescalings Rk>0 , which correspond to (W, U ) 7→ (W A, A−1 U ) for A = diag(a1 , . . . , ak ) a diagonal matrix with positive diagonal elements ai > 0 (Du et al., 2018; Tanaka and Kunin, 2021). A more tractable expression follows by letting wi be the i-th column of W , and u⊤ i be the i-th row of U . Then, k X f (x) = wi σ(u⊤ (E.12) i x) i=1

and the action of the symmetry group is that, for each i, (wi , ui ) 7→ (ai wi , a−1 i ui ). This formulation also makes it obvious that there are k Noether charges corresponding to this symmetry: qi = ∥wi ∥2 − ∥ui ∥2 . (E.13) The P charge related to global rescaling is then simply the sum of these per-neuron charges: q = ki=1 qi . The proof that these charges are individually conserved under gradient flow parallels the proof for homogeneous rescalings considered above: for any i ∈ [k], ∂ ℓ(x) = ⟨∇wi ℓ(x), wi ⟩ − ⟨∇ui ℓ(x), ui ⟩, (E.14) 0= ∂ai A=Ik so ∂t qi = 2(⟨wi , ∂t wi ⟩ − ⟨ui , ∂t ui ⟩) = 2E[⟨∇wi ℓ(x), wi ⟩ − ⟨∇ui ℓ(x), ui ⟩] = 0. (E.15) Under SGD, we have the general update qt+1,i − qt,i = η 2 (∥∇wi L̂t ∥2 − ∥∇ui L̂t ∥2 ). (E.16) We show in Figure S2 that, empirically, SGD training leads to decoder-heavy imbalance for these networks.

Appendix F. Experimental details The simulations were performed with two-layer autoencoders of the form x̂ = W ϕ(U x), with zero bias, input dimension n = 20, and either a contractive hidden layer k = 10 < n or an expansive hidden layer k = 40 > n. The data were generated online from a diagonal Gaussian distribution, x ∼ N (0, Σ), where Σ = diag(λ1 , . . . , λn ) is the covariance eigen32

Broken scale symmetries in undercomplete linear autoencoders

i−1

values were geometrically spaced from λ1 = 12 to λn = 0.18 (concretely, λi = λ1 ( λλn1 ) n−1 ). Training used minibatch SGD with batch size B, learning rate η, mean-squared reconstruction loss, and fresh Gaussian samples at every update. Unless specified, all networks had linear activation ϕ(zi ) = zi . The specific parameters for each plot are as follows. • Figure 1A/Figure S1: n = 20, k = 10 (contractive) and k = 40 (expansive), B = 48, η = 0.003. • Figure S2: n = 20, k = 10 (contractive) and k = 40 (expansive), B = 48, η = 0.003. ϕ(zi ) = max(0, zi ) (ReLU activation). • Figure 1B: n = 20, k = 10, variable B and η, as shown on the plot. • Figure 1C: n = 20, k = 10, B = 32, η = 0.005. • Figure 1D-E: n = 20, k = 10, B = 128, η = 0.02. • Figure 1F: n = 20, k = 10, B = 128, η ∈ {0.014, 0.016, . . . , 0.026}. Networks in Figure 1A/ Figure S1 and Figure S2 were initialized with i.i.d. zero-mean Gaussian weights with standard deviation 0.01, while the contractive networks in the other plots were initialized at the PCA solution with the identity scale. Specifically, for those networks,    
Ik W0 = , U0 = Ik 0k×(n−k) , W0 U0 = diag Ik , 0(n−k)×(n−k) , 0(n−k)×k resulting in an initial scale matrix S0 = W0⊤ W0 = Ik .

33

Pashakhanloo Zavatone-Veth

Figure S1: Continual minibatch training drives scale imbalance in linear autoencoders. (A) Validation loss for contractive and expansive linear autoencoders during online training with random initialization at the onset. (B) Scale imbalance is measured by q = ∥W ∥2F − ∥U ∥2F . (C,D) Frobenius norms of decoder and encoder weights, ∥W ∥F and ∥U ∥F , for contractive and expansive autoencoders, respectively (see §F for experimental details).

34

Broken scale symmetries in undercomplete linear autoencoders

Figure S2: Continual minibatch training drives scale imbalance in homogeneous autoencoders. (A) Validation loss for contractive and expansive ReLU autoencoders during online training with random initialization at the onset. (B) Scale imbalance is measured by q as defined in Eq. (E.3). In this case we have a decoder weight matrix W and an encoder weight matrix U , so q = ∥W ∥2F − ∥U ∥2F . (C,D) Frobenius norms of decoder and encoder weights, ∥W ∥F and ∥U ∥F , for expansive and contractive ReLU autoencoders, respectively (see §F for experimental details).

35

Related documents

• Broken scale symmetries in undercomplete linear autoencoders#1188351
• Scale-invariant Gaussian derivative residual networks#1104974
• Path Regularization: A Near-Complete and Optimal Nonasymptotic Generalization Theory for Multilayer Neural Networks and Double Descent Phenomenon#1062760
• Look But Don't Touch with Sparse Autoencoders for Unlearning in Diffusion Models#324950
• Stiefel-AdamW: Geometry-Aware AdamW for Linear Factorization Blocks#1035009