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

Tianxiao Zhang,Wenju Xu,Bo Luo,Guanghui Wang

2024-08-02

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 <a class="link-external link-https" href="https://github.com/ZTX-100/Efficient_ViT_with_DW" rel="external noopener nofollow">this https URL</a>.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper aims to address the following issues:

1. **Training efficiency of Vision Transformers (ViT) on small datasets**: Existing Vision Transformer models, due to their global self-attention mechanism, overlook the local relationships between adjacent pixels when processing images. This results in slower convergence and poor performance when training on small datasets.
2. **Combining the advantages of Convolution and Transformer**: Convolutional Neural Networks (CNNs) have a natural inductive bias that allows them to quickly capture local features, while ViTs excel at capturing global information. The focus of the research is on how to efficiently combine the strengths of these two models to improve performance on small datasets.

To address these issues, the authors propose a lightweight Depth-Wise Convolution module as a shortcut path in the ViT model. This allows the model to capture local information while retaining global context, significantly improving performance on small datasets without substantially increasing the number of parameters and computational complexity. Experimental results show that this method achieves significant performance improvements on multiple benchmark datasets (such as CIFAR-10, CIFAR-100, Tiny-ImageNet, etc.).