Which Tokens to Merge? Diffusion Dynamics for Efficient Image Generation

AI summary

Overview

Research area: Computer Vision — efficient generative modeling, specifically training-free token reduction (token merging) for diffusion transformers (DiTs) used in text-to-image synthesis.

Technical level: Intermediate. The paper assumes familiarity with diffusion models, the score function, classifier-free guidance, and transformer token merging, but its core idea is explained geometrically and can be followed with that background.

Scope: The paper proposes DaTMe, a training-free token merging framework that decides which tokens to merge inside a diffusion transformer by measuring, at each denoising step, how unresolved a token is (temporal signal) and how much structure it carries (spatial signal), both derived from the Tweedie denoised estimate.

What This Paper Is About

Diffusion models repeatedly process a large set of image tokens at every denoising step, which is expensive; token merging speeds this up by combining similar tokens, but it risks destroying the image if the wrong tokens are merged. The authors argue that the existing state-of-the-art criterion, which uses classifier-free guidance (CFG) to protect prompt-relevant tokens, is biased toward text-mentioned foreground objects and therefore discards background structures, boundaries, and other details that the prompt never names. Their goal is to redefine token importance using the model's own denoising dynamics so that merge decisions protect everything still being refined, not just what the text describes.

Key Contributions

1. A dynamics-aware importance map from the Tweedie estimate. The authors build token importance from two axes computed on the Tweedie denoised estimate x̂₀(t): a temporal axis equal to the inter-step difference |x̂₀(t) − x̂₀(t−1)|, and a spatial axis equal to the residual between x̂₀(t) and its downsampled-then-upsampled coarse version, which acts as a high-pass filter for boundaries and fine texture.

2. A theoretical and empirical argument for Tweedie difference over raw score slope. They decompose the score slope into a "difference term" (reflecting actual model prediction changes) and a "scale term" (arising purely from the predefined noise schedule σ_t), show the scale term dominates in late denoising steps, and show the Tweedie difference removes that term while retaining only the difference term, scaled by σ_{t−1}/α_{t−1}.

3. The DaTMe framework. A training-free pipeline that assigns tokens into a constraint set (intersection of the top (1−r) fractions on both axes, excluded from merging), a source set (intersection of the lowest-r fractions on both axes), grid-based source completion (P grid cells, (n_tgt − n_src)/P tokens sampled per cell) to hit the compression budget without local collapse, and locality-constrained bipartite matching that pairs each source token with its spatially nearest destination.

4. Evaluation across two architectures. Validation on PixArt-α (a DiT) and FLUX.1-dev (an MM-DiT), reporting text alignment, preference, FID, fidelity, aesthetic score, latency, and TFLOPs, plus ablations and a human preference study.

Main Findings

• DaTMe gives the best quality–efficiency trade-off on PixArt-α, at a small latency cost. At compression ratio 0.7, DaTMe reaches FID 47.0519 versus 70.7062 for the CFG-based method and 123.0139 for ToMeSD, with CLIP-T 30.6259 versus 29.4543 (CFG) and 28.2805 (ToMeSD) and HPSv2 25.2393 versus 22.6156 and 18.9797. Latency is 1.6426 s versus 1.5761 s (CFG) and 1.5356 s (ToMeSD), and TFLOPs are 165.193 versus 157.897 for both baselines. The uncompressed PixArt-α base is listed at CLIP-T 31.2485, HPSv2 29.2780, FID 42.1001, latency 1.9417 s, TFLOPs 183.45.

• Structural fidelity improves substantially on PixArt-α. At ratio 0.3, DaTMe obtains SSIM 0.8112, PSNR 20.5136, LPIPS 0.1501, and Aesthetic 6.2014, compared with SSIM 0.7137, PSNR 17.8157, LPIPS 0.2369, Aesthetic 6.1521 for the CFG method and SSIM 0.7306, PSNR 18.9748, LPIPS 0.2433, Aesthetic 5.9906 for ToMeSD. The fidelity advantage persists at ratios 0.5 and 0.7.

• DaTMe transfers to FLUX, where the CFG baseline does not apply. Because FLUX does not employ CFG, the prior CFG-based token selection cannot be used, whereas DaTMe applies directly. At ratio 0.3 DaTMe records CLIP-T 31.3524, HPSv2 29.5042, FID 30.9258, latency 13.772 s, and 2033 TFLOPs, versus the FLUX.1-dev base of 31.2274, 29.6927, 30.6141, 14.601 s, and 2098 TFLOPs. DaTMe's FID is slightly above the uncompressed base at all three ratios (30.9258, 31.3423, 31.1585 at 0.3, 0.5, 0.7), while its fidelity metrics are the strongest among compressed variants.

• CFG importance is static; Tweedie-based signals track evolving regions. In the analysis, the CFG-based signal shows consistently high temporal autocorrelation across all timesteps, while both score-based signals show noticeably lower autocorrelation early on. The Tweedie difference maintains higher correlation with the pseudo-GT change magnitude throughout denoising, whereas the raw score slope's correlation drops sharply around normalized timestep 0.4–0.6, matching the range where the scale term overtakes the difference term.

• Both axes are needed. In the ablation at ratio 0.3, using only the temporal axis gives CLIP-T 31.0355, SSIM 0.8109, PSNR 20.5021, FID 42.5269; using only the spatial axis gives 31.1318, 0.8110, 20.5029, 42.5736; using both gives 31.1375, 0.8112, 20.5136, 42.4474. Substituting the raw score slope for the Tweedie difference degrades all metrics (31.1166, 0.8028, 20.3944, 42.9554) and introduces visible artifacts.

• The constraint set should be the intersection of high temporal and high spatial responses. At ratio 0.3, that definition yields SSIM 0.8112, PSNR 20.5136, LPIPS 0.1501, FID 42.4474; using temporal-high ∩ spatial-low drops to SSIM 0.8005, PSNR 20.3459, LPIPS 0.1641, FID 43.8062.

• The importance signal alone helps, independent of the merging machinery. Keeping the prior merging pipeline but swapping in the proposed importance signal (combined via rank-sum) improves CLIP-T from 30.8801 to 31.0545, HPSv2 from 27.7375 to 27.9366, SSIM from 0.7137 to 0.7644, PSNR from 17.8157 to 20.0917, and LPIPS from 0.2369 to 0.1915; the full DaTMe pipeline reaches 31.1375, 28.5591, 0.8112, 20.5136, and 0.1501.

• Humans prefer DaTMe. In a study at ratio 0.3 with 20 randomly sampled PixArt-α cases and 32 evaluators comparing ToMeSD, the CFG method, and DaTMe without access to the uncompressed reference, DaTMe received the highest ratings for text alignment and visual quality and was the most preferred overall.

Methodology in Plain English

The starting point is that a diffusion model, at every step, produces a guess of what the final clean image looks like — the Tweedie denoised estimate. The authors treat this guess as a readable "state of progress" map.

Two measurements are taken from this guess. The first is change over time: compare the guess at this step with the guess at the previous step; where the difference is large, the model is still actively resolving that region, so those tokens should not be averaged away. The second is local structure: blur the guess by downsampling and upsampling it, then subtract the blur from the original; large residuals mark edges, boundaries, and fine textures, which must survive even if the prompt never mentions them.

These two maps are used to hand out roles each step under a target compression ratio r. Tokens ranked in the bottom r fraction on both maps are considered safe to merge; tokens in the top (1−r) fraction on both maps are locked as anchors. Because the strict intersection of two conditions usually yields fewer candidates than the budget requires, extra merge candidates are drawn evenly from a P-by-P grid over the token layout so no single region gets over-merged. Merging itself is restricted to spatial neighbors: each source token is matched to the nearest destination token rather than to whatever token happens to have the most similar feature vector, which matters because transformers have no built-in spatial bias. The method requires no training or fine-tuning and is applied to the middle-range self-attention layers for PixArt-α and from the middle block to the final block of the joint-attention blocks for FLUX.1-dev.

Why This Matters

For research, the paper reframes token reduction in diffusion models as a question about generative dynamics rather than about semantic relevance to the text prompt. It offers a concrete, cheap signal (a difference of Tweedie estimates, plus a blurred-residual filter) that requires no attention-map computation and no extra training, and it supplies a decomposition showing why the more obvious "score slope" signal degrades late in sampling. The finding that a CFG-derived criterion cannot even be applied to FLUX, while a dynamics-derived one can, is a useful argument for architecture-agnostic criteria.

Real-world applications:

• On-device and edge image generation, where latency and FLOP budgets are tight and a training-free speedup that preserves background quality is more deployable than a model that must be retrained.
• High-throughput creative pipelines (marketing assets, game art, product mockups) that generate many images per request and can trade a modest amount of compute for large batches.
• Interactive design tools where users iterate on prompts and need fast previews that do not lose scene structure outside the named subject.
• Cost reduction for hosted generation APIs, since the savings come from skipping token computation rather than from shortening the sampling trajectory.

Industry relevance: the method is a drop-in, training-free layer, and the reported trade-off is explicit — DaTMe is slightly slower than the baselines it beats (for example 1.8030 s versus 1.7458 s at ratio 0.3 on PixArt-α) but degrades far less as compression increases, which is the regime where serving costs actually fall.

Future Directions

• Closing the latency gap. DaTMe's FID and fidelity advantages come with slightly higher latency and TFLOPs than the baselines at the same nominal ratio (175.631 versus 172.514 TFLOPs at ratio 0.3), so making the role-assignment overhead negligible is an open engineering question.
• Restoring FID parity with the uncompressed model on FLUX. DaTMe's FLUX FID stays just above the uncompressed base at every tested ratio, so the residual gap between compressed and uncompressed generation is unresolved.
• Extending beyond images. The paper evaluates only text-to-image DiTs and MM-DiTs; whether dynamics-aware merging transfers to video diffusion transformers, which the introduction cites as a major application area, is not tested.
• Tuning the two free hyperparameters. The spatial scale s = 4 and grid count P = 4 are fixed across experiments; how sensitive results are to these choices, and whether they should adapt to resolution or timestep, is not reported in the provided content.

Target Audience

Graduate students and applied researchers working on diffusion model inference efficiency, and engineers deploying text-to-image models who need training-free acceleration. It is also relevant to readers interested in token reduction for transformers more broadly, and to practitioners who care about preserving image regions that the text prompt does not describe. Readers without background in diffusion sampling or the Tweedie formula will need to consult the preliminaries section and cited references first.

Authors’ abstract

Token merging has emerged as an effective training-free strategy for accelerating diffusion models by reducing redundant token computation while minimizing generation degradation. A key challenge is determining which tokens can be safely merged and which should be preserved for subsequent computation. Existing method addresses this by preserving prompt-aligned foreground regions, thereby protecting semantically relevant structures. However, this text-relevance criterion overlooks background objects, local details that are not explicitly mentioned in the prompt, leading to degraded structural coherence. In this work, we propose Dynamics-aware Token Merging (DaTMe), a training-free framework that grounds token merging in the underlying diffusion dynamics. Our key insight is that token importance should reflect whether a token remains unresolved during denoising and whether it carries structural information. To this end, DaTMe constructs a dynamics- aware importance map from the Tweedie denoised estimate. The map combines a temporal signal, measuring inter-step changes in the predicted clean image, with a structural signal, capturing local high-frequency content such as boundaries and fine details. Using this map, DaTMe adaptively assigns tokens into merging roles under a given compression ratio: important tokens are preserved as anchors, while tokens that are both temporally stabilized and structurally homogeneous are selected as safe merge candidates. Experiments on PixArt-a and FLUX demonstrate that DaTMe maintains visual fidelity across the entire image, highlighting diffusion dynamics as an effective token-level merging decision in efficient generation. Our code is available at https://github.com/SeungJuCha/DaTMe.