Gradient descent for deep equilibrium single-index models
This work provides incremental theoretical insights into DEQ training dynamics, addressing a gap in the literature for researchers in machine learning theory.
The authors tackled the problem of theoretically understanding gradient descent dynamics for deep equilibrium models (DEQs) by analyzing linear and single-index models, proving linear convergence to a global minimizer under specific conditions.
Deep equilibrium models (DEQs) have recently emerged as a powerful paradigm for training infinitely deep weight-tied neural networks that achieve state of the art performance across many modern machine learning tasks. Despite their practical success, theoretically understanding the gradient descent dynamics for training DEQs remains an area of active research. In this work, we rigorously study the gradient descent dynamics for DEQs in the simple setting of linear models and single-index models, filling several gaps in the literature. We prove a conservation law for linear DEQs which implies that the parameters remain trapped on spheres during training and use this property to show that gradient flow remains well-conditioned for all time. We then prove linear convergence of gradient descent to a global minimizer for linear DEQs and deep equilibrium single-index models under appropriate initialization and with a sufficiently small step size. Finally, we validate our theoretical findings through experiments.