Diffusion language models offer a promising path to efficient agents through parallel token generation. However, efficient interaction requires adapting both where computation is spent and how much text is generated. We introduce a metacognitive controller that guides diffusion generation according to the needs of the current response and the agent's interaction with the environment. It learns when to finalize tokens and how far generation should extend, reducing overall computation within responses. Our controller, a lightweight plug-in for the diffusion backbone, is trained on agent trajectories to balance correctness and generation cost. This approach connects the agent's generation process to the interaction the agent serves. Experiments across five benchmarks demonstrate that our approach reduces inference FLOPs by 20.6% on average compared with the most efficient baseline, while maintaining competitive task accuracy. Ablation studies further validate the effectiveness of our approach.
Figures & tables
Figure 1: Metacognitive control of environment-aware diffusion decoding. Left: fixed-range decoding computes positions beyond the response consumed by the environment. Right: our controller guides content commitment and adjusts the active computation range during interaction. Orange denotes controller-guided revelation, gray denotes computed tail positions, and dashed cells lie outside the active range.
Figure 2: Overview of our approach. The metacognitive controller monitors the generation state and interaction progress, then guides backbone predictions through the confidence alignment policy and flux policy. Tool calls provide environmental feedback for subsequent generation.
Method
BFCL Live
BFCL Non-live
HotpotQA
2Wiki
MuSiQue
Acc. ↑
FLOPs ↓
Acc. ↑
FLOPs ↓
Acc. ↑
FLOPs ↓
Acc. ↑
FLOPs ↓
Acc. ↑
FLOPs ↓
Fast-dLLM
57.46
57.58
52.66
54.98
43.20
324.16
67.70
340.49
21.20
342.32
LLaDA-8B-Tools
50.00
2889.56
37.27
1375.23
16.50
4310.92
22.80
4185.11
2.20
2644.99
DLLM-Searcher
40.85
547.10
16.09
433.45
37.10
1277.25
50.40
1261.22
17.80
1237.25
LLaDA-8B (SFT)
66.52
32.49
63.60
19.28
45.70
175.00
67.10
204.51
20.70
197.38
LLaDA2.0-mini
69.72
43.79
79.42
14.12
36.30
142.09
41.00
199.93
10.80
168.37
Table 1: Task accuracy and inference cost of our method and baseline methods. FLOPs accumulate model computation over the complete interaction trajectory, including controller computation.
Setting
BFCL Live
BFCL Non-live
HotpotQA
2Wiki
MuSiQue
Acc. ↑
FLOPs ↓
Acc. ↑
FLOPs ↓
Acc. ↑
FLOPs ↓
Acc. ↑
FLOPs ↓
Acc. ↑
FLOPs ↓
w/o flux policy
70.07
26.95
64.46
17.39
47.60
149.71
66.70
162.46
19.60
162.93
w/o conf. align.
69.63
23.11
64.46
14.55
48.90
145.49
68.90
166.48
19.10
167.32
Controller (Ours)
70.96
18.48
65.47
12.51
46.60
131.91
67.90
147.68
20.30
142.83
Table 2: Component ablations of the metacognitive controller. Each ablation disables one component of the full controller. Conf. align. denotes the confidence alignment policy. Accuracy (%) and FLOPs ( 1012 ) follow Table 1 .
Figure 3: Accuracy versus FLOPs reduction on BFCL and multi-hop tasks. Gray points denote evaluated checkpoints; blue points and lines highlight their empirical Pareto frontier. The upper-right direction is better.
Dataset
LLaDA
Dream
Δ Acc.
FLOPs ↓
Δ Acc.
FLOPs ↓
BFCL Live
+4.44
43.1%
+0.00
17.2%
BFCL Non-live
+1.87
35.1%
+1.72
29.0%
HotpotQA
+0.90
24.6%
+1.10
52.9%
2Wiki
+0.80
27.8%
−0.60
58.0%
MuSiQue
−0.40
27.6%
−2.20
67.5%
Table 3: Effect of the metacognitive controller on LLaDA and Dream. Δ Acc. reports the change in accuracy (percentage points), and FLOPs ↓ reports the reduction in inference FLOPs (%), relative to the same backbone without the controller. LLaDA results compare Ours with LLaDA-8B (SFT) in Table 1 .
Figure 4: An interaction comparison using the same backbone and search queries. Green highlights the tokens in three structural segments; Red marks post-call generation; gray denotes padding. Blank blue cells have no saved token text. Flux decoding reduces the active range (gray). With the metacognitive controller, the agent uses 24.3% fewer FLOPs over the complete interaction and answers correctly in this example.
Diffusion language models can generate many tokens in parallel, but they still require repeated denoising steps during inference. This makes generation costly, especially when the model continues to recompute tokens that are already stable. To address these limitations, we propose CAI-DLLM, a training-free inference method that uses first-step confidence to guide denoising and reduce inference time. Specifically, CAI-DLLM commits easy tokens earlier, allocates more denoising steps to harder tokens, and adjusts decoding schedules across output blocks. As it relies only on first-step confidence signals, it does not require retraining, extra predictors, or weight updates. We evaluate CAI-DLLM on LLaDA-8B-Instruct and Dream-7B-Instruct across math, code, reasoning, commonsense, and long-context tasks. CAI-DLLM achieves up to 18.2x wall clock inference speedup on LLaDA GSM8K while improving accuracy from 76.27% to 77.41%, and up to 13.1x speedup on Dream HumanEval while achieving higher pass@1 than no-cache inference, 48.17% compared with 46.95%. On harder reasoning tasks, speedups reach 44.8x, with a largest accuracy drop of 4.4 points, while energy consumption is reduced by up to 95.3%.
Farhana Amin, Sabiha Afroz, Dimitrios S. Nikolopoulos
Diffusion large language models (dLLMs) can generate multiple tokens in parallel, but often require many decoding rounds to complete a response. Efficient generation depends not only on which tokens are easy to predict, but also on how committing to them helps predict the remaining positions. A highly predictable token may provide little useful context, while revealing a less predictable token can help several other positions become predictable sooner. We propose dSpark, a new schedule distillation-based framework for learning when to commit tokens for faster diffusion language generation. To achieve this, we construct unmasking schedules that prioritize tokens whose reveal makes other positions easier to predict, reducing uncertainty earlier in decoding. We then distill these schedules into the model to enable faster parallel sampling, with a KL-divergence regularizer to help preserve generation quality. We extensively evaluate our approach across mathematical reasoning and code generation tasks. On GSM8K, our method achieves remarkable 10.12× and 5.12× speedups over the various baselines on LLaDA-8B-Instruct and Dream-7B-Instruct, respectively, while maintaining superior task performance. Our code is available at https://github.com/VILA-Lab/dSpark.
Diffusion language models generate text through iterative refinement, a process that is often computationally inefficient because many tokens reach stability long before the final denoising step. We introduce a training-free, token-level early stopping approach that identifies convergence independently at each position. Our method leverages lightweight signals derived from the model's predictions and local context to dynamically determine when individual tokens can be finalized. This yields adaptive per-token freezing without task-specific fine-tuning, substantially reducing the total number of diffusion steps required. Across diverse benchmarks, spanning mathematical reasoning, general question answering, and scientific understanding, our approach achieves substantial efficiency gains while preserving generation quality.