Lipschitz Multiscale Deep Equilibrium Models: A Theoretically Guaranteed and Accelerated Approach

TL;DR

This work tackles the slow training and inference of Deep Equilibrium Models by ensuring fixed-point convergence through a contractive forward mapping. It introduces Lipschitz MDEQ, a multiscale DEQ architecture whose components—Mean-Only Group Normalization, Scaled-ReLU, spectral-norm constrained convolutions, convex residuals, and adaptive fusion—collectively bound the Lipschitz constant $L$ below 1, guaranteeing convergence for both forward and backward passes. Theoretical analysis decomposes the overall contraction bound into modular terms and identifies the hyperparameters that most influence $L$, while extensive experiments on CIFAR-10 (and ImageNet) demonstrate up to $4.75\times$ speedups with manageable accuracy loss, along with insightful ablations. The approach paves the way for fast, memory-efficient implicit models suitable for high-resolution vision tasks, balancing speed and accuracy through principled control of fixed-point dynamics.

Abstract

Deep equilibrium models (DEQs) achieve infinitely deep network representations without stacking layers by exploring fixed points of layer transformations in neural networks. Such models constitute an innovative approach that achieves performance comparable to state-of-the-art methods in many large-scale numerical experiments, despite requiring significantly less memory. However, DEQs face the challenge of requiring vastly more computational time for training and inference than conventional methods, as they repeatedly perform fixed-point iterations with no convergence guarantee upon each input. Therefore, this study explored an approach to improve fixed-point convergence and consequently reduce computational time by restructuring the model architecture to guarantee fixed-point convergence. Our proposed approach for image classification, Lipschitz multiscale DEQ, has theoretically guaranteed fixed-point convergence for both forward and backward passes by hyperparameter adjustment, achieving up to a 4.75$\times$ speed-up in numerical experiments on CIFAR-10 at the cost of a minor drop in accuracy.

Figures (5)

• Figure 1: Architecture of MDEQ and Lipschitz MDEQ. Lipschitz MDEQ inherits the general structure of MDEQ, with each operation modified to prevent the Lipschitz constant from becoming too large. Note that the layout and notation of the figure are taken from the prior work Bai2020Mul with only minor modifications made for clarity in explaining our results. The sample image is from the AFHQ dataset Choi2020Sta.
• Figure 2: Effect of hyperparameters constituting $f_{\bm{\theta}}$ on Lipschitz bound $L$. Each graph is plotted by varying only one parameter while holding all others constant. Fixed parameters are set as $\alpha_1=0.5, \alpha_2=0.3, L_{\text{Conv}^\star}=1.5, L_{\text{MGN}}=1.0, L_{\text{SReLU}}=0.4, n=4$, and $p=0.3$. Note that changing the values of the fixed parameters will also change the plot.
• Figure 3: Comparison of fixed-point convergence in forward and backward passes of training.
• Figure 4: Accuracy and speed trade-off in Lipschitz MDEQ. The region where convergence is guaranteed, $L < 1$, is shaded green.
• Figure 5: Comparison of fixed-point convergence in the forward pass of testing.

Theorems & Definitions (18)

• Definition 3.1: $L$-Lipschitz mapping
• Lemma 3.1
• Lemma B.1
• proof
• Lemma B.2
• proof
• Lemma B.3
• proof
• Lemma B.4
• proof