Source-linked AI summary
Token Merging for Fast Stable Diffusion
Daniel Bolya, Judy Hoffman
TL;DR
Diffusion generation is slow because transformer computation processes redundant tokens and scales quadratically with token count. The paper adapts training-free Token Merging to Stable Diffusion, merging and later unmerging tokens during U-Net computation. The resulting method minimally impacts visual quality while providing up to 2× faster evaluation and 5.6× lower memory use, with further speedups when combined with xFormers.
Problem
Transformer-based diffusion models process every token even though generated images contain substantial redundancy, making large-image generation computationally expensive.
Method
The paper adapts Token Merging to Stable Diffusion by merging similar tokens during reduced U-Net computation and unmerging them afterward without additional training.
Results
ToMe for Stable Diffusion minimally impacts visual quality while offering up to 2× faster evaluation using 5.6× less memory.
Takeaways & Limitations
The results demonstrate a training-free way to apply token reduction to Stable Diffusion and motivate further exploration of ToMe for dense prediction tasks.
Takeaways & Limitations
The simplest unmerging strategy loses information, and the authors identify better unmerging strategies as a likely avenue for improvement.
Abstract
from arXiv · showhide
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.
1. Introduction
Diffusion models generate high-quality images but are expensive because transformer computation scales quadratically with token count, while existing speedups still process every token. This paper applies training-free Token Merging to exploit image redundancy and improve Stable Diffusion speed while preserving quality.
- Motivation: Transformer computation scales with the square of token count, making large-image diffusion especially expensive.The scaling also follows the square of the number of pixels.
- Motivation: Existing optimizations improve transformer implementation efficiency but still evaluate every token.Flash Attention, xFormers, and PyTorch 2.0 optimizations reduce implementation overhead without reducing the amount of work.
- Contribution: 2048 × 2048 generation takes 28 seconds with ToMe and xFormers on a 4090, 5.4× faster than the original model.The caption describes this as a combined benefit of ToMe and xFormers.
- Motivation: Diffusion-generated images contain substantial redundancy, motivating token pruning and merging to reduce computation with small accuracy loss.Token Merging is especially useful because it does not require additional training.
- Contribution: Applying ToMe to Stable Diffusion introduces diffusion-specific token partitioning and application choices to retain speed and improve memory efficiency while keeping images close to the original model.The paper reports that the resulting speedup also stacks with xFormers.
2. Background
The paper targets an off-the-shelf Stable Diffusion model, whose U-Net contains transformer blocks operating on image tokens and prompt-conditioned cross-attention. It applies Token Merging by reducing similar tokens progressively within those blocks.
- Stable Diffusion: The goal is to speed up an off-the-shelf Stable Diffusion model without training, using ToMe.This establishes the paper’s training-free deployment setting.
- Token Merging: The paper applies ToMe around each U-Net block component, merging tokens before computation and unmerging afterward to reduce compute costs.This preserves the block’s token structure while reducing the tokens processed by its components.
- Stable Diffusion: Stable Diffusion repeatedly denoises noise through a U-Net with transformer-based blocks during image generation.The U-Net encodes the current noised image as tokens before processing them through transformer blocks.
- Stable Diffusion: Each transformer block combines self-attention, cross-attention for prompt conditioning, and multi-layer perception modules.These components are applied to token representations inside the U-Net.
- Token Merging: ToMe gradually reduces token count by merging r similar source tokens into destination tokens in each transformer block.The source-destination partition enables efficient similarity-based merging.
3. Token Merging for Stable Diffusion
Stable Diffusion requires predictions for every image token, so ToMe must merge redundant tokens during computation and unmerge them afterward. Naive reduction preserves coherence but can alter content and degrade quality, motivating the paper’s unmerging and partitioning choices.
- Defining Unmerging: Dense diffusion prediction requires retaining information for every token, unlike classification, which can rely on a single prediction token.This makes unmerging necessary after reduced computation.
- Defining Unmerging: ToMe merges tokens rather than removing them, allowing merged tokens to be unmerged after computation so dense outputs remain token-complete.This distinguishes ToMe from pruning for dense prediction.
- Initial Naive Approach: Naive ToMe preserves overall image coherence but can drastically change image content at high reduction, whereas pruning tokens quickly degrades the resulting images.These observations motivate merging and improved token partitioning rather than direct token replacement.
- Defining Unmerging: The simplest unmerging strategy reconstructs merged tokens from their combined representation when the original tokens are similar.The method assumes similarity makes the reconstruction error small.
- Defining Unmerging: The unmerging strategy loses information, although the authors report that the resulting error is small and note that better strategies may improve it.This is an explicit limitation of the simplest reconstruction method.
- Initial Naive Approach: Naively applying ToMe before each block component can speed diffusion by up to 2× and reduce memory by approximately 4× for 512×512 images, but substantially raises FID.The quality degradation motivates further refinement.
4. Further Exploration
The paper refines ToMe for Stable Diffusion by improving token partitioning and selecting where and when merging should occur. These choices preserve quality better than naïve reduction while retaining substantial speed and memory benefits.
- Motivation: Naïve ToMe can run up to 2× faster with up to 4× less memory, but it substantially increases FID and changes image content.The authors therefore investigate diffusion-specific design choices rather than applying token merging unchanged.
- Partitioning: Alternating source and destination tokens creates columns that effectively halve image resolution along rows when 50% of tokens are merged.A 2d stride improves results but still places destination tokens on a regular grid.
- Partitioning: Fixing random destination-token selection across the classifier-free-guidance batch avoids the massive FID increase caused by independently random sampling.The fixed-randomness method improves beyond a 2d stride, with FID values of 36.00 for fixed random sampling and 35.66 for fixed random 2 × 2 sampling.
- Partitioning: Random methods perform best when randomness is fixed across the batch, and one randomly selected destination token per 2 × 2 region becomes the default partition.All partition experiments merge 50% of tokens.
- Design Experiments: Applying ToMe only to self-attention provides the clearest speed-versus-FID trade-off, while blocks with the most tokens provide most of the speed-up.Merging more tokens earlier and fewer later is slightly better, but not enough to justify the added schedule.
- Results: At high token reduction, the method retains image content and handles complex scenes, although slight background detail can be lost.The quantitative results report similar or better FID than baseline with up to 60% of tokens reduced.
5. Putting It All Together
The authors combine their partitioning and design choices into ToMe for Stable Diffusion. The resulting method minimally impacts visual quality while accelerating evaluation and reducing memory, with additional speed gains when combined with xFormers.
- Combined Method: ToMe for Stable Diffusion combines the techniques developed in the preceding design experiments.The method is evaluated visually and quantitatively after integrating those choices.
- ToMe + xFormers: Combining ToMe with xFormers produces substantial speed benefits for 2048 × 2048 images, with still more speed available when sacrificing visual quality.The speed-up is less pronounced for smaller images, and memory benefits do not stack with xFormers.
6. Conclusion and Future Directions
The paper concludes that ToMe can make Stable Diffusion significantly faster while preserving high-quality images without training. It identifies better unmerging, proportional attention, key-based similarity, and dense prediction as directions for further study.
- Conclusion: ToMe successfully accelerates Stable Diffusion and generates high-quality images without additional training.The authors present the method as both a practical tool and a starting point for research on token merging.
- Future Directions: Better unmerging strategies and the usefulness of proportional attention or key-based similarity remain open questions for diffusion.The authors also motivate further exploration of ToMe for dense prediction tasks.
- Conclusion: At 60% of tokens merged, ToMe for Stable Diffusion keeps the image the same most of the time.The paper illustrates this behavior with additional examples.