Depth-Wise Convolutions in Vision Transformers for Efficient Training on Small Datasets

TL;DR

This paper addresses the inefficiency of Vision Transformers (ViTs) on small datasets due to the lack of local inductive bias. It introduces a lightweight Depth-Wise Convolution (DWConv) module that bypasses Transformer blocks to inject local details, enabling simultaneous local and global representation with minimal overhead. The authors present architecture variants, including bypassing multiple blocks and parallel kernels, and demonstrate substantial accuracy gains across ViT, CaiT, and Swin on CIFAR-10/100, Tiny-ImageNet, ImageNet-1K, and COCO, with faster convergence and negligible parameter increases. The approach enables small ViT models to outperform larger counterparts trained on limited data, offering a practical, plug-and-play fusion of convolutions with Transformers that broadens the applicability of ViTs in data-constrained scenarios.

Abstract

The Vision Transformer (ViT) leverages the Transformer's encoder to capture global information by dividing images into patches and achieves superior performance across various computer vision tasks. However, the self-attention mechanism of ViT captures the global context from the outset, overlooking the inherent relationships between neighboring pixels in images or videos. Transformers mainly focus on global information while ignoring the fine-grained local details. Consequently, ViT lacks inductive bias during image or video dataset training. In contrast, convolutional neural networks (CNNs), with their reliance on local filters, possess an inherent inductive bias, making them more efficient and quicker to converge than ViT with less data. In this paper, we present a lightweight Depth-Wise Convolution module as a shortcut in ViT models, bypassing entire Transformer blocks to ensure the models capture both local and global information with minimal overhead. Additionally, we introduce two architecture variants, allowing the Depth-Wise Convolution modules to be applied to multiple Transformer blocks for parameter savings, and incorporating independent parallel Depth-Wise Convolution modules with different kernels to enhance the acquisition of local information. The proposed approach significantly boosts the performance of ViT models on image classification, object detection, and instance segmentation by a large margin, especially on small datasets, as evaluated on CIFAR-10, CIFAR-100, Tiny-ImageNet and ImageNet for image classification, and COCO for object detection and instance segmentation. The source code can be accessed at https://github.com/ZTX-100/Efficient_ViT_with_DW.

Figures (7)

• Figure 1: The architecture of our proposed method. The Depth-Wise Convolution module bypasses the entire Transformer block so that the local details can be attained and added to the output of the Transformer block. In the DWConv module, the 1D image patch tokens are first reshaped to 2D feature maps. If the class token exists, it would not be involved in the DWConv module and only image patch tokens are utilized to reconstruct the feature maps. Batch normalization and GELU activation are employed before the Depth-Wise Convolution. Finally, the feature maps would be reshaped to 1D tokens and added to the output of the Transformer block. The DWConv module is exploited in all Transformer blocks.
• Figure 2: The architecture variants of our proposed approach involve bypassing multiple Transformer blocks. Structures (a), (b), and (c) represent the Depth-Wise module bypassing 2, 3, and 4 Transformer blocks, respectively. For Vision Transformer models with deeper layers, bypassing additional blocks may be a beneficial strategy to reduce both parameters and computational costs.
• Figure 3: The architecture variants of our method involve multiple DWConv modules operating in parallel. These independent DWConv modules, each with different kernel sizes, run concurrently on Transformer blocks to capture local details simultaneously. This structure can be combined with previous variants shown in Fig. \ref{['fig:2']} to include N Transformer blocks in the DWConv modules.
• Figure 4: Some Grad-CAM visualization with ViT-Tiny and Swin-Tiny models. The vanilla Transformer models tend to capture the global information, as illustrated in the CAM visualizations. With our method, the models are able to capture both local details and global perspectives, particularly when dealing with smaller objects. Please note that the original images in the figure are from the Tiny-ImageNet dataset with a low resolution of 64$\times$64 pixels. Thus, they appear blurred when enlarged.
• Figure 5: The accuracy for val set during the training on Tiny-ImageNet for 300 epochs. The blue curves indicate our method and the red curves are from the original models. The accuracy for val set is recorded for each epoch. The convergence of the models with our approach is much faster than the original models.