PaperlayerView exact revision on arXiv

Source-linked AI summary

Vision-LSTM: xLSTM as Generic Vision Backbone

Benedikt Alkin, Maximilian Beck, Korbinian Pöppel, Sepp Hochreiter, Johannes Brandstetter

arXiv:2406.04303v3cs.CVcs.AIcs.LG

TL;DR

High-resolution vision tasks motivate alternatives to architectures whose self-attention cost grows quadratically with sequence length. This paper introduces ViL, which adapts xLSTM using alternating mLSTM blocks, and reports that it outperforms several vision architectures across three benchmarks.

  • Problem

    For high-resolution vision tasks, ViT’s self-attention has quadratic computational complexity, while ViL is presented as a linear-complexity alternative.

    Source: §1 · p. 2

  • Method

    ViL adapts xLSTM to vision by processing image patch tokens with alternating mLSTM blocks that reverse traversal direction on even blocks.

    Source: §1 · p. 2, §2.2 · p. 4

  • Results

    ViL outperforms SSM-based vision architectures, other autoregressive vision architectures, and optimized ViT models on ImageNet-1K classification, VTAB-1K transfer classification, and ADE20K semantic segmentation.

    Source: §7 · p. 10

  • Takeaways & Limitations

    The paper sees potential for ViL in high-resolution applications such as semantic segmentation and medical imaging, where its linear complexity allows compute-efficient processing of long sequences.

    Source: §7 · p. 10

    The lack of optimized mLSTM hardware implementations lengthens runtimes and makes scaling models, tuning hyperparameters, and exploring larger datasets costly.

    Source: §5 · pp. 8–9

Abstract

from arXiv · show

Transformers are widely used as generic backbones in computer vision, despite initially introduced for natural language processing. Recently, the Long Short-Term Memory (LSTM) has been extended to a scalable and performant architecture - the xLSTM - which overcomes long-standing LSTM limitations via exponential gating and parallelizable matrix memory structure. In this report, we introduce Vision-LSTM (ViL), an adaption of the xLSTM building blocks to computer vision. ViL comprises a stack of xLSTM blocks where odd blocks process the sequence of patch tokens from top to bottom while even blocks go from bottom to top. Experiments show that ViL holds promise to be further deployed as new generic backbone for computer vision architectures.

1 Introduction

Vision-LSTM (ViL) adapts xLSTM to vision by alternating the direction in which mLSTM blocks process image patches. The introduction motivates this design for efficient modeling and outlines its evaluation and future work.

1 Introduction

  • Vision-language-modeling architectures, including Transformers and State Space Models, are commonly adapted to vision, where images can be represented as projected patch tokens with positional information.The introduction describes ViT’s approach of grouping an image into non-overlapping patches, projecting them into tokens, and adding positional information.Source: §1 · pp. 1–2
  • ViL adapts mLSTM to computer vision, with odd blocks processing patches top-left to bottom-right and even blocks processing them in reverse.The paper describes the alternating traversal as a way to process non-sequential image inputs without additional computations.Source: §1 · p. 2
  • Linear sequence-length complexity makes ViL appealing for high-resolution tasks, whereas ViT self-attention scales quadratically.The paper names medical imaging, segmentation, and physics simulations as examples of tasks that benefit from high-resolution images.Source: §1 · p. 2
  • The paper evaluates ImageNet-1K pretraining followed by transfer classification and semantic segmentation, and ablates architectural design choices.It also identifies future directions and current limitations as topics for discussion.Source: §1 · p. 2

2 Method

The method turns images into patch-token sequences and processes them with mLSTM blocks, drawing on the cell’s matrix memory, exponential gates, and parallelizable computation. ViL alternates scan direction across blocks and offers recurrent, parallel, and chunkwise processing modes.

2.1 Preliminaries

  • Images are reshaped into flattened patches, projected into D-dimensional vectors, and refined through a stack of mLSTM blocks that flips the sequence every second block.The sequence length T is the number of patches; each input vector corresponds to one position in the sequence.Source: §2.1 · p. 3
  • The mLSTM uses a matrix memory and exponential input and forget gates, while its lack of memory mixing allows full parallelization.Within ViL’s gated MLP, the weight matrices process patches individually and the mLSTM exchanges information between patches.Source: §2.1 · pp. 3–4

2.2 Vision-LSTM (ViL)

  • ViL is a residual stack of alternating mLSTM blocks: odd blocks scan top-left to bottom-right, and even blocks flip the sequence before and after processing.The block design combines patchwise image processing with a change in traversal direction across blocks.Source: §2.2 · p. 4
  • Chunkwise processing combines parallel computation within chunks and recurrent processing between them, offering high parallelization, minimal operations, and linear scaling with sequence length.The paper contrasts this mode with sequential recurrent processing and fully parallel processing, which have distinct FLOPS and runtime characteristics.Source: §2.2 · p. 4

3 Experiments

ViL is evaluated after ImageNet-1K pre-training on classification, ADE20K segmentation, and VTAB-1K transfer classification. It performs favorably against optimized transformer protocols and Vision Mamba across the compute overview, with strong results on the downstream tasks.

3 Experiments

  • Models are pre-trained on ImageNet-1K and transferred to classification, semantic segmentation, and few-shot classification benchmarks.The benchmarks cover global understanding, semantic local and global understanding, and diverse VTAB-1K tasks; training schedules and resolutions vary by stage.Source: §3 · p. 4
  • Across classification, segmentation, and transfer classification, ViL performs favorably against heavily optimized transformer protocols and Vision Mamba in relation to pre-training compute.The overview covers ImageNet-1K, ADE20K, and VTAB-1K.Source: Figure 3 · p. 4, §3 · p. 5
  • ViL has seen less hyperparameter tuning and architecture optimization than established ViTs, so future work could improve it further.The authors note that this work is the first to apply xLSTM to computer vision.Source: §3 · p. 5

3.1 ImageNet-1K Classification

  • ViL outperforms heavily optimized ViT protocols and other backbones at tiny and small scales on ImageNet-1K, but not every model at base scale.ViL-B nevertheless learns strong features, particularly for semantic segmentation and structured tasks.Source: §3.1 · p. 5, Table 1 · p. 5

3.2 ADE20K Semantic Segmentation

  • On ADE20K, ViL performs strongly and ViL-B outperforms DeiT-III-B despite lower ImageNet-1K accuracy, with fewer FLOPS than transformer and Mamba-based vision models.At 512x512 resolution, the task uses 1,024 patch tokens; the passage attributes ViL's lower FLOPS to mLSTM's linear complexity and its alternating block design.Source: §3.2 · p. 5, Table 2 · p. 6

3.3 VTAB-1K Transfer Classification

  • On VTAB-1K, ViL outperforms all other models in average accuracy across 19 datasets and performs particularly well on structured datasets.ViL-B outperforms DeiT-III-B on structured datasets despite lower ImageNet-1K accuracy.Source: §3.3 · p. 6, Table 3 · p. 6

4 Ablation Studies

The ablations examine traversal direction, mLSTM block design, positional information, and classification aggregation. They identify bidirectional traversal and first/last-patch aggregation as choices for the final models, while noting technical constraints on broader traversal designs.

4 Ablation Studies

  • The ablations train ViL-T for 100 epochs at 224 × 224 resolution, evaluate ImageNet-1K validation accuracy, and fine-tune on ADE20K with a reduced segmentation pipeline.The segmentation head is linear and trained for 40K updates with batch size 16.Source: §4 · p. 7

4.1 Architectural Design

  • Traversing in at least two directions greatly improves performance; column-wise quad-directional traversal further improves segmentation, while oct-directional traversal adds no benefit, at equal FLOPs.The authors use bi-directional traversal in final models because current technical constraints make training with more than two directions slower.Source: §4.1 · p. 7
  • Quad- and oct-directional traversals are incompatible with torch.compile and take approximately twice the runtime, so the authors train final models with bi-directional traversal.The implementation relies on torch.compile because optimized mLSTM hardware implementations are lacking.Source: §4.1 · pp. 7–8
  • A 2D convolution with a 3 × 3 kernel replaces the causal 1D convolution so the mLSTM can use local image context for storage and retrieval decisions.The original causal 1D convolution reflects language modeling's causal sequence structure, which images do not share.Source: §4.1 · p. 8
  • Providing positional information explicitly improves segmentation results without hurting classification performance, despite the mLSTM's ability to infer position from sequence order.Source: §4.1 · p. 8
  • The study investigates applying both traversal directions in parallel rather than sequentially, halving block depth to keep parameter count and FLOPs constant.Source: §4.1 · p. 8

4.2 Classification Design

  • Concatenating the first and last patches gives strong classification performance; learnable [CLS] tokens show no benefit, while [AVG] aggregation may help segmentation.The authors therefore use first-and-last-patch concatenation and no [CLS] tokens in their final models.Source: §4.2 · p. 8

5 Limitations and Future Work

The authors identify the lack of optimized mLSTM hardware as ViL's main practical constraint, limiting runtime analysis and making several research directions costly. They also report speed comparisons and a classification-versus-segmentation aggregation trade-off.

5 Limitations and Future Work

  • The lack of optimized mLSTM hardware leads to longer runtimes than ViTs and makes runtime or throughput analysis difficult; FLOPs are used as a proxy, though they are imperfect.The authors report comparable FLOPs to ViT on low-resolution tasks and much lower FLOPs on high-resolution tasks, and suggest ViL could be faster on high-resolution tasks with optimized hardware.Source: §5 · pp. 8–9
  • Averaging all patches or the four center patches gives strong segmentation but lackluster classification, while averaging or concatenating the first and last patches performs well for classification.Source: Table 5 · p. 9
  • Scaling model size, tuning hyperparameters, using larger datasets, exploring self-supervised pretraining, and investigating hierarchical architectures are costly without optimized hardware.Source: §5 · p. 9
  • The authors characterize the hardware gap as technical rather than methodical, while noting that implementing fast CUDA or Triton kernels requires specialized expertise and development effort.Source: §5 · p. 9
  • ViL is reported to be up to 70% faster than Vim despite Vim's custom CUDA kernel; the authors attribute xLSTM's language-modeling speed gap with Mamba to mLSTM hardware inefficiency.Source: §5 · p. 9

6 Related Work

The related work traces the development of generic visual backbones from CNNs to transformers and then to autoregressive sequence models. It situates ViL among vision adaptations of sequence-modeling architectures.

6 Related Work

  • CNN inductive biases drove early computer-vision advances, and CNNs were found to learn generic visual features useful across tasks.Subsequent work improved CNN architectures and pretraining strategies.Source: §6 · p. 9
  • Transformers moved from language processing into vision, while ViT's larger patches reduced the costs associated with applying self-attention to pixels or small patches.Later work refined ViT architectures and training procedures.Source: §6 · p. 9
  • Recent autoregressive vision models draw interest for their linear scaling, which permits applications such as medical imaging and video understanding.Source: §6 · pp. 9–10

7 Conclusion

ViL outperforms several vision-backbone families on ImageNet-1K classification, VTAB-1K transfer classification, and ADE20K segmentation. The conclusion identifies high-resolution applications as a promising direction, while a runtime comparison reports faster training than Vim.

7 Conclusion

  • ViL outperforms SSM-based, other autoregressive, and optimized ViT architectures on ImageNet-1K, VTAB-1K, and ADE20K benchmarks.The conclusion also says ViL outperforms ViT training pipelines developed through years of tuning and improvements.Source: §7 · p. 10
  • High-resolution applications such as semantic segmentation and medical imaging are identified as promising settings for ViL's compute-efficient processing of long sequences.The conclusion attributes this potential to ViL's linear complexity, contrasted with transformers' quadratic self-attention costs.Source: §7 · p. 10

A.1 Runtime Comparison of ViL vs Vim

  • Up to 69% faster training than Vim for 10 ImageNet-1K epochs, despite no optimized hardware implementation.Runtimes were extrapolated from short benchmark runs on one A100-80GB-PCIe using float16 and 224x224 images.Source: Table 6 · p. 17

A.2 Impact of Longer Training

The appendix compares tiny models trained for 400 and 800 epochs, with ADE20K mIoU evaluated at single scale.

A.2 Impact of Longer Training

  • Table 7 compares tiny models trained for 400 and 800 epochs.The passage does not state the comparison's outcome.Source: Table 7 · p. 17, §A.2 · p. 17
  • ADE20K mIoU is reported using single-scale evaluation.Source: Table 7 · p. 17
  • The longer-training comparison follows the settings used in the main paper.Source: Table 7 · p. 17

A.3 VTAB-1K Individual Dataset Results

The appendix reports VTAB-1K accuracies for individual datasets, organized under natural, specialized, and structured categories.

A.3 VTAB-1K Individual Dataset Results

  • The section describes accuracies for each individual VTAB-1K dataset.Source: §A.3 · p. 17
  • The displayed category labels are natural, specialized, and structured.Source: §A.3 · p. 17
  • Table 8 reports results for all datasets in the VTAB-1K benchmark.Source: Table 8 · p. 17, §A.3 · p. 17

A.4 Robustness and Domain Generalization

The appendix evaluates robustness and out-of-distribution performance for ImageNet-1K pretrained classifiers across several ImageNet variants and ImageNet-Sketch.

A.4 Robustness and Domain Generalization

  • Table 9 covers robustness and out-of-distribution evaluations of ImageNet-1K pretrained classifiers.Source: Table 9 · p. 18, §A.4 · p. 18
  • The evaluation includes ImageNet-C, ImageNet-A, ImageNet-R, and ImageNet-Sketch.Source: Table 9 · p. 18
  • For ImageNet-C, the reported metric is mean corruption error with AlexNet as the baseline.Source: Table 9 · p. 18

B Implementation Details

Appendix B documents hardware and FLOPS accounting, model-training settings, transfer-evaluation procedures, and the DeiT-III-T reimplementation protocol and results.

B.1 Hardware

  • Training used 8×A100 or 4×A100 nodes; total project compute is estimated at 38K A100 GPU-hours, including exploration and analysis.Source: §B.1 · p. 18

B.2 FLOPS Calculation

  • fvcore FLOPS accounting omits causal upper-triangular QK entries; Vim CUDA calls are replaced by PyTorch references, and ViT uses dropped-block stochastic depth.Source: §B.2 · p. 18

B.3 ViL Hyperparameters

  • ViL uses 16×16 patches, depth 24, and latent dimensions 192/384/768 for tiny/small/base; schedules are 800/400 then 20/5 epochs, with bicubic resizing and ImageNet-1K normalization.Source: §B.3 · p. 19, Table 10 · p. 19

B.4 Fine-tuning on VTAB-1K

  • VTAB-1K tuning searches five learning rates with five seeds on 800 training and 200 validation samples, then averages five test runs using the selected rate.Source: §B.4 · p. 20

B.5 ADE20K Semantic Segmentation Fine-tuning

  • ADE20K fine-tuning uses UperNet at 512×512, interpolates absolute positional embeddings from 224×224, and adds zero-initialized relative biases to ViT attention.Source: §B.5 · p. 21

B.6 DeiT-III Reimplementation Hyperparameters

  • +2.7% ImageNet-1K accuracy over DeiT-II-T; the DeiT-III-T reimplementation is approximately even with the original on ADE20K: 40.1 vs 39.8 single-scale and 41.8 vs 42.2 multi-scale mIoU.Source: §B.6 · p. 22
  • Reducing DeiT-III-T’s learning rate from 3e-3 to 1e-3 greatly improved performance; AdamW and no gradient clipping were stable.Source: Table 13 · p. 22

Related papers

Loading 2406.04303v3…