Source-linked AI summary
Token Fusion: Bridging the Gap between Token Pruning and Token Merging
Minchul Kim, Shangqian Gao, Yen-Chang Hsu, Yilin Shen, Hongxia Jin
TL;DR
ViTs are costly to deploy because self-attention creates substantial computational demand, motivating more efficient token reduction. Token Fusion dynamically combines pruning and merging based on layer-wise functional linearity and introduces MLERP to preserve norm distributions. It outperforms standalone alternatives in reported classification and image-generation evaluations while improving efficiency.
Problem
ViT computational overhead, particularly from self-attention, makes deployment in resource-constrained settings challenging, while pruning and merging offer different efficiency trade-offs.
Method
Token Fusion dynamically combines pruning and merging according to layer-wise functional linearity and introduces MLERP merging to conserve feature-norm distributions.
Results
ToFu outperforms standalone token merging and pruning in ImageNet 1K classification and outperforms ToMe in stable diffusion image generation while improving speed.
Takeaways & Limitations
Token Fusion provides a unified token-reduction framework applicable to ViTs with or without additional training, with reported accuracy and efficiency benefits.
Takeaways & Limitations
ToFu depends on tuning hyperparameter d for the transition from pruned to average merging, and its benefits are strongest when functional linearity varies across depths.
Abstract
from arXiv · showhide
Vision Transformers (ViTs) have emerged as powerful backbones in computer vision, outperforming many traditional CNNs. However, their computational overhead, largely attributed to the self-attention mechanism, makes deployment on resource-constrained edge devices challenging. Multiple solutions rely on token pruning or token merging. In this paper, we introduce "Token Fusion" (ToFu), a method that amalgamates the benefits of both token pruning and token merging. Token pruning proves advantageous when the model exhibits sensitivity to input interpolations, while token merging is effective when the model manifests close to linear responses to inputs. We combine this to propose a new scheme called Token Fusion. Moreover, we tackle the limitations of average merging, which doesn't preserve the intrinsic feature norm, resulting in distributional shifts. To mitigate this, we introduce MLERP merging, a variant of the SLERP technique, tailored to merge multiple tokens while maintaining the norm distribution. ToFu is versatile, applicable to ViTs with or without additional training. Our empirical evaluations indicate that ToFu establishes new benchmarks in both classification and image generation tasks concerning computational efficiency and model accuracy.
1. Introduction
Token Fusion addresses ViT efficiency by combining pruning and merging according to layer-wise functional linearity, while MLERP preserves feature-norm distributions during merging. The method reports higher accuracy and speed than standalone approaches across classification and image generation.
- 1. Introduction: Token Fusion combines pruning and merging according to each layer’s functional linearity with respect to interpolated inputs.Pruning is favored when interpolation causes nonlinear output deviations, whereas merging is beneficial when responses are more linear.
- 1. Introduction: Average merging can alter feature distributions because it does not preserve original feature norms, unlike pruning.MLERP addresses this limitation by merging multiple tokens while conserving the norm distribution.
- 1. Introduction: ToFu outperforms ToMe in ImageNet 1K classification while operating faster, and also produces structurally closer class-conditional images than ToMe in stable diffusion.The reported benefits span both accuracy and computational efficiency across classification and image generation.
- 1. Introduction: The paper contributes Token Fusion, MLERP merging, and empirical validation against token pruning and merging on classification and image-generation tasks.The contribution summary describes dual gains in accuracy and speed.
2. Related Works
Related work reduces ViT cost through attention approximations, learned token reduction, and heuristic pruning or merging. Token Fusion differs by unifying pruning and merging while addressing average merging’s limitations without auxiliary-module fine-tuning.
- 2. Related Works: Prior efficiency methods accelerate attention through hashing or sparsity, low-rank approximations, or combinations of both.These approaches primarily target the attention module rather than reducing the total number of tokens.
- 2. Related Works: Learned token reduction identifies redundant tokens using auxiliary models or losses, but commonly requires additional training or fine-tuning.Examples include DynamicViT’s learned pruning masks and A-ViT’s halting probabilities.
- 2. Related Works: Heuristic token reduction can operate on pretrained ViTs without further fine-tuning, including token pooling, adaptive sampling, and token merging.Token Merging averages similar tokens using efficient bipartite matching.
- 2. Related Works: Token Fusion dynamically combines pruning and merging while using MLERP to address limitations of average merging.This distinguishes it from methods that rely exclusively on one reduction strategy or conventional averaging.
3. Proposed Method
Token Fusion reduces ViT tokens by selecting pruning or merging according to functional linearity, while MLERP preserves feature-norm distributions during merging.
- Token Fusion: Token Fusion unifies pruning and merging through Bipartite Soft Matching, selecting similar token pairs before applying a reduction strategy.BSM partitions tokens into SRC and DST sets, weights cross-set similarities, and selects the top r pairs for reduction.
- Average Merging: Average merging combines source tokens into destination positions using gather and scatter-reduce mean operations.The scattered source tokens are averaged with destination tokens at their assigned indices.
- Pruned Merging: Pruned merging discards similar source tokens, avoiding averaging operations and offering faster execution than average merging.The method treats similarity as evidence of redundancy and retains the corresponding destination tokens.
- Hybrid Merging: Hybrid merging uses pruning in early, less functionally linear layers and average merging in deeper layers with more linear behavior.The transition is controlled by a depth hyper-parameter d, and the paper reports that early pruned merging followed by later averaging improves performance.
- Functional Linearity Analysis: Functional linearity compares direct output distance with the length of the interpolated output path, with perfect linearity equal to 1.The metric is estimated through finite differences along interpolated inputs; higher values indicate more linear behavior.
- MLERP: MLERP normalizes the average feature vector and rescales it by the maximum norm among the merged vectors, improving performance over simple averaging.This design addresses the norm reduction associated with average merging and is reported to enhance results in Tables 1 and 2.
4. Experiments
Experiments evaluate ToFu for ImageNet classification and Stable Diffusion image generation, comparing accuracy, speed, token-reduction strategies, and image quality against ToMe and full-token models.
- Classification Task: ToFu MLERP offers the strongest classification efficiency–accuracy trade-off, with ToFu AVG exceeding ToMe in both speed and accuracy and MLERP improving results with slight speed reduction.The comparison covers ViT-B and ViT-L models and identifies ToFu MLERP as the preferred trade-off.
- Classification Task: PRUNE merging in early layers followed by AVG merging in later layers performs consistently better than reversing or uniformly applying the strategies, with d = 6 selected for subsequent experiments.Table 3 evaluates configurations ranging from all PRUNE to all AVG and supports the early-to-late transition.
- Image Generation Task: For Stable Diffusion, ToFu reduces tokens by 50% before attention and later restores the original token length because subsequent ResNet blocks require full tokens.The reduction targets the attention module while leaving ResNet blocks and MLP layers unchanged.
- Image Generation Task: ToFu image generation more closely matches full-token images than ToMe in distribution, structural quality, and perceptual quality, while providing an approximate 32% speed enhancement over the full model.The evaluation uses FID, LPIPS, MS-SSIM, and consistent random seeds across generated-image comparisons.
5. Conclusion
Token Fusion combines token pruning and merging in Vision Transformers, with MLERP merging designed to preserve feature norm distributions. Its benefits depend on layer-wise variation in functional linearity and require tuning the switching hyperparameter d in some settings.
- Token Fusion integrates token pruning and token merging within Vision Transformers.
- MLERP merging preserves feature norm distributions while merging multiple tokens, addressing distribution shifts associated with average merging.
- ToFu requires tuning hyperparameter d for the layer where PRUNED merging switches to AVG merging, especially across datasets or architectures.