[Paper Review] Multiscale Deep Equilibrium Models

Shaojie Bai, Vladlen Koltun+1 more|arXiv (Cornell University)|Jun 15, 2020

Markov Chains and Monte Carlo Methods63 references53 citations

TL;DR

MDEQ introduces a shallow, implicit, multiscale equilibrium model that solves for synchronized equilibria across multiple resolutions, enabling competitive ImageNet classification and Cityscapes segmentation with constant memory during training.

ABSTRACT

We propose a new class of implicit networks, the multiscale deep equilibrium model (MDEQ), suited to large-scale and highly hierarchical pattern recognition domains. An MDEQ directly solves for and backpropagates through the equilibrium points of multiple feature resolutions simultaneously, using implicit differentiation to avoid storing intermediate states (and thus requiring only $O(1)$ memory consumption). These simultaneously-learned multi-resolution features allow us to train a single model on a diverse set of tasks and loss functions, such as using a single MDEQ to perform both image classification and semantic segmentation. We illustrate the effectiveness of this approach on two large-scale vision tasks: ImageNet classification and semantic segmentation on high-resolution images from the Cityscapes dataset. In both settings, MDEQs are able to match or exceed the performance of recent competitive computer vision models: the first time such performance and scale have been achieved by an implicit deep learning approach. The code and pre-trained models are at https://github.com/locuslab/mdeq .

Motivation & Objective

• Motivate a model class that can handle multiscale structure in vision without explicit deep stacks.
• Develop an implicit, single-stage architecture that maintains multiple resolutions in equilibrium simultaneously.
• Enable joint training for multiple tasks (e.g., classification and segmentation) via auxiliary losses at different scales.
• Demonstrate scalability of implicit models to large-scale vision tasks while managing memory efficiently.
• Investigate practical training dynamics and compatibility of common DL techniques with MDEQ.

Proposed method

• Define a multiscale transformation f_theta that preserves several resolutions side-by-side and drives them to a joint equilibrium.
• Inject input x only at the highest-resolution stream to couple the flow across scales.
• Use a fixed-point solver (limited-memory Broyden) to find the equilibrium z* across all scales.
• Employ a residual-block based per-scale module with group normalization and a multiscale fusion step that mixes features across resolutions.
• Provide a differentiable backward pass through the equilibrium via Jacobian-vector products (implicit differentiation).
• Allow auxiliary losses at multiple scales to enable multi-task training and transfer (e.g., pretraining on ImageNet and finetuning on segmentation).

Experimental results

Research questions

• RQ1 Can implicit models that forego explicit layer stacks achieve competitive accuracy in tasks with rich multiscale structure like vision?
• RQ2 Can simultaneous multiscale equilibria support both high-level labels and dense predictions in a single model?
• RQ3 How do memory-efficient implicit solvers perform on large-scale vision tasks compared to explicit architectures?
• RQ4 What practical adaptations (normalization, dropout, initialization) are needed to train MDEQs stably on high-resolution data?
• RQ5 To what extent can a single MDEQ model pretrained on one task be transferred to related tasks via auxiliary losses at different scales?

Key findings

• MDEQ matches or exceeds performance of competitive explicit models on ImageNet classification with similar parameter counts.
• On Cityscapes, MDEQ achieves high mIoU, with larger models approaching state-of-the-art performance while using less memory than explicit backbones.
• A small MDEQ (18M) outperforms several explicit models with similar size on ImageNet and CIFAR-10, illustrating effective multiscale implicit modeling.
• Training with limited-memory Broyden’s method maintains memory efficiency and makes training feasible at megapixel scales.
• MDEQ demonstrates that shallow implicit models can scale to large vision tasks and provide competitive accuracy without extensive layer-wise depth.