Token Merging for Fast Stable Diffusion

TL;DR

The paper addresses the high computational cost of diffusion models by introducing Token Merging (ToMe) to reduce token counts without retraining. It adapts ToMe to the dense, token-based outputs of diffusion through unmerging, novel partitioning, and selective application—especially to self-attention—yielding up to 2x speed and 5.6x memory reductions with 60% token reduction, and these gains stack with xFormers for large images. Extensive experiments on Stable Diffusion validate that image quality remains close to the baseline, with targeted design choices balancing speed and fidelity. The work offers a practical, training-free acceleration strategy and opens avenues for applying token merging to dense prediction tasks in diffusion models.

Abstract

The landscape of image generation has been forever changed by open vocabulary diffusion models. However, at their core these models use transformers, which makes generation slow. Better implementations to increase the throughput of these transformers have emerged, but they still evaluate the entire model. In this paper, we instead speed up diffusion models by exploiting natural redundancy in generated images by merging redundant tokens. After making some diffusion-specific improvements to Token Merging (ToMe), our ToMe for Stable Diffusion can reduce the number of tokens in an existing Stable Diffusion model by up to 60% while still producing high quality images without any extra training. In the process, we speed up image generation by up to 2x and reduce memory consumption by up to 5.6x. Furthermore, this speed-up stacks with efficient implementations such as xFormers, minimally impacting quality while being up to 5.4x faster for large images. Code is available at https://github.com/dbolya/tomesd.

Figures (8)

• Figure 1: Token Merging for Stable Diffusion. When applied properly, ToMe tome can significantly increase the speed of image generation without jeopardizing quality. Moreover, its benefits stack with existing methods such as xFormers xformers. With ToMe and xFormers together, this $2048\times2048$ image generated in just 28 seconds on a 4090, which is 5.4$\times$ faster than the original model.
• Figure 2: A U-Net Block with ToMe applied. Stable Diffusion ldm uses a U-Net unet model with transformer-based blocks attnisallyouneed. We apply ToMe tome by merging tokens before each component of the block and unmerging after to reduce compute costs (Sec. \ref{['sec:approach']}).
• Figure 3: ToMe applied naïvely as described in Sec. \ref{['sec:approach']} works to maintain image coherence, but the content inside the image can change drastically with high amounts of reduction.
• Figure 4: Importance of Merging. If we just prune tokens instead (and replace them with 0), the resulting images quickly degrade.
• Figure 5: Partitioning src and dst. ToMe tome merges tokens from srcintodst. (a) By default, ToMe alternates src and dst tokens. In our case, this causes dst to form regular columns which leads to bad outputs (poor fish). (b) We can improve generation by sampling dst with a 2d stride (e.g., $2\times2$), but this still forms a regular grid. (c) We can introduce irregularity by sampling randomly, but this can cause undesirable clumps of dst tokens. (d) Thus, we sample one dst token randomly in each $2\times2$ region.