Poster

Stabilizing Equilibrium Models by Jacobian Regularization

Shaojie Bai · Vladlen Koltun · Zico Kolter

Keywords: [ Deep Learning ]

[ Abstract ]
[ Paper ] [ Visit Poster at Spot C4 in Virtual World ]
Tue 20 Jul 9 p.m. PDT — 11 p.m. PDT
 
Spotlight presentation: Deep Learning 2
Tue 20 Jul 5 p.m. PDT — 6 p.m. PDT

Abstract:

Deep equilibrium networks (DEQs) are a new class of models that eschews traditional depth in favor of finding the fixed point of a single non-linear layer. These models have been shown to achieve performance competitive with the state-of-the-art deep networks while using significantly less memory. Yet they are also slower, brittle to architectural choices, and introduce potential instability to the model. In this paper, we propose a regularization scheme for DEQ models that explicitly regularizes the Jacobian of the fixed-point update equations to stabilize the learning of equilibrium models. We show that this regularization adds only minimal computational cost, significantly stabilizes the fixed-point convergence in both forward and backward passes, and scales well to high-dimensional, realistic domains (e.g., WikiText-103 language modeling and ImageNet classification). Using this method, we demonstrate, for the first time, an implicit-depth model that runs with approximately the same speed and level of performance as popular conventional deep networks such as ResNet-101, while still maintaining the constant memory footprint and architectural simplicity of DEQs. Code is available https://github.com/locuslab/deq.

Chat is not available.