PaperlayerView exact revision on arXiv

Source-linked AI summary

Adaptive Token Sampling For Efficient Vision Transformers

Mohsen Fayyaz, Soroush Abbasi Koohpayegani, Farnoush Rezaei Jafari, Sunando Sengupta, Hamid Reza Vaezi Joze, Eric Sommerlade, Hamed Pirsiavash, Juergen Gall

arXiv:2111.15667v3cs.CV

TL;DR

Vision transformers are computationally expensive because they process many tokens, while fixed token-reduction settings are not optimal for every input. The paper introduces ATS, a differentiable parameter-free module that adaptively scores and samples tokens within existing vision transformers. Across image and video benchmarks, ATS reduces computation substantially while preserving accuracy and improving the cost/accuracy trade-off.

  • Problem

    Vision transformers have high computational costs, and fixed token-reduction settings cannot accommodate the varying relevant information in different images or videos.

  • Method

    ATS is a differentiable parameter-free module that scores tokens from classification-token attention and adaptively samples significant tokens within vision-transformer blocks.

  • Results

    ATS decreases computation cost by 27% to 50.8% with a negligible accuracy drop across ImageNet-1K, Kinetics-400, and Kinetics-600 evaluations.

  • Takeaways & Limitations

    ATS can convert off-the-shelf image and video vision transformers into more computationally efficient models with input-dependent token counts.

Abstract

from arXiv · show

While state-of-the-art vision transformer models achieve promising results in image classification, they are computationally expensive and require many GFLOPs. Although the GFLOPs of a vision transformer can be decreased by reducing the number of tokens in the network, there is no setting that is optimal for all input images. In this work, we therefore introduce a differentiable parameter-free Adaptive Token Sampler (ATS) module, which can be plugged into any existing vision transformer architecture. ATS empowers vision transformers by scoring and adaptively sampling significant tokens. As a result, the number of tokens is not constant anymore and varies for each input image. By integrating ATS as an additional layer within the current transformer blocks, we can convert them into much more efficient vision transformers with an adaptive number of tokens. Since ATS is a parameter-free module, it can be added to the off-the-shelf pre-trained vision transformers as a plug and play module, thus reducing their GFLOPs without any additional training. Moreover, due to its differentiable design, one can also train a vision transformer equipped with ATS. We evaluate the efficiency of our module in both image and video classification tasks by adding it to multiple SOTA vision transformers. Our proposed module improves the SOTA by reducing their computational costs (GFLOPs) by 2X, while preserving their accuracy on the ImageNet, Kinetics-400, and Kinetics-600 datasets.

1 Introduction

Vision transformers offer strong image and video understanding but incur high computational costs because they process a static number of tokens. ATS addresses this by adaptively sampling tokens according to input content and reducing computation across image and video classification tasks.

  • Motivation: Vision transformers have strong representation power but their computational costs grow quadratically with the number of tokens.This high cost can limit deployment on edge devices.
  • Motivation: Input images and videos contain varying amounts of relevant information, so different examples require different numbers of tokens at each stage.Some examples need only a few patches, while detailed or complex inputs require more.
  • Adaptive Token Sampler: ATS is a differentiable, parameter-free module that scores tokens with classification-token attention weights and samples a subset using inverse transform sampling.It then softly downsamples the output tokens to remove redundant information.
  • Adaptive Token Sampler: ATS can be added to off-the-shelf pre-trained vision transformers without further training, while its differentiability also supports fine-tuning.A maximum token limit can be adjusted at inference time for different edge-device resources.
  • Experiments: Experiments on ImageNet and Kinetics show that ATS improves the computational cost/accuracy trade-off across image and video vision transformers.The method is evaluated with DeiT, CvT, PS-ViT, XViT, and TimeSformer, and surpasses existing token-reduction approaches.

2 Related Work

Related work improves transformer efficiency through architectural sparsification, distillation, adaptive attention, token pruning, and pooling. Vision-transformer methods reduce attention or token processing costs, but existing approaches commonly add prediction modules or use token pooling.

  • Vision Transformers: Vision Transformer architectures split images into non-overlapping patches, embed them, add positional information, and process them through transformer blocks.A learnable class embedding is also included in the architecture.
  • Transformer Efficiency: Transformer-efficiency research includes sparse connectivity, knowledge distillation, redundant-token removal, and adaptive attention spans.Examples include Star-Transformer, TinyBERT, PoWER-BERT, and adaptive-attention-span methods.
  • Vision Transformer Efficiency: Vision-transformer efficiency methods reduce dense attention complexity through sparse factorization or attention-matrix sparsification.Sparse factorization is reported with complexity O(n√n) for autoregressive image generation.
  • Token Reduction: DynamicViT adds a prediction module to estimate token importance and discard uninformative tokens for image classification.This approach is presented alongside other token-reduction methods for vision transformers.
  • Token Reduction: Hierarchical Visual Transformer uses token pooling, analogous to feature-map downsampling in convolutional networks.Its approach reduces the number of tokens through hierarchical pooling.

3 Adaptive Token Sampler

ATS dynamically reduces the token count in vision transformers while retaining the classification token and adapting computation to each input. It scores tokens from self-attention, samples them probabilistically but deterministically at inference, and propagates softly downsampled outputs to later stages.

  • Motivation: Vision-transformer computation grows quadratically with a static token count, motivating input-dependent token reduction.Fixed pooling is not straightforward because transformer tokens are permutation invariant, and static downsampling is not always optimal.
  • Token Scoring: ATS scores N input tokens using attention from the classification token and samples significant tokens with a parameter-free differentiable procedure.The classification token is retained, while the sampled-token maximum K controls the GFLOPs upper bound.
  • Token Sampling: Duplicate samples are collapsed, yielding a unique token count K′ that varies across inputs and stages and satisfies K′ ≤ K ≤ N.K′ can range from one dominant token to K when scores are relatively balanced.
  • Token Scoring: Self-attention computes A from Q and K and produces outputs by weighting V, so ATS uses classification-token attention weights together with value magnitudes for token significance.The classification-token self-attention weights represent input-token importance for the output classification token.
  • Token Sampling: Inverse transform sampling maps cumulative significance scores to token indices, using fixed sampling points for deterministic training and inference.The procedure begins with the second token because the first classification token is always kept.
  • Output Downsampling: ATS forms a refined attention matrix by selecting sampled rows and computes output tokens as AsV for the next network stage.This replaces the full attention matrix with attention restricted to the sampled tokens.

4 Experiments

Experiments evaluate ATS across image and video transformers, ablations, and sampling strategies. ATS adaptively varies token counts while reducing computation with little accuracy loss.

  • Experimental setup: Experiments use ImageNet, Kinetics-400, and Kinetics-600 to evaluate ATS across image and video classification.The study tests ATS with DeiT, CvT, PS-ViT, XViT, and TimeSformer backbones.
  • Ablation experiments: Inverse transform sampling retains a more diverse token set than Top-K selection, preserving more tokens for cluttered images and fewer for less detailed images.This produces input-dependent computation rather than a fixed selection rate at every stage.
  • Ablation experiments: The number of selected tokens varies across images and stages, with sampling patterns reflecting image detail and progressively changing through the network.Histograms and visualizations show adaptive token counts, while token locations become focused on informative regions related to target objects.
  • Ablation experiments: Multi-stage DeiT-S+ATS performs better than single-stage DeiT-S+ATS because it gradually decreases GFLOPs by discarding fewer tokens in earlier stages.Fine-tuning improves accuracy, and testing different K values controls the desired GFLOPs level.
  • Comparison with state-of-the-art: 37% GFLOPs reduction for DeiT-S+ATS comes with only a 0.1% top-1 accuracy reduction on ImageNet, without extra backbone parameters.CvT+ATS models achieve about 30% GFLOPs reduction with 0.1–0.2% top-1 accuracy reduction.
  • Comparison with state-of-the-art: 50.8% GFLOPs reduction for TimeSformer-L+ATS comes with only a 0.2% top-1 accuracy drop on Kinetics-400, demonstrating applicability to video transformers.XViT+ATS also reduces GFLOPs by 39% on Kinetics-400 and 38.7% on Kinetics-600 with small accuracy drops.

5 Conclusion

ATS is a differentiable, parameter-free module for making image and video vision transformers more computationally efficient through adaptive token selection. Across evaluated tasks, it reduces computation substantially while preserving accuracy with negligible drops.

  • 5 Conclusion: ATS selects informative tokens within transformer stages, replacing static token counts with an adaptive number based on each image or video clip.The module integrates into attention layers and uses only as many tokens as needed for each input.
  • 5 Conclusion: 27%–50.8% lower GFLOPs with a negligible accuracy drop was reported across ImageNet-1K, Kinetics-400, and Kinetics-600 evaluations.The method was integrated into three image vision transformers and multiple video vision transformers.
  • 5 Conclusion: ATS supports both plug-and-play deployment on pretrained vision transformers without additional training and training of ATS-equipped models.Its parameter-free and differentiable design enables both uses.

A Runtime

ATS can reduce computation and improve runtime while supporting mini-batch inference. Its sampling overhead is small relative to the computation saved by dropping tokens.

  • Runtime: 1.5% sampling overhead for DeiT-S with one ATS stage is negligible compared with the savings from dropped tokens.The reported overhead is measured as a fraction of overall computation.
  • Runtime: ATS throughput speed-up aligns with its GFLOPs reduction on DeiT-S.Throughput was averaged over 20 runs on a single RTX6000 GPU using 224×224 images.
  • Batch Processing: Mini-batch inference processes m = max_i(K′_i + 1) tokens after rearranging each image’s sampled tokens to lower indices.The final tokens can then be removed completely to reduce computation.
  • Batch Processing: A batch of 512 ImageNet validation images retained 129 tokens in DeiT-S+ATS stage 7, discarding at least 68 of 197 tokens.The example demonstrates computation reduction even in a mini-batch setting.
  • Effect of K: Sampling 48 tokens corresponds to 2 GFLOPs and achieved 75% accuracy in the fine-tuned single-stage DeiT-S+ATS model.The result comes from varying K in the ATS module.

C ATS Integration Without Further Training

ATS can be added to pretrained vision transformers without further backbone training while retaining most of their accuracy. Its token selection uses classification-token attention to remove less informative tokens.

  • DeiT-S+ATS: A not fine-tuned multi-stage DeiT-S+ATS model lost only 0.6% top-1 accuracy while improving efficiency by about 1.6 GFLOPs.The model was evaluated without further training of the backbone.
  • XViT+ATS: A not fine-tuned XViT+ATS model lost only 1.1% top-1 accuracy while improving efficiency by about 329 GFLOPs.This result was reported for video data without further backbone training.
  • Token Selection: ATS assigns token significance from classification-token attention weights and samples informative tokens using inverse transform sampling.Tokens contributing less to the final prediction receive lower attention weights and are more suitable for removal.
  • Visualization: Later-stage attention maps become more focused on birds and less on background, matching the sampled-token patterns.The visualization concerns a not fine-tuned multi-stage DeiT-S+ATS model.

E Implementation Details

The implementation evaluates ATS across image and video datasets, transformer backbones, training regimes, and token-retention settings. ATS is inserted into selected attention blocks and can prune attention-matrix rows before the FFN.

  • Datasets: Image classification uses ImageNet with 1.28M training images, 1K classes, and 50K validation images; video experiments use Kinetics-400 and Kinetics-600.The Kinetics datasets contain 400 and 600 classes, respectively.
  • DeiT + ATS: DeiT adaptive models initialize from pretrained weights, train for 30 epochs with AdamW, and use K = 197 during training.The learning rate is 5e-4, batch size is 8×96, and cosine scheduling is used.
  • DeiT + ATS: DeiT evaluation targets average GFLOPs of 3, 2.5, and 2 using multi-stage ρ values of 1, 0.87, and 0.72, or single-stage K values of 108, 78, and 48.The multi-stage setting uses Kn = max(⌊ρ × #InputTokensn⌉, 8).
  • CvT + ATS: ATS is integrated into blocks 1–9 of stage 3 in CvT-13 and CvT-21, using pretrained initialization and 30-epoch AdamW training.The CvT experiments omit convolutional projection layers in stage-3 transformer blocks.
  • PS-ViT + ATS: PS-ViT models use pretrained initialization, 30 epochs of AdamW training, and cosine scheduling under the PS-ViT training settings.The reported learning rate is 5e-4 with batch size 8×96.
  • Video Transformers: ATS is integrated into stages 3–11 of XViT and stages 3–5 of TimeSformer, with pretrained initialization and SGD-based training.XViT trains for 10 epochs, whereas TimeSformer trains for 5 epochs.
  • ATS Operation: Within self-attention, ATS scores tokens, inverse-samples rows of the attention matrix to form As, computes O = AsV, and forwards O to the FFN.This describes the module’s insertion point and attention-output transformation.

F Ablation

The ablation analyzes how significance scores should be calculated and experimentally supports retaining the classification token.

  • Significance scores can be calculated by summing attention weights across all attention-matrix rows to identify important tokens.The analysis compares token choices for calculating S and uses this sum to represent the classification token’s importance.

F.2 Candidate Token Selection

The study compares inverse transform sampling with top-K selection for choosing tokens. Inverse transform sampling performs better because it preserves more diverse token selections.

  • The comparison evaluates trained multi-stage DeiT-S+ATS models using either top-K selection or inverse transform sampling.The top-K variant selects tokens with the highest significance scores S.
  • Inverse transform sampling outperforms top-K selection both with and without training.The comparison uses multi-stage DeiT-S+ATS models and reports top-1 accuracy on ImageNet validation.
  • Inverse transform sampling avoids discarding nearly all low-significance tokens, producing a more diverse set for subsequent layers.The authors associate this diversified selection with better performance after training.

F.3 ATS Placement

ATS placement affects accuracy and computational trade-offs in DeiT-S. Multi-stage integration supports progressive token sampling, while early single-stage placement performs poorly.

  • The ablation evaluates ATS placement across different DeiT-S stages, including single-stage and multi-stage configurations.The corresponding comparison is reported in Table A.4.
  • 78.5% top-1 accuracy is achieved when ATS is integrated into stage 3, compared with 73.1% in the first stage.These experiments use approximately 3 GFLOPs without finetuning.
  • Multi-stage ATS integration obtains the highest accuracy by progressively reducing tokens across the network.It discards fewer tokens in earlier stages than a single-stage model targeting the same GFLOPs.
  • At a 3 GFLOPs target, K is set to 130, 108, and 56 for stages 0, 3, and 6, respectively.Later placement requires discarding more tokens to reach the same computation level.

F.4 Adding ATS to Models with Other Token Pruning Approaches

ATS can be added to EViT-DeiT-S, further reducing computation beyond an existing token-pruning approach. Visualizations show adaptive sampling increases with image complexity.

  • ATS further decreases the GFLOPs of the EViT-DeiT-S model after its token reorganization and pruning process.The combined model is fine-tuned for 10 epochs on ImageNet using ATS in stages 4, 5, 7, 8, 10, and 11.
  • Combining ATS with EViT demonstrates the approach’s compatibility with another token-pruning method.The authors frame the results as evidence favoring adaptive sampling over static token pruning.
  • More cluttered images receive more sampled tokens, while detailed images show more scattered sampled regions than plain images.The visualization follows progressive sampling in a multi-stage DeiT-S+ATS model, whose output-token count is bounded by 197 input tokens.
  • The visual analysis orders ImageNet validation images by complexity and displays discarded tokens as masks over the inputs.It visualizes sampling across stages 3 through 11 in multi-stage DeiT-S+ATS.
Loading 2111.15667v3…