Gradient descent dynamics for deep equilibrium models

Abstract

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 matrix conservation law for linear DEQs which implies that the parameters remain trapped on spheres along gradient flow 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. In these simple settings, we also study the effect of common Neumann approximations used to reduce the cost of backpropagation; for linear DEQs, we show that increasing the number of terms in the Neumann approximation can prevent convergence to zero risk, even with exact forward solves. For nonlinear single-index DEQs, we prove linear convergence for fixed Neumann approximations, including Jacobian-free backpropagation under suitable assumptions. Finally, we validate our theoretical findings through experiments.