[KongchangAI]
· 2 min read· 1,430 words

Diffusion-Dynamics-Driven Token Merging: DaTMe Accelerates Image Generation

Diffusion-Dynamics-Driven Token Merging: DaTMe Accelerates Image Generation

DaTMe uses diffusion dynamics instead of text relevance to judge token importance, enabling training-free diffusion model acceleration.

Iterative denoising in diffusion models incurs high computational costs, and token merging is a training-free acceleration approach, but determining which tokens can be safely merged has always been difficult. Existing methods rely on text relevance to protect the foreground but neglect background objects and unmentioned local details, causing structural breaks and blurred details. DaTMe proposes defining token importance via the diffusion process's own dynamics: building a dynamics-aware importance map from Tweedie denoising estimates that fuses temporal signals of prediction changes between adjacent steps with structural signals like boundaries and textures. Important tokens are kept as anchors, while temporally stable and structurally uniform tokens become safe merging candidates, with roles adaptively assigned by a given compression ratio. Experiments on PixArt-a and FLUX show the method maintains full-image visual fidelity, and the code is open-sourced.

Introduction: A New Approach to Accelerating Diffusion Models

Diffusion models have achieved remarkable results in image generation, but the computational overhead of their iterative denoising process remains a persistent bottleneck for practical deployment. Token merging, as a training-free acceleration strategy, improves inference speed by reducing the computation of redundant tokens while preserving generation quality as much as possible. However, which tokens can be safely merged and which must be retained has always been the core challenge in this direction.

Existing methods typically rely on text relevance to protect foreground regions, but overlook background objects and local details not mentioned in the prompt, leading to degraded structural coherence. To address this limitation, researchers have proposed Dynamics-aware Token Merging (DaTMe), which grounds token merging decisions in the dynamics of the diffusion process itself.

Overview of the Diffusion-Dynamics Token Merging Method

Diffusion Models are a class of generative models whose core idea is to progressively add noise to data (the forward process) until the data is destroyed into approximate Gaussian noise, then learn a reverse denoising process that starts from pure noise and gradually recovers the target data. In image generation, this process typically requires dozens or even hundreds of sampling steps, each requiring a forward pass through the entire model. Since each step must process all tokens corresponding to the image (i.e., image patches or latent representation units), inference cost grows rapidly with the number of steps and tokens. Therefore, reducing the number of tokens involved in computation at each step has become an important direction for accelerating diffusion inference without retraining the model, which is the direct reason token merging strategies have attracted attention.

Limitations of Existing Methods: Text Relevance Is Not Enough

Current mainstream token merging strategies use prompt alignment as the importance criterion, retaining foreground regions semantically relevant to text while treating remaining tokens as mergeable. This approach is reasonable at the semantic level, but has two clear problems.

First, objects and structures in the background are often not explicitly mentioned in the prompt, yet are crucial to the overall spatial layout of the image. Second, high-frequency local details such as boundaries and textures, while invisible in text, directly affect visual fidelity. Relying solely on text relevance causes these regions to be merged prematurely, resulting in structural breaks or blurred details in the generated output.

This shows that token importance cannot be judged only from the perspective of semantic alignment; signals related to the denoising process itself must also be introduced.

DaTMe's Core Insight: Defining Token Importance via Diffusion Dynamics

DaTMe's key hypothesis is that whether a token is important depends on two things—whether it is still in an unconverged state during the denoising process, and whether it carries structural information.

Based on this hypothesis, DaTMe constructs a dynamics-aware importance map from the Tweedie denoising estimate. This map fuses two types of signals:

  • Temporal signal: measures the degree of change in the predicted clean image between adjacent steps; tokens with large changes indicate they have not yet stabilized;
  • Structural signal: captures local high-frequency content such as boundaries and fine details.

After combining these two types of signals, token importance no longer depends on text annotations but is determined by the dynamic behavior of the diffusion process itself. This perspective breaks away from the habitual thinking that "text alignment is sufficient" and binds merging decisions to the physical meaning of the generation process.

Tweedie denoising estimation is a classic result in diffusion model theory: given the current noisy image and noise level, one can infer the conditional expectation of the clean image from the model's predicted noise—that is, an estimate of the original image. This estimate is naturally available during sampling, requiring no additional training or extra network. DaTMe leverages the Tweedie estimate to observe the dynamics of the denoising process: if the clean image estimate for a token changes drastically between adjacent sampling steps, it indicates the token has not yet converged and is still participating in structure formation; otherwise, it is stabilizing. By comparing the estimation differences between adjacent steps, one can compute an "unconvergence degree" signal for each token without text annotations, which is one of the key sources of the dynamics-aware importance map.

Adaptive Merging Strategy: Anchors and Safe Candidates

Given a compression ratio, DaTMe adaptively assigns tokens to different roles based on the importance map. Important tokens are retained as anchors and participate in subsequent computation, while tokens that are temporally stable and structurally uniform are selected as safe merging candidates.

The advantage of this role division is that it accounts for both denoising progress and spatial structure. Stable tokens will not disrupt structures being formed due to merging, while redundant computation in homogeneous regions is effectively reduced. Compared to static merging rules, this dynamic allocation better adapts to the needs of different sampling steps and different image regions.

Token merging refers to combining multiple tokens that are similar or of lower importance into one according to some rule, thereby reducing the number of tokens that subsequent attention or feed-forward networks need to process. Merging itself loses some information, so the key to the strategy lies in choosing which tokens can be merged. Common merging criteria include feature similarity, attention scores, or text relevance. The "anchor" mechanism adopted by DaTMe draws on the idea of representative points in graph clustering: tokens selected as anchors serve as the representatives retained after merging, and the remaining safe candidates are assigned to anchors. Since anchors are dynamically determined by the dynamics importance map, different steps and different regions within the same sampling run may select different anchor sets, thereby preserving structural information being formed as much as possible while compressing.

Experimental Results: Validation on PixArt-a and FLUX

Researchers conducted experiments on two diffusion models, PixArt-a and FLUX. Results show that DaTMe can maintain visual fidelity across the entire image, rather than only protecting foreground regions. This means background structures and local details are better preserved, and the degradation in generation quality is kept within a small range.

This result also supports the paper's core argument: diffusion dynamics can serve as an effective basis for token-level merging decisions. When the merging strategy aligns with the intrinsic dynamics of the denoising process, the trade-off between acceleration and quality can be improved.

For developers concerned with inference efficiency, DaTMe provides an acceleration path that requires no retraining and can be directly inserted into existing diffusion pipelines. Its code has been open-sourced on GitHub, facilitating community reproduction and further extension.

PixArt-a and FLUX are two representative diffusion generative models. PixArt-a belongs to the series of Transformer-based diffusion models and achieves high-quality image generation at lower training cost; FLUX adopts a larger Transformer architecture with a flow-matching-style training objective and performs outstandingly in details and text alignment. The two differ in architectural scale, sampler design, and latent space representation, so they are often used to test whether acceleration methods have cross-model generality. DaTMe can maintain full-image visual fidelity on both, indicating that the diffusion dynamics signal it relies on is not a peculiar phenomenon of a single model but has a certain universality, which is also an important prerequisite for this work to be directly inserted into existing diffusion pipelines without retraining.

Conclusion

DaTMe's value lies not only in proposing a new merging criterion, but more importantly in redefining the source of "token importance"—shifting from text relevance to diffusion dynamics. This idea may inspire more acceleration methods centered on the dynamic characteristics of the generation process, making efficient image generation no longer come at the cost of sacrificing details.

One-line Summary

DaTMe replaces text relevance with diffusion dynamics to judge token importance, enabling training-free acceleration of diffusion models.

Paragraph Summary

Iterative denoising in diffusion models incurs high computational costs, and token merging is a training-free acceleration approach, but determining which tokens can be safely merged has always been difficult. Existing methods rely on text relevance to protect the foreground but neglect background objects and unmentioned local details, leading to structural breaks and blurred details. DaTMe proposes defining token importance via the dynamics of the diffusion process itself: constructing a dynamics-aware importance map from Tweedie denoising estimates, fusing temporal signals of prediction changes between adjacent steps with structural signals such as boundaries and textures. Important tokens are retained as anchors, while temporally stable and structurally uniform tokens serve as safe merging candidates, with roles adaptively assigned according to a given compression ratio. Experiments on PixArt-a and FLUX show that the method maintains visual fidelity across the entire image, and the code has been open-sourced.

#DaTMe#diffusion model acceleration#token merging#Tweedie denoising#PixArt-a#FLUX#training-free acceleration
Share:

Related articles