用矩阵运算实现高效视觉模型压缩与恢复,提速不丢精度。
MaMe & MaRe: Matrix-Based Token Merging and Restoration for Efficient Visual Perception and Synthesis

- 基于矩阵操作的训练无关压缩方法,避免显卡低效计算
- ViT-B模型吞吐量翻倍,仅损失2%准确率
- 适用于图像生成与视频任务,兼顾速度与质量
令牌压缩对缓解视觉变换器(ViTs)中自注意力机制的二次复杂性至关重要,但现有方法如ToMe依赖于GPU效率低的操作(如排序、散写),引入开销限制其效果。本文提出MaMe,一种完全基于矩阵运算的训练无关、可微分令牌合并方法,具有良好的GPU友好性,可加速ViTs。此外,我们提出其逆操作MaRe,用于令牌恢复,形成用于图像合成的MaMe+MaRe流水线。应用于预训练模型时,MaMe使ViT-B吞吐量提升一倍,仅下降2%准确率。值得注意的是,仅微调最后一层即可在1.1倍速度下提升ViT-B准确率1.0%。在SigLIP2-B@512零样本分类中,MaMe实现1.3倍加速且性能几乎无损。在视频任务中,MaMe使VideoMAE-L在Kinetics-400上提速48.5%,准确率仅下降0.84%。此外,部分任务中实现了性能与速度的同步提升。在图像生成中,MaMe+MaRe流水线在降低Stable Diffusion v2.1生成延迟31%的同时提升了生成质量。总体表明MaMe与MaRe在加速视觉模型方面具有显著效果。代码已开源。
原文摘要 · Abstract (English)
Token compression is crucial for mitigating the quadratic complexity of self-attention mechanisms in Vision Transformers (ViTs), which often involve numerous input tokens. Existing methods, such as ToMe, rely on GPU-inefficient operations (e.g., sorting, scattered writes), introducing overheads that limit their effectiveness. We introduce MaMe, a training-free, differentiable token merging method based entirely on matrix operations, which is GPU-friendly to accelerate ViTs. Additionally, we present MaRe, its inverse operation, for token restoration, forming a MaMe+MaRe pipeline for image synthesis. When applied to pre-trained models, MaMe doubles ViT-B throughput with a 2% accuracy drop. Notably, fine-tuning the last layer with MaMe boosts ViT-B accuracy by 1.0% at 1.1x speed. In SigLIP2-B@512 zero-shot classification, MaMe provides 1.3x acceleration with negligible performance degradation. In video tasks, MaMe accelerates VideoMAE-L by 48.5% on Kinetics-400 with only a 0.84% accuracy loss. Furthermore, MaMe achieves simultaneous improvements in both performance and speed on some tasks. In image synthesis, the MaMe+MaRe pipeline enhances quality while reducing Stable Diffusion v2.1 generation latency by 31%. Collectively, these results demonstrate MaMe's and MaRe's effectiveness in accelerating vision models. The code is available at https://github.com/cominder/mame}{https://github.com/cominder/mame.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。