Stabilizing Equilibrium Models by Jacobian Regularization
This addresses the brittleness and inefficiency of DEQs for researchers and practitioners in deep learning, offering a more practical alternative to traditional deep networks.
The paper tackles the instability and slowness of deep equilibrium networks (DEQs) by proposing a Jacobian regularization scheme, which stabilizes fixed-point convergence and enables implicit-depth models to achieve similar speed and performance as ResNet-101 while maintaining constant memory usage.
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 .