Deep Equilibrium Models: How Implicit Layers and Root-Finding Eliminate Activation Memory in Deep Learning

Deep neural networks derive their expressive capacity from composition. In conventional architectures, depth is explicit: an input vector is transformed sequentially through a chain of $L$ discrete layers, each parameterized by dedicated weights. While effective, this paradigm creates a fundamental computational dependency where training memory scales linearly with network depth ($O(L)$) because standard backpropagation requires caching intermediate activations across every layer.

Deep Equilibrium Models (DEQs), introduced by Bai, Kolter, and Koltun (2019), present an alternative formulation known as implicit deep learning. Instead of stacking discrete layers, a DEQ formulates the forward pass as finding the fixed-point equilibrium of a single non-linear transformation. By decoupling the definition of a model's state from the algorithmic process used to compute it, DEQs compute exact analytical gradients via the Implicit Function Theorem (IFT). This mechanism eliminates intermediate activation caching entirely, enabling constant $O(1)$ memory consumption regardless of effective receptive field or representation depth.

1. The Mathematical Limitation of Explicit Depth

In a standard feedforward neural network or stacked Transformer, the hidden representations evolve according to a discrete recurrence relation:

$$z_{i+1} = f_{\theta_i}(z_i, x)$$

where $x$ represents the input context, $z_i \in \mathbb{R}^d$ denotes the feature representation at layer $i$, and $\theta_i$ represents layer-specific parameters.

When training via backpropagation, the loss gradient with respect to parameter tensor $\theta_i$ depends directly on the activation tensor $z_i$:

$$\frac{\partial \mathcal{L}}{\partial \theta_i} = \frac{\partial \mathcal{L}}{\partial z_{i+1}} \frac{\partial f_{\theta_i}(z_i, x)}{\partial \theta_i}$$

To compute this product during the reverse pass, the automatic differentiation engine must retain $\{z_0, z_1, \dots, z_L\}$ in GPU High Bandwidth Memory (HBM). For modern sequence models operating at long context lengths, activation memory dwarfs parameter memory, necessitating engineering mitigations such as activation checkpointing, sequence parallelism, and offloading.

Weight-Tied Limits and Fixed-Point Convergence

Empirical research in weight-tied models—such as ALBERT (Lan et al., 2019) and TrellisNet (Bai et al., 2019)—demonstrated that repeatedly applying the same layer transformation causes hidden activations to stabilize:

$$z_{i+1} \approx z_i \quad \text{as } i \to \infty$$

Deep Equilibrium Models formalize this stationary behavior mathematically. Rather than computing an arbitrary number of discrete steps, a DEQ defines the network output directly as the infinite-depth equilibrium state $z^\star$:

$$z^\star = f_\theta(z^\star, x)$$

Under this formulation, the model does not specify how many layers to execute; it specifies the equilibrium conditions that the final representation must satisfy.

2. Forward Pass: Root-Finding in Non-Linear Vector Fields

Computing the forward activation of a DEQ requires finding a state vector $z^\star$ that satisfies the fixed-point equation. This can be framed equivalently as finding the root of an objective residual function $g_\theta$:

$$g_\theta(z; x) = f_\theta(z, x) - z = 0$$

+-------------------------------------------------------------+
|                     DEQ Forward Pass                        |
|                                                             |
|  Input x ────> [ Root-Finding Algorithm ] ────> Output z*   |
|                      ▲               │                      |
|                      │   Iterative   │                      |
|                      └── Update ◄────┘                      |
|             (Broyden / Anderson Acceleration)               |
|                                                             |
|         Satisfies:  g_θ(z*; x) = f_θ(z*; x) - z* = 0        |
+-------------------------------------------------------------+

Limitations of Picard Iteration

The simplest approach to finding $z^\star$ is Picard fixed-point iteration (repeated forward evaluation):

$$z^{(k+1)} = f_\theta(z^{(k)}, x)$$

By the Banach Fixed-Point Theorem, Picard iteration is guaranteed to converge to a unique fixed point if $f_\theta$ is a strict contraction mapping (its Lipschitz constant $\text{Lip}(f_\theta) < 1$). However, enforcing strict contraction constraints restricts parameter expressivity. Furthermore, Picard iteration exhibits linear convergence rates, often requiring dozens of iterations to reach numerical tolerance.

Quasi-Newton Solvers: Broyden's Method and Anderson Acceleration

To achieve practical efficiency, DEQs decouple the model architecture from the numerical solver. Production DEQ implementations employ accelerated root-finding algorithms:

1. Broyden's Method (Broyden, 1965): A quasi-Newton technique that iteratively updates a low-rank approximation of the inverse Jacobian $B_k \approx \left(\frac{\partial g_\theta}{\partial z}\right)^{-1}$. Using the Sherman-Morrison formula, Broyden's method achieves superlinear convergence without computing or factorizing the full $d \times d$ Jacobian matrix:

$$\Delta z_k = -B_k g_\theta(z_k)$$ $$z_{k+1} = z_k + \Delta z_k$$ $$B_{k+1} = B_k + \frac{\Delta z_k - B_k \Delta g_k}{\Delta z_k^\top B_k \Delta g_k} \Delta z_k^\top B_k$$

1. Anderson Acceleration (Anderson, 1965): Computes $z_{k+1}$ as an optimal linear combination of the previous $m$ iterates $\{z_{k-m}, \dots, z_k\}$ that minimizes the norm of the linear residual, accelerating convergence in multi-dimensional vector spaces.

Because forward evaluation is treated as a black-box numerical search, the forward solver operates entirely under torch.no_grad(). Intermediate solver states are discarded immediately after convergence.

3. Backward Pass: Implicit Differentiation via the Implicit Function Theorem

If training a DEQ required unrolling the numerical root-finder and backpropagating through time (BPTT), the memory footprint would scale linearly with solver iterations ($O(K)$), negating the efficiency benefits of the implicit formulation.

DEQs resolve this by deriving analytical gradients directly at the equilibrium point using the Implicit Function Theorem (IFT) (Krantz & Parks, 2002).

+-------------------------------------------------------------+
|                     DEQ Backward Pass                       |
|                                                             |
|  Loss Gradient ∂L/∂z* ───> [ Linear Solver ] ───> Vector v  |
|                                                             |
|             Solves:  v^T = v^T (∂f_θ/∂z*) + ∂L/∂z*          |
|                                                             |
|                                                             |
|  Parameter Update:                                          |
|  ∂L/∂θ = v^T (∂f_θ(z*, x)/∂θ)                               |
|                                                             |
|  Activation Memory Stored: O(1) (Only z* and x retained)    |
+-------------------------------------------------------------+

Derivation of Equilibrium Gradients

Let $z^\star(\theta)$ denote the implicit function mapping parameters $\theta$ to the unique equilibrium point satisfying $g_\theta(z^\star, x) = 0$. Differentiating both sides of the identity $f_\theta(z^\star(\theta), x) - z^\star(\theta) = 0$ with respect to $\theta$ via the multivariate chain rule yields:

$$\frac{\partial f_\theta(z^\star, x)}{\partial z^\star} \frac{d z^\star}{d \theta} + \frac{\partial f_\theta(z^\star, x)}{\partial \theta} - \frac{d z^\star}{d \theta} = 0$$

Factoring out the total derivative $\frac{d z^\star}{d \theta}$:

$$\left( I - \frac{\partial f_\theta(z^\star, x)}{\partial z^\star} \right) \frac{d z^\star}{d \theta} = \frac{\partial f_\theta(z^\star, x)}{\partial \theta}$$

Assuming the Jacobian matrix $\left( I - \frac{\partial f_\theta(z^\star, x)}{\partial z^\star} \right)$ is non-singular, inverting yields:

$$\frac{d z^\star}{d \theta} = \left( I - \frac{\partial f_\theta(z^\star, x)}{\partial z^\star} \right)^{-1} \frac{\partial f_\theta(z^\star, x)}{\partial \theta}$$

The Adjoint Fixed-Point System

For a scalar loss function $\mathcal{L}(z^\star)$, applying the chain rule gives the total gradient:

$$\frac{\partial \mathcal{L}}{\partial \theta} = \frac{\partial \mathcal{L}}{\partial z^\star} \frac{d z^\star}{d \theta} = \underbrace{\frac{\partial \mathcal{L}}{\partial z^\star} \left( I - \frac{\partial f_\theta(z^\star, x)}{\partial z^\star} \right)^{-1}}_{v^\top} \frac{\partial f_\theta(z^\star, x)}{\partial \theta}$$

Computing the matrix inverse explicitly is computationally intractable ($O(d^3)$). Instead, DEQs define the adjoint vector $v^\top$:

$$v^\top = \frac{\partial \mathcal{L}}{\partial z^\star} \left( I - \frac{\partial f_\theta(z^\star, x)}{\partial z^\star} \right)^{-1}$$

Multiplying both sides by $\left( I - \frac{\partial f_\theta(z^\star, x)}{\partial z^\star} \right)$ and rearranging terms yields a linear fixed-point system:

$$v^\top = v^\top \frac{\partial f_\theta(z^\star, x)}{\partial z^\star} + \frac{\partial \mathcal{L}}{\partial z^\star}$$

This equation is a linear system of the exact same form as the forward pass. It can be solved iteratively for $v$ using Broyden's method, GMRES, or Richardson iteration. Crucially, evaluating $v^\top \frac{\partial f_\theta}{\partial z^\star}$ requires only a single vector-Jacobian product (VJP), which standard autodiff frameworks evaluate in $O(d)$ time via backward automatic differentiation through a single layer evaluation at $z^\star$.

Once $v$ is computed, the final parameter gradient is obtained by backpropagating $v$ through $f_\theta(z^\star, x)$:

$$\frac{\partial \mathcal{L}}{\partial \theta} = v^\top \frac{\partial f_\theta(z^\star, x)}{\partial \theta}$$

Because this calculation requires only the converged state $z^\star$, the input $x$, and the model parameters $\theta$, memory consumption during training is strictly $O(1)$.

4. Architectural Paradigms: Vision and Sequence Modeling

The implicit formulation extends beyond single-layer vector representations to complex multi-dimensional modalities.

Explicit Layered Network:
Input ──> [ Layer 1 ] ──> [ Layer 2 ] ──> ... ──> [ Layer L ] ──> Output
Act Mem:     z_1             z_2                     z_L          Total: O(L)

Deep Equilibrium Network:
Input x ──┐
          ▼
   ┌─> [ Non-Linear Transformation f_θ ] ──┐
   │                 ▲                     │
   │                 └─────────────────────┘
   │               Iterative Convergence
   └────────────────> Equilibrium State z* ──> Output
Act Mem:                     z* only                              Total: O(1)

Multiscale Deep Equilibrium Models (MDEQ)

In computer vision tasks such as semantic segmentation and object detection, explicit networks (e.g., HRNet, U-Net) maintain feature representations at multiple spatial resolutions simultaneously.

Bai et al. (2020) extended DEQs to multi-scale feature hierarchies by formulating the equilibrium state as a tuple of feature maps at varying spatial resolutions:

$$\mathbf{z}^\star = \left( z^{(1)\star}, z^{(2)\star}, z^{(3)\star}, z^{(4)\star} \right)$$

where $z^{(i)}$ represents feature maps downsampled by $2^{i-1}$. The transformation function $f_\theta$ executes intra-resolution residual convolutions and inter-resolution cross-scale message passing. Solving for the joint equilibrium $\mathbf{z}^\star = f_\theta(\mathbf{z}^\star, x)$ forces all resolution scales to reach simultaneous consistency, matching state-of-the-art vision backbones while consuming a fraction of activation memory.

Equilibrium Transformers (MDEQ-Transformer)

Standard Transformers stack $N$ self-attention and feed-forward network (FFN) blocks. Bai et al. (2021) formulated the Equilibrium Transformer, where sequence representations converge under an equilibrium multi-head self-attention layer:

$$z^\star = \text{LayerNorm}\left( \text{FFN}(\text{SelfAttn}(z^\star)) + x \right)$$

By driving self-attention updates to fixed points, Equilibrium Transformers capture deep contextual dependencies across long sequence horizons without instantiating multi-layer KV-cache checkpoints during training.

Monotone Operator Equilibrium Networks (MonDEQ)

A theoretical vulnerability of unconstrained DEQs is that a fixed point is not mathematically guaranteed to exist or remain unique for arbitrary parameter configurations.

To guarantee well-posedness, Winston and Kolter (2020) introduced MonDEQ, parameterizing $f_\theta$ via strongly monotone operators. By structuring the layer Jacobian to satisfy a negative definite constraint:

$$\langle u - v, f_\theta(u) - f_\theta(v) \rangle \le -m \|u - v\|^2 \quad (m > 0)$$

MonDEQ guarantees that the root-finding problem possesses a unique, globally stable equilibrium that can be found via forward-backward operator splitting (e.g., Peaceman-Rachford or Douglas-Rachford splitting) without numerical divergence.

5. Architectural Comparison

Explicit Deep Networks (ResNets / Standard Transformers)

• Mathematical Formulation: Discrete layer composition: $z_{i+1} = f_{\theta_i}(z_i, x)$ across $L$ physical layers.
• Depth Parameterization: Explicit integer depth $L$ fixed prior to training and inference.
• Training Memory Complexity: $O(L)$ scaling linearly with layer count, requiring intermediate activation retention for backpropagation.
• Backward Pass Mechanism: Standard reverse-mode automatic differentiation (Backpropagation Through Time) unrolled across all $L$ layers.
• Inference Execution: Fixed deterministic computation requiring exactly $L$ sequential layer evaluations.

Neural Ordinary Differential Equations (Neural ODEs)

• Mathematical Formulation: Continuous parameterization: $\frac{dz(t)}{dt} = f_\theta(z(t), t)$ over continuous depth $t \in [t_0, t_1]$.
• Depth Parameterization: Continuous integration time horizon.
• Training Memory Complexity: $O(1)$ memory achieved via the continuous Adjoint Sensitivity Method integrating backward in time.
• Backward Pass Mechanism: Reverse-time numerical integration of the joint adjoint state differential equation.
• Inference Execution: Adaptive step size numerical ODE solvers (such as Dormand-Prince / RK45).

Deep Equilibrium Models (DEQs)

• Mathematical Formulation: Non-linear root-finding: finding equilibrium state $z^\star$ satisfying $z^\star = f_\theta(z^\star, x)$.
• Depth Parameterization: Implicit infinite depth ($L \to \infty$) defined by stationary fixed-point convergence.
• Training Memory Complexity: $O(1)$ constant memory via the Implicit Function Theorem, retaining only input $x$ and equilibrium state $z^\star$.
• Backward Pass Mechanism: Direct linear fixed-point solver: $v^\top = v^\top \left(\frac{\partial f_\theta}{\partial z^\star}\right) + \frac{\partial \mathcal{L}}{\partial z^\star}$.
• Inference Execution: Accelerated quasi-Newton root-finding (Broyden's method or Anderson acceleration) until $\|f_\theta(z, x) - z\| < \epsilon$.

6. Engineering Challenges and Practical Bottlenecks

While DEQs resolve the activation memory bottleneck, deploying implicit models in production environments introduces distinct engineering trade-offs:

1. Computational Overhead (Compute vs. Memory Trade-Off): While explicit models evaluate a layer exactly once, DEQ root-finding typically requires 15 to 35 solver iterations during the forward pass and 10 to 20 iterations during the backward adjoint solve. On compute-bound workloads, DEQs run 2x to 4x slower than equivalent-parameter explicit networks.
2. Jacobian Conditioning and Spectral Norm Drift: If the Jacobian matrix $\frac{\partial f_\theta}{\partial z^\star}$ develops eigenvalues with magnitude $|\lambda| \ge 1$, the linear system $(I - J)$ becomes ill-conditioned, leading to numerical divergence in Broyden's solver. Mitigating this requires explicit Jacobian regularization (Geng et al., 2021), penalizing the Frobenius norm $\|J_f\|_F$ or applying weight normalization to constrain the spectral radius $\rho(J_f) < 1$.
3. Hardware Vectorization and Dynamic Halting: In batched GPU inference, different batch items converge to numerical tolerance in different numbers of iterations. Naive batching forces all samples to wait for the slowest converging sample (the tail-latency problem). Efficient deployment requires asynchronous execution kernels or fixed-iteration warm-starting.

Summary

Deep Equilibrium Models demonstrate that architectural depth does not require stacking discrete physical layers. By reframing feature extraction as numerical root-finding and computing gradients via the Implicit Function Theorem, DEQs achieve the representational capacity of infinite-depth networks with constant $O(1)$ activation memory. As memory bandwidth remains a primary constraint in large-scale machine learning hardware, implicit formulations offer a mathematically rigorous path toward training ultra-deep architectures within fixed hardware budgets.

Sources

• Bai, Kolter, & Koltun (2019) - Deep Equilibrium Models
• Bai, Koltun, & Kolter (2020) - Multiscale Deep Equilibrium Models
• Bai, Koltun, & Kolter (2021) - Transformers are Deep Equilibrium Models
• Winston & Kolter (2020) - Monotone Operator Equilibrium Networks
• Broyden (1965) - A Class of Methods for Solving Nonlinear Simultaneous Equations
• Anderson (1965) - Iterative Procedures for Nonlinear Integral Equations
• Krantz & Parks (2002) - The Implicit Function Theorem: History, Theory, and Applications
• Geng et al. (2021) - On Training Implicit Models
• Chen et al. (2018) - Neural Ordinary Differential Equations