Source-linked AI summary
Vision Transformer for Small-Size Datasets
Seung Hoon Lee, Seunghyun Lee, Byung Cheol Song
TL;DR
ViT’s limited locality inductive bias creates dependence on large-scale pre-training, making small-dataset learning from scratch difficult. The paper proposes SPT and LSA as broadly applicable add-ons, and reports improved ViT performance, including a 2.96% average gain on Tiny-ImageNet and a 4.08% gain for Swin Transformer.
Problem
ViT has lower locality inductive bias than CNNs and depends on large datasets such as JFT-300M, motivating methods for learning small-size datasets from scratch.
Method
SPT embeds neighboring-pixel information through spatially shifted images, while LSA promotes local attention using diagonal masking and learnable softmax temperature.
Results
2.96% average and 4.08% maximum classification improvements were reported on Tiny-ImageNet when SPT and LSA were applied to ViTs.
Takeaways & Limitations
SPT and LSA are presented as generic add-on modules applicable to various ViTs and capable of improving performance independently or together.
Abstract
from arXiv · showhide
Recently, the Vision Transformer (ViT), which applied the transformer structure to the image classification task, has outperformed convolutional neural networks. However, the high performance of the ViT results from pre-training using a large-size dataset such as JFT-300M, and its dependence on a large dataset is interpreted as due to low locality inductive bias. This paper proposes Shifted Patch Tokenization (SPT) and Locality Self-Attention (LSA), which effectively solve the lack of locality inductive bias and enable it to learn from scratch even on small-size datasets. Moreover, SPT and LSA are generic and effective add-on modules that are easily applicable to various ViTs. Experimental results show that when both SPT and LSA were applied to the ViTs, the performance improved by an average of 2.96% in Tiny-ImageNet, which is a representative small-size dataset. Especially, Swin Transformer achieved an overwhelming performance improvement of 4.08% thanks to the proposed SPT and LSA.
1. INTRODUCTION
ViT depends on large-scale pre-training for strong visual representation, partly because it lacks CNN-like locality inductive bias. The paper introduces SPT and LSA to improve locality and enable learning small datasets from scratch.
- ViT structurally has less locality inductive bias than CNNs and may require large-scale pre-training such as JFT-300M.
- Learning small-size datasets from scratch is challenging because of the trade-off between dataset capacity and performance.
- Two identified problems are poor tokenization, which limits receptive fields, and poor attention, which smooths attention-score distributions.
- SPT uses spatially shifted images to embed neighboring-pixel relations and widen visual-token receptive fields.
- LSA induces local attention by excluding self-tokens and learning the softmax temperature.
- 2.96% average and 4.08% maximum classification improvements were reported on Tiny-ImageNet when SPT and LSA were applied to ViTs.
2. RELATED WORK
Prior data-efficient ViTs addressed large-data dependence through augmentation, regularization, distillation, overlapping patches, or hierarchical receptive fields. This paper targets small-size datasets from scratch with more spatially informed tokenization and attention.
- DeiT improved ViT data efficiency through data augmentation, regularization, and knowledge distillation using a distillation token.
- T2T flattened overlapping patches before applying a transformer to capture local structure around tokens.
- The paper states that ViTs for small-size datasets had not yet been reported and proposes more spatially informed tokenization and attention.
- The proposed approach is intended to allow ViTs to learn small-size datasets from scratch.
3. PROPOSED METHOD
The proposed method combines SPT, which embeds neighboring spatial information into tokens, with LSA, which sharpens and localizes attention through temperature learning and diagonal masking.
- SPT: SPT spatially shifts an input image, concatenates shifted images with it, and then applies standard patch partitioning and token embedding.
- SPT: SPT embeds more spatial information into visual tokens and increases ViT locality inductive bias.
- LSA: LSA sharpens attention distributions by learning softmax temperature parameters.
- LSA: LSA applies diagonal masking to suppress self-token similarity and relatively increase attention between different tokens.
3.1. Preliminary
Standard ViT tokenizes non-overlapping image patches into visual tokens before self-attention, with token receptive fields fixed by patch size. Its attention scores are formed from Query–Key relations, but scaling and self-token relations can smooth the distribution and degrade performance on small datasets.
- Tokenization: ViT divides an image into non-overlapping patches, flattens them, and linearly projects each patch into a visual token.Each patch embedding becomes an input token to the transformer encoder.
- Tokenization: Token receptive fields are determined during tokenization and cannot be adjusted by the subsequent transformer encoder.For standard ViT, the token receptive field equals the patch size.
- Tokenization: A ViT token with patch size 16 has a receptive field of 16, compared with 483 for the last ResNet50 feature on 224 × 224 ImageNet images.The ViT token receptive field is therefore about 30 times smaller in this comparison.
- Self-attention: Self-attention projects tokens into Query, Key, and Value representations, applies softmax to scaled Query–Key similarities, and combines the resulting scores with Value.The scaling divides the similarity matrix by the square root of the Key dimension.
- Self-attention: Self-token relations tend to exceed inter-token relations, while division by √dk can act as a high softmax temperature and smooth attention scores.The paper reports that this smoothing degrades ViT performance and motivates LSA.
3.2. Shifted Patch Tokenization
Shifted Patch Tokenization (SPT) expands visual tokens’ spatial context by shifting and concatenating images before patch embedding. It can also be applied to pooling layers to reduce token count while retaining spatial information.
- Patch embedding layer: SPT shifts the input image in four diagonal directions, concatenates the shifted features with the original, and then partitions the result into patches.The shifted features are cropped, concatenated, flattened, normalized, and linearly projected into visual tokens.
- Patch embedding layer: SPT applies patch flattening, layer normalization, and linear projection to produce visual tokens from the concatenated spatial features.When used as a patch-embedding layer, class and positional embeddings are added afterward when applicable.
- Pooling layer: SPT can serve as a pooling layer by reshaping visual tokens into a spatial tensor, applying SPT, and embedding a reduced number of tokens.Class tokens are separated and reattached after the spatial tokenization process when present.
- Pooling layer: SPT embeds richer spatial information into visual tokens by increasing the receptive field according to the spatial shifts.The pooling formulation uses a learnable linear projection for the next stage.
3.3. Locality Self-Attention Mechanism
Locality Self-Attention (LSA) combines diagonal masking with learnable temperature scaling to make ViT attention less smooth and more focused on inter-token relations.
- Overview: LSA consists of diagonal masking and learnable temperature scaling.These are presented as the two core components of the locality attention mechanism.
- Diagonal masking: Diagonal masking excludes self-token relations from softmax by forcing the similarity matrix’s diagonal components to −∞.This directs attention toward relations between different tokens.
- Learnable temperature scaling: Learnable temperature scaling lets ViT determine the softmax temperature during training, with lower temperatures sharpening the attention score distribution.The learned temperature is lower than the constant temperature used by standard ViT.
- Effect of LSA: 0.5 larger average D_total_KL than standard ViT indicates that fully applied LSA attenuates smoothing of the attention score distribution.Lower D_total_KL corresponds to a flatter attention distribution.
4. EXPERIMENT
The experiments evaluate the proposed method across small- and mid-size image-classification datasets using several ViT architectures and standardized training configurations.
- Evaluation: The experiment section measures quantitative performance improvements and reduction of the gap with CNNs, alongside qualitative attention visualizations.The supplied implementation detail reports Tiny-ImageNet throughput measured with batch size 128 on an RTX 2080 Ti.
- Datasets: CIFAR-10, CIFAR-100, Tiny-ImageNet, and SVHN were used for small-size dataset experiments, while ImageNet was used for the mid-size experiment.The study evaluates the method across multiple dataset scales.
- Model configurations: The experiments cover ViT, PiT, T2T, Swin, and CaiT model families.Small-dataset ViT uses depth 9, hidden dimension 192, and 12 heads; ImageNet uses DeiT’s ViT-Tiny configuration.
- Training settings: All models use data augmentation, regularization, AdamW optimization, weight decay, batch sizes, and warm-up schedules adapted from DeiT.The listed techniques include CutMix, Mixup, AutoAugment, repeated augmentation, label smoothing, stochastic depth, and random erasing.
4.2. QUANTITATIVE RESULT
Across small- and mid-size datasets, SPT and LSA improve ViT performance, with gains reported for individual components, their combination, and multiple ViT architectures.
- CIFAR100 performance increased by +3.43% for CaiT and +4.01% for PiT, while Tiny-ImageNet performance increased by up to +4.00% for ViT and +4.08% for Swin.
- The proposed method generally improved ViTs with limited inference-latency overhead, including 1.12% for ViT, 1.15% for T2T, and 1.06% for CaiT.
- Learnable temperature scaling and diagonal masking improved Tiny-ImageNet performance by +0.88% and +1.22%, while both together achieved +1.43%.
- In Tiny-ImageNet, SPT alone improved performance by +1.43%, LSA alone by +3.60%, and both together by +4.00%.
4.3. QUALITATIVE RESULT
Attention visualizations show that applying SPT and LSA concentrates final class-token attention on target objects and captures object shapes more clearly across several ViTs.
- SPT and LSA concentrate attention from the background onto target classes, producing clearer object-shape capture in multiple ViTs.The effect is noted particularly for CaiT, T2T, ViT, and PiT.
5. CONCLUSION
The paper proposes SPT and LSA to increase ViT locality inductive bias, enabling small-size datasets to be learned from scratch and improving multiple ViTs.
- SPT embeds spatial information into visual tokens, while LSA promotes local attention through softmax with learnable parameters.
- Both techniques improve performance independently and are applicable to various ViT architectures.
- The study reports that ViTs can learn small-size datasets from scratch using these locality-enhancing techniques.
Supplementary
Supplementary experiments examine how SPT’s shift direction and ratio affect performance, finding that optimal settings vary across datasets.
- Three SPT shift-direction schemes use four cardinal, four diagonal, or eight combined directions.
- The best shift direction varies by dataset: four cardinal directions lead on CIFAR-10 and SVHN, four diagonal on CIFAR-100, and eight on Tiny-ImageNet.
- Shift intensity is defined as SHIFT = P ×rshift, where P is patch size and rshift is the shift ratio.
- The best shift ratio is 0.5 for CIFAR-100 and ImageNet, but 0.25 for Tiny-ImageNet.