Accelerating Vision Transformers on Brain Processing Unit

Technical Summary: Accelerating Vision Transformers on Brain Processing Unit

Problem Statement
The paper addresses a critical architectural mismatch between specialized neural processing hardware and modern Vision Transformer (ViT) models. Specifically, the Horizon Brain Processing Unit (BPU) is optimized for four-dimensional convolution operations (INT8 quantized) typical of Convolutional Neural Networks (CNNs). However, ViT models, such as the Data-efficient Image Transformer (DeiT), rely heavily on linear layers and Layer Normalization (LayerNorm) that operate on three-dimensional data streams. This discrepancy prevents the BPU from accelerating the core computations of ViTs, forcing such models to run on general-purpose CPUs and forfeiting the hardware's acceleration capabilities.

Methodology
To resolve this incompatibility without requiring model retraining, the authors propose a structural reformation of the DeiT architecture. The core methodology involves replacing the non-convolutional operators with mathematically equivalent convolutional operators that the BPU can accelerate:

1. Linear Layer Transformation: Linear layers, which typically process 3D data $(B, N, C)$, are transformed into point-wise convolutions. This is achieved by reshaping the weight matrix $W \in \mathbb{R}^{O \times I}$ into a 4D tensor $W' \in \mathbb{R}^{O \times I \times 1 \times 1}$ via unsqueeze operations. The input tensor is similarly transposed from $(B, N, C)$ to $(B, C, 1, N)$ to align with the BPU's 4D data stream requirements.
2. LayerNorm Reformulation: The Layer Normalization operation is decomposed into a sequence of $1 \times 1$ convolution operations. The mean is calculated using a convolution with weights set to $1/H$ (where $H$ is the hidden dimension), and the variance is computed by subtracting the mean, squaring the result, and applying a second convolution. The learnable parameters $\gamma$ and $\beta$ are reshaped to match the input dimensions.
3. Weight Inheritance: A crucial aspect of this approach is that the restructured model does not require retraining or fine-tuning. The original pre-trained weights from standard DeiT checkpoints are directly inherited by the restructured convolutional layers.
4. Quantization: The modified models are processed through the Horizon BPU algorithm toolchain, utilizing Post-Training Quantization (PTQ) with a min-max strategy and a calibration dataset to convert the models to INT8 precision.

Key Contributions

• First Successful Deployment: The authors claim this is the first successful deployment of Vision Transformers on the BPU platform that fully leverages its hardware acceleration capabilities.
• Architecture Adaptation: They propose a novel method to restructure DeiT by replacing linear and LayerNorm layers with equivalent convolutional operators, enabling INT8 acceleration without architectural loss.
• No-Retraining Deployment: The method allows for the direct inheritance of pre-trained weights, eliminating the computational cost and complexity of retraining or fine-tuning the entire model for the new operator structure.
• Extensive Validation: The approach is validated on two distinct datasets: the large-scale ImageNet dataset and a smaller flower classification dataset.

Experimental Results
The paper reports extensive experiments comparing the original DeiT models with the proposed BPU-optimized, quantized versions:

• ImageNet Performance:
  • Accuracy: The quantized DeiT-Base model achieves 80.4% Top-1 accuracy on ImageNet, a drop of only 1.4% from the original 81.8%. Distilled variants show even greater resilience, with the DeiT-Base Dist. maintaining 81.7% accuracy (a 1.7% drop).
  • Anomaly Detection: The authors observed that both the original and quantized models correctly predicted labels for several images in the ImageNet validation set that were known to be mislabeled in the official annotations, suggesting the models' true accuracy may be slightly higher than reported metrics.
• Flower Classification Performance:
  • On a dataset with 5 flower classes, the quantized DeiT-Small model achieved 97.4% accuracy with only a 0.3% drop from the full-precision version.
• Inference Speedup:
  • The quantized models demonstrated significant inference speedups on the BPU hardware compared to floating-point counterparts.
  • DeiT-Base: Achieved a 3.82× speedup.
  • DeiT-Small: Achieved a 2.16× speedup.
  • DeiT-Tiny: Achieved a 1.30× speedup.
  • The results indicate that the benefits of quantization and hardware acceleration are more pronounced in larger model architectures.

Significance and Claims
The paper positions this work as a pivotal step in bridging the gap between Transformer-based computer vision models and specialized embedded hardware. By demonstrating that ViTs can be deployed on BPUs with minimal accuracy degradation and substantial speed improvements, the authors argue that their method provides a practical solution for real-world edge deployment scenarios. The ability to utilize pre-trained weights without retraining is highlighted as a major efficiency gain, making the transition to hardware-accelerated Transformers feasible for resource-constrained environments. The authors modestly claim this represents the first successful attempt to fully utilize BPU acceleration for DeiT models.