PaperlayerView exact revision on arXiv

Source-linked AI summary

HeatViT: Hardware-Efficient Adaptive Token Pruning for Vision Transformers

Peiyan Dong, Mengshu Sun, Alec Lu, Yanyue Xie, Kenneth Liu, Zhenglun Kong, Xin Meng, Zhengang Li, Xue Lin, Zhenman Fang, Yanzhi Wang

arXiv:2211.08110v2cs.ARcs.AIcs.CV

TL;DR

Image-adaptive token pruning lacks efficient hardware support for edge-device ViT inference. HeatViT combines a multi-head token selector with reused ViT hardware and 8-bit quantization. On ImageNet ViTs, it reports better accuracy-computation trade-offs than prior pruning studies.

  • Problem

    Prior studies have not considered efficient edge hardware for adaptive token pruning, which must also accommodate selector overhead, multi-head GEMM tiling, and nonlinear operations.

    Source: D. Static vs. Image-Adaptive Token… · p. 4, A. Challenges · p. 7

  • Method

    HeatViT uses an attention-based multi-head selector to consolidate non-informative tokens, reuses backbone GEMM hardware, and applies polynomial approximations for 8-bit quantization.

    Source: §I · pp. 1–3, B. Overview of HeatViT · p. 5

  • Results

    Against state-of-the-art ViT pruning studies, HeatViT achieves 0.7%–8.9% higher accuracy at similar computation cost, or more than 28.4%–65.3% computation-cost reduction at similar accuracy on ImageNet ViTs.

    Source: §I · p. 2

  • Takeaways & Limitations

    The FPGA proof of concept supports adaptive token pruning by reusing backbone hardware components and approximating nonlinear functions for 8-bit quantization.

    Source: §VIII · p. 11

Abstract

from arXiv · show

While vision transformers (ViTs) have continuously achieved new milestones in the field of computer vision, their sophisticated network architectures with high computation and memory costs have impeded their deployment on resource-limited edge devices. In this paper, we propose a hardware-efficient image-adaptive token pruning framework called HeatViT for efficient yet accurate ViT acceleration on embedded FPGAs. By analyzing the inherent computational patterns in ViTs, we first design an effective attention-based multi-head token selector, which can be progressively inserted before transformer blocks to dynamically identify and consolidate the non-informative tokens from input images. Moreover, we implement the token selector on hardware by adding miniature control logic to heavily reuse existing hardware components built for the backbone ViT. To improve the hardware efficiency, we further employ 8-bit fixed-point quantization, and propose polynomial approximations with regularization effect on quantization error for the frequently used nonlinear functions in ViTs. Finally, we propose a latency-aware multi-stage training strategy to determine the transformer blocks for inserting token selectors and optimize the desired (average) pruning rates for inserted token selectors, in order to improve both the model accuracy and inference latency on hardware. Compared to existing ViT pruning studies, under the similar computation cost, HeatViT can achieve 0.7%$\sim$8.9% higher accuracy; while under the similar model accuracy, HeatViT can achieve more than 28.4%$\sim$65.3% computation reduction, for various widely used ViTs, including DeiT-T, DeiT-S, DeiT-B, LV-ViT-S, and LV-ViT-M, on the ImageNet dataset. Compared to the baseline hardware accelerator, our implementations of HeatViT on the Xilinx ZCU102 FPGA achieve 3.46$\times$$\sim$4.89$\times$ speedup.

I. INTRODUCTION

HeatViT addresses ViTs’ token-related computation and memory costs, the limits of existing pruning approaches, and hardware-unfriendly nonlinear operations. It combines image-adaptive token pruning, hardware-aware implementation, quantization, and latency-aware training for embedded FPGA inference.

I. INTRODUCTION

  • Existing approaches leave a gap: static pruning ignores image-specific redundancy, adaptive methods discard tokens and overlook cross-head redundancy, and prior studies do not support efficient edge hardware implementation.The paper identifies quadratic self-attention time and memory costs as a deployment challenge for resource-constrained devices.Source: §I · p. 1
  • The selector scores tokens using redundancies across attention heads, retains informative tokens, and packages non-informative tokens together for later blocks.The design reflects that token redundancy differs across heads and that tokens deemed non-informative earlier may still hold information for later blocks.Source: §I · p. 2
  • For FPGA efficiency, the selector uses linear layers and dense token inputs to reuse backbone GEMM hardware, while 8-bit quantization and polynomial approximations target nonlinear operations.The approximations cover GELU, Softmax, and Sigmoid, with regularization of quantization error to support more ambitious quantization.Source: §I · p. 2
  • Latency-aware multi-stage fine-tuning inserts selectors from later blocks backward, raises pruning until accuracy loss exceeds a threshold, and consolidates consecutive selectors with similar rates.The authors report that selector insertion requires roughly 90% of the training effort used to train the backbone from scratch.Source: §I · p. 2
  • At similar computation cost, HeatViT achieves 0.7%–8.9% higher accuracy, and at similar accuracy it reduces computation by more than 28.4%–65.3% across five ImageNet ViTs.Compared with state-of-the-art ViT pruning studies, the models include DeiT-T, DeiT-S, DeiT-B, LV-ViT-S, and LV-ViT-M.Source: §I · p. 2
  • On Xilinx ZCU102, HeatViT achieves 3.46×–4.89× speedup over the 16-bit, unpruned baseline, with 8%–11% more DSPs and 5%–8% more LUTs.The reported speedup accompanies a trivial resource-utilization overhead.Source: §I · p. 2

II. BACKGROUND AND RELATED WORK

A Vision Transformer converts an image into patch tokens, processes them through transformer encoders, and uses the final class token for classification. The encoders combine multi-head self-attention and MLP operations.

A. Vision Transformer

  • An image is reshaped into 2D patches, projected into D-dimensional patch embeddings, and combined with a prepended class embedding and learnable positional embeddings.The number of patches N is the effective input sequence length to the transformer encoders.Source: A. Vision Transformer · p. 3
  • The final encoder’s class token is fed to an MLP classification head with two linear layers and an intermediate GELU activation.Source: A. Vision Transformer · p. 3
  • Each transformer encoder block contains a multi-head self-attention module and an MLP feed-forward module, with layer normalization applied before each module.Source: A. Vision Transformer · p. 3
  • Each attention head forms query, key, and value through linear projections, and the outputs from all heads are concatenated and linearly transformed.Source: A. Vision Transformer · p. 3

B. Model Pruning on ViT

ViT pruning targets redundancy in attention heads, token channels, and token counts; pruning token counts can reduce operations linearly or quadratically. HeatViT is compared with prior pruning methods, while static and image-adaptive pruning and FPGA accelerator limitations are discussed.

B. Model Pruning on ViT

  • ViT pruning targets redundancy in attention heads, token channels, and token counts, with token pruning aimed at removing non-informative input data.Head pruning can reduce computation but may harm accuracy; channel pruning can make high pruning rates difficult without accuracy deterioration.Source: B. Model Pruning on ViT · p. 3

C. Computational Complexity Analysis

  • Pruning token number N reduces all operations linearly or quadratically, whereas pruning other dimensions offers fewer compression benefits.The quadratic dependence occurs in layers y and z; token pruning can also be integrated with model quantization.Source: C. Computational Complexity Analysis · p. 3

D. Static vs. Image-Adaptive Token Pruning

  • At similar computation cost, HeatViT achieves 0.7%∼8.9% higher accuracy; at similar accuracy, it reduces computation cost by more than 28.4%∼65.3%.The final HeatViT models are quantized to 8-bit fixed-point format.Source: Figure 2 · p. 4
  • Static pruning applies a fixed token-reduction ratio across images, while image-adaptive pruning varies the rate with image content and can ideally achieve a larger overall rate.Existing adaptive methods may discard tokens completely, have limited overall pruning rates, and lack efficient edge-hardware implementations.Source: D. Static vs. Image-Adaptive Token… · pp. 3–4

E. Transformer Accelerators on FPGAs

  • Prior FPGA accelerators compress model weights, with coarse pruning granularity or omitted FFN computation cited as limitations.The cited accelerators are mainly for NLP rather than computer vision, while adaptive token pruning at the data level is reported to speed inference with negligible accuracy losses and little hardware overhead.Source: E. Transformer Accelerators on FPGAs · p. 4, Figure 5 · p. 5

III. ANALYSIS OF TOKEN REDUNDANCY IN VIT AND OVERVIEW OF HEATVIT

Token redundancy varies across attention heads and transformer blocks: later blocks align more strongly with the final CLS token, while heads capture distinct token-level redundancy. HeatViT uses a multi-head classifier and attention-based head weighting to score tokens, alongside a packager that preserves information from pruned tokens.

A. Token Redundancy in ViT

  • Different attention heads contain distinct token-level redundancy, and token similarity to the final CLS token strengthens from earlier to later blocks.This motivates lower pruning rates in front blocks and a packager that can help recover from pruning mistakes.Source: A. Token Redundancy in ViT · p. 5

A. Attention-Based Multi-Head Token Classifier

  • The token classifier scores each input token separately for each attention head, reflecting that heads extract different image features.An attention-based branch estimates head importance and combines head-level scores into an overall token score.Source: A. Attention-Based Multi-Head Token… · pp. 5–6

IV. ADAPTIVE TOKEN PRUNING MODULE

  • HeatViT inserts token selectors between transformer stages; each selector combines a multi-head token classifier with a token packager.The packager consolidates rather than discards non-informative tokens, preserving information for later blocks.Source: §IV · p. 5, B. Overview of HeatViT · p. 5, Figure 7 · p. 6

B. Token Packager

The token packager summarizes predicted non-informative tokens into one package token instead of discarding them. That token continues through later blocks with informative tokens, allowing the model to correct scoring mistakes while operating on a smaller dense input.

B. Token Packager

  • Predicted non-informative tokens are compressed into one package token by weighted averaging.The weights use the tokens’ keep scores.Source: B. Token Packager · p. 6
  • The package token continues through subsequent blocks, giving the model an opportunity to correct token-scoring mistakes.It is carried alongside the informative tokens.Source: B. Token Packager · p. 6
  • Informative tokens and the package token form a smaller dense matrix for subsequent computation, avoiding sparse hardware computation.The resulting smaller input directly speeds inference in following blocks.Source: B. Token Packager · p. 6

V. VIT HARDWARE ACCELERATOR DESIGN WITH ADAPTIVE TOKEN PRUNING

HeatViT’s FPGA accelerator reuses backbone ViT computation hardware for token selection and accommodates multi-head attention in its GEMM tiling. Its throughput depends on parallelism constrained by available compute and on-chip memory resources.

A. Challenges

  • The design must minimize selector overhead, support an extra GEMM tiling dimension for multi-head parallelism, and implement nonlinear operations efficiently under quantization.These are the three hardware challenges identified for dynamic token pruning.Source: A. Challenges · p. 7

B. ViT Accelerator with Dynamic Token Selection

1) Loop Tiling in GEMM:

  • GEMM loop tiling adds a head dimension for attention computations and tunes Ti, To, and Th to improve throughput.Attention inputs and outputs are divided into h groups; Ti and To tile the input and output dimensions, while Th tiles heads.Source: 1) Loop Tiling in GEMM · p. 7

2) Throughput and Resource Utilization Analysis:

  • Inference throughput depends on parallelism factors Ti, To, and Th, which are bounded by available DSPs and BRAMs.Throughput is inferred from the clock cycles needed to process all ViT layers and the FPGA working frequency.Source: 2) Throughput and Resource Utilization… · p. 7

V. VIT HARDWARE ACCELERATOR DESIGN WITH ADAPTIVE TOKEN PRUNING

  • The accelerator architecture covers both pruned ViTs with token selectors and the GEMM engine’s computation flow.The architecture figure identifies these two design views.Source: §V · p. 6, Figure 8 · p. 7

C. Token Selection Flow

Token selection classifies tokens, concatenates informative ones, and averages the others into a package token. For nonlinear functions, HeatViT uses polynomial approximations with regularization for GELU and Softmax; reported implementations reduce resource use relative to the built-in HLS library without accuracy drops in the tested pruning configurations.

C. Token Selection Flow

  • Token selection classifies tokens by a Gumbel-Softmax threshold, concatenates informative tokens, and averages non-informative ones into a package token.The sparse selection is reorganized as a dense sequence to avoid indexing overhead on hardware.Source: C. Token Selection Flow · p. 7, Figure 9 · p. 8

D. Polynomial Approximation of Nonlinear Functions

  • Polynomial approximations for GELU and Softmax use δ1 and δ2 to regularize quantization error; the token selector’s Sigmoid approximation has no such regularization.The authors identify the erf computation in GELU as resource-intensive with the built-in Xilinx HLS math library.Source: D. Polynomial Approximation of… · pp. 7–8
  • 1.5×∼572× resource improvement is reported for the proposed nonlinear-function implementations over the built-in Xilinx Vitis HLS math library.The authors report no accuracy drops between approximate and original models across the tested token-pruning ratios.Source: D. Polynomial Approximation of… · p. 8, Table III · p. 8

E. Regularization Effect on Quantization Error

HeatViT uses regularization effects on quantization error to support more aggressive approximation of GELU and Softmax. The paper shows that each approximation reduces output error relative to the input quantization error.

E. Regularization Effect on Quantization Error

  • GELU and Softmax approximations exploit regularization effects on quantization error to support more aggressive quantization, including 8-bit fixed-point quantization.The paper examines how input quantization error propagates through both nonlinear functions.Source: E. Regularization Effect on… · p. 8
  • For GELU, the output quantization error is reduced after approximation because its derivative is always below 1.The analysis treats quantization on x as adding a small error to the input.Source: E. Regularization Effect on… · p. 8
  • For Softmax, the total output error after approximation is less than the input error Δe0 under the stated bound.The passage derives this bound using 0 ≤ A0 ≤ 1 and δ2 < 1.Source: E. Regularization Effect on… · pp. 8–9

VI. LATENCY-AWARE MULTI-STAGE TRAINING STRATEGY

The latency-aware multi-stage strategy uses measured hardware latency to select token-selector locations and pruning rates while respecting an accuracy constraint. It progressively inserts and consolidates selectors, with quantization providing additional speedup.

VI. LATENCY-AWARE MULTI-STAGE TRAINING STRATEGY

  • The strategy combines a hardware-aware latency-sparsity loss with a block-to-stage pipeline to learn selector number, locations, and pruning rates.The latency-sparsity table maps keep ratios to measured FPGA block latency, and 8-bit quantization provides 2×–2.4× speedup without accuracy drops.Source: §VI · p. 9, Table IV · p. 9
  • The average pruning rate across images in a batch is the convergence target for image-adaptive pruning under the latency-sparsity loss.The training objective also includes cross-entropy and distillation losses.Source: §VI · p. 9
  • The block-to-stage pipeline inserts selectors between adjacent blocks, adjusts their pruning rates, and returns the pruned ViT when the latency limit is satisfied.If the final latency remains above the target, the strategy relaxes the accuracy or latency constraints and repeats training.Source: §VI · pp. 9–10
  • Selectors are inserted progressively from later blocks toward earlier ones, increasing pruning until accuracy loss exceeds the threshold or the target latency is met.The process fine-tunes the current selector and other network components at each insertion.Source: §VI · p. 9
  • Insertion stops before the fourth block because pruning in the front three blocks causes more severe accuracy drops; adjacent selectors are merged when their pruning rates differ by less than 8.5%.For each insertion, the stated stopping threshold is an accuracy drop above 0.5%.Source: §VI · pp. 9–10

A. Experimental Setup

Experiments cover multiple ViT backbones on ImageNet-1K and evaluate HeatViT on a ZCU102 FPGA. Training uses selector insertions and an 8-bit fixed-point hardware implementation.

A. Experimental Setup

  • Quantization after token pruning preserves accuracy in the reported settings, except for a 1.2% drop on DeiT-T.The experiments apply 8-bit quantization to weights and activations.Source: A. Experimental Setup · p. 10

1) Training Setup for ViT Pruning:

  • Training uses ImageNet-1K and DeiT, LV-ViT backbones; each selector insertion costs 30 epochs on eight A100 GPUs, and most models use 3–4 selectors.The total block-to-stage training effort is reported as equivalent to training the backbone ViT from scratch.Source: 1) Training Setup for ViT Pruning · p. 10, Table V · p. 10

2) Hardware Platform:

  • The hardware evaluation uses a 150 MHz Xilinx ZCU102 FPGA, with model data in 8-bit fixed-point format and accelerators generated using Vitis HLS.The platform includes 2520 DSPs, 912 BRAM blocks, and 274.1k LUTs.Source: 2) Hardware Platform · p. 10

B. Accuracy and GMACs Results

HeatViT reports improved accuracy-computation trade-offs across ViT models, alongside FPGA speedups and modest selector resource overhead. MLP selectors reuse ViT GEMM hardware, while polynomial approximation improves GELU resource utilization.

B. Accuracy and GMACs Results

  • HeatViT reports 16.1%–42.6% pruning for various backbones with no more than 0.75% accuracy degradation, and accuracy gains of 4%, 4.67%, and 0.81% over specified scaled DeiT baselines.The corresponding accuracy pairs are 72.1% vs. 68.1% at 0.9 GMACs, 76.87% vs. 72.20% at 1.3 GMACs, and 79.34% vs. 78.53% at 2.64 GMACs.Source: Figure 12 · p. 10

1) Operations in Token Selector:

  • MLP-based selectors outperform convolution-based selectors at 0.9 GMACs and reuse the original ViT GEMM hardware without additional resources for the new computation design.The paper also reports 35×–572× improved GELU resource utilization with polynomial approximation.Source: 1) Operations in Token Selector · p. 10

C. Hardware Results

  • Selector-equipped HeatViT accelerators use 9%–11% more DSPs and 5%–8% more LUTs than the 16-bit, unpruned baselines, which the paper characterizes as negligible overhead.The resource differences vary by model family.Source: C. Hardware Results · p. 10

1) Comparisons with CPUs and GPUs:

  • Final FPGA speedups are 3.46× for DeiT-T, 4.22× for DeiT-S, 4.59× for LV-ViT-S, and 4.89× for DeiT-B.The corresponding frame rates are 271.2, 109.2, 89.1, and 54.8 FPS, respectively.Source: Figure 13 · p. 11

VIII. CONCLUSION

  • The conclusion reports superior pruning-rate and accuracy results versus state-of-the-art pruning studies, with trivial hardware-resource overhead.The comparison scope is the paper’s evaluated ViT pruning studies and FPGA implementation.Source: §VIII · p. 11
Loading 2211.08110v2…