Reproduce “MaMe & MaRe: Matrix-Based Token Merging and Restoration for Efficient Visual Perception and Synthesis” on OpenResearch

About this paper

Simin Huo, Ning Li

This research introduces MaMe (Matrix-Based Token Merging) and MaRe (Matrix-Based Token Restoration), a framework that employs GPU-friendly matrix operations to efficiently reduce and restore tokens in Vision Transformers. The framework significantly boosts throughput in visual perception tasks, such as a 79% increase for ViT-S, and simultaneously improves both speed and quality in image synthesis, demonstrating a 31.2% latency reduction and enhanced fidelity in Stable Diffusion v2.1.

Key insights

• Constructing token merging and restoration solely with matrix operations bypasses GPU-inefficient computations (e.g., sorting, scattered writes), yielding substantial throughput gains while preserving differentiability.
• The adaptive weight refining and preservation of unmerged source tokens ensure that critical, high-frequency information is retained during compression, which is vital for maintaining model performance.
• The combined MaMe+MaRe pipeline can simultaneously improve inference speed and enhance output quality in generative models, suggesting it functions as a 'high-pass filter' that mitigates attention dilution.

Summary by alphaXiv.