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 (90)

• ALBERT: A Lite BERT for Self-supervised Learning of Language Representations
• You Only Learn One Representation: Unified Network for Multiple Tasks
• A Survey of Contextual Optimization Methods for Decision Making under Uncertainty
• RAFT: Recurrent All-Pairs Field Transforms for Optical Flow
• DynaBERT: Dynamic BERT with Adaptive Width and Depth
• An Adaptive and Robust Deep Learning Framework for THz Ultra-Massive MIMO Channel Estimation
• Pretrained Transformers as Universal Computation Engines
• Dissecting Neural ODEs
• DC3: A learning method for optimization with hard constraints
• Implicit Graph Neural Networks
• Multiscale Deep Equilibrium Models
• Nonlocal Kernel Network (NKN): a Stable and Resolution-Independent Deep Neural Network
• Deep Idempotent Network for Efficient Single Image Blind Deblurring
• Enforcing Policy Feasibility Constraints through Differentiable Projection for Energy Optimization
• Efficient and Modular Implicit Differentiation
• Parameter Efficient Multimodal Transformers for Video Representation Learning
• Addressing Some Limitations of Transformers with Feedback Memory
• Bilevel methods for image reconstruction
• Continuous-in-Depth Neural Networks
• Lipschitz Bounded Equilibrium Networks
• On the Iteration Complexity of Hypergradient Computation
• Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey
• Implicit Feature Pyramid Network for Object Detection
• Hypersolvers: Toward Fast Continuous-Depth Models
• M6-10T: A Sharing-Delinking Paradigm for Efficient Multi-Trillion Parameter Pretraining
• Capsules with Inverted Dot-Product Attention Routing
• Nonsmooth Implicit Differentiation for Machine Learning and Optimization
• Model-based Adversarial Meta-Reinforcement Learning
• Sinkformers: Transformers with Doubly Stochastic Attention
• On the Theory of Implicit Deep Learning: Global Convergence with Implicit Layers
• SPIDE: A Purely Spike-based Method for Training Feedback Spiking Neural Networks
• On Training Implicit Models
• Monotone operator equilibrium networks
• Synthesis of Stabilizing Recurrent Equilibrium Network Controllers
• Memory Efficient Model Based Deep Learning Reconstructions for High Spatial Resolution 3D Non-Cartesian Acquisitions
• Fighting Gradients with Gradients: Dynamic Defenses against Adversarial Attacks
• To The Point: Correspondence-driven monocular 3D category reconstruction
• Almost Surely Stable Deep Dynamics
• Neural SDEs as Infinite-Dimensional GANs
• Training Feedback Spiking Neural Networks by Implicit Differentiation on the Equilibrium State
• Solving Structured Hierarchical Games Using Differential Backward Induction
• SITReg: Multi-resolution architecture for symmetric, inverse consistent, and topology preserving image registration
• A Shooting Formulation of Deep Learning
• Implicit Normalizing Flows
• A global convergence theory for deep ReLU implicit networks via over-parameterization
• Not All Memories are Created Equal: Learning to Forget by Expiring
• Deep Implicit Optimization enables Robust Learnable Features for Deformable Image Registration
• Semialgebraic Representation of Monotone Deep Equilibrium Models and Applications to Certification
• A Recipe for Global Convergence Guarantee in Deep Neural Networks
• Opening the Blackbox: Accelerating Neural Differential Equations by Regularizing Internal Solver Heuristics
• Operator Splitting for Learning to Predict Equilibria in Convex Games
• $\xi$-torch: differentiable scientific computing library
• Learning Graph Structure With A Finite-State Automaton Layer
• A composable autoencoder-based iterative algorithm for accelerating numerical simulations
• Adversarially Robust Learning for Security-Constrained Optimal Power Flow
• Multi-scale Transformer Language Models
• Bi-level Score Matching for Learning Energy-based Latent Variable Models
• Speeding up Deep Model Training by Sharing Weights and Then Unsharing
• Feasibility-based Fixed Point Networks
• Implicit Behavioral Cloning
• Can You Learn an Algorithm? Generalizing from Easy to Hard Problems with Recurrent Networks
• On the Regularity of Attention
• Differentiable Implicit Layers
• Optimization Induced Equilibrium Networks
• Scaling up machine learning-based chemical plant simulation: A method for fine-tuning a model to induce stable fixed points
• Multiset-Equivariant Set Prediction with Approximate Implicit Differentiation
• Understanding Deep Architectures with Reasoning Layer
• Differentiable Forward and Backward Fixed-Point Iteration Layers
• Convergent Graph Solvers
• The Uncanny Similarity of Recurrence and Depth
• Learning MDPs from Features: Predict-Then-Optimize for Sequential Decision Problems by Reinforcement Learning
• Gradients are Not All You Need
• Neural Parameter Allocation Search
• ${\ell }_1$DecNet+: A new architecture framework by ${\ell }_1$ decomposition and iteration unfolding for sparse feature segmentation
• One-step differentiation of iterative algorithms
• Meta-learning Transferable Representations with a Single Target Domain
• Combinatorial Optimization for Panoptic Segmentation: A Fully Differentiable Approach
• Approximate Fixed-Points in Recurrent Neural Networks
• Differentiable Multiple Shooting Layers
• Robust Implicit Networks via Non-Euclidean Contractions
• On the complexity of nonsmooth automatic differentiation
• Subformer: Exploring Weight Sharing for Parameter Efficiency in Generative Transformers
• Recurrent networks improve neural response prediction and provide insights into underlying cortical circuits
• A nonlinear diffusion method for semi-supervised learning on hypergraphs
• Lightweight, Dynamic Graph Convolutional Networks for AMR-to-Text Generation
• Implicit vs Unfolded Graph Neural Networks
• A Survey on Green Deep Learning
• Input Convex Gradient Networks
• Recurrent Stacking of Layers in Neural Networks: An Application to Neural Machine Translation
• Youla-REN: Learning Nonlinear Feedback Policies with Robust Stability Guarantees