PaperlayerView exact revision on arXiv

Source-linked AI summary

Multiscale Deep Equilibrium Models

Shaojie Bai, Vladlen Koltun, J. Zico Kolter

arXiv:2006.08656v2cs.LGcs.CVstat.ML

TL;DR

Vision models rely on hierarchical multiscale processing, while prior implicit networks do not clearly support such structure. MDEQ simultaneously solves and differentiates through equilibria at multiple resolutions with constant memory, and it matches or exceeds competitive explicit models across ImageNet classification and Cityscapes segmentation.

  • Problem

    Prior implicit networks replace explicit layers and stages, leaving unclear whether they can model the multiscale structure required for competitive high-dimensional vision tasks.

  • Method

    MDEQ maintains multiple feature resolutions side by side and simultaneously solves for their joint equilibrium, using implicit differentiation rather than storing intermediate states.

  • Results

    MDEQs match or exceed competitive explicit vision models across ImageNet classification and Cityscapes segmentation, including 77.5% ImageNet top-1 accuracy and over 80% Cityscapes mIoU.

  • Takeaways & Limitations

    A single shallow implicit model can support multiple vision tasks and scale to practical large-scale computer vision while retaining O(1) memory training.

  • Takeaways & Limitations

    The experiments focus on competitiveness rather than setting a new state of the art on ImageNet or Cityscapes, and note that explicit structure within the equilibrium function remains empirically helpful.

Abstract

from arXiv · show

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 .

1 Introduction

MDEQ addresses the challenge of applying implicit networks to multiscale, high-dimensional vision by maintaining multiple resolutions at equilibrium. It provides constant-memory training and achieves competitive results on ImageNet classification and Cityscapes segmentation.

  • Vision systems commonly use sequential hierarchical feature extractors because images have multiscale structure and extremely high signal dimensionality.A typical image contains millions of pixels that must be processed coherently.
  • Implicit networks offer effectively infinite depth within a constant memory footprint but traditionally lack flexible layers and stages for modeling multiscale structure.This limitation motivates testing whether implicit models can attain competitive accuracy in computer vision.
  • MDEQ simultaneously solves for multiple-resolution equilibria, keeping feature scales side by side in one shallow model rather than processing resolutions successively.The model directly optimizes stable representations across feature scales at the same time.
  • MDEQ uses O(1) training memory and exposes equilibrium features at multiple scales for auxiliary losses and diverse task heads.A single model can support losses defined at different scales and tasks such as classification and segmentation.
  • 77.5% top-1 accuracy is achieved by MDEQs on ImageNet, outperforming baseline ResNets with similar parameter counts.The cited comparison includes ResNet-101 as an example baseline.
  • MDEQs match recent explicit models on Cityscapes segmentation while using much less memory, with the largest model surpassing 80% mIoU.Cityscapes involves dense labeling of 2-megapixel images.

2 Background

The paper places MDEQ within implicit deep learning and multiscale computer vision. It extends equilibrium modeling from single feature streams to synchronized representations at multiple resolutions for realistic visual tasks.

  • Implicit models define analytical conditions for hidden states instead of prescribing explicit computation graphs, enabling black-box forward solvers and analytical backward passes.Examples include Neural ODEs and deep equilibrium models.
  • Neural ODEs model continuous residual dynamics, while DEQs find fixed points of sequence models using black-box root-finding methods.Both represent effectively deep computations without explicit layer stacks.
  • Prior implicit models generally operate on a single feature stream, whereas MDEQ jointly optimizes features for different tasks and dimensionalities.The paper gives image segmentation and classification as examples.
  • Computer vision outputs range from whole-image labels to pixel-level labels, and established models address this range through sequential stages at different resolutions.These stages use combinations of upsampling and downsampling.
  • MDEQ keeps different resolutions side by side in one implicit stage, injecting input at the highest resolution and optimizing all scales simultaneously toward a joint equilibrium.A black-box solver drives the representations to satisfy the equilibrium condition.
  • MDEQ represents an effectively infinitely deep network with constant memory cost while extending implicit modeling to megapixel-level visual tasks.The paper contrasts this scope with earlier ODE-based applications to relatively low-dimensional signals.

3 Multiscale Deep Equilibrium Models

MDEQ extends deep equilibrium modeling to multiple feature resolutions, solving and backpropagating through their equilibria simultaneously rather than processing scales successively. Its multiscale interfaces support different supervision types while retaining equilibrium-based training and one-block memory usage.

  • Multiscale equilibrium formulation: The DEQ foundation replaces infinitely many weight-tied layers with a fixed-point or root-finding problem for the equilibrium state.Repeated application of fθ converges toward z⋆, which can be solved directly with Newton or quasi-Newton methods.
  • Implicit training: Implicit differentiation computes gradients through the equilibrium state without tracing the forward root-finding process.The backward pass uses the Jacobian at the equilibrium and a cheaper vector-Jacobian-product linear solve, while memory matches one block rather than infinitely many.
  • MDEQ architecture: The transformation uses one shallow residual-block layer per resolution, with the precomputed input injected into the highest-resolution stream.Group normalization is used in the residual block, and the model contains only one such residual layer before fusion.
  • MDEQ architecture: MDEQ keeps feature scales side by side and couples them through multi-resolution fusion, producing outputs that mix transformed features from incoming scales.Upsampling and downsampling connect each scale to the others, driving the representations toward a coordinated equilibrium.
  • Multiscale equilibrium formulation: MDEQ maintains feature states at multiple resolutions and solves for their equilibrium points simultaneously.The state is a collection of tensors with different dimensionalities, resolutions, and semantics; each resolution induces its own loss.
  • Supervision and scalability: MDEQ exposes multiple-resolution equilibrium states as interfaces for auxiliary losses, classification, segmentation, joint training, and pretraining followed by fine-tuning.The highest resolution supports dense prediction, while the lowest can collapse spatial dimensions for image-level labeling.

4 Experiments

Experiments evaluate MDEQs on CIFAR-10, ImageNet, and Cityscapes, comparing accuracy, convergence, runtime, and memory with implicit and explicit models. MDEQs achieve competitive or better performance while substantially reducing training memory, though they require more computation time.

  • CIFAR-10: MDEQs improve CIFAR-10 accuracy over comparable NODEs, Augmented NODEs, and single-stream DEQs, while larger models match or exceed ResNet-18 at equal capacity.The small MDEQ improves accuracy by more than 20 percentage points over NODEs and Augmented NODEs, reducing error by more than a factor of 2.
  • ImageNet Classification: An 18M-parameter MDEQ outperforms ResNet-18, ResNet-34, and ResNet-50, while a 64M-parameter MDEQ matches ResNet-101 on ImageNet.These comparisons use top-1 and top-5 accuracy after 100 training epochs.
  • Cityscapes Semantic Segmentation: A 7.8M-parameter MDEQ achieves 75.1 mIoU on Cityscapes, while 53.5M- and 70.9M-parameter models reach 77.8 mIoU and surpass recent explicit models, respectively.The models use the highest-resolution equilibrium output for segmentation after ImageNet pretraining.
  • Runtime and Memory: MDEQ saves more than 60% of GPU memory during training compared with explicit models while maintaining competitive accuracy.Training a large MDEQ on ImageNet uses about 6GB of memory, primarily for Broyden’s method.
  • Runtime and Memory: MDEQ is generally slower than explicit networks, with a 2.7× slowdown relative to ResNet-101, but a similarly sized MDEQ is 3× faster than ANODEs with a 3× error reduction.Maintaining features at all resolutions contributes to the slowdown relative to progressively downsampling explicit networks.
  • Convergence Analysis: Higher-resolution inputs require more equilibrium-solving steps, while limited-memory Broyden’s method and multiscale fusion help stabilize convergence.The convergence comparison spans CIFAR-10, ImageNet, and Cityscapes and measures residual change across function evaluations.

5 Conclusion

The paper concludes that MDEQs extend implicit modeling to high-dimensional, multiscale vision tasks by synchronizing equilibria across resolutions. A single shallow implicit model can support classification and segmentation with performance competitive with deeply stacked explicit architectures.

  • Conclusion: MDEQ synchronizes equilibrium solving and backpropagation across multiple feature representations at different resolutions.This distinguishes MDEQ from prior implicit models designed without explicit multiscale equilibrium structure.
  • Conclusion: A single MDEQ supports different tasks, including image classification and semantic segmentation.The same model family is used across ImageNet classification and Cityscapes segmentation.
  • Conclusion: Experiments show that shallow implicit models can scale to practical computer vision tasks and match explicit architectures using sequentially stacked layers.The conclusion frames this as evidence that implicit modeling can operate at realistic visual scale.

A Task Descriptions and Training Settings

The paper evaluates MDEQ across CIFAR-10, ImageNet, and Cityscapes using task-specific training and evaluation protocols. ImageNet-trained MDEQs are reused for Cityscapes segmentation, with shared model configurations across tasks.

  • Image Classification on CIFAR-10: CIFAR-10 contains 60,000 32 × 32 color images across 10 classes, split into 50K training and 10K testing images.The experiments use both unaugmented and augmented training settings.
  • Image Classification on ImageNet: ImageNet provides 1.2 million labeled training images across 1,000 classes, with 224 × 224 crops used as model inputs.MDEQs are trained for 100 epochs and later serve as backbones for Cityscapes fine-tuning.
  • Semantic Segmentation on Cityscapes: Cityscapes contains 5,000 pixel-annotated urban images from 50 cities, evaluated with 19-way per-pixel classification.The dataset is divided into 2,975 training, 500 validation, and 1,525 test images.
  • Semantic Segmentation on Cityscapes: Cityscapes training uses random 1024 × 512 crops and horizontal flips, while validation uses original 2048 × 1024 resolution at a single scale without flipping.The identical MDEQ models from ImageNet are used with the high-resolution prediction head.
  • Shared Training Settings: A single MDEQ configuration is shared between ImageNet classification and Cityscapes segmentation.Cityscapes models are initialized from pretrained ImageNet MDEQs.
  • Hardware: MDEQ-Large models use 4 RTX-2080 Ti GPUs, MDEQ-XL models use 8 Quadro RTX 8000 GPUs, and CIFAR-10 models use 1 GPU.These hardware settings apply to the reported experiments.

B Equilibrium Solving and Convergence Analysis

The equilibrium-solving analysis introduces limited-memory Broyden’s method as the root solver used for MDEQ computations. The method is part of the paper’s approach to obtaining equilibrium states and enabling implicit differentiation.

  • Equilibrium Solving: Limited-memory Broyden’s method is used to perform root-solving for MDEQ equilibrium states.The broader equilibrium framework uses quasi-Newton methods to find roots and differentiates through the resulting equilibrium.

B.1 (Limited-memory) Broyden’s Method

The section explains why direct Newton updates are impractical for large implicit vision models and motivates a limited-memory Broyden solver using low-rank updates.

  • Challenges: Newton’s method is impractical because realistic network Jacobians are prohibitively large to store and expensive to invert.A 32 × 32 × 80 tensor induces an 81920 × 81920 Jacobian requiring 25GB of memory, while inversion has cubic complexity.
  • Broyden’s Method: MDEQ therefore uses a variant of Broyden’s method to approximate the inverse Jacobian during equilibrium solving.The method uses an adjustable step size and a low-rank approximation to J−1.
  • Low-rank Representation: The Broyden matrix need not be formed explicitly because it can be represented as a sum of low-rank updates.The updates are derived using the Sherman-Morrison formula.
  • Limited Memory: The limited-memory solver stores at most m low-rank updates and discards the oldest update when that storage limit is reached.The matrix is initialized as B[0] = −I, while update sets U and V retain only the most recent updates.

B.2 Discussions

MDEQ runtime depends on root-finding iterations, and multiscale fusion with Broyden’s method supports stable convergence while higher-resolution inputs require more iterations.

  • Runtime: MDEQ runtime is assessed by root-finding iterations because each Broyden iteration evaluates fθ once.Iterations may stop at a threshold such as 22 before reaching the exact equilibrium, yet the resulting estimates can still support competitive training.
  • High-resolution Inputs: Higher-resolution inputs make equilibrium solving harder because the Jacobian grows and small convolutional receptive fields require broader effective coverage.More low-rank updates are needed for the larger Jacobian, while repeated local convolutions may need enough iterations to cover the entire image.
  • High-resolution Inputs: Over 100 steps are used on Cityscapes images, whereas MDEQ typically reaches equilibrium with good accuracy within 30 steps on CIFAR-10.The number of required Broyden iterations increases as input images become larger.
  • Solver Comparison: Broyden’s method converges more stably and efficiently than simply iterating fθ, which often converges poorly or not at all.This pattern holds across CIFAR-10, ImageNet, and Cityscapes convergence experiments.
  • Convergence Mechanism: Simultaneous multiscale fusion effectively stabilizes equilibrium convergence.This observation is reported alongside the cross-resolution convergence behavior in MDEQ.
  • Multiscale Convergence: All MDEQ resolution streams converge in parallel, with lower-resolution streams converging faster than higher-resolution streams.The comparison is made on CIFAR-10 against a single-stream DEQ maintaining only the highest-resolution stream.
  • Multiscale Convergence: High-resolution convergence is much faster in the multiscale setting than in the single-stream setting.Figure 7 identifies the multiscale high-resolution curve as pink and the single-stream curve as orange.
  • Convergence Mechanism: Broyden’s global location mixing and multiscale up- and downsampling are hypothesized to expand the receptive field faster than simply stacking local transformations.The proposed explanation links solver behavior and multiscale fusion to broader feature interaction on high-resolution streams.

C Qualitative Segmentation Results on Cityscapes

The section presents qualitative examples of MDEQ-large semantic segmentation outputs on Cityscapes validation images.

  • Qualitative Results: Figure 8 shows examples of MDEQ-large segmentation results on Cityscapes validation images.The displayed images have resolution 2048 × 1024.
Loading 2006.08656v2…