Source-linked AI summary
Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup
Sebastian Goldt, Madhu S. Advani, Andrew M. Saxe, Florent Krzakala, Lenka Zdeborová
TL;DR
The paper studies how SGD generalises in over-parameterised two-layer teacher-student networks. It derives and validates differential-equation dynamics, then uses their fixed points to analyse how activation functions, trainable layers, and network size affect asymptotic generalisation.
Problem
The paper examines how over-parameterised student networks generalise relative to teacher networks, including when teacher outputs contain Gaussian noise.
Method
The authors derive ordinary differential equations for order-parameter dynamics, prove their asymptotic validity, and analyse fixed points numerically to predict generalisation error.
Results
Training only the first layer can make generalisation error increase with student size, whereas training both layers yields better performance and can avoid this growth depending on the activation function.
Takeaways & Limitations
Generalisation depends not only on SGD but also on activation functions, parameter arrangement, and which weights are trainable.
Takeaways & Limitations
For large students, SGD may converge slowly or remain on plateaus and fail to reach the lowest-generalisation-error solution in finite time.
Abstract
from arXiv · showhide
Deep neural networks achieve stellar generalisation even when they have enough parameters to easily fit all their training data. We study this phenomenon by analysing the dynamics and the performance of over-parameterised two-layer neural networks in the teacher-student setup, where one network, the student, is trained on data generated by another network, called the teacher. We show how the dynamics of stochastic gradient descent (SGD) is captured by a set of differential equations and prove that this description is asymptotically exact in the limit of large inputs. Using this framework, we calculate the final generalisation error of student networks that have more parameters than their teachers. We find that the final generalisation error of the student increases with network size when training only the first layer, but stays constant or even decreases with size when training both layers. We show that these different behaviours have their root in the different solutions SGD finds for different activation functions. Our results indicate that achieving good generalisation in neural networks goes beyond the properties of SGD alone and depends on the interplay of at least the algorithm, the model architecture, and the data set.
1 Online learning in teacher-student neural networks
The paper models online SGD in over-parameterised two-layer teacher-student networks through macroscopic order parameters and coupled ODEs. Under stated assumptions, the ODE description is asymptotically exact and agrees closely with SGD experiments.
- Setup: The teacher generates noisy regression labels from a two-layer network with M hidden units, while the student has K ≥ M hidden units and trains by SGD.The student is therefore more expressive than the generative teacher function.
- Setup: The student’s generalisation error is the input-averaged squared difference between teacher and student outputs.The analysis expresses this error through macroscopic order parameters, including student-student and teacher-student overlaps.
- ODE description: The order parameters obey a closed set of coupled ODEs that describes the generalisation dynamics for any K ≥ M.The ODE system extends earlier analyses beyond the matched-width case K = M.
- ODE description: Theorem 1.1 establishes that the macroscopic SGD state converges to the unique ODE solution with an error controlled by a constant independent of N.The result holds under the paper’s assumptions on the data, activation function, and initial macroscopic state.
- Validation: Numerically integrated ODE dynamics and a single SGD run show very good agreement for sigmoidal students with varying hidden-layer size K.Figure 1 compares the generalisation error trajectories and teacher-student overlaps.
- Implication: The fixed points found by the ODEs yield different test-error dependence on over-parameterisation across activation functions and training setups.This motivates analysing asymptotic performance rather than SGD dynamics alone.
2 Asymptotic generalisation error of Soft Committee machines
For Soft Committee Machines trained only in the first layer, asymptotic generalisation error worsens as over-parameterisation grows. The analysis attributes this to redundant student units learning output-noise fluctuations, with related scaling also observed for linear and ReLU networks.
- Analytical reduction: The reduced dynamics contain eight coupled differential equations for arbitrary K and M, enabling closed-form asymptotic generalisation-error expressions.The analysis studies the fixed points reached after convergence, with L = K − M denoting excess student units.
- Noise and fixed points: In the absence of output noise, a student with K ≥ M can achieve zero asymptotic generalisation error as α → ∞.With small noise, the zero-error fixed point becomes unstable and SGD approaches a nearby positive-error fixed point.
- Sigmoidal networks: The asymptotic generalisation error of Soft Committee Machines increases linearly with excess width L when only the first layer is trained.The result is supported by theoretical predictions and experiments across learning rates, noise levels, and L.
- Robustness: Reducing the learning rate does not remove the width-dependent error increase before a plateau and instead requires longer training and more online-learning data.This trade-off persists even when choosing η = 1/K.
- Mechanism: The extra L student units specialise toward nonexistent teacher units, acquire output-noise fluctuations, and increase the final error.These units effectively learn from teacher hidden units with zero weights.
- Other activations: Linear and ReLU students also show increasing asymptotic error with added hidden units, despite differing internal representations.Linear networks add redundant parameters, while ReLU networks distribute finite overlaps across student and teacher nodes.
3 Training both layers: Asymptotic generalisation error of a neural network
Training both layers changes how over-parameterisation affects generalisation: sigmoidal students improve as K grows, whereas linear and ReLU students remain approximately constant. The improvement in sigmoidal networks is associated with a denoising solution, but SGD may not reach the best fixed point from random initialization.
- Sigmoidal networks: Increasing K decreases the asymptotic generalisation error of sigmoidal students trained on both layers.The observed solution specialises multiple student units to each teacher unit, with outgoing weights approximately summing to the teacher weight.
- Sigmoidal networks: The sigmoidal denoising solution averages several student estimates of each teacher node through specialised hidden units.For K = 5, two student units specialise to one teacher node and three to the other, with corresponding outgoing weights approximately summing to the teacher weights.
- Optimization and fixed points: SGD does not always reach the lowest-error solution: large students can remain on long plateaus, and some better fixed points are not reached from small random weights.For sigmoidal networks, convergence to the denoising solution can be difficult; for linear and ReLU networks, better fixed points may exist but remain unreachable from the stated initialization.
- Linear and ReLU networks: For linear and ReLU networks trained on both layers, the asymptotic generalisation error remains constant as K increases.This improves over the increase observed when only the first layer is trained, but does not produce the 1/K decay found for sigmoidal networks.
- Linear and ReLU networks: Normalised first-layer-only training decreases the asymptotic error with K because setting v_k = 1/K effectively reduces the learning rate.The paper relates this behaviour to the linear-network scaling ϵ_g ∼ η.
A Proof of Theorem 1.1
The proof establishes that the stochastic order-parameter dynamics concentrate around their expected updates and converge to the proposed ODE description as the input dimension grows. It combines moment bounds, induction, and a coupling argument with finite-difference convergence.
- Proof strategy: The proof shows that the mean order-parameter updates match the equations of motion while their variances are bounded at order N^-2.This makes the stochastic component negligible in the thermodynamic limit N →∞.
- Proof strategy: The time-dependent order parameters include teacher-student overlaps R, student-student overlaps Q, and second-layer weights v.The proof treats these variables through their update equations and conditional expectations.
- Moment bounds: Induction bounds the variance increments by C(T)N^-2 over the training horizon.The argument uses bounded activation functions and derivatives, deterministic initial macroscopic state, and Taylor expansion of the update terms.
- ODE convergence: The coupling trick compares the stochastic Markov process with a finite-difference process for the ODEs.An Euler argument then controls the finite-difference approximation to the equations of motion.
- Assumptions: Large input dimension and one-pass sampling are essential assumptions for the concentration and local-field arguments.The proof uses N →∞ and a sufficiently large dataset so each sample is visited only once before convergence.
B.1 Expressing the generalisation error in terms of order parameters
The generalisation error can be reduced to a function of macroscopic order parameters that encode student-teacher and within-network weight overlaps. This reduction enables closed ODE dynamics and asymptotic analysis across over-parameterised architectures.
- Order parameters: The order parameters R_in, Q_ik, and T_nm measure student-teacher overlaps and within-network self-overlaps.They summarise the relevant geometric relationships among student and teacher weight vectors.
- Order parameters: The order parameters are sufficient to determine the generalisation error ϵ_g.The high-dimensional input average is replaced by an average over K + M jointly Gaussian local fields in the large-N limit.
- ODE construction: The required local-field averages include I2, I3, I4, and J2-type integrals over multivariate normal distributions.For sigmoidal activations, the relevant I2 average can be evaluated analytically for g(x) = erf(x/2).
- ODE construction: The equations of motion for Q and R follow by squaring the first-layer weight update and taking its inner product with teacher weights.Closing these equations requires evaluating averages over local fields.
- Over-parameterised analysis: Symmetry reduction re-parameterises the overlap matrices and produces eight ODEs for analysing over-parameterised students.The reduction exploits permutation symmetry among student hidden units and makes the dependence on K and M more tractable.
- Over-parameterised analysis: For noiseless teachers, online learning reaches a fixed point with R = Q = 1, C = D = E = F = 0 and perfect generalisation ϵ_g = 0.The result is stated for T_nm = δ_nm and σ = 0.
C.2 Perturbative solution of the ODEs
The perturbative ODE analysis characterizes the asymptotic generalisation error of noisy sigmoidal networks and shows that over-parameterisation and learning rate can worsen it. With η = 1/K, the error still rises with added units before plateauing, while the perturbative result diverges beyond a learning-rate boundary.
- Perturbative calculation: The asymptotic generalisation error is obtained by expanding the ODEs to first order in the teacher-output noise variance σ2 around a noiseless fixed point.The resulting fixed point is evaluated through the generalisation-error expression.
- Perturbative calculation: The resulting function f(M, L, η) is an unwieldy rational function, so its full form is supplied with the accompanying notebook and source code.
- Dependence on L and η: The generalisation error increases monotonically with both the number of superfluous units L and the learning rate η for fixed values of the other variable.This trend is shown for teachers with M = 2 and M = 16.
- Dependence on L and η: Reducing η for larger students lengthens training and requires more data for convergence, while η ∼1/K does not eliminate the increase with L.At η = 1/K, the error eventually plateaus at a constant value independent of M.
- Validity range: For K = M, the divergence rate ηdiv equals the maximum learning rate ηmax guaranteeing exponential convergence to the noiseless optimum.The perturbative result diverges for sufficiently large L or learning rate.
- Small-learning-rate expansion: At small η, the generalisation error is linear in η, with second-order corrections quadratic in L.The first-order expression also matches the sum of asymptotic errors from M learning perceptrons and L non-learning perceptrons.
D Asymptotic generalisation error of a noisy continuous perceptron
This section analyzes noisy continuous perceptrons as the K = 1 case of the teacher-student setting. Its equations of motion have a perfect-generalisation fixed point without output noise, and a perturbative calculation gives the asymptotic noisy error.
- Problem setup: The continuous perceptron problem asks for the asymptotic generalisation error when a K = 1 student learns from a teacher with additive Gaussian output noise.
- Problem setup: The teacher noise is Gaussian with mean 0 and variance σ2, and the student and teacher weights are represented by w and w∗.
- Comparison with ReLU networks: The ReLU-network error scales as ϵg ∼ησ2L, matching the analytic scaling found for sigmoidal networks.The scaling is reported for the learning rate, teacher-noise variance, and additional hidden units.
- Dynamics: The equations of motion have a fixed point at Q = R = T that yields perfect generalisation when σ = 0.
- Asymptotic error: A perturbative ansatz in σ2 is used to derive the asymptotic generalisation error.The resulting expression is compared with the corresponding full soft committee machine result.
E Calculation of the asymptotic generalisation error in two-layer sigmoidal networks
For two-layer sigmoidal networks trained in both layers, the authors use a symmetry-aware ansatz and reduce the dynamics to coupled ODEs. The resulting fixed point yields the asymptotic generalisation-error formula plotted for comparison with linear and ReLU networks.
- Numerical comparison: Figure S4 compares asymptotic performance for linear and ReLU two-layer networks using the same y-axis and five-run standard-deviation error bars.The displayed parameters are N = 500, M = 2, v∗ = 4, η = 0.01, and σ = 0.01.
- Ansatz: The standard soft committee machine ansatz is unsuitable because all student units participate and multiple student units can overlap with one teacher unit.
- Ansatz: When K = ZM, the authors use an overlap ansatz adapted to students whose hidden-unit count is a multiple of the teacher’s.
- ODE calculation: The adapted ansatz produces reduced coupled ODEs for Q, C, R, and S, expanded around the noiseless fixed point R = 1, S = 0, Q = 1, C = 0.
- ODE calculation: Substituting the resulting fixed point into the generalisation-error expression yields the formula plotted in Fig. 3c.
F Unbalanced weights rescale effective learning rate in two layer linear networks
In linear two-layer networks, the product of the layer weights determines the output, but unequal layer norms change the effective learning dynamics. The analysis connects this imbalance to noisy updates and suggests that more balanced norms can reduce online-learning generalisation error.
- Network and updates: A linear two-layer network is parameterized by output weights v, first-layer weights w, and input x, with online SGD updating both weight layers.
- Network and updates: Defining the product of the student weights as a vector u provides a representation for analyzing the network’s output and updates.
- Effective learning rate: Rescaling equivalent networks as ṽ = av and w̃ = 1/a w changes the scaling of their update steps despite preserving the same output function.
- Effective learning rate: Layer-norm imbalance can affect noisy update fluctuations at late training times by changing the effective learning rate.
- Implication: More balanced norms across layers tend to produce lower generalisation error during online learning.
- Weight decay: Weight decay does not improve asymptotic generalisation error in the reported setup and fails to prevent its increase with L.The illustrative experiment varies the weight-decay constant κ.
G.2 SGD with mini-batches
Mini-batch SGD reduces asymptotic generalisation error up to a point, while preserving the qualitative dependence on student size seen in online learning.
- Figure S6 compares sigmoidal and ReLU students learning from a teacher with M = 4 across varying K and batch sizes.The plotted parameters include N = 500, η = 0.2, σ = 0.1, and κ = 0.
- Increasing the mini-batch size decreases asymptotic generalisation error until further gains become minimal.The comparison spans batch sizes from 1 to 20,000.
- The shape of the ϵ∗g−L curve is the same across mini-batch sizes.
- The minimal generalisation error is attained by a student with K = M.
G.3 Using MNIST images for training and testing
The analysis tests its input-distribution assumptions by replacing Gaussian inputs with MNIST images in the same teacher-student regression setup.
- The MNIST experiment replaces i.i.d. standard-normal inputs with images having non-trivial higher-order moments.
- The experiment retains the supervised teacher-student regression setup while changing the input distribution.
G.4 The scaling of ϵ∗ g with L for finite training sets
With finite training sets, the relationship between student size and final generalisation depends on activation function and dataset size, while early stopping shows qualitatively similar scaling.
- For nonlinear networks with P = 50, best generalisation occurs at K = M, whereas with P = 4 it occurs at K > M.
- Training-set size relative to network size strongly influences the scaling of ϵ∗g with L.
- Early-stopping generalisation error shows qualitatively the same dependence on L as final generalisation error.ReLU networks show little to no over-training, while sigmoidal networks show more significant over-training.
- For P ≫ K, increasing network size hurts generalisation, whereas for small P ∼ K it improves generalisation or has little effect.