Which Tokens to Merge? Diffusion Dynamics for Efficient Image Generation

The gist The authors propose Dynamics-aware Token Merging (DaTMe), a training-free framework that grounds token merging in diffusion dynamics by constructing a dynamics-aware importance map from the Tweedie denoised estimate to adaptively assign tokens into merging roles, thereby preserving structurally critical regions while minimizing generation degradation.

Why it matters

Token merging has emerged as an effective training-free strategy for accelerating diffusion models by reducing redundant token computation while minimizing generation degradation The authors address the challenge that existing methods overlook background objects and local details not explicitly mentioned in the prompt, leading to degraded structural coherence. Their key insight is that token importance should reflect whether a token remains unresolved during denoising and whether it carries structural information.

How it works

The DaTMe framework constructs a dynamics-aware importance map from the Tweedie denoised estimate, which 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. This map is used to adaptively assign tokens into merging roles under a given compression ratio, where important tokens are preserved as anchors, while those that are both temporally stabilized and structurally homogeneous are selected as safe merge candidates.

The process involves defining two complementary signals for token importance:

1. Temporal dynamics: This is quantified by the inter-step residual in predicted clean-image space, defined as M(t)temp = xˆ0(t) − xˆ0(t − 1), which quantifies how much the predicted clean structure is still evolving at each spatial location.

2. Structural saliency: This is quantified by a spatial axis derived from the Tweedie estimate, defined as M(t)spat = xˆ0(t) − UpSample×s DownSample×s xˆ0(t), which assigns high responses to contours, boundaries, and fine textures.

DaTMe Pipeline

DaTMe performs three main steps at each denoising step t given a target compression ratio r:

1. Role Assignment: It defines a conservative source set Ssrc as the intersection of the lowest-r fractions of tokens ranked by Mtemp and Mspat, and a constraint set Scon as the intersection of the top-(1 − r) fractions on both axes, preserving tokens that are both temporally unstable and structurally salient.

2. Spatially Distributed Source Completion: To meet the target compression budget ntgt = ⌊rN⌋, DaTMe selects an additional set Saddsrc from the remaining unassigned tokens and samples them into P non-overlapping grid cells to avoid local information loss after merging. The final source set is then S∗src = Ssrc ∪ Saddsrc.

3. Locality-Constrained Routing: The remaining tokens form the destination pool D, and S∗src is merged into D via bipartite soft matching using spatial adjacency as the matching criterion to ensure that merging does not collapse spatially distant tokens and distort local spatial structure.

Experimental Validation

Experiments on PixArt-α 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. Quantitative results show that DaTMe achieves strong preservation of image-wide details and state-of-the-art performance across most metrics when compared to prior methods. Ablation studies confirm that jointly using temporal and spatial signals yields the best overall performance, while replacing the temporal signal with the raw score slope consistently degrades all evaluation metrics.

Conclusion

DaTMe successfully grounds token merging in underlying generation dynamics by leveraging a spatiotemporal importance map to identify safely mergeable tokens while preserving structurally critical regions, achieving strong image fidelity and favorable quality-efficiency trade-offs across DiT-based diffusion models. The work shows that grounding token merging in the underlying generation dynamics is a key step toward faithful and efficient diffusion inference.

--- Page 1 ---

Which Tokens to Merge? Diffusion Dynamics for Efficient Image Generation

SeungJu Cha Ye-Chan Kim HyunGee Kim Sungho Koh Dong-Jin Kim† Hanyang University, South Korea

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 dynamicsaware 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-α 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.

--- Page 2 ---

CFG Score Slope Tweedie Diff Pseudo GT

T (a) Results on different compression ratio (b) Spatial distribution of top-ranked tokens (c) Blind spots at non-prompt region Figure 1: (a) Comparison under different compression ratio. CFG-based merging [9] often overlooks background details and structures, whereas our method better preserves background integrity. (b) Spatial distribution of top-ranked tokens over denoising timesteps. CFG-based importance remains concentrated on the foreground object with limited temporal variation. In contrast, our score-based signals (Score Slope and Tweedie Diff) progressively shift from object regions to fine-grained spatial details, better matching the pseudo-GT pattern. Here, pseudo-GT is defined as the difference between the final denoised latent and the Tweedie estimate at each timestep. (c) Non-prompt region blind spot. Low-ranked tokens under the CFG-based signal consistently overlap with structurally salient regions in text-irrelevant areas (red). However, as shown in Fig. 1(a), CFGbased merging tends to underprotect non-prompt regions of the image, losing fine details and structural coherence. We attribute this to CFG’s prompt-centric bias: by focusing exclusively on text-conditioned influence, CFG overlooks two key aspects of how diffusion models generate images. First, CFG fails to capture time-varying token importance. Diffusion generation follows a coarse-tofine progression [15, 16], where tokens undergoing large inter-step changes represent regions in which the model is actively injecting details [17]. In this context, prematurely merging such still-evolving tokens directly attenuates their ongoing refinement by averaging their features with neighboring tokens. Crucially, the spatial distribution of these actively evolving tokens is time-varying: while it remains concentrated on foreground objects in early steps, it progressively expands toward broader spatial regions, including non-prompt areas, at later timesteps (Fig. 1(b)). High-ranked tokens in CFG importance signal, however, remain anchored to foreground objects throughout the entire timestep, confining its search for important tokens only to prompt-aligned regions. As a result, CFG fails to identify the very tokens that should be protected: actively evolving tokens outside the foreground are prematurely merged into nearby or similar tokens, suppressing their distinct updates. Second, CFG is insensitive to prompt-independent structural importance. To produce natural and coherent images, diffusion models actively generate scene elements that are not directly mentioned in the prompt such as background objects and supporting structures [1, 2]. Although these promptindependent structures contribute little to text guidance, they carry essential structural information such as boundaries between homogeneous regions, edges of background objects, and fine textures that anchor visual coherence. However, because CFG only measures text-conditioned discrepancy, structural boundaries in such non-prompt regions tend to receive low scores across all timesteps (Fig. 1(c)), making them vulnerable to merging into neighboring homogeneous areas. Taken together, these observations indicate that token importance should not be defined solely by text relevance, but by diffusion dynamics: how unresolved a token remains and how much structural information it carries at the current denoising step [18].

--- Page 3 ---

To faithfully capture diffusion dynamics, we define two complementary signals that quantify the spatiotemporal dynamics of each token during denoising. For temporal dynamics, we start from the score function [19], which quantifies how each token is updated at the current denoising step. In particular, we project this score into clean-image space via the Tweedie denoised estimate [20, 21], and use its inter-step difference as the temporal signal, assigning high values to tokens that are unresolved. For the structural signal, we measure the residual between the Tweedie estimate and its coarse-scale smoothed version, which acts as a high-pass filter highlighting structural components throughout the image.

1. Guarda token merging in diffusion dynamics through DaTMe (Dynamics-aware Token Merging). This framework grounds token merging in the underlying diffusion dynamics by constructing a dynamics-aware importance map from the Tweedie denoised estimate, which combines a temporal signal, measuring inter-step changes in the predicted clean image, with a structural signal. This allows for adaptive assignment of tokens into roles: important tokens are preserved as anchors, while tokens that are both temporally stabilized and structurally homogeneous are selected as safe merge candidates.

2. Enhance structural coherence and background integrity in generated images. The paper demonstrates that DaTMe better preserves background integrity than CFG-based merging, which the authors attribute to CFG’s prompt-centric bias, leading to a method that better preserves the original spatial configuration and object structure while maintaining the intended prompt semantics.

3. Develop a more reliable temporal importance signal for token evolution. By defining the temporal map as the inter-step residual in predicted clean-image space: M(t) temp = xˆ0(t) − xˆ0(t − 1), DaTMe captures how much the predicted clean structure is still evolving at each spatial location, providing a signal that is more stable than the raw score slope, which is prone to being heavily dominated by the noise schedule scaling in later denoising steps.

4. Implement locality-constrained spatial matching for robust token merging in DiT architectures. To prevent structural distortion, DaTMe uses locality-constrained routing, pairing source tokens with their spatially nearest destination to form the merged set, thereby compensating for the weak spatial inductive bias of DiTs.

5. Achieve superior quality-efficiency trade-offs across diverse diffusion model architectures. Experiments show that DaTMe consistently achieves the strongest CLIP-T and HPSv2 scores while maintaining comparable inference efficiency across PixArt-α, FLUX, Stable Diffusion 1.5, and SDXL, demonstrating effectiveness beyond DiT models.