D$^{3}$ToM: Decider-Guided Dynamic Token Merging for Accelerating Diffusion MLLMs

AI summary

Overview

• Research area: Computer Vision / efficient inference for diffusion-based multimodal large language models (Diffusion MLLMs).
• Technical level: Advanced (assumes familiarity with Transformer attention, denoising diffusion decoding, and FLOPs accounting), though the core idea is describable in plain terms.
• Scope: The paper introduces D³ToM, a training-free, plug-and-play module that physically shortens visual token sequences inside a Diffusion MLLM so that denoising inference runs faster while retaining most task accuracy.

What This Paper Is About

Diffusion MLLMs generate text non-autoregressively by iteratively denoising a masked sequence, but each denoising step runs full bidirectional self-attention over the entire sequence, so cost grows cubically with sequence length (O(N³) for T = O(N) steps of O(N²) attention). Because these models process thousands of visual tokens, N is far larger than in text-only settings, making decoding computationally impractical. The paper's goal is to exploit heavy redundancy among visual tokens to cut that cost without retraining the model or altering any of its parameters.

Key Contributions

1. Decider-guided token merging. D³ToM uses decider tokens — the output tokens decoded in the immediately preceding denoising step — to build an importance map over all visual tokens, keeps the top (1 − α) fraction, and merges the rest into their most cosine-similar kept tokens by adding their embedding vectors.
2. A timestep-dependent merge schedule (D³ToM-t). The merge ratio α_t varies linearly with the denoising timestep, retaining more tokens early (when global semantics are being drafted) and merging more aggressively late (when only fine details remain).
3. Training-free, plug-and-play design. The merge happens inside a single transformer layer (l*), physically shortening the visual sequence for all downstream layers with no modification of existing model parameters.
4. Demonstrated orthogonality to KV caching. Combining D³ToM with the Prefix-DLM cache yields complementary savings, addressing a different computational bottleneck than caching.

Main Findings

• Visual tokens are largely redundant. The probability density distribution of attention weights from output tokens to grouped visual tokens is heavily skewed toward zero: because an image contains over one thousand visual tokens while the decoder generates only 64 output tokens, every 16 consecutive visual tokens were averaged into a group, and most groups receive near-zero attention.
• Visual saliency shifts across denoising steps. Attention snapshots at steps 8, 16, 24, and 32 on an image of a small white Chihuahua on a wooden stump show early steps attending to the dog's body (as "white" is decoded) and later steps attending to the background stump (as "stump" is generated) — a step-wise sparsity that static pruning cannot capture.
• Strong accuracy retention across seven benchmarks. On LaViDa, evaluated on MMMU, SQA, MMBench, MathVision, AI2D, LLaVA-Bench, and GQA, D³ToM-t retains 98.12% of baseline average performance when retaining 25% of tokens, and 96.75% when retaining only 10%. D³ToM (fixed ratio) retains 95.23% and 95.31% at those same budgets.
• Clear margin over adapted autoregressive baselines. At 10% retention the baselines score markedly lower: FastV 76.34%, PDrop 86.41%, VisionZip 88.68%, versus D³ToM 95.31% and D³ToM-t 96.75%.
• Efficiency matches inner-LLM pruning baselines. Averaged over LLaVA-Bench with T = 32 and O = 64, the baseline LaViDa costs 262.60 TFLOPs and 1156.65s; at 10% retention D³ToM-t costs 79.37 TFLOPs (30.2% relative) and 493.88s (42.7% relative), while FastV costs 77.14 TFLOPs (29.4%) and 487.19s (42.1%) and PDrop costs 97.16 TFLOPs (37.0%) and 578.31s (50.0%).
• VisionZip is cheaper but weaker. VisionZip reduces raw cost further (51.36 TFLOPs, 19.6% relative at 10% retention) by deleting tokens in the visual encoder, but its one-shot pruning cannot follow evolving visual saliency during diffusion.
• Variable schedule is essentially free. The paper derives that the FLOPs difference between a constant merge ratio and a time-varying linear schedule with the same mean arises only from the quadratic self-attention term, is 4d|V|²[(T−1)Var(α_t)], and is negligible (Δ/(4dN²) < 1%), so FLOPs for D³ToM-t ≃ FLOPs for D³ToM.
• Compatibility with KV caching. With the Prefix-DLM cache, D³ToM* retains 99.7% of average accuracy and D³ToM-t* retains 99.4% at 10% token retention, confirming the two techniques act on distinct parts of the computation.

Methodology in Plain English

The method exploits two observations: most visual tokens barely matter, and which tokens matter changes from one denoising step to the next.

At each denoising step, the model has already decoded some output tokens in the previous iteration. D³ToM calls these the decider tokens and uses the attention weights they assign to every visual token as an importance score — summing the attention each visual token receives from all decider tokens. Given a merge ratio α, it keeps the top (1 − α) fraction of visual tokens by score and sends the rest to be merged.

Merging works on the post-attention hidden states at a chosen layer l*. Each token marked for merging is matched to the kept token with the highest cosine similarity, and its representation is simply added into that kept token's representation. The merged positions are then physically removed from the hidden-state tensor, so the sequence is shorter for every subsequent layer. The model's parameters are untouched.

For the timestep-aware variant, α_t is set by a linear schedule from α_min at the start of decoding to α_max near the end, so more aggressive merging happens as the model moves from building a coarse semantic draft to refining details.

The paper also derives closed-form FLOPs expressions — a per-layer cost of 4nd² + 2n²d (self-attention) plus 3ndm (feed-forward), where n is token count, d the hidden size, and m the FFN width — and matches baseline methods' total effective token budget (cumulative sequence length over all denoising steps) for a fair comparison.

Implementation details reported: the method is built on LaViDa-LLaDA-8B, run on an NVIDIA RTX A6000 (48G) GPU, with output lengths |O| of 16 for MMMU, SQA, MMBench, AI2D and GQA, 64 for LLaVA-Bench, and 100 for MathVision. Results are produced using the official lmms-eval implementation.

Why This Matters

Impact on research. Token compression has been extensively studied for autoregressive MLLMs, but the paper argues that these techniques are fundamentally incompatible with the non-causal denoising process of diffusion MLLMs. D³ToM opens that direction and shows that decider tokens — a signal unique to the diffusion decoding loop — can drive compression decisions that a one-shot or static attention-based criterion cannot. It also positions token merging and KV caching as orthogonal rather than competing optimizations.

Real-world applications (drawn from the paper's evaluation domain):

• Real-time visual question answering and image captioning, where LLaVA-Bench's instruction-following and detailed-description tasks mirror assistant-style interactions.
• Diagrammatic and scientific reasoning, matching AI2D and ScienceQA, useful in education and technical documentation tools.
• Expert-level multimodal reasoning such as MMMU, relevant to professional analysis workflows over charts, chemical structures, and engineering diagrams.
• Mathematical problem solving over images, matching MathVision, relevant to tutoring and automated grading systems.

Industry relevance. Serving Diffusion MLLMs at scale is limited by GPU time and cost; the paper reports reducing cost to about 30% of baseline FLOPs and roughly 42–43% of wall-clock time at 10% token retention while keeping over 96% of performance, and shows the method stacks with existing caching infrastructure rather than requiring new training runs.

Future Directions

• Other Diffusion MLLM backbones. The method is implemented only on LaViDa-LLaDA-8B; whether it transfers to MMaDA, LLaDA-V, or other diffusion MLLMs is not reported.
• Joint optimization with parallel decoding. The paper notes Confidence-Aware Parallel Decoding and SlowFast reduce the number of denoising steps but do not reduce visual token count; combining step reduction with token merging is an untested combination.
• Extending beyond images. HoliTom, cited in related work, performs pruning with spatial-temporal merging for video LLMs; video introduces a temporal token dimension that the current image-based merging does not address.
• Better merge policies. The schedule is a simple linear ramp between α_min and α_max, and the merge layer l* appears fixed; learned or step-adaptive schedules and automatic selection of the merge layer are open questions.

Target Audience

Researchers and practitioners working on efficient multimodal LLM inference — particularly those building or serving Diffusion MLLMs — will benefit most. It is also useful for readers interested in token pruning/merging for vision-language models who want to understand how the non-causal, iterative denoising setting changes the design constraints relative to autoregressive MLLMs. Readers without background in Transformer attention or diffusion decoding will need to consult the preliminaries in Section 3.1.

Authors’ abstract

Diffusion-based multimodal large language models (Diffusion MLLMs) have recently demonstrated impressive non-autoregressive generative capabilities across vision-and-language tasks. However, Diffusion MLLMs exhibit substantially slower inference than autoregressive models: Each denoising step employs full bidirectional self-attention over the entire sequence, resulting in cubic decoding complexity that becomes computationally impractical with thousands of visual tokens. To address this challenge, we propose D$^{3}$ToM, a Decider-guided dynamic token merging method that dynamically merges redundant visual tokens at different denoising steps to accelerate inference in Diffusion MLLMs. At each denoising step, D$^{3}$ToM uses decider tokens-the tokens generated in the previous denoising step-to build an importance map over all visual tokens. Then it maintains a proportion of the most salient tokens and merges the remainder through similarity-based aggregation. This plug-and-play module integrates into a single transformer layer, physically shortening the visual token sequence for all subsequent layers without altering model parameters. Moreover, D$^{3}$ToM employs a merge ratio that dynamically varies with each denoising step, aligns with the native decoding process of Diffusion MLLMs, achieving superior performance under equivalent computational budgets. Extensive experiments show that D$^{3}$ToM accelerates inference while preserving competitive performance. The code is released at https://github.com/bcmi/D3ToM-Diffusion-MLLM.