Gradient descent dynamics for deep equilibrium models

arXiv:2511.16976v3 Announce Type: replace 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.

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.

Machine Learning (cs.LG); Statistics Theory (math.ST); Machine Learning (stat.ML)

arXiv:2511.16976 [cs.LG]

(or arXiv:2511.16976v3 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2511.16976 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.

analiza aiqbench

Gradient descent dynamics for deep equilibrium models

Subiecte conexe