Long-Short Transformer: Efficient Transformers for Language and Vision

Chen Zhu; Wei Ping; Chaowei Xiao; Mohammad Shoeybi; Tom Goldstein; Anima Anandkumar; Bryan Catanzaro

Source-linked AI summary

Long-Short Transformer: Efficient Transformers for Language and Vision

Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, Bryan Catanzaro

arXiv:2107.02192v3Jul 5, 2021cs.CVcs.CLcs.LGcs.MM

TL;DR

Long sequences are expensive because self-attention scales quadratically in time and memory, while existing efficient mechanisms do not uniformly cover language and vision or autoregressive and bidirectional settings. Transformer-LS combines dynamic-projection long-range attention with local window attention and dual normalization, achieving state-of-the-art results across language and vision tasks with linear complexity.

Problem
Quadratic self-attention makes long documents and high-resolution images expensive, and efficient attention mechanisms have not been broadly effective across language, vision, autoregressive, and bidirectional models.
Method
Transformer-LS combines dynamic low-rank long-range attention, local window attention, and dual normalization in a linear-complexity Transformer.
Results
Transformer-LS achieves state-of-the-art results on Long Range Arena, character-level language modeling, and ImageNet classification.
Takeaways & Limitations
The method provides one efficient attention mechanism for long-sequence language and vision modeling in both autoregressive and bidirectional settings.
Takeaways & Limitations
Sliding-window attention alone does not guarantee that long-range correlations are captured in every layer.

Abstract

from arXiv · show

Transformers have achieved success in both language and vision domains. However, it is prohibitively expensive to scale them to long sequences such as long documents or high-resolution images, because self-attention mechanism has quadratic time and memory complexities with respect to the input sequence length. In this paper, we propose Long-Short Transformer (Transformer-LS), an efficient self-attention mechanism for modeling long sequences with linear complexity for both language and vision tasks. It aggregates a novel long-range attention with dynamic projection to model distant correlations and a short-term attention to capture fine-grained local correlations. We propose a dual normalization strategy to account for the scale mismatch between the two attention mechanisms. Transformer-LS can be applied to both autoregressive and bidirectional models without additional complexity. Our method outperforms the state-of-the-art models on multiple tasks in language and vision domains, including the Long Range Arena benchmark, autoregressive language modeling, and ImageNet classification. For instance, Transformer-LS achieves 0.97 test BPC on enwik8 using half the number of parameters than previous method, while being faster and is able to handle 3x as long sequences compared to its full-attention version on the same hardware. On ImageNet, it can obtain the state-of-the-art results (e.g., a moderate size of 55.8M model solely trained on 224x224 ImageNet-1K can obtain Top-1 accuracy 84.1%), while being more scalable on high-resolution images. The source code and models are released at https://github.com/NVIDIA/transformer-ls .

1 Introduction

Transformer-LS addresses the quadratic cost of self-attention for long sequences by combining local and long-range attention with linear complexity. It applies across autoregressive and bidirectional language and vision models and reports state-of-the-art results across several benchmarks.

Self-attention has quadratic time and memory costs, making long documents and high-resolution image sequences expensive to process.
Transformer-LS combines local window attention for fine-grained correlations with dynamic-projection attention for long-range correlations.
The method uses linear time and memory complexity and supports both autoregressive and bidirectional models.
Dynamic low-rank projection depends on input content and is designed to handle semantic-preserving positional variations more flexibly than previous low-rank methods.
DualLN addresses scale mismatch between long-range and short-term attention embeddings to improve their aggregation.
Transformer-LS outperforms state-of-the-art models across Long Range Arena, autoregressive language modeling, and vision tasks.

2 Related Work

Efficient Transformer research includes sparse, low-rank, and hybrid attention methods for reducing the quadratic cost of full attention. Transformer-LS combines local windows with dynamic low-rank projections to retain both short-term and long-range modeling.

Sparse attention reduces cost through predefined patterns such as sliding windows or random connections, but restricts which tokens can attend to one another.
Low-rank methods project full key and value sequences into smaller representations to reduce attention complexity.
Transformer-LS integrates local window attention with dynamic low-rank projection to model short-term and long-range correlations.
Vision-efficient Transformers reduce complexity through local windows, blocked attention, or cross-attention with latent arrays.

3 Long-Short Transformer

Transformer-LS approximates full attention by combining local sliding-window attention with dynamic low-rank long-range attention, preserving linear complexity while modeling both nearby and distant correlations. DualLN aligns the scales of the two attention streams during aggregation.

Short-term and long-range attention: Sliding-window attention gives each query access to a fixed neighborhood spanning 2w key-value pairs, with memory scaling linearly in sequence length.Segments include their home tokens and w/2 neighboring tokens on each side.
Short-term and long-range attention: Dynamic projection maps sequence-dependent key and value embeddings to r-dimensional representations, replacing full attention with an implicit low-rank product of complexity O(rn).The projection depends on the input sequence and is intended to be more flexible than Linformer’s fixed projection.
Short-term and long-range attention: The global attention lets every query attend to all token embeddings within one layer, unlike sparse mechanisms that require multiple layers to build such correlations.The effective attention weights over all tokens still sum to 1.
Autoregressive application: Autoregressive models use equal-length segments so dynamic projections are computed once per segment in parallel, preserving linear complexity and training speed.Recomputing projections for every token would instead require O(rn^2) computation.
Short-term and long-range attention: Transformer-LS combines sliding-window attention for fine-grained local correlations with dynamic low-rank attention for long-range correlations.Each query can attend to both local keys and values and globally projected low-rank keys and values.
Dual normalization: DualLN applies separate Layer Normalizations to local and global key-value paths, aligning their scales and yielding consistently lower validation loss than omitting it.The normalization addresses a mismatch that initially biases attention toward the local window.

4 Experiments

Experiments evaluate Transformer-LS across long-sequence language and vision tasks, comparing accuracy, efficiency, robustness, and scalability against efficient and full-attention baselines.

Long Range Arena: Transformer-LS outperforms other efficient Transformers on the evaluated Long Range Arena tasks while using about 50% to 70% of their computation.The advantage is strongest on ListOps, while global attention is more effective than window attention on Retrieval.
Long Range Arena: 66.28 vs. 56.12 test accuracy shows dynamic projection outperforming Linformer on Text.The comparison uses the Text task, which has the highest variance in sequence length.
Robustness: Dynamic projection is designed to remain robust to test-time insertion and deletion changes that preserve labels in most cases.The robustness comparison evaluates models trained on clean data and perturbed only at test time.
Autoregressive Language Modeling: 0.97 test BPC is achieved by the larger Transformer-LS model on enwik8, matching Compressive Transformer with 2× parameters.The smaller model reaches 0.99 on enwik8 and 1.09 on text8; Transformer-LS is also reported as more memory- and computation-efficient than full attention.
ImageNet Classification: CvT*-LS-21 at 448^2 achieves 0.3% higher accuracy than the best reported CvT result using the same parameters and 76% of its FLOPs.The experiments use ImageNet-1K and replace attention in CvT and ViL backbones with long-short attention.
ImageNet Classification: Transformer-LS improves ViL-Medium and ViL-Base from 83.5 and 83.7 to 83.8 and 84.1 without increasing FLOPs.Increasing ViL-LS-Medium training resolution from 224^2 to 384^2 improves accuracy by 0.6% while FLOPs increase approximately linearly.
Robustness: The vision evaluation also examines robustness across ImageNet-Real, ImageNet-V2, ImageNet-C, ImageNet-R, ImageNet-9, and ImageNet-A.These datasets cover distribution shifts, corruptions, semantic shifts, background dependence, and natural adversarial examples.

5 Conclusion

Long-Short Transformer models long sequences across language and vision, supporting both bidirectional and autoregressive settings. Its dynamic global attention has linear complexity, while DualLN addresses scale mismatch when aggregating local and global attention.

Transformer-LS targets long-sequence modeling in both language and vision, including bidirectional and autoregressive models.
Linear-complexity dynamic projection provides global attention for long-range correlations.The design uses linear computational and memory complexity in sequence length.
Transformer-LS obtains state-of-the-art results on Long Range Arena, character-level language modeling, and ImageNet classification.
DualLN addresses the initialization-scale mismatch between long-term and short-term attention embeddings.The short-term embeddings have larger norms because long-term embeddings are weighted means with smaller variance.

B Details for Experiments on Long Range Arena

The Long Range Arena experiments evaluate Transformer-LS on structured, sentiment, and other long-sequence tasks using compact two-layer models and reported multi-seed averages. The setup follows prior work while varying training and configuration details across tasks.

The tasks: Long Range Arena includes ListOps for hierarchical parsing and Text for binary sentiment classification over long character-level sequences.ListOps instances contain 500–2000 tokens, while Text uses IMDb reviews.
Architecture: The evaluated models use two layers, embedding dimension d = 64, two attention heads, and FFN hidden dimension 128.A CLS token provides a global representation for classification.
The appendix provides best-result configurations and training hyperparameters for the Long Range Arena experiments.
Hyperparameters for Training: Training uses Adam with learning rate 10^-4, no weight decay, and batch size 32 across tasks.Warmup and total training steps vary with the number of training samples.
Results are averaged over 4 runs with different random seeds, with standard deviations reported separately.

C.1 Results on the image-based tasks of LRA

The image-based Long Range Arena results compare Transformer-LS with prior methods using reported image-task scores. Table 10 presents the PyTorch comparison and selects the highest prior scores from published sources.

Table 10 compares Transformer-LS with other methods on the image-based Long Range Arena tasks.
Prior-model results in the comparison are taken from the highest scores reported in earlier work.

C.2 Compare models implemented in JAX

The JAX comparison re-implements Transformer-LS alongside prior Long Range Arena methods and evaluates both accuracy and per-batch latency on A100 GPUs. Transformer-LS improves results while remaining efficient in this evaluation.

Transformer-LS is re-implemented in JAX to compare its Long Range Arena results with methods from the original benchmark paper.The other methods’ accuracies come from the Long Range Arena paper.
Per-batch latency is evaluated for all models on A100 GPUs using the benchmark’s official JAX implementation.
Transformer-LS achieves improvements while remaining efficient enough in the JAX comparison.

D Details for Autoregressive Language Modeling

Transformer-LS combines segment-wise dynamic projection with sliding-window attention for efficient autoregressive modeling. The window component restores recent-token coverage omitted by the long-range mechanism while preserving linear complexity.

Autoregressive attention: Segment-wise dynamic projection computes low-rank projections in parallel while preventing queries from attending to future tokens, maintaining O(n) complexity.Each query attends only to permitted earlier segments and excludes its own future context.
Complementary local attention: Sliding-window attention supplies local coverage for recent tokens that the low-rank attention can omit.The window becomes indispensable when recent tokens are missing from the low-rank attention span.
Experimental setup: The autoregressive language-modeling setup uses window size w = 512, segment length l = 16, and dynamic projection dimension r = 1.These settings achieved a better efficiency-BPC trade-off than the alternative configurations tested.
Sliding-window pattern: With group size w = 2, each token attends to at most 2w tokens in the illustrated one-dimensional window scheme.Bidirectional attention includes the home segment and neighboring tokens, while autoregressive attention restricts access to the left context and home segment.

E Details for ImageNet Classification

The ImageNet experiments replace attention in established vision Transformer backbones with long-short attention and explore higher-resolution architectures. The resulting models retain or improve accuracy while reducing computational or parameter costs in selected comparisons.

Architecture: CvT∗-LS models replace the attention mechanisms of CvT backbones with long-short term attention while preserving the backbone’s staged feature-map structure.The implementation also adapts dynamic projection with depth-wise separable convolution and uses no positional encoding.
Architecture scaling: CvT∗-LS-17 uses 25% fewer parameters and fewer FLOPs than CvT-21 while achieving the same level of accuracy.This comparison supports scaling the efficient attention mechanism across deeper or wider higher-resolution stages.
Architecture scaling: CvT∗-LS-21S achieves 0.4% higher accuracy than CvT∗-LS-21 with fewer parameters but more FLOPs.The architecture concentrates additional computation on higher-resolution feature maps.
Normalization: Removing DualLN lowers CvT∗-LS-13 test accuracy from 81.9 to 81.3.The comparison measures the contribution of the dual-normalization strategy in the vision model.
Memory scalability: Figure 7 compares running memory consumption of full self-attention and Long-Short Transformer as sequence resolution increases until a 32GB V100 runs out of memory.The figure evaluates memory scalability across different tasks and sequence lengths.

F Evaluate the robustness of models trained on ImageNet-1k.

The ImageNet robustness evaluation tests Transformer-LS against distribution shifts, corruptions, semantic changes, background changes, and natural adversarial examples. The reported results show improvements over ResNet across several robustness benchmarks, while the authors identify deeper analysis as future work.

Comparison protocol: The robustness comparison selects CNN- and Transformer-based models with similar parameter counts, including ResNet and DeiT.Detailed corruption results are reported for different ImageNet-C corruption types.
Evaluation scope: The evaluation covers ImageNet-C corruptions, ImageNet-A natural adversarial examples, ImageNet-R semantic shifts, and ImageNet-9 background dependence.ImageNet-C includes 15 corruption types across five severity levels, with lower corruption error indicating higher robustness.
Results: Transformer-LS improves accuracy over ResNet by 23.6% on ImageNet-C, 22.1% on ImageNet-A, and 9.7% on ImageNet-R.These comparisons use the reported results from the robustness evaluation.
Results: Transformer-LS improves ImageNet-9 accuracy by 4.3% on average across the Mixed-same and Mixed-rand settings.The authors interpret this result as indicating reduced sensitivity to background changes.
Scope boundary: The authors leave an in-depth study of the observed robustness behavior as important future work.The stated future-work boundary limits the current study’s explanatory depth rather than its reported comparisons.