On Lazy Training in Differentiable Programming (1812.07956v5)

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.

• The paper presents a formal criterion for lazy training, demonstrating that scaling choices can convert non-linear models into effective linear systems during optimization.
• The methodology leverages gradient flow analysis and kernel theory to establish linear convergence rates under both over- and under-parameterized regimes.
• Empirical results reveal that models trained in the lazy regime may underperform in vision tasks, emphasizing the need for non-lazy training dynamics in practical applications.

An Essay on "On Lazy Training in Differentiable Programming" by Chizat, Oyallon, and Bach

"On Lazy Training in Differentiable Programming," authored by Lénaïc Chizat, Edouard Oyallon, and Francis Bach, presents a rigorous analysis of an intriguing phenomenon in machine learning—termed "lazy training." This work thoroughly investigates the conditions under which differentiable programming models, specifically neural networks, behave as inherently linear systems during training. Such behavior has profound implications for understanding optimization dynamics and practical performance in machine learning tasks.

Lazy Training Explained

Lazy training describes scenarios where neural networks exhibit negligible parameter changes and converge swiftly to zero training loss. Contrary to intuitive expectations, the authors demonstrate that lazy training is not just a feature of over-parameterized neural networks. Instead, it can manifest in various models due to specific scaling choices, thereby causing the models to operate near their linearization around initial parameter values.

The crux of lazy training lies in an implicit scaling that ensures the output remains close to zero at the start. Such scaling transforms the training of complex, non-linear models into the training of linear models with positive-definite kernels. This fundamentally simplifies the optimization landscape, making it resemble that of a convex problem amenable to linear methods.

Theoretical Contributions

The paper makes several important theoretical contributions:

1. General Criterion for Lazy Training: The authors derive a formal criterion that quantifies when lazy training occurs. This criterion depends on the initial scale of the model's output, the Lipschitz constants of the model and its differential, and the second-order derivatives of the model's output. For instance, they establish that for $q$-homogeneous models (such as deep neural networks), lazy training is bound to emerge with increasing variance of initialization.
2. Dynamical Analysis: By examining gradient flows, Chizat et al. show that lazy training dynamics closely adhere to those of the corresponding linearized models, over a finite time horizon. In the case of strongly convex loss functions, they prove that these dynamics guarantee linear convergence rates to global minimizers in the over-parameterized setting, or local minima in under-parameterized scenarios.
3. Differentiable Models: The paper extends the understanding to broad instances of differentiable programming, beyond just neural networks. They also handle the high-dimensional complexities by linking the phenomenon to kernel theory, particularly relating to the neural tangent kernel (NTK) and random feature models.
4. Numerical Validation: Crucially, the authors balance their theoretical findings with extensive numerical experiments. They demonstrate that in practical settings, like computer vision tasks using convolutional neural networks (CNNs), models trained in the lazy regime tend to perform poorly compared to their non-lazy counterparts. This finding is pivotal, suggesting that the impressive empirical performance of neural networks in high-dimensional tasks cannot solely be attributed to lazy training dynamics.

Implications and Future Directions

The practical implications of this research lie in:

• Model Initialization and Scaling:

Understanding lazy training helps guide better initialization and scaling practices. This can have significant impacts on training efficiency and model performance, especially for deep and wide neural networks.

• Interpretation of Over-parameterization:

This work challenges the naive interpretation that over-parameterization alone suffices for good performance. It makes clear that effective training dynamics, which steer away from laziness, are critical for achieving desirable generalization properties.

• Kernel Methods and Neural Networks:

The linkage to kernel methods provides a theoretical scaffold for interpreting neural network performance and extends the toolkit available for theoretical analysis of deep learning models.

Conclusion

"On Lazy Training in Differentiable Programming" offers a comprehensive exploration of a subtle yet significant aspect of machine learning. By merging theoretical insights with empirical observations, the paper provides a nuanced understanding of the conditions and implications of lazy training. This work is a stepping stone for future research to explore the dynamics of neural networks, optimize training methodologies, and further reconcile empirical success with theoretical foundations in machine learning.