Source-linked AI summary
AutoFormer: Searching Transformers for Visual Recognition
Minghao Chen, Houwen Peng, Jianlong Fu, Haibin Ling
TL;DR
Vision transformer architecture design is difficult because key dimensions are manually configured and their best combination is unclear. AutoFormer performs one-shot architecture search using weight entanglement to train a supernet whose subnets inherit effective weights. The searched AutoFormers achieve state-of-the-art ImageNet results among vision transformers and transfer well to downstream classification tasks, while remaining less efficient than inverted-residual models.
Problem
Vision transformer architecture design remains difficult because key dimensions are manually chosen, while prior automatic transformer searches focus on natural language tasks.
Method
AutoFormer performs one-shot vision transformer architecture search with weight entanglement, sharing weights among common parts of different blocks during supernet training.
Results
AutoFormers achieve state-of-the-art ImageNet results among vision transformers, transfer well to downstream classification tasks, and produce subnets whose inherited-weight performance is comparable to retraining from scratch.
Takeaways & Limitations
The trained supernet supports efficient search for diverse transformer architectures under different resource constraints without extra retraining or finetuning.
Takeaways & Limitations
Vision transformers including AutoFormer remain inferior to inverted-residual models optimized for edge devices, which have much smaller model sizes and FLOPs.
Abstract
from arXiv · showhide
Recently, pure transformer-based models have shown great potentials for vision tasks such as image classification and detection. However, the design of transformer networks is challenging. It has been observed that the depth, embedding dimension, and number of heads can largely affect the performance of vision transformers. Previous models configure these dimensions based upon manual crafting. In this work, we propose a new one-shot architecture search framework, namely AutoFormer, dedicated to vision transformer search. AutoFormer entangles the weights of different blocks in the same layers during supernet training. Benefiting from the strategy, the trained supernet allows thousands of subnets to be very well-trained. Specifically, the performance of these subnets with weights inherited from the supernet is comparable to those retrained from scratch. Besides, the searched models, which we refer to AutoFormers, surpass the recent state-of-the-arts such as ViT and DeiT. In particular, AutoFormer-tiny/small/base achieve 74.7%/81.7%/82.4% top-1 accuracy on ImageNet with 5.7M/22.9M/53.7M parameters, respectively. Lastly, we verify the transferability of AutoFormer by providing the performance on downstream benchmarks and distillation experiments. Code and models are available at https://github.com/microsoft/AutoML.
1. Introduction
AutoFormer addresses the difficulty of manually designing vision transformer architectures by searching across key structural dimensions with a one-shot supernet. Its weight-entanglement strategy enables well-trained subnets and supports diverse resource-constrained models that outperform handcrafted transformer baselines.
- Design Challenge: Vision transformer design is difficult because depth, embedding dimension, head number, and MLP ratio jointly affect accuracy and model capacity.Increasing some dimensions can initially improve accuracy but later overfit, while larger embedding dimensions eventually yield diminishing gains.
- Design Challenge: Prior vision transformers rely mainly on manual crafting, while existing automatic transformer searches focus on natural language tasks and do not readily generalize to vision.This leaves vision transformer architecture design dependent on human expertise and trial-and-error.
- AutoFormer: AutoFormer searches a large space spanning embedding dimensions, heads, query/key/value dimensions, MLP ratios, depth, and varied transformer block structures.The search space is designed to cover diverse architectures and model complexities rather than requiring identical structures for all blocks.
- AutoFormer: Weight entanglement shares weights among common parts of different transformer blocks within each layer during supernet training.Unlike classical one-shot NAS, updates to one block affect other blocks sharing those weights.
- Results: Thousands of subnets inherit weights from the trained supernet and achieve performance comparable to independently retrained models.This enables efficient search for architectures meeting different resource constraints without extra finetuning or retraining.
- Results: 81.7% top-1 accuracy: AutoFormer-S with 22.9M parameters exceeds DeiT-S by 1.8% and ViT-S/16 by 2.9% on ImageNet.The result is reported for a searched model obtained through evolutionary search over well-trained supernets.
2. Background
Vision Transformers convert images into patch-token sequences processed by alternating self-attention and MLP blocks, followed by classification. One-shot NAS encodes many candidate architectures in a weight-sharing supernet and searches them using inherited validation performance.
- 2.1. Vision Transformer: Images are split into patches, projected into D-dimensional embeddings, augmented with class and positional embeddings, and classified after transformer encoding.The encoder receives the combined token sequence, while a final linear layer performs classification.
- 2.1. Vision Transformer: Transformer encoders alternate multihead self-attention and MLP blocks, with LayerNorm before and residual connections after each block.The MLP uses two fully connected layers and an activation function, typically GELU.
- 2.1. Vision Transformer: Self-attention linearly maps token embeddings into queries, keys, and values, computes similarity-based weighted sums, and projects concatenated multihead outputs.The Q-K-V dimension determines the dimensionality of each transformed representation, while heads process attention in parallel.
- 2.2. One-Shot NAS: One-shot NAS represents an architecture space A as a supernet N(A, W), allowing candidate subnets to share weights instead of training independently.The supernet weight W is shared across architecture candidates.
- 2.2. One-Shot NAS: The first NAS stage optimizes supernet weights on training loss, while the second ranks candidate subnets by validation accuracy using inherited weights.Sampling subnets during optimization reduces memory usage, and search procedures approximate evaluation across the large architecture space.
3. AutoFormer
AutoFormer addresses transformer-search difficulties with weight entanglement, a broad elastic search space, and supernet-based evolution search under resource constraints.
- One-Shot NAS with Weight Entanglement: Classical weight sharing causes slow supernet convergence because independent transformer blocks receive limited weight updates.The paper motivates weight entanglement as a response to slow convergence and unsatisfactory subnet performance.
- One-Shot NAS with Weight Entanglement: Weight entanglement shares weights across common parts of candidate blocks within each layer, so updating one block affects the others.The strategy targets homogeneous blocks such as attention modules with different head counts and MLPs with different hidden dimensions.
- One-Shot NAS with Weight Entanglement: The strategy reduces storage by retaining only the largest homogeneous building block’s parameters in each layer.This contrasts with storing independent parameters for every candidate block.
- Results and Discussion: Weight entanglement produces well-trained subnets whose inherited-weight performance is comparable to independently retrained models.The paper attributes this capability to improved training and discusses regularization as a possible underlying reason.
- Search Space: The search space varies embedding dimension, Q-K-V dimension, head number, MLP ratio, and network depth across transformer architectures.The space permits different building-block structures and allows attention heads, hidden dimensions, and other factors to vary across layers.
- Search Space: All candidate architectures are encoded in supernets whose uniformly sampled subnets update corresponding shared weights during training.Three independent supernets partition the space by parameter constraints, enabling searches specialized to resource ranges.
- Search Pipeline: Evolution search evaluates trained-supernet subnets to maximize classification accuracy while minimizing model size.Random seeds generate successive populations through parent selection, crossover, and mutation.
4. Experiments
Experiments show that weight entanglement enables efficient supernet-based search, producing strong AutoFormer models without retraining and transferring effectively across classification benchmarks.
- Ablation Study and Analysis: Weight entanglement trains many high-quality subnets whose inherited-weight performance remains close to retraining from scratch.The strategy supports direct subnet evaluation without extra finetuning or retraining.
- Ablation Study and Analysis: 80.1%–82.0% top-1 accuracy is achieved by 1000 sampled high-performing subnets from the supernet-small without retraining.These subnets exceed recent DeiT and RegNetY results according to the reported experiment.
- Results on ImageNet: 81.7% top-1 accuracy is achieved by AutoFormer-S with approximately 23M parameters, outperforming DeiT-S by 1.8% and ViT-S/16 by 2.9%.The searched models inherit weights directly from the supernet without extra retraining or postprocessing.
- Results on ImageNet: AutoFormers outperform manually designed ResNet, ResNeXt, and DenseNet models while remaining inferior to MobileNetV3 and EfficientNet.The paper attributes the latter gap to inverted-residual models having much smaller model sizes and FLOPs for edge devices.
- Transfer Learning Results: On downstream classification datasets, AutoFormer achieves comparable or better accuracy with fewer parameters than existing convolutional and transformer models.The reported comparison includes CIFAR, fine-grained classification datasets, and roughly fourfold fewer parameters than transformer-based models.
- Transfer Learning Results: 75.7%/82.4%/82.9% top-1 accuracy is reached by AutoFormer-tiny/small/base after combining AutoFormer with knowledge distillation.The corresponding undistilled results are 74.7%/81.7%/82.4%.
5. Related Work
Related work spans vision transformers and one-shot neural architecture search, while AutoFormer targets pure vision transformers with a once-for-all supernet that avoids post-search retraining.
- Vision Transformer: ViT introduced pure transformer visual recognition using stacked transformer blocks over projected non-overlapping image patches.DeiT showed that ImageNet-scale data can support vision transformers without the unavailable JFT-300M dataset.
- Neural Architecture Search: NAS methods evolved from reinforcement learning and evolutionary search toward one-shot weight sharing with over-parameterized supernets.SPOS is described as a simple representative of this one-shot approach.
- Neural Architecture Search: OFA, BigNAS, and slimmable networks train once-for-all supernets for convolutional networks, whereas AutoFormer addresses transformer-specific architecture choices.AutoFormer searches dimensions involving multihead self-attention and MLP components.
- Neural Architecture Search: Prior transformer NAS studies mainly target natural language tasks, and HAT requires retraining or finetuning after search unlike AutoFormer.HAT also searches an encoder-decoder structure, while AutoFormer searches a pure encoder.
- Summary: AutoFormer differs by providing a vision-transformer-specific search method whose subnets inherit weights directly from a once-for-all supernet without extra finetuning or retraining.The paper presents this as its main distinction from previous work.
6. Conclusion
AutoFormer introduces a one-shot transformer search method with Weight Entanglement, finding promising architectures and achieving strong transfer and distillation results.
- AutoFormer uses Weight Entanglement to make subnets in the search space almost fully trained.The strategy is presented as a component of the one-shot architecture search method.
- Extensive experiments show that AutoFormer improves supernet training and finds promising transformer architectures.
- AutoFormers achieve state-of-the-art results on ImageNet among vision transformers.
- AutoFormers transfer well to several downstream classification tasks and can be further improved by distillation.
- Future work includes adding convolutions to the search space, applying Weight Entanglement to convolution search, and theoretically analyzing the strategy.