Block Selective Reprogramming for On-device Training of Vision Transformers

TL;DR

The paper addresses the challenge of on-device fine-tuning for Vision Transformers by focusing on activation-memory as the key bottleneck. It introduces Block Selective Reprogramming (BSR), which selectively tunes a small subset of transformer blocks and applies token dropping based on self-attention scores from frozen layers to reduce activation memory and compute while preserving accuracy. Empirical results show substantial gains, including up to $6\times$ activation-memory reduction and up to $2\times$ compute savings on ViT-B and DeiT-S across multiple datasets, with modest accuracy losses; the approach also extends to Mixture-of-Experts models for multitask learning. The work demonstrates that on-device ViT fine-tuning can be practical and efficient, and its concepts are orthogonal to activation-quantization techniques for further reductions.

Abstract

The ubiquity of vision transformers (ViTs) for various edge applications, including personalized learning, has created the demand for on-device fine-tuning. However, training with the limited memory and computation power of edge devices remains a significant challenge. In particular, the memory required for training is much higher than that needed for inference, primarily due to the need to store activations across all layers in order to compute the gradients needed for weight updates. Previous works have explored reducing this memory requirement via frozen-weight training as well storing the activations in a compressed format. However, these methods are deemed inefficient due to their inability to provide training or inference speedup. In this paper, we first investigate the limitations of existing on-device training methods aimed at reducing memory and compute requirements. We then present block selective reprogramming (BSR) in which we fine-tune only a fraction of total blocks of a pre-trained model and selectively drop tokens based on self-attention scores of the frozen layers. To show the efficacy of BSR, we present extensive evaluations on ViT-B and DeiT-S with five different datasets. Compared to the existing alternatives, our approach simultaneously reduces training memory by up to 1.4x and compute cost by up to 2x while maintaining similar accuracy. We also showcase results for Mixture-of-Expert (MoE) models, demonstrating the effectiveness of our approach in multitask learning scenarios.

Figures (5)

• Figure 1: Test accuracy, training time, and memory comparison for a ViT-B on CIFAR-10 with a batch size of 32. In particular, we achieve a 5.47$\times$ memory reduction and 2.43$\times$ training speedup while yielding similar test accuracy. Our benefits are even more significant for higher batch sizes.
• Figure 2: Block selective reprogramming framework for a ViT model with 12 layers. The token drop locations are set at $4^{th}$, $7^{th}$, and $10^{th}$ blocks, where 50% of incoming tokens are dropped based on token importance calculated by the MHSA module. The $4^{th}$, $8^{th}$, and $12^{th}$ blocks along with the classification head are trainable. The gradient propagation does not occur beyond the last trainable block.
• Figure 3: Left: Gradient flow in a residual learning architecture like yang2022repRight: Gradient flow in our block selective reprogramming approach
• Figure 4: Distribution of patch embedding, positional embedding and classification token of two DeiT-S: a pre-trained ImageNet model and a fully fine-tuned model on CIFAR-10
• Figure 5: Accuracy and on-device memory for DeiT-S and ViT-B models on CIFAR-100