Broken scale symmetries in undercomplete linear autoencoders

AI summary

Overview

• Research area: Deep learning theory — specifically the geometry of neural network loss landscapes, continuous parameter symmetries, and the implicit bias of stochastic gradient descent (SGD).
• Technical level: Advanced. The paper uses matrix calculus, Hessian spectral analysis, Noether charges, the Lambert W function, and stability theory for discrete dynamical systems.
• Scope: A solvable toy model — the undercomplete linear autoencoder — used to determine the direction, timescale, and endpoint of SGD drift along a continuous symmetry manifold.

What This Paper Is About

Neural networks contain symmetries: transformations of parameters that leave the network's input-output function unchanged. Gradient flow respects these symmetries exactly, but finite-stepsize SGD does not. The paper asks what selects the direction of the resulting motion along such a symmetry orbit, and where that motion stops.

The authors answer this in a linear autoencoder with encoder U ∈ ℝ^{k×d}, decoder W ∈ ℝ^{d×k}, output f(x) = WUx, trained with an L2 reconstruction loss in the undercomplete regime k < d. In this setting the loss minima form a manifold of functionally equivalent PCA solutions related by the GL(k) symmetry (W, U) ↦ (WA, A^{-1}U).

Key Contributions

1. A geometric account of directed symmetry breaking. The authors show that on the PCA solution manifold of an undercomplete linear autoencoder, the minibatch gradient with respect to the encoder is identically zero (∇_U L̂_t = 0), while the decoder gradient is not. Substituting into the update rule for the Noether charge Q = W^T W − U U^T gives Q_{t+1} − Q_t = η² (∇_W L̂_t)^T(∇_W L̂_t) ⪰ 0, forcing drift in a preferred direction within the cone of positive-semidefinite matrices.

2. An analytically tractable effective theory of scale drift. Exploiting a separation of timescales between fast relaxation onto the solution manifold (T_fast ∼ 1/η) and slow drift along it (T_slow ∼ B/η²), the authors derive a matrix flow for the scale degrees of freedom S = A^T A and solve it in closed form using the principal branch of the Lambert W function.

3. Prediction of an anisotropic, finite endpoint set by a stability boundary. Scale growth cannot continue indefinitely; the authors identify the "dangerous" Hessian modes μ_{1a}^∥ = λ_1(s_1^{-1} + s_a) and predict that all scales saturate at a common value s_a = κ/(η λ_1) + O(η), with κ ≃ 2 in simulations.

4. A demonstration that different sharpness measures disagree. The authors show that while the maximum Hessian eigenvalue increases with decoder-heavy scale imbalance, the Hessian trace is minimized at a decoder-heavy scale s_*² = ((1/k)Σ_{j=1}^k λ_j)/((1/d)Σ_{j=1}^d λ_j) ≥ 1, and the "batch sharpness" measure of Andreyev and Beneventano (2025) is decreasing in global scale on the PCA manifold.

Main Findings

• Decoder-heavy drift is geometrically forced. Because I_d − WU = I_d − P^T P ≡ P_⊥ projects onto the complement of the principal subspace at any PCA solution, the residual r = (WU − I_d)x lies entirely outside the principal subspace. This makes ∇_U L̂_t vanish identically and ∇_W L̂_t = −P_⊥ Σ̂_t P^T A_t^{-T} nonzero, so the charge Q can only grow. On the scalar charge q = tr(Q) = ‖W‖_F² − ‖U‖_F², associated with global rescalings (W,U) ↦ (aW, a^{-1}U) for a > 0, this is a directional bias toward decoder-heavy weights.

• No noise-balanced equilibrium is available. The paper states that in this setting no equilibrium where the two terms in the Q update cancel on average — the mechanism proposed by Ziyin et al. (2025a) — is possible. In contrast, an expansive autoencoder can reconstruct perfectly, so there is no fundamental geometric reason for directed drift there.

• Scale drift is anisotropic and follows a closed-form flow. With S = A^T A and C = E[‖x_⊥‖² x_∥ x_∥^T] defined via x_∥ = Px and x_⊥ = (I_d − P^T P)x, the effective flow ∂_τ S + S^{-1}(∂_τ S)S^{-1} = S^{-1/2} C S^{-1/2} solves analytically when S(τ=0) and C commute, giving s_a(τ) = √(W[s_a(0)² exp(s_a(0)² + 2c_a τ)]). For Gaussian data the growth rates are c_a = λ_a Σ_{j=k+1}^{d} λ_j — the variance along direction a multiplied by the total variance outside the principal subspace.

• Drift couples within- and out-of-subspace fluctuations. The authors note this mirrors Pashakhanloo (2025)'s finding that the same coupling governs drift rates along rotation-type symmetry directions.

• The endpoint is a stability boundary, akin to the edge of stability. Assuming ordered scales s_1 > s_2 > ⋯ > s_k, the population-loss Hessian has dangerous modes μ_{1a}^∥ = λ_1(s_1^{-1} + s_a) for a ∈ [k]. As scales grow, curvature along each mode rises until η μ_{1a}^∥ = κ, at which point the discrete dynamics destabilize. The prediction s_a = κ/(η λ_1) + O(η) with κ ≃ 2 matches simulations (Figure 1 D–F), and the authors report sequential saturation of eigenvalues and saturation scaling with the inverse learning rate.

• Sharpness measures can move in opposite directions. The Hessian trace tr(H) = (d/s)Σ_{j=1}^k λ_j + ks Σ_{j=1}^d λ_j is minimized at s_*² = ((1/k)Σ_{j=1}^k λ_j)/((1/d)Σ_{j=1}^d λ_j) ≥ 1, so decoder-heavy solutions can be flatter by this measure than balanced or encoder-heavy ones. The paper also shows the "batch sharpness" measure of Andreyev and Beneventano (2025) decreases with global scale on the PCA manifold, and proves that for Gaussian data with balanced initialization their "modified batch sharpness" measure strictly decreases under the effective dynamics even as the population Hessian sharpens.

• Preliminary extension to ReLU autoencoders. Nonlinear homogeneous activations break the GL(k) symmetry down to a neuron-wise scaling symmetry with A diagonal and positive. The authors report preliminary evidence that undercomplete ReLU autoencoders show similar directed drift toward decoder-heavy weights (Fig. S2), but leave deep investigation to future work.

• Progressive sharpening without functional learning. The paper emphasizes that the edge of stability here is reached purely through slow drift along a manifold of functionally equivalent solutions, after the PCA manifold is already reached — so progressive sharpening does not necessarily reflect continued learning of the function.

Methodology in Plain English

The authors deliberately choose the simplest model that still has a continuous parameter symmetry: a linear autoencoder with a bottleneck, trained on a reconstruction loss. Because the model is linear and the loss is quadratic, the loss landscape can be analyzed exactly with linear algebra rather than approximated numerically.

The argument proceeds in three steps. First, they identify the exact set of global minima — the PCA solution manifold, parameterized by an invertible matrix A encoding the symmetry degrees of freedom. Second, they evaluate the minibatch gradients at that manifold and find a structural asymmetry: the encoder gradient vanishes identically, so the balance between incoming and outgoing weights can only move one way under SGD. Third, they separate the timescales of fast relaxation back onto the manifold and slow drift along it, derive an effective equation for the scale matrix S, and solve it in closed form.

They then compute the population Hessian and find the specific directions whose curvature grows fastest as scale increases. Combining the growth law with the finite-stepsize stability condition gives a predicted saturation scale, which they compare to simulations. Finally, they check alternative definitions of sharpness — the Hessian trace, and two batch-sharpness measures proposed elsewhere — to show that the sign of the effect depends on the measure.

The paper also reports simulation figures and cites appendices for derivations: Appendix A for linear autoencoder basics, Appendix B for the Hessian spectrum, Appendix C for the effective dynamics and stability analysis, Appendix D for the batch sharpness comparison, Appendix E for the nonlinear/homogeneous case, and Appendix F and Figure S1 for additional training details.

Why This Matters

• Impact on research. The paper gives a concrete counterexample to the idea that symmetry breaking under SGD is generically explained by noise balancing or a statistical equilibrium. Instead, the endpoint of drift is selected dynamically, by the geometry of the loss rising until the solution becomes unstable. It connects three otherwise separate literatures: parameter symmetries and conserved quantities, the implicit bias of SGD, and progressive sharpening / the edge of stability. It also reinforces the parameterization-dependence of sharpness measures (as argued by Dinh et al., 2017), showing that "sharpness" is not a single well-defined quantity.

• Real-world applications (of the underlying phenomena):

  • Training of autoencoders and other bottleneck representation-learning models, where the relative scale of encoder and decoder weights affects numerics, quantization, and regularization.
  • Weight-balancing and reparameterization tricks used in practice (for example, balancing layers to stabilize training or to control effective learning rates).
  • Model compression and quantization, where scale imbalance between layers determines which weights dominate the representation and how gracefully precision can be reduced.
  • Pruning and sparsity analyses of autoencoder-style architectures, where the geometry of functionally equivalent solutions determines what can be removed without changing outputs.

• Industry relevance. Large-scale training pipelines routinely rely on heuristics about sharpness and stability (learning-rate schedules, warmup, batch-size scaling). This paper suggests that some sharpening observed during training may be functionally inert — motion along a manifold of equivalent solutions rather than genuine learning — which matters for anyone interpreting loss curves or sharpness diagnostics as evidence of model progress.

Future Directions

• Extending beyond linear models. The authors show only preliminary evidence (Fig. S2) that undercomplete ReLU autoencoders exhibit the same decoder-heavy drift, and explicitly leave the geometric basis of this bias "to future work."
• Defining the edge of stability in the stochastic setting. The paper notes this remains an active area of research, citing Andreyev and Beneventano (2025) and Emmanouilidis et al. (2026), and its own analysis sits at the interface of this question.
• Transferring the phenomenology to other optimization settings. The authors pose as an open question whether progressive sharpening driven purely by symmetry-manifold drift occurs in optimizers and architectures beyond this toy model.
• Reconciling conflicting sharpness measures. Since the Hessian trace, the maximum Hessian eigenvalue, and batch sharpness move in different directions on the same solution manifold, a principled account of which measure matters for generalization is unresolved here.

Target Audience

This paper is best suited to machine learning theorists and graduate students working on optimization dynamics, loss-landscape geometry, and the implicit regularization of SGD. Readers interested in parameter symmetries, Noether charges, and the edge-of-stability literature will find the most value. Practitioners who tune learning rates, batch sizes, and weight-balancing schemes may benefit from the qualitative picture — that scale imbalance is a directed, geometrically forced process with a stability-determined endpoint — but the paper is written at an advanced mathematical level and assumes comfort with matrix differential equations and Hessian spectral analysis.

Authors’ 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.