Source-linked AI summary
Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss
Lenaic Chizat, Francis Bach
TL;DR
The paper asks why gradient-based logistic-loss training selects effective classifiers in over-parameterized two-layer networks. It analyzes infinitely wide homogeneous networks through gradient-flow limits and max-margin characterizations, showing favorable statistical behavior under hidden low-dimensional structure while leaving runtime and some assumptions unresolved.
Problem
The paper addresses the unclear mechanism behind good generalization when over-parameterized neural networks fit training data with gradient-based methods and no explicit regularization.
Method
The paper analyzes infinite-width gradient flow for homogeneous two-layer networks and characterizes its limits using functional max-margin problems, with simplified dynamics related to online mirror descent.
Results
Gradient flow converges asymptotically to an F1-max-margin classifier under stated assumptions, with corresponding finite-width asymptotic guarantees.
Takeaways & Limitations
Hidden low-dimensional structure can yield margins and generalization bounds that are independent of the ambient dimension, while training both layers differs fundamentally from training only the output layer.
Takeaways & Limitations
The analysis excludes ReLU networks from its differentiability assumption, and the runtime or convergence rate of the full non-convex gradient flow remains unknown.
Abstract
from arXiv · showhide
Neural networks trained to minimize the logistic (a.k.a. cross-entropy) loss with gradient-based methods are observed to perform well in many supervised classification tasks. Towards understanding this phenomenon, we analyze the training and generalization behavior of infinitely wide two-layer neural networks with homogeneous activations. We show that the limits of the gradient flow on exponentially tailed losses can be fully characterized as a max-margin classifier in a certain non-Hilbertian space of functions. In presence of hidden low-dimensional structures, the resulting margin is independent of the ambiant dimension, which leads to strong generalization bounds. In contrast, training only the output layer implicitly solves a kernel support vector machine, which a priori does not enjoy such an adaptivity. Our analysis of training is non-quantitative in terms of running time but we prove computational guarantees in simplified settings by showing equivalences with online mirror descent. Finally, numerical experiments suggest that our analysis describes well the practical behavior of two-layer neural networks with ReLU activation and confirm the statistical benefits of this implicit bias.
1. Introduction
The paper studies why gradient-based training can select well-generalizing classifiers among many zero-training-error solutions in over-parameterized wide two-layer networks. It connects this implicit bias to max-margin problems, dimension-independent generalization under hidden structure, computational guarantees in simplified settings, and empirical ReLU behavior.
- 1. Introduction: The paper investigates the implicit bias induced by gradient descent on logistic or exponentially tailed losses in over-parameterized two-layer networks.The setting concerns models with more parameters than data, trained to zero error without explicit regularization.
- 1. Introduction: Prior work established max-margin behavior for linear classifiers, while this paper characterizes the two-layer case through a convex max-margin problem.The characterization is precise enough to support statistical analysis.
- 1. Introduction: In the infinite-width setting, exponentially tailed losses lead non-convex gradient flow to a max-margin classifier for the variation norm.The analysis uses infinite-width dynamics to obtain a functional-space characterization.
- 1. Introduction: Fixing neuron directions or training only the output layer yields online mirror-ascent interpretations and convergence guarantees of O(log(t)/t) in simplified settings.These computational guarantees apply to the simplified dynamics described in the contributions.
- 1. Introduction: The resulting margins support dimension-independent generalization bounds under hidden linear structures, and ReLU experiments confirm the statistical efficiency of the implicit bias.The paper also notes that mean-field analyses are generally non-quantitative in neuron count and iteration count.
2. Preliminaries on infinitely wide two-layer networks
The paper formulates wide two-layer networks through homogeneous feature functions and functional norms, then connects their infinite-width training dynamics to Wasserstein gradient flow. It distinguishes variation-norm max-margin classifiers from RKHS max-margin classifiers and identifies assumptions and quantitative limitations of the analysis.
- 2.1. 2-homogeneous neural networks: Two-layer networks use m hidden units with a 1/m scaling, and the over-parameterized regime considers large m.The feature function φ is assumed to be positively 2-homogeneous and balanced; ReLU and S-ReLU networks are examples.
- 2.3. Max-margins and functional norms: The variation norm is the minimum representing-measure mass, and its max-margin problem defines the F1-max-margin γ1.Functions are represented by measures on the unit sphere, with balanced features allowing signed representations to be replaced by nonnegative measures.
- 2.3. Max-margins and functional norms: The RKHS norm defines a second function space F2, whose max-margin problem is denoted γ2 and corresponds to a separable kernel support vector machine.F2 is a reproducing kernel Hilbert space, unlike the variation-norm space F1.
- 2.3. Max-margins and functional norms: When hidden low-dimensional structure is present, γ1 can remain large in high dimension and yield strong generalization guarantees, whereas this is not a priori true for γ2.Training only the output layer can find F2-max-margin classifiers by convex optimization, while over-parameterized two-layer training is shown to find F1-max-margin classifiers.
- 2.4. Training dynamics in the infinite width limit: The main training theorem’s differentiability assumption excludes ReLU networks, although ReLU satisfies the other stated structural conditions.Subanalyticity is used in a technical proof step, and piecewise-polynomial activations such as ReLU and S-ReLU satisfy it.
- 2.4. Training dynamics in the infinite width limit: Under exponential-tail loss and regularity assumptions, infinitely wide finite networks converge on bounded time intervals to a unique Wasserstein gradient flow.The result assumes convergence of the empirical initialization measures in P2; discrete-time analysis is treated only in simpler settings.
- 2.4. Training dynamics in the infinite width limit: The infinite-width analysis is largely non-quantitative, with quantitative bounds potentially carrying exponential dependence on time.The framework includes lazy training as a short-time regime, but the paper’s runtime remains unknown.
3. Main result: implicit bias of gradient flow
Under stated convergence and initialization assumptions, infinitely wide two-layer networks trained with exponentially tailed losses converge to an F1 max-margin solution. The result gives a global characterization of the non-convex dynamics, while convergence itself and some ReLU conditions remain open or restrictive.
- Main theorem: Theorem 3 identifies the limiting normalized measure as a maximizer of the F1 max-margin problem.This conclusion assumes full support at initialization and convergence of the relevant quantities.
- Interpretation: The non-convex dynamics reaches a global minimizer of the convex max-margin problem, despite retaining dependence on initialization.A Dirac initialization can remain Dirac and typically fails to reach a global minimizer.
- Open questions: Convergence of the gradient and normalized measure is still an open question in this setting.The unnormalized measure does not converge, so a global convergence result from prior mean-field analysis does not apply.
- Finite-width implication: Theorems 2 and 3 provide asymptotic guarantees for finite-width neural-network training.The finite-width statement is formulated through convergence of the empirical initialization measures.
- ReLU scope: The ReLU extension relies on a strong Morse-Sard-type assumption and bounded-support, bounded-density input conditions.The authors describe this assumption as arguably too strong and use the result to identify technical open questions.
4. Insights on the convergence rate and choice of step-size
A simplified fixed-direction dynamics permits quantitative convergence analysis by reducing normalized iterates to online mirror ascent. This yields an asymptotic rate and shows similar convergence behavior across substantially different step-size choices.
- Simplified dynamics: The simplified model fixes each hidden-unit direction and trains only its nonnegative magnitude.The resulting signed features are fixed, and the analysis focuses on the exponential loss.
- Convergence rate: The best iterate converges to max-ℓ1-margin maximizers at an asymptotic rate log(t)/t.The proof exploits an analogy with online mirror ascent.
- Quantitative limitation: The bound on B grows to infinity as γ(m) approaches zero.This identifies a limitation of the quantitative guarantee in the relevant regime.
- Optimization reduction: The normalized dynamics approximately follows online mirror ascent over a sequence of smooth-margin objectives.A smoothing bound relates the soft minimum to the hard minimum through log(n)/β.
- Rate refinement: The analysis notes that a logarithmic dependence on m in Proposition 6 could be removed with a finer argument.This observation concerns the quantitative bound rather than the underlying limiting classifier.
5. Training only the output layer
Training only the output layer produces a max-ℓ2-margin classifier through a random-feature or kernel formulation. For larger, non-vanishing step sizes, the dynamics converges at O(log(t)/t) and exhibits self-regularizing behavior.
- Output-layer training: With the input layer fixed at random initialization, the model reduces to optimizing the output-layer objective over random features.The signed neuron outputs form the features used by the training dynamics.
- Implicit bias: The output-layer dynamics converges to a max-ℓ2-margin classifier.This is the kernel support-vector-machine-type implicit bias associated with the fixed random features.
- Convergence: For larger, non-vanishing step sizes, convergence occurs at O(log(t)/t).The result uses an online projected-gradient-ascent reduction with a step size scaled by the evolving norm.
- Step-size effect: Similar convergence speeds across very different step sizes indicate an advantageous self-regularizing property.The comparison is made within this simplified output-layer dynamics.
- Random-feature interpretation: Random features can approximate the F2 max-margin problem with feature counts typically of order 1/ε^2 for approximation error ε.This contrasts with the stated high-dimensional approximability properties of the corresponding function space.
6. Dimension independent generalization bounds
The paper defines a low-dimensional interclass-distance measure and uses it to derive strong non-asymptotic generalization bounds for the F1-max-margin classifier.
- Dimension-independent complexity: The quantity Δr identifies the r-dimensional subspace maximizing the distance between the two classes.When Δr is close to Δd for r much smaller than d, the data contain a hidden low-dimensional structure.
- Generalization bound: Theorem 8 gives a high-probability generalization bound under bounded inputs and empirical and population interclass-distance conditions.The bound applies for any ϵ ∈ (0,1) and r ∈ [d].
- Generalization bound: The same bound applies to the F2-max-margin classifier when r = d.
- Proof strategy: The proof lower-bounds the margins using Δr(Sn) and then applies margin-based generalization and Rademacher-complexity bounds.
7. Numerical experiments
Experiments compare training both layers with output-layer-only training in wide ReLU networks, showing distinct implicit biases, favorable performance for both-layer training, and a transition between kernel and F1 biases.
- Experimental setup: The experiments use a wide ReLU network with m = 1000 hidden units and compare training both layers against training only the output layer.The former leads to an F1 max-margin classifier, while the latter leads to an F2 max-margin classifier.
- Implicit-bias comparison: In d = 2, the F1-max-margin classifier has a nonsmooth piecewise-affine boundary, whereas the F2 classifier has a smooth boundary.The passage relates these shapes to sparse mass-constrained solutions and RKHS learning, respectively.
- Performance: Higher-dimensional experiments show superior test performance for training both layers as test error varies with m or d.The experiments average results from 20 independent trials with k = 3.
- Performance: Figure 2(d) evaluates the converged F1-margin as a function of width m using 30 fresh-sample experiments.
- Dynamics: With large initialization variance and a suitably small initial step-size, training first follows the tangent-kernel max-margin solution and later converges to the F1-max-margin classifier.The intermediate lazy-regime behavior requires an initial step-size inversely proportional to initialization scale.
8. Conclusion
The conclusion distinguishes the implicit biases of training both layers and only the output layer, emphasizing favorable statistics for the former while identifying unresolved computational questions.
- Conclusion: Training both layers converges to a max-margin classifier for a non-Hilbertian norm with favorable statistical properties.
- Conclusion: Training both layers and training only the output layer therefore produce very different implicit biases.
- Open questions: The non-Hilbertian max-margin problem does not seem directly solvable with known convex methods in high dimension.
- Open questions: A convergence rate under reasonable conditions remains unknown even for infinite-width continuous-time dynamics.
Appendix B. Equivalence of two variation norms
For ReLU networks, the paper relates two variation-norm representations through explicit measure transformations and proves that their functional spaces coincide up to a factor of two.
- Definitions: The two ReLU feature representations use φ(θ,z) = c(a · z)+ and ˜φ(a,z) = (a · z)+, inducing spaces F1 and ˜F1.
- Equivalence result: Proposition 9 states that F1 = ˜F1 and that ∥f∥F1 = 2∥f∥˜F1 for every f in F1.It also characterizes where norm-minimizing measures are concentrated.
- Constructive proof: The proof constructs a measure Π(ν) feasible for one norm from a feasible measure ν for the other, with norm changing by at most a factor of two.
- Consequence: The equivalence implies that weight decay regularization and the total variation regularization known as the path-norm are equivalent for empirical risk minimization.
- Constructive proof: The reverse construction uses pushforward measures and yields a norm multiplied exactly by two.
C.2. Proof of Theorem 2
The proof verifies the regularity and initialization assumptions needed to apply the Wasserstein gradient-flow result, including non-compactly supported initializations.
- The smooth-margin objective satisfies the required differentiability and local Lipschitz conditions under the stated assumptions.The feature map has a locally Lipschitz gradient and at most linear growth due to 2-homogeneity.
- Positive 2-homogeneity allows compact-support restrictions to be removed by characterizing the dynamics through projection onto the sphere.Every nonzero second-moment initialization can be replaced by a compactly supported measure with the same projected dynamics.
- Existence and uniqueness for compactly supported projected dynamics therefore extend to arbitrary initializations in P2(Rp).The velocity field is determined by the projected measure, so matching projections yields the same dynamics.
D.1. Proof of the main theorem
The proof shows that normalized gradient-flow measures converge toward max-margin optimality by establishing unbounded mass growth and complementary slackness.
- D.1. Proof of the main theorem: Theorem 10 states that any weak limit of the normalized path is a maximizer of the max-margin problem under the convergence assumptions.The result assumes full support initially, convergence of the loss gradient, and weak convergence of the normalized measures.
- D.1. Proof of the main theorem: The normalized dynamics are analyzed through the limiting function J′∞ and its maximizers on the sphere.The proof uses convergence of J′ associated with the normalized projected measures and identifies its positive maximum M.
- D.1. Proof of the main theorem: The normalized limit concentrates on the maximizer set of J′∞, yielding one complementary slackness condition.Mass outside increasingly narrow super-level sets vanishes under the weak-convergence argument.
- D.1. Proof of the main theorem: The loss-gradient soft-argmin property eliminates weights on nonminimal-margin examples, supplying the second complementary slackness condition.For exponentially tailed losses, gradient coordinates away from the limiting minimum margin converge to zero.
- D.1. Proof of the main theorem: The two complementary slackness conditions imply max-margin optimality through the stated optimality characterization.Proposition 12 provides necessary and sufficient conditions for global maximizers.
Appendix E. Appendix for Section 4
The appendix establishes computational and statistical properties of the simplified dynamics, relating normalized iterates to online mirror ascent and deriving low-dimensional generalization bounds.
- Online mirror ascent: The normalized dynamics follow perturbed online mirror ascent updates with entropy geometry and second-order step-size corrections.The perturbation appears because the multiplicative update includes corrections beyond the plain exponential mirror-ascent term.
- Limitations: The lower bound on β(t) is not optimized, so the computational growth guarantee is explicitly non-tight.The proof notes that constants are not tracked and that the lower bound was not made tight.
- Online mirror ascent: The mirror-ascent analysis proves that β(t) diverges and that normalized iterates approach the ℓ1-max-margin solution.The analysis tracks a bounded cumulative step-size quantity while establishing β(t) →∞.
- Online projected gradient ascent: A projected-gradient analysis gives analogous growth guarantees for the ℓ2-normalized dynamics.The iterates are online projected-gradient ascent steps on a Euclidean constraint set, with β(t) bounded below by a √t-scale expression.
- Generalization bounds: The F1-max-margin classifier admits a generalization bound whose exponent depends on the intrinsic dimension r rather than the ambient dimension d.The bound applies when the distribution has low-dimensional structure measured by ∆r(P); the F2 result is stated for r = d.
- Generalization bounds: The F1 class contains projected F2 functions, while its Rademacher complexity is no larger than the corresponding F2 complexity.These inclusions and complexity inequalities support the comparison between F1- and F2-based margin bounds.
Appendix H. Proof for the ReLU case
The ReLU appendix establishes well-defined Wasserstein dynamics under continuous-input assumptions and shows that suitable normalized limits solve the F1-max-margin problem.
- Regularity and assumptions: ReLU’s nondifferentiability is handled by restricting initialization and assuming an input distribution without atoms.The proof addresses parameter singularities at zero and hyperplane boundaries separately.
- Regularity and assumptions: For continuous functions on compact input support, the soft-margin functional is Fréchet differentiable with a probability-measure gradient.For exponential loss, its derivative is represented using an exponentially weighted measure.
- Regularity and assumptions: The ReLU feature differential is valid away from zero, and its restricted gradient is Lipschitz on the sphere under bounded second moments.The bounded-density assumption controls mass near activation boundaries.
- Existence of the flow: Initialization on the set D yields a well-defined Wasserstein gradient flow whose support remains in D for all time.The invariant relation between the parameter components preserves the initialization manifold.
- Implicit bias for ReLU networks: Under the ReLU theorem’s assumptions, a normalized limiting measure produces a classifier maximizing the F1 margin.The result relies on existence of the flow and convergence or regularity conditions specified in the theorem.