Source-linked AI summary
Training data-efficient image transformers & distillation through attention
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, Hervé Jégou
TL;DR
Vision transformers had required large datasets and expensive training, limiting data-efficient adoption. This paper trains convolution-free DeiT models on ImageNet alone and introduces attention-based token distillation, achieving results competitive with convnets.
Problem
Vision transformers had required large curated datasets and extensive computing resources to achieve effective image classification.
Method
DeiT trains convolution-free vision transformers on ImageNet alone and uses a teacher-guided distillation token that interacts with the class token through attention.
Results
DeiT achieves competitive ImageNet results against convolutional networks without external data and transfers competitively to other tasks.
Takeaways & Limitations
Image transformers can be on par with convnets and may become a method of choice because of their lower memory footprint for a given accuracy.
Takeaways & Limitations
The distillation approach assumes access to a strong image-classifier teacher, such as a convnet or classifier mixture.
Abstract
from arXiv · showhide
Recently, neural networks purely based on attention were shown to address image understanding tasks such as image classification. However, these visual transformers are pre-trained with hundreds of millions of images using an expensive infrastructure, thereby limiting their adoption. In this work, we produce a competitive convolution-free transformer by training on Imagenet only. We train them on a single computer in less than 3 days. Our reference vision transformer (86M parameters) achieves top-1 accuracy of 83.1% (single-crop evaluation) on ImageNet with no external data. More importantly, we introduce a teacher-student strategy specific to transformers. It relies on a distillation token ensuring that the student learns from the teacher through attention. We show the interest of this token-based distillation, especially when using a convnet as a teacher. This leads us to report results competitive with convnets for both Imagenet (where we obtain up to 85.2% accuracy) and when transferring to other tasks. We share our code and models.
1 Introduction
The paper presents data-efficient, convolution-free vision transformers trained on ImageNet alone, including a transformer-specific distillation-token strategy. These models achieve competitive ImageNet and downstream-task performance while requiring substantially less training data and infrastructure than earlier vision transformers.
- Data-efficient training: 53 hours of pre-training on a single 8-GPU node produces DeiT models competitive with similarly sized and efficient convnets using ImageNet as the sole training set.Optional fine-tuning takes 20 hours.
- Transformer distillation: The transformer-specific DeiT⚗ procedure introduces a distillation token that interacts with the class token through attention to reproduce the teacher’s label.The distillation token plays the same role as the class token, but targets the teacher’s estimated label.
- Data-efficient training: Convolution-free networks achieve competitive ImageNet results with no external data, including smaller DeiT-S and DeiT-Ti models corresponding to ResNet-50 and ResNet-18.The models are learned on a single 4-GPU node in three days.
- Transformer distillation: The token-based distillation strategy outperforms vanilla distillation by a significant margin.The paper describes this strategy as specific to transformers.
- Transformer distillation: With the proposed distillation, image transformers learn more from a convnet than from another transformer with comparable performance.The comparison concerns teachers with comparable performance.
- Transfer performance: ImageNet-pretrained models remain competitive when transferred to fine-grained classification across CIFAR-10, CIFAR-100, Oxford-102 flowers, Stanford Cars, and iNaturalist-18/19.These are listed as popular public benchmarks.
2 Related work
Image classification has been dominated by convolutional networks, while prior vision transformers lagged behind or required large-scale pre-training. Related work also connects transformer-inspired attention mechanisms in vision with knowledge distillation from stronger teacher networks.
- Image classification: Since AlexNet, convnets have dominated ImageNet image classification and become the de facto standard for measuring progress in image understanding.Image classification progress often translates to improvements in related tasks such as detection and segmentation.
- Vision transformers: Earlier vision transformers underperformed convnets, whereas hybrid convnet-transformer architectures achieved competitive results across image, video, object-discovery, and text-vision tasks.The cited applications include classification, detection, video processing, unsupervised object discovery, and unified text-vision tasks.
- Vision transformers: Vision transformers closed the ImageNet state-of-the-art gap without convolutions, but prior work required pre-training on a large volume of curated data.This requirement contrasts with the paper’s stated goal of achieving strong performance without large training data.
- Transformer-inspired architectures: Transformer self-attention has influenced image-classification convnet designs, including Squeeze and Excitation, Selective Kernel, and Split-Attention Networks.The Transformer architecture was introduced for machine translation and is described as a reference model for NLP tasks.
- Knowledge distillation: Knowledge distillation trains a student with softmax outputs from a strong teacher, improving student performance or compressing the teacher into a smaller model.Soft labels preserve the teacher’s output vector rather than only its highest-scoring class.
3 Vision transformer: overview
The vision transformer represents images as sequences of fixed-size patch tokens processed by self-attention and feed-forward transformer blocks. A trainable class token produces predictions, while positional encoding supports changes in input resolution.
- Multi-head Self Attention layers (MSA): Self-attention matches queries with keys, applies scaled softmax weights, and returns weighted sums of the corresponding value vectors.For sequences, queries, keys, and values are computed through linear transformations of the input vectors.
- Transformer block for images: Each transformer block combines multi-head self-attention with a two-layer FFN that expands dimension D to 4D and reduces it back to D.Both components use residual skip-connections and layer normalization; the FFN uses a GeLu activation.
- Transformer block for images: A ViT decomposes each RGB image into N = 14 × 14 patches of 16 × 16 pixels, each linearly projected from dimension 3 × 16 × 16 = 768.The resulting patch sequence is processed as transformer tokens.
- Class token: A trainable class token is appended to patch tokens, traverses the transformer layers, and alone is projected to predict the class.The architecture therefore processes N + 1 tokens of dimension D while using only the class vector for output prediction.
- Fixing the positional encoding across resolutions: Keeping the patch size fixed while increasing image resolution changes N, so positional encoding must accommodate the altered number of input patches.The passage also describes training at lower resolution and fine-tuning at larger resolution.
4 Distillation through attention
This section develops transformer distillation from a strong image-classifier teacher, covering soft- and hard-label objectives and introducing a distillation token that learns through self-attention. At inference, the method combines predictions from separate class and distillation heads by late fusion.
- Distillation objectives: The framework distills a transformer from a strong teacher, including convnets or classifier mixtures, through soft- and hard-label alternatives.Soft distillation minimizes the KL divergence between teacher and student softmax outputs; hard distillation uses the teacher’s argmax prediction as a label.
- Distillation objectives: The hard teacher label can vary with data augmentation and is reported as better than traditional distillation while remaining parameter-free and simpler.The teacher prediction plays the same role as the ground-truth label.
- Distillation token: The distillation token is added alongside patch and class tokens, interacts with them through self-attention, and learns to reproduce the teacher’s hard label.Both the class and distillation tokens are learned by back-propagation.
- Distillation token: 0.06 is the average cosine similarity between learned class and distillation tokens initially, increasing to 0.93 at the final layer while remaining below 1.The different vectors reflect their similar but nonidentical target objectives.
- Distillation token: cos=0.999 is reached by two independently initialized class tokens with the same target, and the added token provides no classification benefit.This comparison supports the distinct contribution of the teacher-supervised distillation token.
- Classification with our approach: joint classifiers: At test time, the reference method uses late fusion, adding the softmax outputs of separate class-token and distillation-token classifiers.Both embeddings can independently infer the image label, and the three classification options are evaluated in Section 5.
5 Experiments
The experiments show that DeiT’s transformer-specific distillation, especially from a convolutional teacher, improves accuracy and efficiency while narrowing or surpassing convnet comparisons. Distilled DeiT models also transfer competitively beyond ImageNet and benefit from longer training.
- Accuracy and efficiency: 85.2% top-1 accuracy on ImageNet-1k makes the best distilled model outperform the best ViT-B model at 84.15%.The distilled model also outperforms EfficientNet in the throughput–accuracy trade-off, while DeiT nearly closes the transformer–convnet gap using ImageNet-only training.
- Teacher architecture: A convnet teacher gives better performance than a transformer teacher, with RegNetY-16GF serving as the default teacher in subsequent experiments.The authors attribute this advantage to inductive bias inherited through distillation.
- Distillation strategy: Hard distillation reaches 83.0% at resolution 224×224, compared with 81.8% for soft distillation in transformers.The proposed two-token strategy further improves performance by combining complementary class and distillation information.
- Training duration: With 300 epochs, the distilled DeiT-B⚗ is already better than DeiT-B, and longer training continues improving the distilled network.In contrast, DeiT-B performance saturates with longer schedules.
- Accuracy and efficiency: +6.3% top-1 in a comparable setting improves over previous ViT models trained only on ImageNet1k.DeiT-B⚗ and DeiT-B⚗↑384 also outperform the state of the art on the trade-off between accuracy and GPU inference time.
- Transfer learning: DeiT is on par with competitive convnet models on transfer-learning tasks beyond ImageNet.On CIFAR-10-only training, performance reaches 98.5% versus 99.1% with ImageNet pre-training, while still demonstrating reasonable transformer learning.
6 Training details & ablation
DeiT uses strong augmentation, AdamW-based optimization, stochastic regularization, and resolution fine-tuning to train transformers efficiently. The ablations identify augmentation and repeated augmentation as key ingredients, while EMA provides only a transient gain.
- Data-Augmentation: Strong augmentation improves results: Auto-Augment, Rand-Augment, and random erasing help, while dropout is excluded after ablation.Rand-Augment is selected instead of AutoAugment, and almost all evaluated augmentation methods prove useful.
- Regularization & Optimizers: AdamW performs best with ViT’s learning rates but much smaller weight decay, because the reported ViT weight decay harms convergence.The study cross-validates three learning rates and three weight-decay values, scaling learning rate with batch size.
- Regularization & Optimizers: Repeated augmentation provides a significant performance boost and is one of the key ingredients of the training procedure.Stochastic depth facilitates transformer convergence, while Mixup and Cutmix improve performance.
- Exponential Moving Average (EMA): 0.1 accuracy points is the EMA model’s edge before fine-tuning, but both EMA and non-EMA models reach the same improved performance afterward.The EMA gain vanishes after fine-tuning.
- Fine-tuning at different resolution: 224 is the default DeiT training resolution and 384 is the default fine-tuning resolution, with positional embeddings interpolated during resolution changes.Direct bilinear interpolation can reduce positional-vector ℓ2-norms and cause a significant accuracy drop.
- Training time: 37 hours is the typical 300-epoch DeiT-B training time with 2 nodes, versus 53 hours on a single node; DeiT-S and DeiT-Ti train in less than 3 days on 4 GPU.Fine-tuning DeiT-B at 384×384 takes 20 hours on a single node with 8 GPU for 25 epochs.
7 Conclusion
DeiT introduces data-efficient image transformers trained with improved training and a novel distillation procedure. The reported results place image transformers on par with convnets, while the method offers a lower memory footprint for a given accuracy.
- Contributions: DeiT introduces image transformers that do not require very large amounts of training data, enabled by improved training and a novel distillation procedure.The conclusion identifies improved training and distillation as the basis for reducing data requirements.
- Contributions: The method builds on data augmentation and regularization strategies developed for convnets while adding no significant architectural changes beyond the novel distillation token.The conclusion contrasts these established strategies with the new distillation token.
- Results: Image transformers are reported to be on par with convnets already, supporting their potential to become a method of choice.This conclusion is explicitly based on the paper’s reported results.
- Results: For a given accuracy, image transformers are argued to have a lower memory footprint than convnets.The conclusion presents lower memory usage as a reason for their expected adoption.
- Availability: The authors provide an open-source implementation of DeiT through the project’s GitHub repository.The repository is identified as https://github.com/facebookresearch/deit.