[Paper Review] On Lazy Training in Differentiable Programming

Lénaïc Chizat, Edouard Oyallon+1 more|arXiv (Cornell University)|Dec 19, 2018

Stochastic Gradient Optimization TechniquesComputer Science287 citations

TL;DR

The paper argues that lazy training, where models behave like their linearizations, emerges from scaling choices rather than over-parameterization alone, analyzes when this occurs, and shows that in practice lazy training can degrade CNN performance compared to non-lazy regimes.

ABSTRACT

In a series of recent theoretical works, it was shown that strongly over-parameterized neural networks trained with gradient-based methods could converge exponentially fast to zero training loss, with their parameters hardly varying. In this work, we show that this "lazy training" phenomenon is not specific to over-parameterized neural networks, and is due to a choice of scaling, often implicit, that makes the model behave as its linearization around the initialization, thus yielding a model equivalent to learning with positive-definite kernels. Through a theoretical analysis, we exhibit various situations where this phenomenon arises in non-convex optimization and we provide bounds on the distance between the lazy and linearized optimization paths. Our numerical experiments bring a critical note, as we observe that the performance of commonly used non-linear deep convolutional neural networks in computer vision degrades when trained in the lazy regime. This makes it unlikely that "lazy training" is behind the many successes of neural networks in difficult high dimensional tasks.

Motivation & Objective

• Motivate and define the lazy training phenomenon in differentiable programming.
• Develop a general criterion for when lazy training occurs via scaling and initialization.
• Analyze gradient flow dynamics under scaled models and compare to their linearizations.
• Provide theoretical bounds and convergence results in lazy regimes across over- and under-parameterized settings.
• Evaluate the practical implications of lazy training through synthetic and CNN experiments.

Proposed method

• Introduce a scaling factor alpha and study the objective F_alpha(w) = (1/alpha^2) R(alpha h(w)).
• Define the linearized model bar{h}(w) around initialization and compare F_alpha with its linearization bar{F}_alpha.
• Derive a general lazy training criterion kappa_h(w0) = ||h(w0)-y*|| * ||D^2 h(w0)|| / ||Dh(w0)||^2 and relate to lazy dynamics.
• Prove finite-horizon lazy training bounds showing w_alpha(t) approaches w0 and remains close to the linearized path as alpha grows (Theorem 2.2).
• Provide a square-loss quantitative bound (Theorem 2.3) and analyze over-parameterized and under-parameterized regimes (Theorems 2.4 and 2.5).
• Extend analysis to homogeneous models and two-layer networks, linking to random feature and mean-field limits.

Experimental results

Research questions

• RQ1 Under what conditions does gradient-based optimization of a scaled model behave like training the linearized model around initialization?
• RQ2 How do initialization, scaling, and network architecture influence the emergence of lazy training?
• RQ3 What are the convergence properties and generalization implications of lazy training in over- and under-parameterized regimes?
• RQ4 Do practical neural networks (e.g., CNNs) exhibit degraded performance in lazy regimes compared to non-lazy regimes?
• RQ5 Can we bound the distance between lazy training dynamics and the linearized dynamics over time?

Key findings

• Lazy training can arise not only in over-parameterized networks but in essentially any parametric model with output near zero at initialization, due to implicit scaling.
• For large alpha, the training dynamics of F_alpha are close to those of the linearized model bar{F}_alpha, making learning effectively linear.
• In the square loss case, explicit bounds show the distance between the nonlinear and linearized outputs shrinks as alpha increases, under mild smoothness assumptions.
• Over-parameterized lazy training converges to a global minimizer under suitable conditions on the Jacobian Dh(w0) and the loss (Theorem 2.4).
• Under-parameterized lazy training converges to a local minimum for large alpha, indicating potential stagnation away from global optima in finite-dimensional settings (Theorem 2.5).
• Numerical experiments indicate that CNNs trained in the lazy regime can underperform compared to non-lazy training and can be ill-conditioned, challenging the view that lazy training explains NN success.