[Paper Review] Deep Equilibrium Models

Shaojie Bai, J. Zico Kolter+1 more|arXiv (Cornell University)|Sep 3, 2019

Topic ModelingComputer Science245 citations

TL;DR

Introduces the Deep Equilibrium Model (DEQ), which finds fixed points of a weight-tied nonlinear transformation to model sequences with constant memory, and demonstrates competitive performance on language modeling with substantial memory savings.

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 .

Motivation & Objective

• Motivate a memory-efficient approach to deep sequence modeling by seeking equilibrium points of weight-tied transformations.
• Propose a general DEQ framework that directly solves for sequence-level fixed points and differentiates through them.
• Demonstrate instantiations of DEQ with TrellisNet and weight-tied transformers on large-scale language tasks.

Proposed method

• Define the equilibrium condition z* = fθ(z*; x) and solve for z* with black-box root-finding methods.
• Backpropagate through the equilibrium using implicit differentiation via Theorem 1, avoiding storage of intermediate activations.
• Accelerate training with Broyden’s quasi-Newton updates to approximate the inverse Jacobian in forward and backward passes.
• Show universality: stacking multiple DEQs does not exceed the representational power of a single DEQ (Theorem 2).
• Instantiate DEQ with TrellisNet (weight-tied TCN) and memory-augmented self-attention transformers to cover major sequence model families.

Experimental results

Research questions

• RQ1 Can a fixed-point formulation of a weight-tied deep sequence model match or exceed the performance of layer-stacked architectures on language modeling tasks?
• RQ2 Does implicit differentiation through the equilibrium enable constant-memory backpropagation for DEQ models?
• RQ3 What are practical instantiations of DEQ for TrellisNet and transformers, and how do they fare on benchmarks like WikiText-103 and PTB?
• RQ4 How does the memory footprint and training efficiency of DEQ compare to traditional deep networks and gradient checkpointing?
• RQ5 Is stacking multiple DEQs necessary or beneficial for additional representational power?

Key findings

• DEQ can achieve competitive or better perplexities on PTB and WikiText-103 compared to state-of-the-art models of similar size.
• DEQ provides over 80% memory savings (up to 88%) during training relative to layer-based equivalents, and can outperform gradient checkpointing in memory usage.
• Backward gradients are computed via an implicit fixed-point formulation, enabling constant-memory backpropagation.
• Two instantiations, DEQ-TrellisNet and DEQ-Transformer, demonstrate compatibility with convnet, RNN, and transformer architectures while using similar or fewer resources.
• Forward and backward passes converge to equilibrium with quasi-Newton methods, and runtime is modestly higher than fixed-depth models but memory efficiency is substantially improved.