PaperlayerView exact revision on arXiv

Source-linked AI summary

DiffRate : Differentiable Compression Rate for Efficient Vision Transformers

Mengzhao Chen, Wenqi Shao, Peng Xu, Mingbao Lin, Kaipeng Zhang, Fei Chao, Rongrong Ji, Yu Qiao, Ping Luo

arXiv:2305.17997v1cs.CV

TL;DR

Token compression can make vision transformers more efficient, but existing approaches rely on manually chosen compression rates. DiffRate learns compression rates differentiably while combining pruning and merging; on an off-the-shelf ViT-H (MAE), it reduces FLOPs 40% with a 0.16% accuracy drop.

  • Problem

    Handcrafted token-compression rates are tedious to set across layers and can lead to suboptimal performance, while ViTs' self-attention cost grows quadratically with token length.

    Source: §1 · pp. 1–2

  • Method

    DiffRate uses a Differentiable Discrete Proxy to learn compression rates with gradient back-propagation and integrates token pruning and merging.

    Source: §1 · p. 2

  • Results

    40% FLOPs reduction and 0.16% accuracy drop on an off-the-shelf ViT-H (MAE), with 50% throughput improvement and no fine-tuning.

    Source: §1 · p. 2

  • Takeaways & Limitations

    DiffRate identifies an appropriate compression rate using only 1,000 images and performs comparably or better than prior state-of-the-art token-compression methods without fine-tuning.

    Source: §7 · p. 9

Abstract

from arXiv · show

Token compression aims to speed up large-scale vision transformers (e.g. ViTs) by pruning (dropping) or merging tokens. It is an important but challenging task. Although recent advanced approaches achieved great success, they need to carefully handcraft a compression rate (i.e. number of tokens to remove), which is tedious and leads to sub-optimal performance. To tackle this problem, we propose Differentiable Compression Rate (DiffRate), a novel token compression method that has several appealing properties prior arts do not have. First, DiffRate enables propagating the loss function's gradient onto the compression ratio, which is considered as a non-differentiable hyperparameter in previous work. In this case, different layers can automatically learn different compression rates layer-wisely without extra overhead. Second, token pruning and merging can be naturally performed simultaneously in DiffRate, while they were isolated in previous works. Third, extensive experiments demonstrate that DiffRate achieves state-of-the-art performance. For example, by applying the learned layer-wise compression rates to an off-the-shelf ViT-H (MAE) model, we achieve a 40% FLOPs reduction and a 1.5x throughput improvement, with a minor accuracy drop of 0.16% on ImageNet without fine-tuning, even outperforming previous methods with fine-tuning. Codes and models are available at https://github.com/OpenGVLab/DiffRate.

1. Introduction

DiffRate addresses the difficulty of hand-setting token-compression rates by learning them differentiably while combining pruning and merging. The paper reports improved efficiency with a small accuracy drop on an off-the-shelf ViT-H model.

1. Introduction

  • ViTs' self-attention complexity grows quadratically with token length, motivating efficiency improvements for practical applications.Token compression can reduce redundancy without modifying network structure and complements other compression methods.Source: §1 · pp. 1–2
  • Handcrafted layer-wise compression rates make it tedious to meet complexity constraints while retaining performance, and can discard or merge informative foreground tokens.Token pruning and merging each have distinct token-selection behavior, but both use manually selected rates.Source: §1 · p. 2
  • The Differentiable Discrete Proxy sorts tokens by importance and uses reparameterization to select top-K tokens with gradient back-propagation.This lets images retain the same number of top-K tokens for parallel batch computation; optimization converges within 3 epochs (2.7 GPU hours for ViT-B).Source: §1 · p. 2
  • DiffRate combines token pruning and merging in one forward pass, preserving detailed foreground information while pruning irrelevant background tokens.The paper describes this as a trade-off between efficiency and performance.Source: §1 · p. 2
  • 40% FLOPs reduction and 50% throughput improvement with a 0.16% accuracy drop on off-the-shelf ViT-H (MAE), outperforming previous methods that require tuning network parameters.The model is compressed without fine-tuning.Source: §1 · p. 2

2. Related Work

Prior token-compression work emphasizes identifying redundant tokens but generally leaves block-wise compression schedules to hand-picked settings. DiffRate learns the schedule differentiably and combines pruning with merging.

2. Related Work

  • Most token-compression methods design metrics to distinguish redundant tokens but do not address the compression schedule in each block.VTC-LFC and ViT-Slim use threshold-based pruning choices influenced by hand-picked hyperparameters.Source: §2 · p. 2
  • DiffRate differs by learning token-compression schedules differentiably and including both pruning and merging to trade off speed and accuracy.Source: §2 · p. 2
  • Differentiable architecture-search methods such as DARTS and ProxylessNAS learn probabilities over candidate operations and select the highest-probability operation.DMCP searches layer-wise convolution channel counts, but channel-pruning methods do not directly transfer to token compression because token positions are agnostic and tokens discarded earlier cannot return.Source: §2 · p. 3

3. Differentiable Compression Rate

Token pruning and merging reduce redundant ViT tokens, but existing approaches handcraft their compression rates. DiffRate learns per-block rates differentiably and can combine pruning with merging.

3. Differentiable Compression Rate

  • Existing pruning and merging methods handcraft compression rates, a tedious choice associated with sub-optimal performance.Pruning preserves informative tokens; merging combines semantically similar tokens, including less informative foreground tokens.Source: §3 · p. 3
  • DiffRate formulates compression as minimizing classification loss on training data under a target FLOPs budget, with learnable pruning and merging rates for each block.The rates are represented across all blocks, and the corresponding FLOPs are expressed as a differentiable function of them.Source: §3 · p. 3
  • DiffRate can represent pruning or merging alone, and the paper combines them by pruning tokens first and then merging the remaining tokens.Direct gradient-based optimization is difficult because compression rates are non-differentiable; per-image masks also complicate parallel computation and require manually chosen rates to preserve equal token counts.Source: §3 · p. 3

4. Differentiable Discrete Proxy

DiffRate’s Differentiable Discrete Proxy (DDP) sorts tokens by importance and reparameterizes compression rates so top-K selection can receive gradient updates. Its pipeline uses class attention for ranking, then combines token pruning, merging, and attention masking.

4. Differentiable Discrete Proxy

  • DDP uses token sorting and a reparameterization trick to select top-K important tokens while allowing gradient back-propagation.The method is designed to retain a common number of tokens across images for parallel batch computation.Source: §4 · p. 4

4.1. Token Sorting

  • Class attention ranks tokens by their contribution to the class token; DiffRate prunes low-importance tokens, then cosine-similarity matches and averages the next unimportant tokens with remaining tokens.The learned compression rates determine the numbers pruned and merged in each block.Source: §4.1 · p. 4

4.2. Compression Rate Re-parameterization

  • DDP represents compression rates as learnable probabilities over discrete candidate rates, with the most important token always retained and less important tokens having greater compression probability.Each block uses separate reparameterization modules for pruning and merging, producing token-level masks.Source: §4.2 · pp. 4–5
  • Attention masking simulates token dropping while preserving gradient propagation from the loss to the mask and, through the learnable probabilities, to the compression rate.A token removed in the preceding block must also be compressed in the current block.Source: §4.2 · p. 5

5. Training Objective

DiffRate trains its compression rates by combining classification loss with a FLOPs constraint and back-propagating through the learnable probabilities. The learned rates can also be supervised by latency or power constraints in a co-search with accelerator parameters.

5. Training Objective

  • The training objective adds a weighted squared FLOPs-budget error to classification loss, with λf balancing the two terms.The paper sets λf to 5 by default in its experiments.Source: §5 · p. 5
  • Straight-through estimation and the chain rule allow gradients to update ρk, so compression rates can be optimized by back-propagation and applied to off-the-shelf models.The algorithm updates ρk after computing the classification and optimization losses.Source: §5 · p. 5
  • DiffRate can optimize compression rates under latency or power constraints while co-searching accelerator parameters for implementation on an FPGA.The extension uses Gemmini to search compression ratios and accelerator parameters together.Source: §5 · pp. 5–6

6. Experiments

Experiments on ImageNet-1k evaluate DiffRate with DeiT, MAE, and LV-ViT backbones, including off-the-shelf compression, fine-tuning, complexity constraints, ablations, and visualizations. The results report gains over prior methods, near-optimal schedule search, and effective joint pruning and merging.

6.2. Comparison with state-of-the-art

  • DiffRate’s searched schedules perform almost optimally across compression settings and FLOPs constraints, with a more prominent advantage at lower FLOPs constraints.The schedule comparison tests pruning, merging, and their combination against constant schedules and 10,000 randomly sampled schedules.Source: §6.2 · p. 6, Figure 4 · p. 7
  • 0.14%–1.14% over ToMe and 0.39%–4.93% over EViT across various models: DiffRate consistently outperforms both on off-the-shelf models.The pre-trained network parameters are not updated.Source: Table 1 · p. 6, §6.2 · p. 7
  • 79.58% on off-the-shelf DeiT-S compares with 79.50% for trained EViT and 79.49% for trained ToMe; DiffRate also maintains an advantage over training-required methods.ATS reaches 79.70% on DeiT-S, close to DiffRate’s 79.83%, but cannot be used for batch inference, unlike DiffRate.Source: §6.2 · p. 7, Table 2 · p. 7

6.4. Visualization

  • With only 34 tokens left in block 12, visualizations show DiffRate pruning redundant background tokens and merging less-discriminative foreground tokens.The examples include merging dog-hair and butterfly-wing tokens while preserving salient information across image regions.Source: Figure 5 · p. 8, §6.4 · p. 8

6.2. Comparison with state-of-the-art

  • 2.5 ms lower latency and 8 mW lower power consumption accompany only a 0.03% accuracy decline with multiple rather than single FLOPs constraints.The comparison is on ViT-S (DeiT).Source: §6.2 · pp. 7–8, Table 3 · p. 7

6.3. Ablation Study

  • 1,000 images suffice to find an appropriate compression rate, and DiffRate requires only three epochs for rate-search convergence.The ablations also find pruning followed by merging performs best among the tested compression options, while class attention is the default sorting metric because it adds no computation.Source: §6.3 · p. 8, Table 4 · p. 8

7. Conclusion

The conclusion presents DiffRate as a unified framework for differentiable compression-rate optimization using both pruning and merging. It reports comparable or superior performance without fine-tuning and rate selection from only 1,000 images.

7. Conclusion

  • DiffRate combines token pruning and merging in a unified framework that optimizes compression rates differentiably.Its Differentiable Discrete Proxy determines compression rates using gradient back-propagation.Source: §7 · pp. 8–9
  • DiffRate performs comparably or better than prior state-of-the-art token-compression methods, including without model fine-tuning.Source: §7 · p. 9
  • Only 1,000 images are needed to identify an appropriate compression rate, demonstrating data efficiency in the reported experiments.Source: §7 · p. 9

A. Training Details

The appendix describes the rate-search and fine-tuning training settings and details the computational-cost constraints. It covers FLOPs estimation and hardware-aware latency and power constraints, including iterative optimization of compression rates and hardware parameters.

A. Training Details

  • Rate searching initializes from official pre-trained backbones, fixes their weights, and trains the Differentiable Discrete Proxy for three epochs; compressed models can then be fine-tuned for 30 epochs.Attention masking simulates token dropping during rate search, while fine-tuning directly drops redundant tokens.Source: A. Training Details · p. 11

B.1. FLOPs Calculation

  • The FLOPs algorithm accumulates attention and MLP costs across blocks using the compression rates, with patch embedding and classifier costs added to the final calculation.A straight-through estimator is used for backpropagation through the max operation.Source: §B.1 · p. 11

B.2. Latency and Power Constraint

  • Latency and power constraints use a hardware-performance loss based on the measured metric and its target, with Gemmini providing candidate accelerator designs.Source: §B.2 · pp. 11–12
  • Compression rates and hardware parameters are optimized iteratively, alternating between rate optimization for fixed hardware and hardware optimization for fixed rates.The differentiable hardware-cost formulation supports updates toward target latency and power.Source: §B.2 · p. 12

C. More Results

The appendix reports expanded DeiT and MAE results, examines the searched schedules and their transferability, and adds evidence on training speed, visualizations, and computational overhead.

C. More Results

  • The appendix covers additional FLOPs constraints for DeiT and MAE, searched schedules, schedule transfer across models, and training compressed models from scratch.The overview identifies these as the contents of the additional-results section.Source: C. More Results · p. 12

C.1. DeiT Models; C.2. MAE Models

  • The appendix provides complete off-the-shelf and fine-tuned results for DeiT and MAE models; ViT-H is not fine-tuned because of computational-resource constraints.The supplied passages do not include the tables’ cell values.Source: §C.2 · p. 12, §C.1 · p. 12, Table 8 · p. 12

C.3. Seached Compression Schedule

  • The searched schedules retain more tokens in deep blocks for MAE than for DeiT, which the authors relate to their different classification token usage.DeiT classifies using the class token, whereas MAE uses the average of all image tokens.Source: §C.3 · p. 12

C.4. Compression Schedule Transfer

  • ViT-T’s compression rates transfer to ViT-S with a similar outcome, whereas transferring ViT-B’s rates leads to significantly poorer performance.The authors conclude that rates suit different network structures based on their features, while transfer can work among similar structures.Source: §C.4 · pp. 12–13, Figure 6 · p. 13

C.5. Train from Scratch

  • Using a searched compression rate, DiffRate’s compressed model can be trained from scratch; eliminating redundant tokens yields a 1.4× increase in training speed with only a −0.06% performance degradation.Source: §C.5 · p. 13

D. More Visualization

  • The visualizations show DiffRate identifying semantic objects, removing irrelevant background tokens, and merging less-discriminative foreground tokens.The visualizations map compressed tokens back to input patches, coloring pruned patches black and merged patches by their average color.Source: D. More Visualization · p. 13

E. Overhead of DiffRate

  • For DeiT-S, DiffRate adds 4.7k parameters and 0.24M FLOPs overall; one block’s module uses 1.36M FLOPs in the stated example.The authors characterize the framework overhead as negligible relative to DeiT-S’s 22M parameters and 4.6G FLOPs.Source: E. Overhead of DiffRate · p. 13

F. Extending to Hierarchical Architecture

DiffRate is adapted to hierarchical architectures by restoring compressed token sequences between stages, addressing the spatial-integrity constraint of hierarchical models. On CAFormer-S36, the reported ImageNet results show FLOPs reductions with small accuracy drops.

F. Extending to Hierarchical Architecture

  • Hierarchical models use staged resolution reduction and depend on spatial integrity, which traditional unstructured token compression disrupts.The passage names convolution and pooling as downsampling operations.Source: F. Extending to Hierarchical… · pp. 13, 15
  • A token uncompression module restores the compressed sequence by copying tokens based on their relationships, enabling DiffRate’s use in hierarchical architectures.The module is placed at the end of each stage.Source: F. Extending to Hierarchical… · pp. 15–16, Figure 8 · p. 16
  • 25% FLOPs reduction comes with a 0.24% accuracy sacrifice without fine-tuning, while fine-tuning yields 35% FLOPs reduction with a 0.42% accuracy drop on CAFormer-S36.DiffRate is applied only to the third stage; the first two use convolution blocks, and the final stage accounts for 6% of CAFormer-S36’s computational cost.Source: F. Extending to Hierarchical… · p. 16, Table 12 · p. 16

G. Extending to Downstream Tasks

DiffRate with the uncompression module is evaluated on ADE20K semantic segmentation using a CAFormer-S36 backbone, without further fine-tuning.

G. Extending to Downstream Tasks

  • On ADE20K, applying DiffRate to the third stage of the trained segmentation model derived from uncompressed pre-trained CAFormer-S36, without further fine-tuning, yields 30% fewer FLOPs with negligible performance degradation.Source: G. Extending to Downstream Tasks · p. 16

Related papers

Loading 2305.17997v1…