Gradient descent for deep equilibrium single-index models

TL;DR

This paper provides a rigorous theoretical analysis of gradient descent dynamics in deep equilibrium models, focusing on linear and single-index models, and demonstrates convergence properties and training stability.

Contribution

It offers the first theoretical results on gradient descent behavior in DEQs, including convergence and stability analysis in simplified models.

Findings

Gradient flow remains well-conditioned during training.

Gradient descent converges linearly to a global minimum.

Parameters stay on spheres during training.

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 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.