PaperlayerView exact revision on arXiv

Source-linked AI summary

Three things everyone should know about Vision Transformers

Hugo Touvron, Matthieu Cord, Alaaeldin El-Nouby, Jakob Verbeek, Hervé Jégou

arXiv:2203.09795v1cs.CV

TL;DR

Vision Transformers still leave open practical questions about efficient computation, transfer, and compatibility between patch preprocessing and masked self-supervised learning. The paper studies parallel residual blocks, attention-only fine-tuning, and an MLP-based patch stem, finding these variants viable within the tested settings while identifying hardware and architecture scope boundaries.

  • Problem

    Existing Vision Transformer design and training procedures have been explored only to a limited extent, including uncertainty about patch-preprocessing compatibility with patch-masked self-supervision.

  • Method

    The paper evaluates parallel residual-block arrangements, attention-only fine-tuning, and a hierarchical MLP patch stem using ImageNet-based experiments and transfer settings.

  • Results

    The three variants provide viable alternatives: parallelization can preserve accuracy, attention-only fine-tuning is sufficient in the studied adaptations, and the hMLP stem supports masked pretraining while remaining effective for supervised ViT training.

  • Takeaways & Limitations

    These design choices can reduce latency, compute, and peak memory or enable patch preprocessing without patch communication, within the evaluated ViT settings.

  • Takeaways & Limitations

    Parallel-throughput benefits depend on suitable hardware or kernels, and whether the parallel design applies to other architectures remains future work.

Abstract

from arXiv · show

After their initial success in natural language processing, transformer architectures have rapidly gained traction in computer vision, providing state-of-the-art results for tasks such as image classification, detection, segmentation, and video analysis. We offer three insights based on simple and easy to implement variants of vision transformers. (1) The residual layers of vision transformers, which are usually processed sequentially, can to some extent be processed efficiently in parallel without noticeably affecting the accuracy. (2) Fine-tuning the weights of the attention layers is sufficient to adapt vision transformers to a higher resolution and to other classification tasks. This saves compute, reduces the peak memory consumption at fine-tuning time, and allows sharing the majority of weights across tasks. (3) Adding MLP-based patch pre-processing layers improves Bert-like self-supervised training based on patch masking. We evaluate the impact of these design choices using the ImageNet-1k dataset, and confirm our findings on the ImageNet-v2 test set. Transfer performance is measured across six smaller datasets.

1 Introduction

The paper identifies three practical design choices for Vision Transformers: parallelizing residual blocks, fine-tuning attention layers selectively, and using patch preprocessing compatible with masked self-supervision.

  • Parallel vision transformers: Parallelizing MHSA and FFN residual blocks preserves parameter count and compute while producing wider, shallower ViTs with more parallel processing.The authors report that this can ease optimization and reduce latency depending on the implementation.
  • Parallel vision transformers: Parallel ViTs can maintain comparable accuracy to sequential baselines, with some improvements when networks are sufficiently deep.The reported latency reductions occur on GPUs particularly for small batch sizes.
  • Fine-tuning attention is all you need: Fine-tuning only multi-head attention layers is mostly sufficient for adapting ViTs to higher resolutions and other classification tasks.Freezing FFN layers saves compute and peak training memory while allowing parameter-heavy FFN weights to be shared across tasks.
  • Parallel vision transformers: The parallelization analysis leaves its applicability to other architectures for future work.The paper specifically notes that it has not found prior work analyzing width versus depth for ViTs on common GPUs and CPUs.
  • Patch preprocessing with masked self-supervised learning: Existing convolutional stems can propagate information across patches, making them incompatible with masked patch prediction in methods such as BeiT and MAE.This motivates preprocessing that preserves patch independence during masking.
  • Patch preprocessing with masked self-supervised learning: The hierarchical MLP stem interleaves MLP layers with patch aggregation while prohibiting communication between patches, supporting masked pretraining and supervised ViT training.The masking can be applied after preprocessing, and the stem is reported to be on par with the best compared convolutional stem in supervised training.

2 Background

The paper builds on vanilla Vision Transformers and evaluates them with standardized ImageNet training and validation procedures. Its baseline uses common ViT variants, 16×16 patches, and controlled training configurations including LayerScale and epoch count.

  • Training procedure: The training procedure adapts the A2 setting with ViT-specific modifications, including a reduced learning rate and model-dependent stochastic-depth rates.The paper reports learning rates of 4.10^-3 for ViT-Ti and ViT-S and 3.10^-3 for ViT-B and ViT-L.
  • Training procedure: LayerScale improves performance for large models, and longer training is beneficial in that setting.The study compares a 300-epoch baseline without LayerScale against a 400-epoch configuration with LayerScale.
  • Evaluation: The adapted training procedure outperforms existing supervised-training procedures for the model sizes considered.The baseline models use 16×16 patches and are evaluated at 224×224 resolution.

3 Depth vs Width: Parallel ViT

The paper examines depth–width trade-offs in ViTs and proposes parallelizing residual blocks to reduce effective depth without changing parameter and compute complexity. Parallelization is most beneficial for deeper, higher-capacity models that are harder to optimize, while proper optimization can make sequential and parallel designs roughly equivalent.

  • Depth versus width: ViT capacity depends on width and depth, while increasing width raises parameter count, compute, and peak memory quadratically.Latency may improve with wider architectures because they are more parallel, but actual speedups depend on implementation and hardware.
  • Parallel ViT construction: Parallel ViTs reorganize sequential MHSA and FFN residual blocks into parallel branches, reducing layers while preserving the number of blocks, parameters, and compute.The approximation becomes more accurate as networks deepen because each residual block contributes less relative to the overall function.
  • Experimental findings: Two parallel branches achieve the best performance across tested model capacities, whereas more than two branches are generally unfavorable for accuracy.The comparison uses fixed total block counts, so sequential and parallel variants have matched parameter and FLOP budgets.
  • Experimental findings: Parallelization helps most for deeper and higher-capacity models that are difficult to optimize, while smaller 24-block Ti and S models favor sequential processing.For the largest tested model, ViT-L12×2 is stronger than its sequential counterpart despite LayerScale, whereas B24×1 and B12×2 are comparable.
  • Optimization: With LayerScale and sufficiently strong optimization, sequential and parallel ViTs become roughly equivalent for large models.LayerScale improves both variants, bringing their performance approximately on par.
  • Depth versus width: Compared with increasing working dimensionality, using more blocks in sequential or parallel models yields substantially higher accuracy at approximately matched parameter and FLOP complexity.The parallel model is better than the sequential model in the 48-block setting because the sequential model becomes deeper and harder to optimize.
  • Efficiency: Measured latency benefits are strongest for per-sample processing, while compelling throughput gains require hardware or kernels specialized for the parallel architecture.The reported implementation is suboptimal because it lacks a dedicated CUDA kernel.

4 Fine-tuning attention is all you need

Fine-tuning only the MHSA weights adapts ViT models to higher resolutions with accuracy comparable to full fine-tuning while reducing memory, computation, and storage. For downstream classification, this strategy is especially effective on smaller datasets and with larger models, but can lag on larger datasets for smaller models.

  • Fine-tuning at different resolutions: Solely fine-tuning MHSA weights at 384×384 matches full fine-tuning within ±0.1 standard deviation on ImageNet-val and ImageNet-V2.FFN-only fine-tuning does not achieve the same result despite having twice as many parameters as MHSA.
  • Fine-tuning at different resolutions: 10% less GPU memory and 10% faster training result from fine-tuning MHSA instead of all weights at higher resolution.The approach computes fewer gradients during fine-tuning.
  • Fine-tuning at different resolutions: 66% of storage is saved for each additional model fine-tuned at a different input resolution.Attention weights comprise approximately one third of the model weights.
  • Fine-tuning on different datasets: On smaller transfer datasets such as CARS and Flower, MHSA-only fine-tuning can outperform full fine-tuning, whereas ViT-S shows a significant gap on iNaturalist.The limitation tends to disappear for ViT-L, whose MHSA capacity is larger.
  • Fine-tuning on different datasets: The transfer experiments compare full, MHSA-only, and FFN-only fine-tuning across six downstream classification datasets using ViT models pretrained on ImageNet-1k.The method is evaluated as a transfer-learning strategy across multiple model sizes and datasets.

5 Patch preprocessing for Bert-like self-supervised learning

The paper addresses the compatibility of patch preprocessing with masked self-supervised learning by introducing an hMLP stem that processes patches independently. The hMLP stem improves BeiT performance with minimal added computation, unlike convolutional stems.

  • Motivation: Existing patch-preprocessing designs lacked clear evidence of their effectiveness with vanilla transformers and compatibility with patch-masked self-supervised methods such as BeiT.Convolutions can propagate information across patches and impede masked prediction.
  • hMLP stem: The hMLP stem interleaves linear layers, nonlinearities, normalization, and patch aggregation while preventing communication between patches.Its independent processing allows masking before or after patch processing without convolutional overlap artifacts.
  • Complexity: The hMLP stem adds less than 1% of compute compared with the usual linear projection stem.For ViT-B, the design requires 17.73 GFLOPS.
  • Stem comparison: The hMLP stem improves accuracy over linear projection without significantly increasing complexity, whereas the convolutional stem improves only supervised performance and adds +7.5% FLOPS.The hMLP comparison reports +0.8% FLOPS relative to ViT-B12.
  • Stem comparison: The hMLP stem obtains comparable performance with lower complexity and no interaction between the 16×16 patches.This preserves patch independence during preprocessing.
  • Results with BeiT training: +0.3/+0.4 top1 accuracy over the baseline is achieved with BeiT, while existing stems add compute without improving the baseline.The improvement is reported as significant relative to measurement uncertainty.

6 Conclusion

The paper presents three design insights for Vision Transformers: parallelizing residual blocks, fine-tuning attention weights, and using an independent hMLP patch-preprocessing stem. These choices target capacity, efficient adaptation, and masked self-supervised learning.

  • Parallel vision transformers: Parallel residual blocks provide an alternative for increasing capacity without significantly increasing working dimensionality.The paper leaves application of this design principle to other architectures for future work.
  • Fine-tuning attention: Fine-tuning self-attention weights is sufficient for resolution adaptation and can be useful for downstream transfer, especially with large models or few training images.The conclusion specifically highlights resolution fine-tuning and selected transfer settings.
  • Patch preprocessing: The hMLP stem processes patches independently across linear layers, nonlinearities, and patch aggregation, and is especially useful with BeiT-style masked self-supervised learning.Its stated scope is patch preprocessing for mask-based self-supervised training.

A Baselines

The baseline comparison includes vanilla ViT models trained on ImageNet-1k at 224×224 resolution without distillation, spanning ViT-B, ViT-L, ViT-Ti, and ViT-S configurations.

  • Baseline models: The baseline table compares vanilla ViT results for ViT-B, ViT-L, ViT-Ti, and ViT-S models trained on ImageNet-1k at 224×224 without distillation.The table includes models from the original ViT work and Touvron et al., with one result marked as 200 epochs.

B Transfer Learning Datasets

Table 9 lists the datasets used in the transfer experiments and provides their corresponding references.

  • Table 9 lists the datasets used in transfer experiments and their corresponding references.

C Pytorch code of our hMLP Stem

The hMLP stem converts input images into patch embeddings using a sequence of convolutional, normalization, and activation layers, then flattens the spatial features into tokens.

  • The number of patches is computed from the image and patch dimensions before projection.
  • The hMLP stem uses three convolutional stages with intermediate SyncBatchNorm and GELU operations.
  • The projected feature map is flattened across spatial dimensions and transposed to produce patch-token representations.
Loading 2203.09795v1…