Learning the identity: a case study of how SGD selects among functional decompositions

Andy Arditi; Weian Xie; David Bau; Liu Ziyin

cs.LGSep 30, 2026

Learning the identity: a case study of how SGD selects among functional decompositions

Authors: Andy Arditi, Weian Xie, David Bau, Liu Ziyin

Organizations: Northeastern University · Massachusetts Institute of Technology

Abstract

One might think that learning the identity function with a deep linear residual network is trivial - the path along residual connections already implements the identity, and so the network need only drive its weights to zero. However, this zero-weight solution is just one point on an entire manifold of population-loss minimizers, each corresponding to a different decomposition of the identity across the network's layers. Although the population loss does not distinguish among these solutions, stochastic gradient descent (SGD) reproducibly favors particular ones. For instance, under anisotropic label noise, the learned layers exhibit a noise-dependent spectrum; even with weight decay, SGD does not generally recover the zero-weight solution. Changing only the parametrization, while leaving the set of realizable functions unchanged, yields different behavior: factoring each weight matrix as a product of two matrices causes the weights to collapse to zero, even without explicit weight decay. While perhaps mysterious and unintuitive at first, these phenomena can be understood through the lens of entropic loss, which augments the population loss with a term proportional to the expected squared norm of the minibatch gradient (Ziyin et al., 2025). On the identity manifold, the population loss is constant, while the entropic term distinguishes among these decompositions. We characterize its minimizers analytically and use them to derive predictions for the structure of solutions favored by SGD. Networks trained with SGD closely match these predictions. Overall, the identity learning task studied here serves as a clean and simple case study of how the lens of entropic loss can clarify why SGD favors particular decompositions of the same input-output function.

Figures & tables

Figure 1: The one-dimensional case. SGD quickly reaches the identity manifold

\widetilde{w}_{2}\widetilde{w}_{1}=1

, where the population loss is minimized. It then drifts along the manifold toward one of the two factorizations,

\widetilde{w}_{1}=\widetilde{w}_{2}=1

\widetilde{w}_{1}=\widetilde{w}_{2}=-1

, that minimize the expected squared gradient norm. Left: four SGD trajectories through parameter space, colored by training step. Population loss is indicated by shading, and the identity manifold is represented by the dotted hyperbola. Right: population loss

\mathcal{L}

and the expected squared gradient norm

\mathbb{E}\|\nabla\ell\|^{2}

as functions of training step; dotted lines mark their respective minima on the identity manifold. All runs use step size

\eta=0.01

and batch size

B=1

Figure 2: Isotropic noise selects orthogonal factorizations, while weight decay favors the identity. (a) Frobenius distance of the learned effective layer

\widetilde{W}_{1}

to its nearest orthogonal matrix and to the identity matrix. Each point represents one of 16 seeds; horizontal bars show medians. With or without weight decay, the learned layers are nearly orthogonal, while weight decay (

\gamma=10^{-2}

) additionally brings them close to the identity. (b) Eigenvalues of the same learned matrices in the complex plane (dashed: unit circle), with all 16 eigenvalues from each of 16 seeds overlaid in each subplot. Without weight decay, the eigenvalues spread along an arc near the unit circle; with weight decay, they concentrate near

+1

. Here

d=16

\eta=0.02

B=64

, and

(W_{\ell})_{ij}\sim\mathcal{N}(0,1/d)

. We train for

2\times 10^{6}

steps and average the weights over the final

1\%

Figure 3: Anisotropic noise induces a fourth-root singular-value spectrum, while weight decay pulls the singular values toward one. Left: Without weight decay, SGD closely follows the fourth-root law $s_{i}=\lambda_{i}^{1/4}$ . Points show mean learned singular values over five runs. Right: With weight decay, the singular-value spectrum is compressed toward one, closely matching theoretical predictions. Points show mean learned singular values over five runs, and lines show the corresponding predictions of Equation 13 ; colors indicate the weight-decay strength $\gamma$ . The predictions retain the correction $\gamma\mapsto\gamma(1+\eta\gamma)$ described in Appendix C.3 . In both panels, singular values are sorted within each run and paired with the ordered noise variances; error bars show one standard deviation and are smaller than the markers. Here $d=16$ , $\eta=0.005$ , and $B=16$ . Singular values are computed from weights averaged over the final $1\%$ of training.

Figure 4: In deeper networks, noise anisotropy is confined to the first and last layers. Singular values of the three effective layers in a depth

D=3

network, plotted against the noise variances. The first layer follows a fourth-root law, the middle layer has a flat spectrum, and the last layer follows an inverse fourth-root law. Points show mean singular values over six runs; error bars show one standard deviation. Dashed lines show the corresponding predictions

c^{-1/2}\lambda_{i}^{1/4}

c

, and

c^{-1/2}\lambda_{i}^{-1/4}

. Within each run, the first- and middle-layer singular values are sorted in ascending order, and the last-layer singular values in descending order, before pairing them with ascending noise variances. Here

d=16

\eta=0.02

B=64

\gamma=0

, and

1.5\times 10^{6}

SGD steps; weights are averaged over the final

1\%

of training.

Figure 5: Factoring the residual weights causes the residual branch to collapse to zero. We compare direct residual weights

W_{\ell}

with factored residual weights

W_{\ell}=A_{\ell}B_{\ell}

, starting from matched nonzero residual maps on the identity manifold. The y-axis shows the total residual map norm

\sum_{\ell}\|R_{\ell}\|_{F}

, where

R_{\ell}=W_{\ell}

for the direct parametrization and

R_{\ell}=A_{\ell}B_{\ell}

for the factored parametrization. Gray curves show direct weights and blue curves show factored weights; solid and dashed lines denote isotropic and anisotropic label noise, respectively. In both noise settings, the factored residual collapses toward zero while the direct residual remains nonzero. Lines show means over six runs and shaded regions show their ranges. Here

d=16

\eta=0.02

B=64

, and

\gamma=0

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Anisotropy remains confined to the boundary layers at greater depths. Left and middle: Mean boundary singular values for

D=4

and

D=5

, respectively, divided by

a_{D}

and plotted against the noise variance. Error bars show one standard deviation across runs. Dashed lines show

\lambda_{i}^{1/4}

and

\lambda_{i}^{-1/4}

. Within each run, the first layer’s singular values are sorted in ascending order and the last layer’s in descending order, then paired with ascending noise variances. Right: Interior singular values divided by

c_{D}

; the dashed line marks one. Points show individual singular values across layers, noise directions, and runs; boxes show the median and interquartile range. Each depth uses six runs with

d=8

\eta=0.02

B=64

\gamma=0

, and

1.5\times 10^{6}

SGD steps; singular values are computed from weights averaged over the final

1\%

of training.

parametrization	activation	noise	relative branch RMS	relative function error
direct	linear	iso	$0.9146$	$0.0080$
direct	ReLU	iso	$0.8677$	$0.0087$
direct	linear	aniso	$1.2027$	$0.0080$
direct	ReLU	aniso	$0.8049$	$0.0093$
factored	linear	iso	$0.0016$	$0.0014$
factored	ReLU	iso	$0.0028$	$0.0022$

Appendix

Table 1: Residual branches after training from matched matrix products. The branch RMS and function error columns report

\mathcal{R}

and

\mathcal{E}

, respectively. Entries are means over six runs, using weights averaged over the final

1\%

8\times 10^{5}

SGD steps. Each run is evaluated on 8192 Gaussian inputs independent of training. All runs use

d=16

\eta=0.02

B=64

, and

\gamma=0

Figure 7: Functional measurements of ReLU branch collapse. Top panels show relative branch RMS

\mathcal{R}

; bottom panels show clean function error

\mathcal{E}

. Columns correspond to isotropic and anisotropic label noise. Gray curves show the direct model and blue curves the factored model. Lines are means across six runs and shaded bands show their ranges. At each plotted training step, we evaluate the current weights on 2048 Gaussian inputs per run, independent of training, without averaging weights across steps. Training uses

d=16

\eta=0.02

B=64

\gamma=0

, and

8\times 10^{5}

steps.

Figure 8: Covariance orientation and the accuracy of the entropic prediction. Each panel fixes the covariance eigenvalues and varies their relative orientation. Blue circles and gray triangles show the commuting controls (

0^{\circ}

and

90^{\circ}

); orange squares show the noncommuting case (

45^{\circ}

). Points show

\|\overline{G}-G^{\star}\|_{F}/\|G^{\star}\|_{F}

, averaged across eight runs, with one standard error; lines connect measurements. Here

\overline{G}

averages

G=\widetilde{W}_{1}^{\top}\widetilde{W}_{1}

sampled at regular intervals over the final quarter of each run, and

G^{\star}

is computed for each covariance pair. All cases start at

\widetilde{W}_{1}=\widetilde{W}_{2}=I

. Training uses batch size

B=8

, and no weight decay. Learning rates

0.02

0.01

, and

0.005

use

400{,}000

1{,}600{,}000

, and

6{,}400{,}000

steps, respectively, keeping

\eta^{2}t/B=20

at the endpoint.

Figure 9: Eigenvalue distribution at initialization. Each panel shows one draw of

\widetilde{W}_{1}=I+(s/\sqrt{d})G

d=400

, where the entries of

G

are independent standard Gaussians. Dashed circles have center

1

and radius

s

; the dotted vertical line marks zero real part. Blue points are nonreal eigenvalues and orange points are real eigenvalues. The disk reaches the origin at

s=1

and extends into the left half-plane for

s>1

Figure 10: Learned eigenvalues across initialization scales. Each panel overlays the eigenvalues of

\widetilde{W}_{1}

from

16

runs with

d=16

and isotropic noise. The dashed curve is the unit circle. Initial entries of each

W_{\ell}

are independent

\mathcal{N}(0,s^{2}/d)

, with

s\in\{\sqrt{0.1},1,\sqrt{2},2\}

from left to right; the second column uses the initialization of Figure 2 (b). The rows share initial weights and use

\gamma=0

(top) and

\gamma=10^{-2}

(bottom). We train for

2\times 10^{6}

steps with

\eta=0.02

B=64

, then average weights over the final

1\%

of training.

Explore similar work

Sep 30, 2026cs.LG

Exact information accounting for SGD methods

As an alternative to the standard geometric analyses, we give an exact, information-theoretic analysis of stochastic gradient descent (SGD) and its variants. We show that a preconditioned SGD step is the posterior-mean update of a Gaussian Bayes model, and that its one-step regret splits into an intrinsic-time cost and a change in comparator information. The split extends to an identity for the objective itself. Convex convergence, strict-saddle-point escape, the link between flatness and generalization, the standard learning-rate schedules, adaptive optimizers, and the noisy, momentum, heavy-tailed, and gradient-free variants of SGD each correspond to a term or a special case of this identity. We measure its terms on synthetic and real training runs. On real networks it attributes the slack of classical convergence bounds to the terms their derivations drop and separates optimizers that reach the same training loss. That separation follows the number and consistency of their steps. Its relation to which of them generalizes better differs between networks. For gradient-free SGD the identity determines how a curvature preconditioner should enter the update. The sharpness-based generalization certificate it yields, with a data-independent isotropic prior, is vacuous at network scale unless the curvature spectrum is nearly flat across all parameters.

Akshay Balsubramani

Stochastic Gradient Descent Gradient-Based Optimization Generalization Bounds Information-Theoretic Limits

Ask PDF

Jul 3, 2026cs.LG

Implicit Bias of SGD in Multivariate ReLU Networks: Effective Width Collapse

We study the implicit bias of noisy stochastic gradient descent in training wide two-layer ReLU networks for multivariate regression. In a mean-field regime, the training dynamics are approximated by a Wasserstein gradient flow that converges to a unique stationary measure. We characterize the structure of this stationary measure and the predictor it represents. We show that, despite the network being infinitely overparameterized, the learned predictor admits an effectively finite representation: the input weights and biases align along finitely many directions, leading to an effective width collapse. In particular, the solution function is continuous piecewise affine, with affine regions determined by the cells of a finite hyperplane arrangement. The number of learned directions, and hence hyperplanes, is bounded above by

2\mathcal{P}-1

, where

\mathcal{P}

denotes the number of linear dichotomies realizable on the training inputs. We further establish a non-redundancy property of the learned representation by proving that each learned direction induces a unique ternary activation pattern on the training data. Consequently, the complexity of the learned predictor is governed by the combinatorial geometry of the training data.

Shuang Liang, Tom Jacobs, Guido Montúfar

Department of Statistics & Data Science University of California, Los Angeles Los Angeles, CA 90095, USA · CISPA Helmholtz Center for Information Security · Departments of Mathematics and Statistics & Data Science University of California, Los Angeles; and Max Planck Institute for Mathematics in the Sciences

Rectified Linear Unit Networks Stochastic Gradient Descent Gradient Mean-Field Limit Representation Collapse

Ask PDF

Sep 30, 2026cs.LG

Awakening of the Buddha: Subspace Learning During Population-Loss Plateaus

Population loss can remain nearly constant while a neural network learns a substantially more predictive representation. We establish this separation for two-layer ReLU and leaky-ReLU networks trained on Gaussian inputs by simultaneous fixed-step population gradient descent on all parameters. For structured additive teachers whose links are positive mixtures of Gaussian-damped cubics in

H^1(γ)

, we give explicit conditions under which small IID Gaussian initialization yields a high-probability guarantee: at a checkpoint during a high-loss plateau, minimum alignment between the rank-

r

teacher subspace and the leading

r

-dimensional eigenspace of the predictor's average gradient outer product (AGOP) increases by at least

1/2

, and the minimum refit MSE under unchanged coefficient budgets decreases by more than

0.399

, both relative to initialization. The same trajectory subsequently attains a trained loss below every value in the plateau window. A complementary result treats unequal-weight cubic teachers and small additive Sobolev perturbations using projected-feature refits. For SwiGLU networks with an exactly fitted intercept, we prove leading-AGOP alignment during a loss plateau at fixed width and dimension as Gaussian initialization vanishes, for square-integrable teachers with nonzero Hermite content of degree one, two, or three. A rank-one cubic specialization also gives simultaneous unrestricted-refit gains at a prescribed width. An approximation lower bound further shows that certain interaction targets retain nonzero error when ridge neurons are restricted to shared orthogonal axes within the teacher subspace. Population-moment experiments with ReLU students across 21 teachers and 50 initializations per teacher complement the analysis.

Akash Kumar

Department of Computer Science & Engineering University of California-San Diego

Anisotropic Loss Landscapes Singular Learning Theory Subspace Teacher Population Dynamics Bodhi

Ask PDF