Accelerating Vision Transformers on Brain Processing Unit

With the advancement of deep learning technologies, specialized neural processing hardware such as Brain Processing Units (BPUs) have emerged as dedicated platforms for CNN acceleration, offering optimized INT8 computation capabilities for convolutional operations. Meanwhile, Vision Transformer (ViT) models, such as the Data-efficient Image Transformer (DeiT), have demonstrated superior performance and play increasingly crucial roles in computer vision tasks. However, due to the architectural mismatch between CNN-optimized hardware and Vision Transformer computation characteristics–namely, that linear layers in Transformers operate on three-dimensional data while BPU acceleration is designed for four-dimensional convolution operations-it is difficult or even impossible to leverage BPU’s advantages when deploying Vision Transformers. To address this challenge, we propose a novel approach that restructures the Vision Transformer by replacing linear layers and layer normalization operations with carefully designed convolutional operators. This enables DeiT to fully utilize the acceleration capabilities of BPUs, while allowing the original weight parameters to be inherited by the restructured models without retraining or fine-tuning. To the best of our knowledge, this is the first successful deployment of Vision Transformers that fully leverages BPU classification datasets demonstrate the effectiveness of our approach. Specifically, the quantized DeiT-Base model achieves 80.4% accuracy on ImageNet, compared to the original 81.8%, while obtaining up to a 3.8* inference speedup. Our finetuned DeiT model on the flower classification dataset also achieves excellent performance, with only a 0.5% accuracy drop for the DeiT-Base model, further demonstrating the effectiveness of our method.

💡 Research Summary

The paper addresses the fundamental mismatch between Vision Transformers (ViTs), specifically the Data‑efficient Image Transformer (DeiT), and Brain Processing Units (BPUs), which are specialized hardware accelerators designed for INT8 2‑D convolution on four‑dimensional tensors. While BPUs excel at accelerating CNNs, ViTs rely heavily on linear (fully‑connected) layers and Layer Normalization (LayerNorm) that operate on three‑dimensional data, making it impossible to exploit BPU’s hardware‑level optimizations for transformer inference.

To bridge this gap, the authors propose a systematic conversion of the core transformer operations into BPU‑friendly equivalents. Linear layers are replaced with point‑wise 1×1 convolutions: the weight matrix W (output‑dim × input‑dim) is reshaped into a 4‑D tensor of shape