Deep Equilibrium Models

Shaojie Bai; J. Zico Kolter; Vladlen Koltun

Source-linked AI summary

Deep Equilibrium Models

Shaojie Bai, J. Zico Kolter, Vladlen Koltun

arXiv:1909.01377v2Sep 3, 2019cs.LGstat.ML

TL;DR

Deep sequence models can converge toward equilibrium states, but conventional depth-based computation stores intermediate activations and limits memory efficiency. DEQ directly solves for these equilibria and differentiates through them implicitly, yielding constant-memory training while supporting multiple sequence architectures. Experiments show competitive or slightly better performance on realistic large-scale sequence tasks, with applicability constrained by the need for stable transformations.

Problem
Conventional deep sequence models rely on layer stacks and stored intermediate activations, while prior continuous approaches faced challenges scaling to realistic high-dimensional datasets.
Method
DEQ directly solves for the fixed point of a weight-tied sequence transformation and differentiates through that equilibrium using implicit differentiation.
Results
DEQs achieve competitive or slightly better performance than corresponding layer-based and state-of-the-art sequence models while using constant memory and scaling to large-scale sequence tasks.
Takeaways & Limitations
DEQ provides a practical implicit-depth perspective that applies across substantially different sequence architectures and supports constant-memory optimization.
Takeaways & Limitations
Reliable and efficient equilibrium finding generally requires the transformation fθ to be stable and constrained.

Abstract

from arXiv · show

We present a new approach to modeling sequential data: the deep equilibrium model (DEQ). Motivated by an observation that the hidden layers of many existing deep sequence models converge towards some fixed point, we propose the DEQ approach that directly finds these equilibrium points via root-finding. Such a method is equivalent to running an infinite depth (weight-tied) feedforward network, but has the notable advantage that we can analytically backpropagate through the equilibrium point using implicit differentiation. Using this approach, training and prediction in these networks require only constant memory, regardless of the effective "depth" of the network. We demonstrate how DEQs can be applied to two state-of-the-art deep sequence models: self-attention transformers and trellis networks. On large-scale language modeling tasks, such as the WikiText-103 benchmark, we show that DEQs 1) often improve performance over these state-of-the-art models (for similar parameter counts); 2) have similar computational requirements to existing models; and 3) vastly reduce memory consumption (often the bottleneck for training large sequence models), demonstrating an up-to 88% memory reduction in our experiments. The code is available at https://github.com/locuslab/deq .

1 Introduction

Deep sequence models often use repeated layers, and weight-tied networks raise the question of their infinite-depth limit. DEQ answers by directly solving and differentiating through the resulting equilibrium, reducing memory while retaining competitive sequence-modeling performance.

Motivation: Deep feedforward sequence models apply layer-indexed nonlinear transformations repeatedly to hidden sequences, often enforcing causality.The hidden sequence is updated across layers, with future time points prevented from influencing past ones.
Motivation: Weight tying applies the same transformation at every layer, prompting the question of what repeated application converges to and how to model that limit.Prior weight-tied models achieved state-of-the-art-competitive results despite reusing the transformation.
DEQ approach: DEQ directly solves for the fixed point representing the eventual hidden values of an infinite-depth network, rather than iterating through all layers.The equilibrium is obtained with black-box root-finding methods.
DEQ approach: Implicit differentiation through the equilibrium avoids storing intermediate activations, enabling constant-memory backpropagation through the effective infinite-depth network.The backward pass uses memory equivalent to a single layer’s activations.
Experiments: DEQ instantiates both trellis networks and memory-augmented universal transformers, achieving accuracy on par with or slightly better than corresponding layer-based networks.The experiments include large-scale sequence modeling benchmarks such as WikiText-103.

2 Background

The paper situates DEQ among sequence-modeling and memory-efficient deep-learning approaches. Unlike methods tied to particular architectures or latent trajectories, DEQ directly solves for sequence-level equilibria and applies to large-scale sequence tasks with constant memory.

Deep sequence models: Deep sequence modeling includes recurrent networks, temporal convolutions, and self-attention transformers, with transformers reaching state-of-the-art results on several NLP benchmarks.Trellis Networks connect recurrent and convolutional perspectives by combining their benefits.
Memory-efficient deep networks: Gradient checkpointing and reversible networks reduce memory for deep models, but checkpointing adds forward computation and reversible methods reconstruct activations during backpropagation.These approaches address the storage of intermediate layer activations through different mechanisms.
Memory-efficient deep networks: DEQs reduce memory consumption to a constant independent of network depth by differentiating directly through the equilibrium instead of maintaining layers.This avoids constructing and storing the intermediate activations of a conventional deep network.
Continuous views: Unlike ODE-based methods that backpropagate through an entire latent trajectory, DEQ solves directly for sequence-level equilibria with a quasi-Newton method.DEQ therefore does not depend on the solution path that leads to the equilibrium.
Positioning: DEQ is not restricted to a specific interlayer transformation and is demonstrated with two different sequence architectures on realistic high-dimensional tasks.The paper contrasts this scope with prior work focused on residual architectures, small symmetric-weight energy models, or limited numerical evaluations.
Implicit layers: An implicit layer defines its output as the solution of a nonlinear equation, allowing DEQ to be viewed either as an infinitely deep network or as a single implicit layer.The paper presents DEQ as a large-scale practical application of implicit layers.

3 The Deep Equilibrium Sequence Model

DEQ replaces the effectively infinite unrolling of weight-tied sequence models with direct equilibrium computation and implicit differentiation, enabling constant-memory training and prediction. It uses root-finding for the forward pass and Jacobian-based linear solves for the backward pass, while requiring stable transformations and offering no extra representational power from stacking DEQs.

Motivation: Weight-tied sequence models can converge toward a fixed point as depth increases, motivating equilibrium as a replacement for predetermined unrolling.The model family includes trellis networks and universal transformers, whose activations are repeatedly refined by the same transformation.
Forward Pass: DEQ directly solves for the equilibrium state instead of iteratively stacking the transformation.Its output is the equilibrium point itself, and the forward pass may use a black-box root-finding procedure such as Newton or quasi-Newton methods.
Backward Pass: Implicit differentiation computes gradients through the equilibrium without storing intermediate hidden states or depending on the root-finding algorithm.The backward gradient is represented using a matrix operation involving the Jacobian at equilibrium, and linear systems can be solved indirectly with vector-Jacobian products.
Memory Cost: Constant memory is achieved by retaining the equilibrium sequence, layer-independent inputs, and the transformation rather than the full sequence of activations.This avoids explicitly constructing the potentially huge Jacobian and can let models fit on a single GPU instead of requiring multiple GPUs or checkpointing.
Design Constraints: Reliable and efficient equilibrium finding generally requires stable, constrained transformations, such as gated activation in TrellisNet or layer normalization in transformers.The analysis itself is independent of the transformation choice, but convergence is not.
Representational Power: Stacking multiple DEQs does not provide additional representational power over a single DEQ.This follows from the paper’s universality theorem for single-layer DEQs.

4 Instantiations of DEQ

The paper instantiates DEQ with TrellisNet and weight-tied self-attention transformers, showing how distinct sequence architectures can be expressed through equilibrium transformations. The comparison emphasizes DEQ’s constant-memory requirement relative to conventional weight-tied networks.

Instantiations: DEQ is demonstrated with TrellisNet and self-attention, illustrating compatibility with distinct deep sequence model architectures.Both instantiations are presented as examples of transformations fθ within the generic DEQ framework.
Trellis Networks: TrellisNet is a weight-tied temporal convolutional network with input injection at every layer.Its DEQ form uses convolutional kernel size k, dilation s, nonlinearity ψ, and a linearly transformed input sequence.
Trellis Networks: The TrellisNet transformation uses history padding or zero-padding and an LSTM gated activation.The previous sequence output may supply the history padding when that option is used.
Weight-Tied Transformers: The weight-tied transformer applies multi-head self-attention with input injection, positional embeddings, layer normalization, and a positionwise feedforward block.The implementation described includes a memory-augmented transformer.
Comparison: Figure 1 compares conventional weight-tied deep networks with DEQ and highlights the latter’s constant memory requirements.The comparison concerns memory storage needed at training time.

5 Experiments

Experiments show that DEQ models retain long-range information and achieve competitive or better language-modeling performance while substantially reducing memory use. Their runtime depends on Broyden iterations, whose convergence and accuracy create practical trade-offs.

DEQ retains memory over sequences of length T = 400, substantially outperforming LSTM/GRU architectures on the copy-memory task.
Penn Treebank: DEQ-TrellisNet matches the original deeply supervised TrellisNet on Penn Treebank test perplexity when trained from scratch.
WikiText-103: DEQ-TrellisNet achieves better test perplexity than the original deep TrellisNet on the roughly 100x larger WikiText-103 corpus.
WikiText-103: DEQ models achieve competitive performance with same-size state-of-the-art models on WikiText-103 while consuming much less training memory.
Memory footprint: Over 80% and up to 88% of model memory consumption is reduced by DEQ, excluding word embeddings; reductions exceed 60% against non-weight-tied baselines.
Memory footprint: DEQ training memory is O(1), whereas conventional L-layer networks require O(L) memory for stored intermediate activations.
Convergence and runtime: Higher tolerance ε or lower Broyden iteration limits can accelerate DEQ, but inaccurate equilibrium estimates can reduce performance.
Convergence and runtime: Broyden iterations generally increase across training epochs, while backward passes require fewer iterations than forward passes and DEQs converge more efficiently than weight-tied transformers.

6 Conclusion

The conclusion presents DEQ as a constant-memory approach that replaces depth with sequence-level equilibrium computation. Experiments support its temporal memory retention, scalability, and competitive or slightly superior performance on sequence tasks.

DEQ models temporal data by directly solving for sequence-level fixed points and optimizing the resulting equilibrium representations.
DEQ requires O(1) training memory, is agnostic to the forward root solver, and supports substantially different architectural choices.
Experiments report good temporal memory retention, scalability to realistic large-scale sequence tasks, and competitive or slightly better performance than state-of-the-art methods.

A Backward Pass of the Deep Equilibrium Model

The backward pass differentiates the equilibrium condition implicitly rather than storing the trajectory used to reach the fixed point. This yields an analytical loss gradient through the equilibrium state.

The gradient formulation uses the inverse Jacobian of gθ evaluated at the equilibrium state.
The equilibrium hidden sequence satisfies fθ(z⋆_1:T; x_1:T) = z⋆_1:T.
Implicit differentiation of the equilibrium condition provides the loss gradient with respect to model parameters or inputs.

B Sufficiency of a Single DEQ “Layer”

The universality result argues that a single DEQ module can represent the behavior of multiple stacked DEQ modules by enlarging the equilibrium state and transformation.

The paper formalizes universality for a single-layer DEQ using two stable transformations and their parameter sets.
The construction combines the first module’s equilibrium with additional features into an augmented state for one transformation.
The theorem is presented as showing that exactly one DEQ layer is sufficient in the stated construction.

C Universality of Weight-tied, Input-injected Networks

Weight-tied, input-injected networks can represent traditional deep networks without reducing representational capacity. The construction preserves depth while increasing hidden width only linearly.

C Universality of Weight-tied, Input-injected Networks: DEQ's weight-tied, input-injected restriction does not reduce representational capacity mathematically.Any traditional deep network can be represented in this form.
C Universality of Weight-tied, Input-injected Networks: A traditional layer applies depth-specific weights, biases, and nonlinearities to hidden features, while the equivalent network uses constant parameters across layers and injects the input.The equivalent update is ˜z[i+1] = σ(Wz˜z[i] + Wxx + ˜b).
C Universality of Weight-tied, Input-injected Networks: The constructive proof builds shared parameters using a shift operation that places the original layer transformations into a larger hidden representation.After repeated applications, the shifted hidden vector reproduces the original network's terms.
C Universality of Weight-tied, Input-injected Networks: The equivalent weight-tied network uses the original depth and a hidden size equal to the sum of the original hidden-unit sizes.Thus the width increase is linear in the original network's layer widths.

D Empirical Convergence of Weight-tied Deep Nets

Experiments show that trained weight-tied TrellisNets and universal transformers tend toward fixed points as depth increases. Transformers may oscillate around the fixed point, which root-finding can locate more easily.

D Empirical Convergence of Weight-tied Deep Nets: Trained weight-tied TrellisNets and universal transformers show activation convergence across sequence lengths T = 100, 200, 400, and 800.The observations use randomly selected test samples.
D Empirical Convergence of Weight-tied Deep Nets: Transformers usually exhibit oscillatory stacked iterations, with activations fluctuating around the actual fixed point.The paper reports this behavior in Figure 4.
D Empirical Convergence of Weight-tied Deep Nets: Newton or quasi-Newton methods empirically find the transformer fixed point more easily than continued stacked iterations.This motivates solving directly for the equilibrium.
D Empirical Convergence of Weight-tied Deep Nets: In practice, computation limits usually restrict models to a predetermined depth, such as 60 layers, rather than the deeper inference depths analyzed.For transformers, auxiliary losses have also stabilized training of 64-layer models.

E More Remarks on DEQ

The paper discusses implementation constraints and practical techniques for training DEQs. These include handling unavailable layer-wise regularization, minibatch convergence imbalance, shallow-network warmup, and very long sequences.

E More Remarks on DEQ: Auxiliary losses are difficult to apply to DEQs because the architecture has no explicit layers.Layer normalization and weight normalization remain compatible with the DEQ function.
E More Remarks on DEQ: Minibatch sequences may require different numbers of equilibrium iterations, causing faster-converging sequences to wait and sometimes reducing GPU utilization.The paper finds this imbalance relatively small empirically.
E More Remarks on DEQ: Pretraining a shallow weight-tied stack can accelerate DEQ training and provide a memory- and computation-efficient initialization.A two-layer stack is given as an example, although shallow models plateau at lower accuracy.
E More Remarks on DEQ: For extremely long sequences T > 1000, the forward fixed point can be difficult to solve accurately, so the paper suggests splitting sequences into subsequences.The subsequence technique can also be used in the backward pass.

F Task Descriptions

The paper evaluates DEQs on synthetic memory retention and language-modeling benchmarks. These datasets vary in corpus size, vocabulary, sequence realism, and long-range memory demands.

F Task Descriptions: The evaluation uses three sequence-prediction tasks or datasets: copy memory, Penn Treebank, and WikiText-103.These are introduced as the evaluation settings for the DEQ approach.
F Task Descriptions: The copy memory task tests whether a model can recall ten initial values and reproduce them after a delimiter and a long intervening gap.The dataset contains 20K training samples and 2K testing samples.
F Task Descriptions: Penn Treebank is a comparatively small corpus with about 888K training words and a 10,000-word vocabulary.Prior work reports that strong regularization is needed for best results on this corpus.
F Task Descriptions: WikiText-103 is about 110 times larger than Penn Treebank, has a vocabulary exceeding 260K, and preserves rare words, punctuation, numbers, and capitalization.It evaluates scaling to long sequences from a large vocabulary.