Timezone: »

 
Poster
Stabilizing Equilibrium Models by Jacobian Regularization
Shaojie Bai · Vladlen Koltun · Zico Kolter

Tue Jul 20 09:00 PM -- 11:00 PM (PDT) @
in Poster Session 2 »

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.

Keywords: Deep Learning

Author Information

Shaojie Bai (Carnegie Mellon University)
Vladlen Koltun (Intel Labs)
Zico Kolter (Carnegie Mellon University / Bosch Center for AI)

Related Events (a corresponding poster, oral, or spotlight)

More from the Same Authors