[Paper Review] Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup

Sebastian Goldt, Madhu Advani+3 more|arXiv (Cornell University)|Jun 18, 2019

Stochastic Gradient Optimization Techniques33 citations

TL;DR

The paper analyzes online SGD dynamics of over-parameterised two-layer networks in a teacher-student framework, derives ODEs for macroscopic order parameters, and shows how generalisation scales with over-parameterisation depending on which layers are trained and activation functions.

ABSTRACT

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.

Motivation & Objective

• Motivate and understand why heavily over-parameterised networks generalise well in practice.
• Develop a rigorous macroscopic description (ODEs) of online SGD dynamics in the teacher-student setting.
• Analyze asymptotic generalisation error for over-parameterised students when only the first layer is trained.
• Analyze how training both layers changes generalisation and identify activation-function dependent behaviours.
• Provide analytic and numerical validation of the ODE framework against SGD simulations.

Proposed method

• Model inputs as i.i.d. Gaussian with teacher and student two-layer networks.
• Define order parameters m = (R, Q, T, v*, v) capturing teacher-student and student overlaps.
• Derive coupled ODEs for dR/dα, dQ/dα, dv/dα and show they close in terms of m(α).
• Prove a rigorous convergence result: SGD macroscopic state follows the unique solution of the ODEs in the large-N limit.
• Compute asymptotic generalisation error ε_g* for various activation functions (sigmoidal, linear, ReLU) and training setups.
• Validate analytical predictions with SGD simulations and finite-size experiments.

Experimental results

Research questions

• RQ1 How does online SGD dynamics evolve in the teacher-student two-layer setting as network size grows?
• RQ2 How does over-parameterisation (K > M) affect the final generalisation error when only the first layer is trained?
• RQ3 How does training both layers alter the asymptotic generalisation error and what solutions does SGD converge to for different activations?
• RQ4 What role do activation functions play in the fixed points and generalisation performance under SGD dynamics?
• RQ5 Can the ODE framework accurately predict SGD outcomes across various architectures and data regimes?

Key findings

• For Soft Committee Machines, the final generalisation error increases with the number of extra hidden units L when only the first layer is trained.
• With sigmoidal and linear activations, ε_g* scales as a function of η, σ^2 and L, showing larger over-parameterisation can degrade generalisation in the single-layer training regime.
• Training both layers yields improved generalisation for sigmoidal networks due to a denoising solution where multiple student units specialise and effectively average teacher outputs.
• ReLU and linear networks show constant ε_g* with increasing K when training both layers, indicating less pronounced benefits from over-parameterisation under SGD in these cases.
• Analytical expressions and numerics indicate the implicit regularisation from SGD depends critically on algorithm, architecture, and data, not SGD alone.
• The authors provide a reproducible pipeline, including an ODE integrator and experiments, at a public repository.