Deep Equilibrium Models

By Shaojie Bai, J. Kolter, Vladlen Koltun. 2019

942 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

Citations in the corpus

ALBERT: A Lite BERT for Self-supervised Learning of Language Representations (2019, methodology): “Very recently, Bai et al. (2019) propose a Deep Equilibrium Model (DQE) for transformer networks and show that DQE can reach an equilibrium point for which the input embedding and the output embedding of a certain layer stay the same.”
Multiset-Equivariant Set Prediction with Approximate Implicit Differentiation (2021, background): “Recent works in deep learning have successfully applied implicit differentiation to large-scale language modelling (Bai et al., 2019), semantic segmentation (Bai et al., 2020), meta-learning (Rajeswaran et al., 2019), neural architecture search (Zhang et al., 2021b), and more (Agrawal et al., 2019).”
Of Non-Linearity and Commutativity in BERT (2021, background): “Recently there has been a trend towards transformers that feature recurrence across the layer-dimension [1], [5], [17].”
Implicit Behavioral Cloning (2021, background): “Many other works have investigated using the notion of implicit functions in 225 learning, including works that investigate implicit layers [39, 40, 41, 42].”
DynaBERT: Dynamic BERT with Adaptive Width and Depth (2020, background): “k. Universal Transformer [7] shares parameters across layers. Compared to the standard Transformer, it gets better performance on language modeling and subject-verb agreement. Deep Equilibrium Models [1] reach an equilibrium point where the input and output of a certain layer stay the same. ALBERT [”
Implicit Deep Learning (2019, background): “icit deep learning. Recent works have considered versions of implicit models, and demonstrated their potential in deep learning. In particular, recent ground-breaking work by Kolter and collaborators [8,29] demonstrated success of an entirely implicit framework, which they call Deep Equilibrium Mode”
Can You Learn an Algorithm? Generalizing from Easy to Hard Problems with Recurrent Networks (2021, methodology): “Also, the basic neural architectures we use are not new and build upon prior studies of weight sharing and recurrence [Pinheiro and Collobert, 2014, Liang and Hu, 2015, Alom et al., 2018, Bai et al., 2018, 2019, Lan et al., 2020, Jaegle et al., 2021].”
Parameter Efficient Multimodal Transformers for Video Representation Learning (2020, background): “Trellis network (Bai et al., 2019b) is a temporal convolutional architecture with weight-tying across time and depth.”
M6-10T: A Sharing-Delinking Paradigm for Efficient Multi-Trillion Parameter Pretraining (2021, background): “[1] both illustrated that it is effective for vanilla transformer with encoder-decoder work.”
JFB: Jacobian-Free Backpropagation for Implicit Networks (2021, methodology): “Moreover, JFB is competitive with Multiscale DEQs (Bai, Koltun, and Kolter 2020) despite having less than a tenth as many parameters.”
Training Graph Neural Networks with 1000 Layers (2021, methodology): “Similar to DEQ (Bai et al., 2019), the forward pass of DEQ-GNN is implemented with a root-finding algorithm (e.”
Training Large Language Models to Reason in a Continuous Latent Space (2024, positive): “Recently, looped transformers (Giannou et al., 2023; Fan et al., 2024) have been proposed to solve algorithmic tasks, which have some similarities to the computing process of continuous thoughts, but we focus on common reasoning tasks and aim at investigating latent reasoning in comparison to…”
Implicit Normalizing Flows (2021, methodology): “We use the similar technique in Bai et al. (2019) to compute ∂ L ∂ z J − 1 G ( z ) : solving the linear system of variable y : where the left hand side is a vector-Jacobian product and it can be efﬁciently computed by autograd packages foy any y without computing the Jacobian matrix.”
Gradients are Not All You Need (2021, background): “…that we are trying to optimize is known and can be phrased in terms of the ﬁxed point of some dynamical system, there can often be signiﬁcant beneﬁt to computing derivatives implicitly rather than unrolling optimization explicitly Rajeswaran et al. [2019], Bai et al. [2019], Blondel et al. [2021].”
Scale down Transformer by Grouping Features for a Lightweight Character-level Language Model (2020, background): “To compensate for the losses, knowledge distillation (Sun et al., 2019) tries to optimize a scale-downed model with a large teacher network, and weight-sharing (Bai et al., 2019) stacks a uniﬁed layer multiple times.”
Efficient and Modular Implicit Differentiation (2021, background): “Another line of applications is “optimization as a layer” [42, 6, 51, 27, 36, 7], which usually includes regularization or constraints in the optimization problem in order to impose desirable structure on the layer output.”
GFlowNet Foundations (2021, neutral): “, for implicit layers in fixed-point iteration layers and Deep Equilibrium Models (Bai et al., 2019).”
Multiscale Deep Equilibrium Models (2020, methodology): “[5] for the proof, which is based on the implicit function theorem [30].”
Optimizing Millions of Hyperparameters by Implicit Differentiation (2019, methodology): “In DEQ [31], implicit diﬀerentiation is used to add diﬀerentiable ﬁxed-point methods into NN architectures.”