Compress image to patches for Vision Transformer
Xinfeng Zhao, Yaoru Sun
TL;DR
This work tackles the high computational cost of Vision Transformers (ViTs) on large images by introducing CI2P-ViT, a CNN-ViT hybrid that replaces ViT's Patch Embedding with a CNN-based image compression pipeline (CI2P). The CI2P-Encoder, based on CompressAI, compresses the input and outputs a downsampled latent that is reshaped into ViT-compatible tokens, with the encoder frozen during ViT training to reduce compute. A dual-scale variant, CI2P-ViT$^{ds}$, further improves efficiency by performing early attention at a larger scale and then downsampling, reducing parameters and FLOPs while improving ImageNet performance. On Animals-10, CI2P-ViT achieves 92.37% accuracy with 63.35% fewer FLOPs and faster training, while the dual-scale variant reduces FLOPs by ~72% and reaches 77% Top-1 on ImageNet, illustrating a practical path to more efficient ViT-based vision systems. Overall, the approach retains ViT architecture, enhances accuracy through CNN inductive biases, and enables scalable deployment on resource-constrained settings.
Abstract
The Vision Transformer (ViT) has made significant strides in the field of computer vision. However, as the depth of the model and the resolution of the input images increase, the computational cost associated with training and running ViT models has surged dramatically. This paper proposes a hybrid model based on CNN and Vision Transformer, named CI2P-ViT. The model incorporates a module called CI2P, which utilizes the CompressAI encoder to compress images and subsequently generates a sequence of patches through a series of convolutions. CI2P can replace the Patch Embedding component in the ViT model, enabling seamless integration into existing ViT models. Compared to ViT-B/16, CI2P-ViT has the number of patches input to the self-attention layer reduced to a quarter of the original. This design not only significantly reduces the computational cost of the ViT model but also effectively enhances the model's accuracy by introducing the inductive bias properties of CNN. The ViT model's precision is markedly enhanced. When trained from the ground up on the Animals-10 dataset, CI2P-ViT achieved an accuracy rate of 92.37%, representing a 3.3% improvement over the ViT-B/16 baseline. Additionally, the model's computational operations, measured in floating-point operations per second (FLOPs), were diminished by 63.35%, and it exhibited a 2-fold increase in training velocity on identical hardware configurations.