Search papers, labs, and topics across Lattice.
This paper introduces MaMe, a training-free, differentiable token merging method for Vision Transformers (ViTs) based on matrix operations to improve GPU efficiency and reduce the quadratic complexity of self-attention. They also introduce MaRe, an inverse operation for token restoration, enabling a MaMe+MaRe pipeline for image synthesis. Experiments show that MaMe doubles ViT-B throughput with minimal accuracy loss, accelerates VideoMAE-L by 48.5%, and reduces Stable Diffusion v2.1 generation latency by 31%, demonstrating its effectiveness in accelerating vision models across various tasks.
Token merging can be fast: MaMe accelerates ViTs by up to 2x and Stable Diffusion v2.1 by 31% using only GPU-friendly matrix operations, outperforming prior methods.
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.
