Metacognitive Decoding: Efficient Inference for Diffusion Language Agents

Mingyu Zhao; Shuhang Lin; Ligong Han; Xiaojing Yu; Zhenting Wang; Yifei Pang; Yang Zhou; Dimitris N. Metaxas

cs.CLOct 27, 2026

Metacognitive Decoding: Efficient Inference for Diffusion Language Agents

Authors: Mingyu Zhao, Shuhang Lin, Ligong Han, Xiaojing Yu, Zhenting Wang, Yifei Pang, Yang Zhou, Dimitris N. Metaxas

Organizations: Rutgers · Amazon · UNC

Abstract

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
Method	Acc. $\uparrow$	FLOPs $\downarrow$	Acc. $\uparrow$	FLOPs $\downarrow$	Acc. $\uparrow$	FLOPs $\downarrow$	Acc. $\uparrow$	FLOPs $\downarrow$	Acc. $\uparrow$	FLOPs $\downarrow$
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
Setting	Acc. $\uparrow$	FLOPs $\downarrow$	Acc. $\uparrow$	FLOPs $\downarrow$	Acc. $\uparrow$	FLOPs $\downarrow$	Acc. $\uparrow$	FLOPs $\downarrow$	Acc. $\uparrow$	FLOPs $\downarrow$
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 (

10^{12}

) follow Table 1 .

Panel of Figure 3 — 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
Dataset	$\Delta$ Acc.	FLOPs $\downarrow$	$\Delta$ Acc.	FLOPs $\downarrow$
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.

\Delta

Acc. reports the change in accuracy (percentage points), and FLOPs

\downarrow

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 Model Inference

Diffusion language model inference studies how to generate, guide, and serve text efficiently through iterative denoising and parallel token updates.

See all 179 papers

CAI-DLLM: Convergence Aware Inference for Diffusion Language Models
Aug 23, 2026Farhana Amin, Sabiha Afroz, Dimitrios S. NikolopoulosDiffusion Language Model Inference Diffusion Language Model Decoding
Virginia Tech
Learning When to Commit: Schedule Distillation for Fast Diffusion Language Models
Oct 26, 2026Aidar Myrzakhan, Zhiqiang ShenDiffusion Language Model Inference Diffusion Language Model Decoding
VILA Lab, MBZUAI
Just on Time: Token-Level Early Stopping for Diffusion Language Models
Feb 11, 2026Zakhar Kohut, Severyn Shykula, Mykola Vysotskyi +7Diffusion Language Model Inference Diffusion Language Model Decoding
SoftServe Inc.

Aug 23, 2026J/K move · Enter open · S save

CAI-DLLM: Convergence Aware Inference for Diffusion Language Models

Farhana Amin, Sabiha Afroz, Dimitrios S. Nikolopoulos

Virginia Tech

Diffusion Language Model Inference Diffusion Language Model Decoding Diffusion Model Inference Acceleration Diffusion Language Models

Open paper AskPDF arXiv

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%.

Abstract

Diffusion Language Model Inference

Explore similar work

CAI-DLLM: Convergence Aware Inference for Diffusion Language Models

Learning When to Commit: Schedule Distillation for Fast Diffusion Language Models

Just on Time: Token-Level Early Stopping for Diffusion Language Models