[Paper Review] Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss

Lénaïc Chizat, Francis Bach|arXiv (Cornell University)|Feb 11, 2020

Stochastic Gradient Optimization Techniques49 references90 citations

TL;DR

The paper characterizes the implicit bias of gradient flow for infinitely wide two-layer networks with exponential-tailed losses, showing convergence to a max-margin classifier in a non-Hilbert space, and compares training both layers to training only the output layer with empirical validation.

ABSTRACT

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.

Motivation & Objective

• Motivate understanding of why over-parameterized neural nets trained with gradient methods generalize well.
• Characterize the limit behavior of gradient flow for infinitely wide two-layer networks with 2-homogeneous activations.
• Show that the learned classifier is a max-margin solution in a variation norm space.
• Compare training both layers versus training only the output layer and analyze implications for generalization.
• Provide numerical evidence supporting the theoretical findings for ReLU-based two-layer nets.

Proposed method

• Model the predictor as a finite-width two-layer network with 2-homogeneous activation and balanced structure.
• Use the measure-based convex reformulation to describe the predictor via a variation norm 1 and its max-margin 1()1 objective.
• Characterize the infinite-width limit as a Wasserstein gradient flow on probability measures over parameters.
• Prove that the limit of the training dynamics yields a maximizer of the 1-max-margin problem under suitable assumptions.
• Contrast with the RKHS framework 2 and discuss the computational aspects of training only the output layer.
• Discuss convergence rates and online mirror descent connections in simplified dynamics.

Experimental results

Research questions

• RQ1 Does gradient flow on exponentially tailed losses converge to a global max-margin solution in the variation-norm function space 1?
• RQ2 How does training dynamics differ when jointly training both layers versus only training the output layer in terms of the implicit bias?
• RQ3 Can dimension-independent generalization be established for networks with hidden low-dimensional structure?
• RQ4 Do numerical experiments on wide two-layer nets with ReLU activations align with the theoretical max-margin characterizations?

Key findings

• The gradient flow of infinitely wide two-layer networks with exponential-tailed losses converges to a max-margin classifier in the 1 variation-norm space.
• With hidden low-dimensional structure, the resulting margin is independent of ambient dimension, enabling strong generalization guarantees.
• Training only the output layer implicitly solves a kernel SVM in the 2 RKHS, which may lack the same adaptivity as the 1 margin.
• In simplified settings, the training dynamics are shown to be equivalent to online mirror ascent with convergence rates of the order log(t)/sqrt(t).
• Numerical experiments suggest the theory describes practical behavior for two-layer ReLU networks and support the statistical benefits of the implicit bias.