Deep Equilibrium Models

Every architecture so far has been a recipe: an ordered list of layers, each applied once, the output read off the top of the stack. A residual block, an LSTM cell, a transformer block — all are functions you run. A deep equilibrium model (DEQ) instead defines a layer not by what it computes but by a condition its output must satisfy, and lets a black-box solver find whatever output meets the condition.¹

A layer that calls itself

Take a single weight-tied transformation $f_{\theta }(\cdot ,x)$ (one layer's worth of parameters, conditioned on the input $x$) and apply it over and over to its own output. If the iteration settles, it settles at a point that the layer maps to itself.

This describes an infinite-depth weight-tied network run to convergence. Stack the same block $L$ times and you compute $z_L=f_{\theta }(\cdots f_{\theta }(f_{\theta }(z_0,x),x)\cdots ,x)$; as $L\to \infty$, if the iteration converges, $z_L\to z^{\star }$ independent of $z_0$. The infinitely deep stack and the self-referential layer are two descriptions of one object.

The payoff is structural. A conventional $L$-layer net stores $L$ sets of activations for the backward pass; the DEQ stores one, at any effective depth — provided we can solve and differentiate the fixed-point condition, which is the rest of this lesson.

Forward pass = root-finding

Finding $z^{\star }$ is solving a nonlinear system. Rearrange the fixed-point condition into a residual whose zero is the equilibrium:

Any root-finder applies. The simplest is fixed-point iteration (literally re-running the layer), which converges when $f_{\theta }$ is a contraction.²

The contraction constant is governed by the Jacobian of $f_{\theta }$: the iteration converges locally when the spectral radius $\rho (J_f)<1$, where $J_f=\partial f_{\theta }/\partial z$ evaluated near $z^{\star }$. Each step of plain fixed-point iteration multiplies the error by roughly $\rho (J_f)$, so convergence is linear and can be slow when $\rho (J_f)$ is close to $1$.

For example, take the one-dimensional layer $f(z)=0.5z+1$, a contraction with $K=0.5<1$. Its fixed point solves $z^{\star }=0.5z^{\star }+1$, so $z^{\star }=2$. Start at $z_0=0$ and iterate:

The error halves every step, exactly the predicted geometric decay $K^k=0.5^k$: reaching a tolerance of $10^{-3}$ takes about $11$ steps ($0.5^{11}\cdot 2\approx 10^{-3}$). Now sharpen the map to $f(z)=0.95z+0.1$, still a contraction but with $K=0.95$. Its fixed point is $z^{\star }=2$ as well, but the error now shrinks by only $5\%$ per step, so the same tolerance needs about $150$ iterations — two orders of magnitude slower. This is why the slope of the map at $z^{\star }$, not just the fact that it is below $1$, decides whether plain iteration is practical, and why the Newton and Broyden solvers below justify their extra per-step cost.

Plain iteration is rarely the fastest route. Treating $h_{\theta }(z,x)=0$ as a general root-finding problem licenses quasi-Newton solvers (Newton and Broyden) that use (an approximation of) the Jacobian to take superlinear steps.³ The forward solver is genuinely a black box: the rest of the model needs only a point that satisfies the residual to tolerance, however it was found.

Backward pass = implicit differentiation

To train the DEQ we need $\partial z^{\star }/\partial \theta$ — how the equilibrium moves when the weights move. The naive route is to backprop through every solver iteration, storing each intermediate $z_k$; that destroys the memory advantage and ties the gradient to an arbitrary solver trajectory. The implicit function theorem lets us skip the solver entirely and differentiate the condition instead.⁴

The fixed point satisfies $z^{\star }=f_{\theta }(z^{\star },x)$ as an identity in $\theta$. Differentiate both sides with respect to any quantity $(\cdot )$, whether $\theta$ or the input $x$, and apply the chain rule, remembering that $z^{\star }$ itself depends on $(\cdot )$:

Collect the $\partial z^{\star }/\partial (\cdot )$ terms on the left and factor:

This is a closed form for the layer's Jacobian that mentions only the fixed point $z^{\star }$ and the local Jacobian $J_f$ — not a single solver iterate. The matrix $I-J_f$ is invertible exactly when $\rho (J_f)<1$, the same condition that made the forward pass converge.

For example, return to $f(z)=0.5z+1$ from the forward-pass trace, but read $0.5$ as a parameter: $f_{\theta }(z)=\theta z+1$ with $\theta =0.5$, whose fixed point is $z^{\star }=1/(1-\theta )=2$. Here $J_f=\partial f/\partial z=\theta =0.5$ and $\partial f/\partial \theta =z^{\star }=2$, so the implicit-gradient formula gives

Check it directly: $z^{\star }(\theta )=1/(1-\theta )$, so $\mathrm{d}{z}^{\star }/\mathrm{d}\theta =1/(1-\theta {)}^2=1/0.25=4$. The implicit formula and the explicit derivative agree exactly, and the $(1-J_f{)}^{-1}$ factor is precisely the $1/(1-\theta )$ that a naive one-step derivative would miss — it accounts for the fact that changing $\theta$ moves the whole equilibrium, not just one application of $f$. In many dimensions this scalar reciprocal becomes the matrix inverse $(I-J_f{)}^{-1}$, and the same the equilibrium shifts too correction is what the adjoint solve below recovers without ever forming that inverse.

In practice we never form $(I-J_f{)}^{-1}$: it is a huge dense matrix. We need only a vector–Jacobian product: the upstream gradient $v^{\top }=\partial \mathcal{L}/\partial z^{\star }$ times the inverse. Define the adjoint $u$ by

which is itself a fixed-point equation — solvable with the same black-box solver used in the forward pass, since $J_f^{\top }$ acts on a vector through one reverse-mode autodiff call on $f_{\theta }$.⁵ Once $u$ is found, every parameter gradient is the cheap product $u^{\top }\partial f_{\theta }/\partial \theta$.

Explicit deep net vs DEQ

The two designs compute related functions but the cost of depth falls in different places. The explicit net spends memory (one activation buffer per layer) and parameters (one weight set per layer); the DEQ spends solver iterations at run time and stores only a single equilibrium.

Convergence in practice

Because the forward pass is a solver, training surfaces a diagnostic the explicit net never had: the residual $\Vert z_k-f_{\theta }(z_k,x)\Vert$, which a healthy DEQ drives toward zero each step. A residual that stalls or grows signals $\rho (J_f)\ge 1$ — the layer is not contractive, the fixed point is not being reached, and gradients computed from a non-equilibrium $z_k$ are wrong. Stabilizing $\rho (J_f)<1$ (via spectral normalization or a Jacobian regularizer) is the central engineering concern.

Tradeoffs and uses

The DEQ gains constant memory and infinite effective depth at the cost of solver time and stability:

The natural fit is settings where depth helps but memory is the bottleneck and a single repeated transformation is a reasonable inductive bias: sequence models (a weight-tied recurrence iterated to equilibrium) and graph models (message passing run to a steady state over the graph). In both, a DEQ matches the accuracy of a deep explicit stack while training within a fixed memory budget.

Implicit layers as a family

DEQ is one member of a broader implicit-layer family, all of which replace run these operations with return the output satisfying this condition, and all of which train through the condition rather than the computation. The public literature gives three landmarks.

Deep equilibrium models (Bai, Kolter & Koltun, 2019) introduced the construction above and showed a single weight-tied DEQ block matches deep transformers and weight-tied LSTMs on sequence tasks at constant memory. The follow-up multiscale DEQ (Bai, Koltun & Kolter, 2020) drove a single implicit block to competitive ImageNet accuracy by solving for equilibria at several resolutions at once, evidence that the constant-memory trick scales beyond toy settings.

The open problem noted above — that nothing guarantees $\rho (J_f)<1$, so the solver can diverge — is what monotone DEQ (Winston & Kolter, 2020) addresses. By parameterizing the layer so the residual operator is provably monotone, it guarantees a unique fixed point and a convergent solver by construction, trading a little expressiveness for a stability certificate. This is the DEQ analogue of the certified defenses from adversarial robustness: prove the property rather than hope for it.

The closest cousin is the neural ODE (Chen et al., 2018), which takes the opposite limit of the same idea. Where a DEQ is a weight-tied layer iterated to a fixed point, a neural ODE is a residual block whose step size shrinks to zero, so the forward pass integrates $\mathrm{d}z/\mathrm{d}t=f_{\theta }(z,t)$ and the backward pass solves an adjoint ODE — again $O(1)$ memory, again differentiating a solver rather than storing its trace. Fixed-point equations and differential equations are two routes to the same infinite effective depth.

The last row, the optimization layer (OptNet, Amos & Kolter, 2017), closes the circle: its output is the argmin of an inner objective, differentiated through the optimality (KKT) conditions — the same implicit-function move, applied to a minimizer instead of a fixed point. All four share one design: define the output by a condition, make the condition differentiable, and let a solver and the implicit-function theorem do the rest.

Closing the loop

DEQ postdates Goodfellow (present it as a frontier method), but it is built entirely from machinery the course has already laid down: the autodiff of backpropagation, the fixed-point and Newton ideas from the optimization landscape, and the weight tying of recurrent networks. And it returns us to the very first idea of the course. We opened by saying a deep network is the composition of simple differentiable pieces learned by following the gradient; the DEQ shows that deep need not mean many stacked layers at all — a layer can be a condition a single transformation satisfies, run to convergence, differentiated through its own definition. From the perceptron's weighted sum to an infinite-depth network that fits in constant memory, the same principle applies: define what the output should satisfy, make it differentiable, and let the gradient do the learning.