Stabilizing Equilibrium Models by Jacobian Regularization

TL;DR

This paper introduces a Jacobian regularization technique for deep equilibrium networks (DEQs) that stabilizes training and convergence, enabling DEQs to match the performance of traditional deep networks with less memory.

Contribution

The paper proposes a novel Jacobian regularization scheme for DEQs that improves stability and scalability without significant computational overhead.

Findings

Regularization stabilizes fixed-point convergence in DEQs.

Method scales well to high-dimensional tasks like language modeling and image classification.

Achieves comparable performance to ResNet-101 with constant memory footprint.

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 nonlinear 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 at https://github.com/locuslab/deq .