Deep Equilibrium Models

Abstract

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 .

NeurIPS 2019 Spotlight Oral

Cited by in corpus (30)

• DC3: A learning method for optimization with hard constraints
• Enforcing Policy Feasibility Constraints through Differentiable Projection for Energy Optimization
• Continuous-in-Depth Neural Networks
• Lipschitz Bounded Equilibrium Networks
• On the Iteration Complexity of Hypergradient Computation
• Implicit Feature Pyramid Network for Object Detection
• M6-10T: A Sharing-Delinking Paradigm for Efficient Multi-Trillion Parameter Pretraining
• Capsules with Inverted Dot-Product Attention Routing
• On the Theory of Implicit Deep Learning: Global Convergence with Implicit Layers
• Almost Surely Stable Deep Dynamics
• To The Point: Correspondence-driven monocular 3D category reconstruction
• Fighting Gradients with Gradients: Dynamic Defenses against Adversarial Attacks
• Not All Memories are Created Equal: Learning to Forget by Expiring
• Implicit Normalizing Flows
• A Recipe for Global Convergence Guarantee in Deep Neural Networks
• A composable autoencoder-based iterative algorithm for accelerating numerical simulations
• $\xi$-torch: differentiable scientific computing library
• Semialgebraic Representation of Monotone Deep Equilibrium Models and Applications to Certification
• Speeding up Deep Model Training by Sharing Weights and Then Unsharing
• Can You Learn an Algorithm? Generalizing from Easy to Hard Problems with Recurrent Networks
• Multi-scale Transformer Language Models
• Implicit Behavioral Cloning
• Differentiable Implicit Layers
• Optimization Induced Equilibrium Networks
• On the Regularity of Attention
• Combinatorial Optimization for Panoptic Segmentation: A Fully Differentiable Approach
• Meta-learning Transferable Representations with a Single Target Domain
• Differentiable Multiple Shooting Layers
• Approximate Fixed-Points in Recurrent Neural Networks
• Recurrent Stacking of Layers in Neural Networks: An Application to Neural Machine Translation