Token Merging for Fast Stable Diffusion

In plain words

Stable Diffusion draws an image by running a large transformer many times. Each run looks at every small patch of the picture, called a token. More tokens means more work, and the work grows faster than the number of tokens. Many patches look almost the same, such as sky or a plain background, so the network repeats a lot of work it does not need.

Token Merging (ToMe) finds the most alike tokens inside a block and merges each one into a similar partner. The block then works on fewer tokens. Once the block is done, the merged positions are copied back, so the image keeps its full size. No retraining is needed: the change is applied to an existing model.

Simply deleting tokens breaks the image, which is why the copy-back step matters. The authors report that merging 60% of the tokens makes generation about twice as fast and uses much less memory on Stable Diffusion v1.5, with a close quality score (FID, where lower means closer to real images). A few choices matter: the tokens that host the merges are picked at random, with the same pick for every image in a batch, and the merge is applied only inside attention.

TL;DR

A training-free speed-up for Stable Diffusion. Inside each U-Net transformer block, similar tokens are merged before self-attention and unmerged afterwards. The final design merges only in self-attention, only in the four highest-resolution blocks, and uses a random 2×2 destination partition shared across the batch. Author-reported on SD v1.5 (512 px, 50 PLMS steps) with 60% merged: FID 33.37 against 33.12 for the baseline, 1.52 against 3.09 s per image, and 0.60 against 3.41 GB per image.

Problem

Each denoising step runs a transformer over every token, and attention cost grows with the square of the token count §1. Efficient kernels such as FlashAttention, xFormers and PyTorch 2.0 make each token cheaper but still process every token, so they do not reduce the work §1. Token pruning and merging reduce the work in general but usually need retraining §1. ToMe avoids retraining, but it was shown only for ViT classification §1.

Two specific problems arise for diffusion. A dense output needs a value for every token, so a pruned token has to be restored §3.1. Pruning with zeros destroys the image fig. 4. A naive port of ToMe keeps the image coherent but changes its content a lot at high reduction, and its FID rises to 38.95 at 50% §3.2 tab. 1 fig. 3.

All numbers in the paper are for Stable Diffusion v1.5 §4.

Method

ToMe for SD applies ToMe inside the transformer blocks of an existing Stable Diffusion U-Net. Tokens are merged before a module and unmerged after it, so the block output keeps the full token count §3.2 fig. 2. The method adds no training step §6.

Step 1: merging and unmerging. Two similar tokens x1 and x2 are merged into their average eq. 1. To unmerge, both positions receive that average eq. 2. The information lost is small when the merged tokens were already similar. The paper notes that other unmerging methods might be worth exploring §3.1.

Step 2: naive application. Merge before each of self-attention, cross-attention and the MLP in every block, and unmerge before the skip connection §3.2 fig. 2. The reduction is a percentage r% of the current tokens. Nothing accumulates across blocks, because merged tokens are unmerged again §3.2. The similarity is computed once per block, from the block input x and not from the attention keys. Proportional attention is not used §3.2.

Step 3: destination partition. The default ToMe alternates source and destination tokens, which forms columns and halves the resolution along rows §4.1 fig. 5. The fix used here picks one random destination in each 2×2 region §4.1 fig. 5. The random pick is the same for every image in a batch, because prompted and unprompted images under classifier-free guidance must allocate destinations identically §4.1.

Step 4: final design. Apply ToMe to self-attention only; cross-attention and the MLP are left alone §4.2 tab. 3. Apply it only to blocks with at least 4096 tokens, which at 512 px are the 4 highest-resolution blocks §4.2 tab. 3 patch.py. Keep the reduction constant across diffusion steps. A linearly varying reduction gave only a small gain, not worth the added complexity §4.2 tab. 3.

# ToMe merge and unmerge for one block  [§3.1] [§3.2] [§4.1]
# Input:  features (B, N, C) = block input; grid h x w, N = h * w
# Output: merge and unmerge functions for this block

def make_merge(features, h, w, stride, r):
    # Pick one destination per stride x stride region; the same pick
    # is used for every image in the batch  [§4.1]
    # [code: tomesd/merge.py#L49-L52]
    dst_index, src_index = partition_2d(h, w, stride, shared=True)

    # Cap r at the number of source tokens  [code: tomesd/merge.py#L89]
    r = min(r, src_index.shape[1])

    # Cosine similarity between each source and each destination, on the
    # block input  [§3.2] [code: tomesd/merge.py#L84-L86]
    unit = features / norm(features)
    scores = cosine(unit[src_index], unit[dst_index])  # (B, N_src, N_dst)

    # Each source keeps its best destination. The r sources with the
    # highest best score are merged, the rest are kept  [§3.2]
    # [code: tomesd/merge.py#L92-L97]
    best_score, best_dst = max_over_dst(scores)
    order = argsort(best_score, descending=True)
    merged_src = order[:, :r]
    kept_src = order[:, r:]
    target_dst = best_dst[:, merged_src]

    def merge(tokens):
        # Average each destination with every source merged into it; the
        # destination itself is one of the terms  [eq. 1]
        # [code: tomesd/merge.py#L99-L107]
        src, dst = split(tokens, src_index, dst_index)
        kept = src[:, kept_src]
        dst = mean_into(dst, target_dst, src[:, merged_src])
        return concat(kept, dst)                    # (B, N - r, C)

    def unmerge(tokens):
        # Each merged source copies the value of its destination  [eq. 2]
        # [code: tomesd/merge.py#L109-L122]
        num_kept = kept_src.shape[1]
        kept, dst = tokens[:, :num_kept], tokens[:, num_kept:]
        output = zeros(B, N, C)
        output = place(output, dst_index, dst)
        output = place(output, src_index[kept_src], kept)
        output = place(output, src_index[merged_src], dst[:, target_dst])
        return output                               # (B, N, C)

    return merge, unmerge

# One U-Net transformer block with ToMe on self-attention  [§4.2] [tab. 3]
# Input:  features (B, N, C); context for cross-attention; grid h x w
# Output: features (B, N, C)

def tome_block(features, context, h, w, ratio):
    # Merge count is a share of the tokens in this block  [§3.2]
    # [code: tomesd/patch.py#L20]
    r = int(features.shape[1] * ratio)
    merge, unmerge = make_merge(features, h, w, stride=2, r=r)

    # Self-attention sees the reduced tokens; its output is unmerged
    # before the residual add  [§3.2] [code: tomesd/patch.py#L58-L62]
    attended = self_attention(merge(norm1(features)))
    features = unmerge(attended) + features

    # Cross-attention and the MLP run on every token by default  [§4.2]
    # [code: tomesd/patch.py#L187-L189]
    features = cross_attention(norm2(features), context) + features
    features = feed_forward(norm3(features)) + features
    return features

Hyperparameters

Values are from the paper unless marked "code" (taken from the official repository, not stated in the paper). The paper's experiment settings are in §4.

Results

Author-reported. Setup: Stable Diffusion v1.5, 512 × 512, 2,000 samples, 50 PLMS steps, guidance 7.5, FID against 5,000 ImageNet-1k validation images, one 4090 GPU §4. Baseline is the unmodified model.

Headline numbers (60% merged, the largest setting in the table):

Table 4, ToMe for SD (final design, §5):

Table 1, naive ToMe (§3.2, the baseline the design is built on):

Table 2, destination partition (§4.1; 50% of tokens merged in every row).

Table 3, design choices (§4.2; r = 50%, each experiment independent of the others, random 2 × 2 partition).

The paper's highlighted choices are: self-attention only (a), 4096 minimum tokens (b), and a constant 50/50 schedule (c). The paper does not report one table row with all three choices together, except the 50% row of Table 4.

Qualitative: at 60% merged, images stay close to the original, though some detail is lost, for example in a background fig. 6 fig. 8. Pruning instead of merging destroys the image fig. 4. Speed-up for large images with xFormers is larger than for small ones §5.

Speed limit: with more aggressive merging, speed-up stops at a point where most time goes to modules ToMe does not touch fig. 7.

Paper vs code

Code: official repository dbolya/tomesd, MIT licence (Copyright 2023 Daniel Bolya), commit 27a14a372beecaa85c101fb588631643f379b8ce. The README calls it the official implementation README.md. The repository is a small pure-PyTorch package. It contains the merge logic, the patch for both the LDM and diffusers block classes, and utilities. It contains no training code and no FID or timing scripts.

Consistent with the paper:

• stride 2 × 2 and the self-attention-only default patch.py patch.py;
• one random destination per 2 × 2 region, with one draw shared across the batch merge.py;
• cosine similarity on the block input merge.py patch.py;
• the README's table reproduces Table 4 exactly README.md.

Differences, and details the paper does not give:

• Merge of several sources into one destination. The paper's merge is written for two tokens eq. 1. The code averages each destination with all sources merged into it, counting the destination once merge.py. This matches eq. 1 when a destination receives one source. Inferred from code and PyTorch's default include_self=True.
• Odd batch size turns randomness off. The paper does not state this patch.py CHANGELOG.md.
• Separate random generator. The destination layout uses its own generator, created on first use from the global RNG state. It is re-created only when the device changes patch.py utils.py CHANGELOG.md. The paper describes only fixing randomness across the batch §4.1.
• Threshold is a downsampling factor, not a token count. max_downsample is compared with a linear factor: the square root of the full-resolution token count divided by the block's current count patch.py. Default is 1 patch.py, which for 512 px gives the 4 blocks in Table 3b. At other image sizes the token threshold changes. The docstring says 1 or 2 is recommended patch.py.
• Non-divisible sizes. Image sizes that the stride does not divide are handled merge.py CHANGELOG.md. The paper says nothing about this.
• Rounding and cap. r = int(N × ratio) truncates patch.py and is capped at the source count merge.py.
• Diffusers path. A second block patch covers diffusers' BasicTransformerBlock with the same merge and unmerge, applied around norm1/attn1, attn2 and ff patch.py. The paper does not say which implementation it used §4.
• No schedule parameter. apply_patch has one ratio for all steps patch.py. This matches the paper's chosen constant schedule.

Pitfalls

• Only SD v1.5 is measured. The README says the patch should work for any Stable Diffusion model README.md. The paper reports numbers for v1.5 only §4.
• Quality is FID only. The paper notes that FID does not measure prompt adherence, which may explain why merging cross-attention lowered FID §4.2. FID is computed from 2,000 samples §4.
• Random pick must be shared across guidance pairs. At 25% merged, FID is 46.08 without the shared pick and 36.00 with it tab. 2. Odd batches turn the randomness off in code patch.py.
• Seeds. The destination layout depends on the generator state at its first use, not on the seed alone utils.py. The README says you may need to reset the seed for each batch to get consistent results README.md.
• Similarity cost is paid every block call. The merge is built once per block call, so it is recomputed on every diffusion step patch.py. The paper says the similarity is computed once at the start of each block §3.2.
• Attention ignores merged size. Proportional attention is not used §3.2. The paper says that whether it or key-based similarity carries over to diffusion needs more exploration §6.
• Unmerging is a copy. Merged positions get the same value, so detail can be lost eq. 2 fig. 6.
• Ablations are one factor at a time. Table 3 rows are independent of each other tab. 3. Combining them is not tested.
• Speed numbers depend on hardware and image size. Speed-up is smaller for small images, where the diffusion model is not the bottleneck §5. Memory gains did not stack with xFormers §5.
• Large sampling settings are not tested. The README's extreme example (ratio 0.9, 4 × 4 stride, max_downsample 2) is not in the paper README.md. The paper's 4 × 4 result at 50% is FID 38.97 tab. 2.
• Numbers are rounded. The paper says memory falls by up to 5.6×; Table 4 gives 3.41 / 0.60 ≈ 5.7×. Speed: 3.09 / 1.52 ≈ 2.0×.

Lineage

• Builds on: ToMe (Bolya et al., 2023), which the paper calls its base method and whose ViT setting it extends to diffusion §1 §3.1. See Token Merging: Your ViT but Faster. The paper adds unmerging for dense outputs §3.1, a 2 × 2 random destination partition §4.1 and a choice of which modules and blocks to merge §4.2.
• Compared against: the naive ToMe port tab. 1; pruning, shown as the failing alternative fig. 4; the unmodified Stable Diffusion v1.5 tab. 4. FlashAttention and xFormers are discussed as complementary, with a combined 2048 × 2048 example §1 fig. 1 §5.
• Superseded by: not stated in the content this wiki holds.