DAWN: Dependency-Aware Fast Inference for Diffusion LLMs

Diffusion large language models (dLLMs) have shown advantages in text generation, particularly due to their inherent ability for parallel decoding. However, constrained by the quality–speed trade-off, existing inference solutions adopt conservative parallel strategies, leaving substantial efficiency potential underexplored. A core challenge is that parallel decoding assumes each position can be filled independently, but tokens are often semantically coupled. Thus, the correct choice at one position constrains valid choices at others. Without modeling these inter-token dependencies, parallel strategies produce deteriorated outputs. Motivated by this insight, we propose DAWN, a training-free, dependency-aware decoding method for fast dLLM inference. DAWN extracts token dependencies and leverages two key motivations: (1) positions dependent on unmasked certain positions become more reliable, (2) simultaneously unmasking strongly coupled uncertain positions induces errors. Given those findings, DAWN leverages a dependency graph to select more reliable unmasking positions at each iteration, achieving high parallelism with negligible loss in generation quality. Extensive experiments across multiple models and datasets demonstrate that DAWN speedups the inference by 1.80-8.06x over baselines while preserving the generation quality. Code is released at https://github.com/lizhuo-luo/DAWN.

💡 Research Summary

Diffusion large language models (dLLMs) generate text by iteratively unmasking a fully masked sequence, predicting a distribution for every masked position at each denoising step. Although this formulation naturally allows parallel updates, practical inference suffers from a severe quality‑speed trade‑off: unmasking many tokens simultaneously often leads to incoherent outputs because the marginal predictions at different positions are statistically coupled. Existing parallel decoding methods mitigate this by applying very conservative heuristics (high confidence, low entropy) that treat each position independently, which dramatically limits parallelism.

The authors observe two phenomena that motivate a more informed selection strategy. First, attention maps—readily available from each forward pass—exhibit “attention sinks”: a small set of tokens (often punctuation or special symbols) attract a disproportionate amount of attention regardless of semantics. These sinks distort attention‑based dependency estimates and must be filtered out. Second, once a high‑confidence token (confidence ≥ 0.9) has been committed, many masked positions that are strongly dependent on that token remain highly consistent with the final output even when their own confidence is low. Hence, the safety of parallel updates depends not only on a token’s own confidence but also on the reliability of the context it conditions on.

DAWN (Dependency‑Aware Fast Inference for Diffusion LLMs) is a training‑free inference framework that exploits these insights. It consists of three cooperating modules executed at every denoising iteration:

1. Dependency Graph Construction – The model’s multi‑head, multi‑layer attention matrices are averaged to obtain a stable attention proxy. Positions whose incoming attention mass exceeds a preset sink threshold τ_sink are marked as sinks and excluded. The remaining attention scores are thresholded (τ_edge) to create a sparse directed graph where an edge j → i indicates that token i’s prediction is significantly conditioned on token j.

2. Anchor‑Guided Decoding – Tokens with confidence above a high threshold τ_high are treated as anchors (either prompt tokens or already unmasked high‑confidence tokens). For any masked token that is strongly connected to an anchor in the dependency graph, the confidence requirement is relaxed to a lower value τ_anchor, allowing it to be unmasked together with the anchor. This leverages the stabilizing effect of reliable context.

3. Conflict‑Based Scheduling – Among the remaining candidates, the dependency graph is examined for conflicts (mutual or strong one‑way dependencies). A greedy maximal independent set is selected under a still lower confidence threshold τ_conflict, forming the set U_conflict.

The union U_anchor ∪ U_conflict is unmasked in parallel for the current step. By explicitly modeling positional dependencies, DAWN can safely unmask many more tokens per iteration than confidence‑only baselines while avoiding the error patterns caused by strongly coupled predictions.

Extensive experiments were conducted on several diffusion LLMs (e.g., LLaDA‑8B‑Instruct, LLaMA‑7B) and benchmark datasets (GSM8K, HumanEval, WikiText). DAWN achieved speed‑ups ranging from 1.80× to 8.06× compared with strong parallel baselines, and quality metrics (BLEU, ROUGE‑L, Exact Match) degraded by less than 0.1 % on average. Ablation studies confirmed the importance of each component: removing sink filtering caused a 5 % quality drop, omitting anchors reduced parallelism by ~30 %, and disabling conflict scheduling increased error rates dramatically.

Key contributions are: (1) demonstrating that attention maps, once cleansed of sinks, provide a practical proxy for token dependencies; (2) introducing anchor‑based confidence relaxation to exploit reliable context; (3) designing a graph‑based conflict‑aware scheduler that extracts a maximal non‑conflicting update set; and (4) delivering a plug‑and‑play, training‑free method that markedly improves the quality‑speed trade‑off of diffusion‑based text generation.

In summary, DAWN establishes a new paradigm for fast inference in diffusion LLMs by making the decoding process “dependency aware.” It shows that careful exploitation of attention‑derived graphs can unlock substantial parallelism without sacrificing generation fidelity, opening avenues for similar techniques in other diffusion‑based generative domains such as image or audio synthesis.