[Paper Review] Locality Guidance for Improving Vision Transformers on Tiny Datasets

Kehan Li, Runyi Yu+4 more|arXiv (Cornell University)|Jul 20, 2022

CCD and CMOS Imaging Sensors4 citations

TL;DR

This paper proposes locality guidance via distillation from a lightweight CNN to improve Vision Transformers (VTs) on tiny datasets, where VTs typically underperform due to difficulty learning local image structures. By jointly training the VT to imitate CNN features and learn from labels, the method boosts performance significantly—e.g., +13.07% on DeiT, +8.98% on T2T, and +7.85% on PVT—without requiring large-scale pretraining.

ABSTRACT

While the Vision Transformer (VT) architecture is becoming trendy in computer vision, pure VT models perform poorly on tiny datasets. To address this issue, this paper proposes the locality guidance for improving the performance of VTs on tiny datasets. We first analyze that the local information, which is of great importance for understanding images, is hard to be learned with limited data due to the high flexibility and intrinsic globality of the self-attention mechanism in VTs. To facilitate local information, we realize the locality guidance for VTs by imitating the features of an already trained convolutional neural network (CNN), inspired by the built-in local-to-global hierarchy of CNN. Under our dual-task learning paradigm, the locality guidance provided by a lightweight CNN trained on low-resolution images is adequate to accelerate the convergence and improve the performance of VTs to a large extent. Therefore, our locality guidance approach is very simple and efficient, and can serve as a basic performance enhancement method for VTs on tiny datasets. Extensive experiments demonstrate that our method can significantly improve VTs when training from scratch on tiny datasets and is compatible with different kinds of VTs and datasets. For example, our proposed method can boost the performance of various VTs on tiny datasets (e.g., 13.07% for DeiT, 8.98% for T2T and 7.85% for PVT), and enhance even stronger baseline PVTv2 by 1.86% to 79.30%, showing the potential of VTs on tiny datasets. The code is available at https://github.com/lkhl/tiny-transformers.

Motivation & Objective

• Address the poor performance of Vision Transformers (VTs) on tiny datasets due to their global self-attention mechanism failing to capture local image structures.
• Overcome the limitations of pre-training-fine-tuning, which is costly, inflexible, and infeasible for domains lacking large-scale data like medical imaging.
• Develop a simple, efficient, and general-purpose method to enhance VTs on tiny datasets without requiring large-scale pretraining or architectural changes.
• Enable VTs to learn hierarchical local-to-global representations similar to CNNs, even when training from scratch on limited data.

Proposed method

• Introduce a dual-task learning framework where the VT learns from both supervised labels and feature imitation from a pre-trained lightweight CNN.
• Apply knowledge distillation by aligning hidden-layer features of the VT with those of a CNN trained on low-resolution images of the same dataset.
• Use a learnable linear projection to match channel dimensions between the CNN and VT features, enabling effective feature distillation.
• Balance the imitation loss and self-supervised loss via a hyperparameter β, allowing the VT to learn both guided locality and task-specific representations.
• Select guidance positions from multiple layers of the CNN, with optimal performance when all features are used (R = 1.0), reflecting hierarchical feature alignment.
• Use lightweight CNNs (e.g., ResNet-20, -56, -110) as guidance models, showing that model complexity has minimal impact on performance gain.

Experimental results

Research questions

• RQ1 Can locality guidance from a lightweight CNN improve Vision Transformer performance on tiny datasets without large-scale pretraining?
• RQ2 How does the choice of guidance position and feature layer selection affect the performance gain in VTs?
• RQ3 What is the optimal balance between imitation loss and self-supervised loss in the dual-task training framework?
• RQ4 Does the complexity of the guidance CNN significantly affect the performance improvement of the VT?
• RQ5 Can the proposed method generalize across different VT architectures and datasets, including medical imaging?

Key findings

• The method improves DeiT performance by 13.07% on tiny datasets, T2T-ViT by 8.98%, and PVT by 7.85% when training from scratch.
• Even stronger baseline PVTv2 achieves a 1.86% improvement, reaching 79.30% top-1 accuracy on CIFAR-100.
• Using all features from the CNN (R = 1.0) yields the best performance, indicating that full hierarchical feature alignment is critical.
• The optimal β value lies within a moderate range (e.g., β = 2.0–2.5), where imitation is strong but not overwhelming.
• A learnable linear projection for channel alignment outperforms fixed methods like Attention and Similarity, achieving the highest gain of 77.29% top-1 accuracy.
• The performance gain is robust across different CNN complexities, with ResNet-110 yielding only 76.62% accuracy on VTs, showing that lightweight guidance models are sufficient.