Metacognitive Decoding: Efficient Inference for Diffusion Language Agents

What it claims

The paper adds a trained controller on top of a frozen block-diffusion language model, LLaDA2.0-mini being the named example. The controller decides which predicted tokens to commit at each denoising step. It also decides how many positions in the current block to keep computing. The targets are two sources of waste specific to agents. Predictable formatting such as tool-call templates gets over-processed. Fixed-size blocks keep computing tail tokens past a tool-call boundary, and the agent discards those tokens while it waits for environment feedback. The abstract's headline figure is a 20.6% average FLOPs reduction against the most efficient baseline across five benchmarks. The authors state they release the implementation on GitHub, but the page gives no link.

How it works

The backbone stays frozen. No backbone post-training is needed. Each denoising step has two stages.

1. Backbone prediction. The model predicts as usual, giving an argmax token and its confidence at each masked position.
2. Decoding control. The controller reads three signals. The denoising state covers uncertainty, which tokens are revealed and progress through the block. The structural state is the match against recurring format patterns. The interaction state is the current turn and whether thinking, tool calls and final answers are done. It then applies two policies: - Confidence alignment. Structural segments are mined from sampled trajectories of the task. Candidates that agree with the tokens already generated are scored by a linear head. The best candidate is committed in one step if its score beats a learned threshold ρ. Each remaining token is accepted when its confidence is at least τ0 − δi, where δi is a per-position offset the controller predicts. - Flux. A recoverable "empty token" ε marks the end of the active range, and positions past it skip denoising. A softmax over lengths 1..B picks the range, constrained to keep all committed tokens. If the range fills while the turn is unfinished, it grows and restores masks. Once it has grown, it cannot shrink again within that block.

flowchart TD
  A[Frozen backbone predicts tokens] --> B[Controller reads denoising, structural, interaction state]
  B --> C{Structural candidate score > ρ?}
  C -- yes --> D[Commit whole segment]
  C -- no --> E[Per-token accept if conf ≥ τ0 − δi]
  D --> F{Active range fully decoded?}
  E --> F
  F -- turn incomplete --> G[Extend range, restore masks] --> A
  F -- tool call or answer --> H[Halt, return to environment]
  F -- masks remain --> A

Training uses GRPO with on-policy rollouts over complete trajectories. A wrong trajectory gets reward 0. A correct trajectory gets a reward that falls as relative cost rises, where relative cost is the FLOPs of backbone plus controller divided by the no-controller FLOPs on the same question. The reward reaches zero once relative cost hits a scale parameter η, and α sets how sensitive the reward is to cost. Reference costs are precomputed. Rewards are standardized within each group of trajectories.

Evidence

The paper evaluates on four datasets:

• BFCL v3, live and non-live subsets, each split 1:1 into training and test following Lu et al. (2025).
• HotpotQA, 1,000 sampled questions.
• 2WikiMultiHopQA, 1,000 sampled questions.
• MuSiQue, 1,000 sampled questions.

The multi-hop sets follow the DLLM-Searcher setup (Zhao et al., 2026). Baselines include general diffusion models, task-adapted diffusion agents and diffusion decoding acceleration methods. The fetched text cuts off at the results table header (accuracy and FLOPs per benchmark), so per-benchmark numbers and named baseline figures are not available here.

Limits

• The page text ends before the results table, the ablations and the cross-backbone evaluation, so only the 20.6% average is confirmed.
• The abstract describes accuracy only as "competitive." The size of any accuracy loss is not visible.
• Structural candidates come from trajectories sampled per task, and the controller is trained per setting. Moving to a new tool schema or task may require new trajectories and retraining.
• FLOPs is the only stated cost metric. The visible text reports no wall-clock latency figure.