TraceQuery: learning what to recompute from full-denoising trajectory in diffusion language models

# TraceQuery: Learning What to Recompute from Full-Denoising Trajectory in Diffusion Language Models
Source: [https://arxiv.org/html/2610.33008](https://arxiv.org/html/2610.33008)
Shuoqiu Li,Xu Gao,Shibo Li,Xin Liu,Shangqian GaoFlorida State University Carnegie Mellon University University of Pennsylvania ByteDance Inc\.

###### Abstract

Diffusion large language models \(dLLMs\) enable flexible generation through iterative masked denoising, but incur substantial inference cost because token representations are repeatedly recomputed across denoising steps\. Existing cache\-based methods reduce this cost by reusing intermediate results across denoising steps and recomputing only a subset of positions at each step, chosen either on a fixed schedule or by heuristic signals of change\. Our trajectory experiments instead point to decoding trajectory as the more relevant quantity\. Cached generation remains accurate when it stays close to the full\-denoising trajectory, with errors concentrated around unmask\-order deviations and reduced as longer reference prefixes are recovered\. This suggests a more direct principle:limiting representation error helps, but keeping generation on the full\-denoising trajectory is a more direct route to accurate decoding\. We therefore formulate efficient dLLM inference as*trajectory\-guided querying*\. Rather than broadly repairing representation error, we use full\-denoising trajectories to supervise which unresolved positions should receive limited fresh computation\. Based on this formulation, we introduceTraceQuery, a lightweight sparse decoding framework that predicts a small set of near\-future trajectory\-relevant positions, queries only those candidates, and lets the backbone select which one to commit from their fresh predictions\.TraceQuerydelivers the strongest quality–computation trade\-off among cache\-based approaches, sustaining high accuracy with substantially fewer fresh token computations, and generalizes across tasks, model variants, and families\. Code:[https://github\.com/fish1928/tracequery\.git](https://github.com/fish1928/tracequery.git)

## 1Introduction

Diffusion large language models \(dLLMs\) have emerged as a promising alternative to autoregressive language models by replacing fixed left\-to\-right generation with iterative masked denoising\([Sahoo et al\., 2024](https://arxiv.org/html/2610.33008#bib.bib7);[Ou et al\., 2025](https://arxiv.org/html/2610.33008#bib.bib8);[Nie et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib9)\)\. By repeatedly predicting unresolved tokens under bidirectional attention, dLLMs support flexible generation orders and expose opportunities for parallel decoding\([Ye et al\., 2025](https://arxiv.org/html/2610.33008#bib.bib10);[Song et al\., 2025](https://arxiv.org/html/2610.33008#bib.bib11);[Labs et al\., 2025](https://arxiv.org/html/2610.33008#bib.bib12)\)\. However, these advantages come with substantial inference cost\. At each denoising step, standard dLLMs repeatedly recompute token representations over large portions of the sequence, even though only one or a few tokens may ultimately be committed\([Hu et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib15);[Cheong et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib18)\)\. When the number of denoising steps grows with generation length, this redundant recomputation can lead to cubic decoding cost\([Ma et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib14);[Liu et al\., 2026b](https://arxiv.org/html/2610.33008#bib.bib17)\)\. To reduce this redundancy, recent work maintains an approximate KV cache across denoising steps, reusing stable representations and selectively refreshing only a subset of positions\([Ma et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib14);[Wu et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib16);[Liu et al\., 2026b](https://arxiv.org/html/2610.33008#bib.bib17);[Jiang et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib19)\)\. Despite increasingly fine\-grained refresh policies, these methods share a common objective:limiting the representation error introduced by cache reuse\.

We revisit this objective by examining how cached decoding behaves relative to full denoising\. Our trajectory analysis \(Figure[1](https://arxiv.org/html/2610.33008#S1.F1)\) shows that cached decoding deviates from the full\-denoising unmasking order even when token\-level predictions agree\. Accuracy also drops as the injected full\-denoising prefix shortens, tying quality to keeping generation anchored to the reference trajectory\. These observations motivate the following causal view:

*representation error*→\\rightarrow*decision perturbation*→\\rightarrow*trajectory divergence*→\\rightarrow*quality degradation*\.

Refreshing cache entries intervenes at the first link of this chain by reducing representation error, but it controls decision perturbation and trajectory divergence only indirectly\. The link from representation error to decoding decisions is highly non\-uniform\. Many errors are decision\-irrelevant and leave the next commitment unchanged, while a small subset can alter a consequential decision and redirect the subsequent trajectory\. This makes representation repair inefficient, because computation may be spent correcting harmless errors, while even useful corrections must propagate through later representations and decisions before improving the decoding trajectory\. This suggests a more direct principle:limiting representation error helps, but keeping generation on the full\-denoising trajectory is a more direct route to accurate decoding\.

![Refer to caption](https://arxiv.org/html/2610.33008v1/2_trajectory_test_figure.png)Figure 1:Trajectory agreement and dependence on trajectory guidance\.\(a\)We compare decoding progress under full denoising and cached decoding\. Each underlying observation corresponds to a token position, paired by the step at which it is unmasked under full denoising and its corresponding step under cached decoding\. Rather than plotting individual observations, we show the 80% density contours\. Concentration along the diagonal indicates agreement with the full\-denoising unmask order, while broader deviations reflect trajectory divergence\.\(b\)Accuracy on examples solved by full denoising as the injected full\-denoising prefix is progressively removed\.This view gives the full\-denoising trajectory a new role\. A full\-denoising run provides a sequence of decisions made by the original backbone from fully recomputed representations, without cache\-induced error\. Although multiple decoding trajectories may lead to equally good generations, this trajectory is immediately available from the model itself and requires neither additional trajectory search nor external supervision\. It therefore provides a natural, unapproximated reference for anticipating the decoding decisions that lie in the near future\.

Based on this view, we formulate efficient dLLM inference as*trajectory\-guided querying*\. Rather than searching for an optimal trajectory, we use the already available full\-denoising trajectory to guide where a limited query budget should be spent\. At each denoising step, the objective is to identify unresolved positions that appear soon along the full\-denoising reference trajectory and include them in a small fresh\-query set\. Importantly, the trajectory does not determine the final decoding decision\. It supervises candidate selection, while the backbone queries the selected candidates with fresh computation and retains control over which position to commit\.

We introduceTraceQuery, a lightweight trajectory\-guided sparse decoding framework illustrated in Figure[2](https://arxiv.org/html/2610.33008#S4.F2)\. During training,TraceQuerylearns from full\-denoising trajectories to rank unresolved positions according to their near\-future relevance\. During inference, it selects a small candidate set from inexpensive signals available in the current decoding state\. Only these candidates receive fresh computation, while the remaining positions continue to reuse cached representations\. The backbone dLLM then selects the final commitment from the freshly evaluated candidates\.TraceQuerytherefore splits the two decisions: the router decides where fresh computation goes, and the backbone decides what to commit\.

Our main contributions are summarized as follows:

- •We propose*trajectory\-guided querying*\. Representation error degrades generation through perturbed consequential decisions, so rather than spending limited computation to repair the error, we let the available full\-denoising trajectory guide it directly toward the decisions that lie ahead\.
- •We buildTraceQuery, a lightweight sparse decoding framework that implements this formulation with a router trained on full\-denoising trajectories\. From inexpensive decoding\-state signals, the router decides which few positions receive fresh computation, and the backbone decides what to commit from the freshly evaluated candidates\.
- •A single learned router approaches full\-denoising quality with substantially fewer fresh token computations across tasks, model variants, and dLLM families, improving the quality–computation trade\-off over cache\-based baselines\.

## 2Related Work

Diffusion language models generate text through iterative masked denoising rather than fixed left\-to\-right factorization, enabling bidirectional context and flexible decoding orders\([Austin et al\., 2021](https://arxiv.org/html/2610.33008#bib.bib4);[Li et al\., 2022](https://arxiv.org/html/2610.33008#bib.bib5);[Lou et al\., 2023](https://arxiv.org/html/2610.33008#bib.bib6);[Sahoo et al\., 2024](https://arxiv.org/html/2610.33008#bib.bib7);[Ou et al\., 2025](https://arxiv.org/html/2610.33008#bib.bib8)\)\. Recent large\-scale models such as LLaDA and Dream demonstrate that this paradigm can scale to competitive language modeling and reasoning performance\([Nie et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib9);[Ye et al\., 2025](https://arxiv.org/html/2610.33008#bib.bib10)\), while systems such as Seed and Mercury further highlight its potential for parallel generation\([Song et al\., 2025](https://arxiv.org/html/2610.33008#bib.bib11);[Labs et al\., 2025](https://arxiv.org/html/2610.33008#bib.bib12)\)\.

Learned and Trajectory\-Aware Decoding\.Prior methods learn, search, guide, or verify unmasking decisions and denoising trajectories\([Jazbec et al\., 2025](https://arxiv.org/html/2610.33008#bib.bib1);[Liu et al\., 2026a](https://arxiv.org/html/2610.33008#bib.bib2);[Hu et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib15);[Pan et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib3)\)\.TraceQueryinstead learns which masked positions to*query*: its router allocates fresh computation, while the backbone determines the committed position and token\. Full\-denoising trajectories provide offline supervision, with no external guider, online search, or draft verification at inference\.

KV Reuse and Adaptive Refresh\.Exact KV reuse is difficult in dLLMs because bidirectional attention changes token representations as decoding proceeds\. Early methods reuse states across completed blocks\([Arriola et al\., 2025](https://arxiv.org/html/2610.33008#bib.bib13)\)or stabilized regions, with periodic recomputation to limit staleness\([Hu et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib15);[Wu et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib16)\)\. dKV\([Ma et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib14)\)identifies*local fidelity*: each decoding decision changes representations most strongly at a limited set of positions, making uniform refresh unnecessary\. Building on this principle, dLLM\-Cache\([Liu et al\., 2026b](https://arxiv.org/html/2610.33008#bib.bib17)\)uses V\-state verification to select refresh positions, while d2Cache\([Jiang et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib19)\)combines confidence and local mask density for unresolved positions with attention\-based selection for prompt and decoded positions\. These methods adapt which cached states to refresh;TraceQueryfocuses on which unresolved positions receive fresh queries before the next commitment\.

Different from the above works, our analysis suggests that the downstream effect of cache\-induced error is more closely reflected by trajectory deviation than by representation error alone\. We use the available full\-denoising trajectory to guide fresh computation toward positions that are likely to matter for upcoming decoding decisions\.

## 3Preliminaries

This section introduces masked diffusion decoding, the full\-denoising reference trajectory, and the cache setting used byTraceQuery\. These definitions apply to both single\-block and block\-wise diffusion decoding\.

### 3\.1Masked Diffusion Decoding

A masked diffusion language model operates on a partially masked sequencex\(t\)x^\{\(t\)\}at denoising steptt\. Letℳt\\mathcal\{M\}\_\{t\}and𝒰t\\mathcal\{U\}\_\{t\}denote the unresolved and already\-resolved positions, respectively\. The model predicts token distributions for positions inℳt\\mathcal\{M\}\_\{t\}in parallel, and a decoding policy selects a subset𝒜t⊆ℳt\\mathcal\{A\}\_\{t\}\\subseteq\\mathcal\{M\}\_\{t\}to commit:

pDLLM\(x1:N∣τ\)=∏t=0T−1∏i∈𝒜tpθ\(xi∣x\(t\)\),τ=\(𝒜0,…,𝒜T−1\),p\_\{\\mathrm\{DLLM\}\}\(x\_\{1:N\}\\mid\\tau\)=\\prod\_\{t=0\}^\{T\-1\}\\prod\_\{i\\in\\mathcal\{A\}\_\{t\}\}p\_\{\\theta\}\(x\_\{i\}\\mid x^\{\(t\)\}\),\\qquad\\tau=\(\\mathcal\{A\}\_\{0\},\\ldots,\\mathcal\{A\}\_\{T\-1\}\),\(1\)whereNNis the sequence length,TTis the number of denoising steps, andτ\\taudenotes the decoding trajectory induced by the sequence of commitments\. After committing𝒜t\\mathcal\{A\}\_\{t\}, the resolved and unresolved sets evolve as

𝒰t\+1=𝒰t∪𝒜t,ℳt\+1=ℳt∖𝒜t\.\\mathcal\{U\}\_\{t\+1\}=\\mathcal\{U\}\_\{t\}\\cup\\mathcal\{A\}\_\{t\},\\qquad\\mathcal\{M\}\_\{t\+1\}=\\mathcal\{M\}\_\{t\}\\setminus\\mathcal\{A\}\_\{t\}\.\(2\)Under single\-token decoding,𝒜t=\{it\}\\mathcal\{A\}\_\{t\}=\\\{i\_\{t\}\\\}, whereiti\_\{t\}is the position selected at steptt\. Unlike autoregressive decoding,ttspecifies only temporal progression and is not tied to a fixed sequence position\. Consequently, the ordering of committed positions itself becomes part of the decoding process\.

### 3\.2Full\-Denoising Reference Trajectory

For a given prompt, we fix the backbone model, decoding policy, denoising schedule, and token\-sampling configuration\. At each denoising step, the full\-denoising baseline performs a forward pass over the entire active generation region and recomputes the representations used to make the next commitment\. In particular, it does not reuse approximate representations from earlier denoising steps\. The decoding policy then selects the next position from predictions obtained using these fully recomputed representations\.

Let𝒜tFD\\mathcal\{A\}\_\{t\}^\{\\mathrm\{FD\}\}denote the set of positions committed at stepttby full denoising\. Under single\-token decoding,𝒜tFD=\{itFD\}\\mathcal\{A\}\_\{t\}^\{\\mathrm\{FD\}\}=\\\{i\_\{t\}^\{\\mathrm\{FD\}\}\\\}, and we define the full\-denoising reference trajectory as

τFD=\(𝒜0FD,…,𝒜T−1FD\)=\(i0FD,…,iT−1FD\)\.\\tau^\{\\mathrm\{FD\}\}=\\left\(\\mathcal\{A\}\_\{0\}^\{\\mathrm\{FD\}\},\\ldots,\\mathcal\{A\}\_\{T\-1\}^\{\\mathrm\{FD\}\}\\right\)=\\left\(i\_\{0\}^\{\\mathrm\{FD\}\},\\ldots,i\_\{T\-1\}^\{\\mathrm\{FD\}\}\\right\)\.\(3\)The trajectory records only the order in which positions are unmasked; it does not include the identities of the committed tokens\. The superscriptFD\\mathrm\{FD\}indicates that the corresponding decisions are produced by full denoising\. We useτFD\\tau^\{\\mathrm\{FD\}\}as an operational reference for the behavior of the unmodified backbone and decoding policy in the absence of cache\-induced approximation\. It is neither assumed to be the unique decoding trajectory nor guaranteed to produce a correct or optimal output\. Nevertheless, it can be generated offline directly from the frozen backbone, without external annotations or additional trajectory search, and therefore provides a well\-defined target for supervising and evaluating trajectory agreement under cached decoding\.

### 3\.3Cached Diffusion Inference

We partition the sequence into three regions: the prefix𝒫\\mathcal\{P\}, the active generation region𝒢\\mathcal\{G\}, and the postfix𝒮\\mathcal\{S\}\. Previously completed blocks belong to𝒫\\mathcal\{P\}, while subsequent blocks belong to𝒮\\mathcal\{S\}\. At denoising steptt, the active generation region is partitioned into

𝒢=𝒰t∪ℳt,𝒰t∩ℳt=∅\.\\mathcal\{G\}=\\mathcal\{U\}\_\{t\}\\cup\\mathcal\{M\}\_\{t\},\\qquad\\mathcal\{U\}\_\{t\}\\cap\\mathcal\{M\}\_\{t\}=\\emptyset\.\(4\)Decoding starts with a full forward pass to initialize the persistent KV cache\. Following existing region\-wise caching strategies\([Wu et al\., 2026](https://arxiv.org/html/2610.33008#bib.bib16);[Liu et al\., 2026b](https://arxiv.org/html/2610.33008#bib.bib17)\), the active generation region is periodically recomputed to control accumulated approximation error, while the prefix and postfix are reused over longer intervals\.TraceQueryoperates only within the active generation region𝒢\\mathcal\{G\}\.

## 4Methodology

![Refer to caption](https://arxiv.org/html/2610.33008v1/figures/3_main_process/3_main_process_figure.png)Figure 2:Overview ofTraceQuery\. Full denoising first provides an offline reference trajectory describing the future order of decoding decisions\. The router is trained to recover the near\-future portion of this trajectory from inexpensive signals available during decoding\. At inference time, after stepttcommits𝒜t\\mathcal\{A\}\_\{t\}, the router ranks the remaining unresolved positions and selects a small query set𝒬t\+1\\mathcal\{Q\}\_\{t\+1\}\. Only these positions receive fresh computation, while other positions reuse cached representations\. The backbone dLLM then makes the final commitment from the freshly queried candidates\. Periodic broad recomputation is retained to control accumulated approximation error\.### 4\.1Trajectory Recovery under a Query Budget

LetτFD\\tau^\{\\mathrm\{FD\}\}be the full\-denoising reference trajectory defined in Section[3\.2](https://arxiv.org/html/2610.33008#S3.SS2)\. After steptt, its remaining ordered suffix is

πtFD=\(it\+1FD,…,iT−1FD\)\.\\pi\_\{t\}^\{\\mathrm\{FD\}\}=\\left\(i\_\{t\+1\}^\{\\mathrm\{FD\}\},\\ldots,i\_\{T\-1\}^\{\\mathrm\{FD\}\}\\right\)\.\(5\)Given a maximum query horizonHH, we define

Ht\\displaystyle H\_\{t\}=min⁡\(H,\|ℳt\+1FD\|\),\\displaystyle=\\min\\\!\\left\(H,\\left\|\\mathcal\{M\}\_\{t\+1\}^\{\\mathrm\{FD\}\}\\right\|\\right\),\(6\)πt,HtFD\\displaystyle\\pi\_\{t,H\_\{t\}\}^\{\\mathrm\{FD\}\}=\(it\+1FD,…,it\+HtFD\),\\displaystyle=\\left\(i\_\{t\+1\}^\{\\mathrm\{FD\}\},\\ldots,i\_\{t\+H\_\{t\}\}^\{\\mathrm\{FD\}\}\\right\),𝒬t\+1oracle\\displaystyle\\mathcal\{Q\}\_\{t\+1\}^\{\\mathrm\{oracle\}\}=\{it\+1FD,…,it\+HtFD\},\\displaystyle=\\left\\\{i\_\{t\+1\}^\{\\mathrm\{FD\}\},\\ldots,i\_\{t\+H\_\{t\}\}^\{\\mathrm\{FD\}\}\\right\\\},whereℳt\+1FD\\mathcal\{M\}\_\{t\+1\}^\{\\mathrm\{FD\}\}is the masked set in the corresponding full\-denoising state\. The ordered prefixπt,HtFD\\pi\_\{t,H\_\{t\}\}^\{\\mathrm\{FD\}\}specifies the nextHtH\_\{t\}reference decisions, while𝒬t\+1oracle\\mathcal\{Q\}\_\{t\+1\}^\{\\mathrm\{oracle\}\}is the query set that would be available if the reference trajectory were known at inference time\. Here, “oracle” denotes access toτFD\\tau^\{\\mathrm\{FD\}\}, rather than global optimality\. We formulate computation allocation as predicting this near\-future reference set from the current decoding state\. During training,πt,HtFD\\pi\_\{t,H\_\{t\}\}^\{\\mathrm\{FD\}\}provides ordered supervision; during inference,TraceQueryselects at mostHHpositions using only signals available from the cached decoding state\. Thus, the target is determined by future decoding behavior rather than by cached\-state magnitude, age, or representation error\.

### 4\.2Trajectory\-Supervised Router Training

We train theTraceQueryrouter offline on intermediate states collected from full\-denoising trajectories\. For each reference state, letit,jFDi\_\{t,j\}^\{\\mathrm\{FD\}\}denote thejj\-th element ofπt,HtFD\\pi\_\{t,H\_\{t\}\}^\{\\mathrm\{FD\}\}\.

For every unresolved positioni∈ℳt\+1FDi\\in\\mathcal\{M\}\_\{t\+1\}^\{\\mathrm\{FD\}\}, we construct

𝐳i\(t\)\\displaystyle\\mathbf\{z\}\_\{i\}^\{\(t\)\}=\[ai,last\(t\),ρi\(t\),di\(t\)\],\\displaystyle=\\left\[a\_\{i,\\mathrm\{last\}\}^\{\(t\)\},\\rho\_\{i\}^\{\(t\)\},d\_\{i\}^\{\(t\)\}\\right\],\(7\)si\(t\)\\displaystyle s\_\{i\}^\{\(t\)\}=fϕ​\(𝐳i\(t\)\),\\displaystyle=f\_\{\\phi\}\\\!\\left\(\\mathbf\{z\}\_\{i\}^\{\(t\)\}\\right\),where the features are normalized last\-layer attention from recently committed positions, local mask density, and relative distance to recently committed positions, respectively\. The scorerfϕf\_\{\\phi\}is a two\-layer MLP shared across positions and denoising steps\. Detailed feature definitions are provided in Appendix[B](https://arxiv.org/html/2610.33008#A2)\.

At ranking stagejj, previously selected reference positions are removed:

𝒞t,j=ℳt\+1FD∖\{it,1FD,…,it,j−1FD\}\.\\mathcal\{C\}\_\{t,j\}=\\mathcal\{M\}\_\{t\+1\}^\{\\mathrm\{FD\}\}\\setminus\\left\\\{i\_\{t,1\}^\{\\mathrm\{FD\}\},\\ldots,i\_\{t,j\-1\}^\{\\mathrm\{FD\}\}\\right\\\}\.\(8\)We train the router with a truncated Plackett–Luce objective:

ℒPL=−1\|𝒟\|∑\(x,t\)∈𝒟∑j=1Htlogexp⁡\(sit,jFD\(t\)\)∑i∈𝒞t,jexp⁡\(si\(t\)\),\\mathcal\{L\}\_\{\\mathrm\{PL\}\}=\-\\frac\{1\}\{\|\\mathcal\{D\}\|\}\\sum\_\{\(x,t\)\\in\\mathcal\{D\}\}\\sum\_\{j=1\}^\{H\_\{t\}\}\\log\\frac\{\\exp\\\!\\left\(s\_\{i\_\{t,j\}^\{\\mathrm\{FD\}\}\}^\{\(t\)\}\\right\)\}\{\\displaystyle\\sum\_\{i\\in\\mathcal\{C\}\_\{t,j\}\}\\exp\\\!\\left\(s\_\{i\}^\{\(t\)\}\\right\)\},\(9\)where𝒟\\mathcal\{D\}is the collection of intermediate reference states\. The objective assigns higher scores to positions that occur earlier in the future full\-denoising order\. The reference trajectory is used only during offline training\.

### 4\.3TraceQueryInference

At inference time, the trained router scores every unresolved positioni∈ℳt\+1i\\in\\mathcal\{M\}\_\{t\+1\}and selects theHtH\_\{t\}highest\-scoring positions:

𝒬t\+1=arg​max𝒬⊆ℳt\+1\|𝒬\|=Ht​∑i∈𝒬si\(t\)\.\\mathcal\{Q\}\_\{t\+1\}=\\underset\{\\begin\{subarray\}\{c\}\\mathcal\{Q\}\\subseteq\\mathcal\{M\}\_\{t\+1\}\\\\ \|\\mathcal\{Q\}\|=H\_\{t\}\\end\{subarray\}\}\{\\operatorname\{arg\\,max\}\}\\sum\_\{i\\in\\mathcal\{Q\}\}s\_\{i\}^\{\(t\)\}\.\(10\)The router scores determine only which positions are queried\. For eachi∈𝒬t\+1i\\in\\mathcal\{Q\}\_\{t\+1\}, the backbone recomputes its representation against the persistent KV cache and obtains a query\-updated token distributionp~i\(t\+1\)\\widetilde\{p\}\_\{i\}^\{\(t\+1\)\}\. Because the remaining representations are cached, this sparse query is not equivalent to full recomputation\.

Under confidence\-based greedy single\-token decoding, the backbone computes

v^i\(t\+1\)\\displaystyle\\widehat\{v\}\_\{i\}^\{\(t\+1\)\}=arg​maxv∈𝒱​p~i\(t\+1\)​\(v\),\\displaystyle=\\underset\{v\\in\\mathcal\{V\}\}\{\\operatorname\{arg\\,max\}\}\\;\\widetilde\{p\}\_\{i\}^\{\(t\+1\)\}\(v\),\(11\)ci\(t\+1\)\\displaystyle c\_\{i\}^\{\(t\+1\)\}=p~i\(t\+1\)​\(v^i\(t\+1\)\),\\displaystyle=\\widetilde\{p\}\_\{i\}^\{\(t\+1\)\}\\left\(\\widehat\{v\}\_\{i\}^\{\(t\+1\)\}\\right\),it\+1\\displaystyle i\_\{t\+1\}=arg​maxi∈𝒬t\+1ci\(t\+1\),𝒜t\+1=\{it\+1\}\.\\displaystyle=\\underset\{i\\in\\mathcal\{Q\}\_\{t\+1\}\}\{\\operatorname\{arg\\,max\}\}\\;c\_\{i\}^\{\(t\+1\)\},\\qquad\\mathcal\{A\}\_\{t\+1\}=\\\{i\_\{t\+1\}\\\}\.It replaces the mask atit\+1i\_\{t\+1\}withv^it\+1\(t\+1\)\\widehat\{v\}\_\{i\_\{t\+1\}\}^\{\(t\+1\)\}and updates

ℳt\+2=ℳt\+1∖\{it\+1\}\.\\mathcal\{M\}\_\{t\+2\}=\\mathcal\{M\}\_\{t\+1\}\\setminus\\\{i\_\{t\+1\}\\\}\.\(12\)All other positions in𝒬t\+1∖\{it\+1\}\\mathcal\{Q\}\_\{t\+1\}\\setminus\\\{i\_\{t\+1\}\\\}remain masked, while positions outside𝒬t\+1\\mathcal\{Q\}\_\{t\+1\}continue to reuse their cached representations\. The loop repeats between scheduled broad\-recomputation steps until the active generation region is resolved\. No reference trajectory or oracle computation is used at inference time\.

## 5Experiments

We have shown in Figure[1](https://arxiv.org/html/2610.33008#S1.F1)that trajectory fidelity is strongly associated with downstream performance\. Based on this observation, we ask three questions: whether trajectory\-guided querying yields end\-to\-end gains across models and tasks, how efficiently it allocates computation relative to existing cache\-based methods, and which router design choices account for its performance\.

### 5\.1Experimental Setup

Models\.We evaluate four representative diffusion language models: LLaDA\-Base, LLaDA\-Instruct, Dream\-Base, and Dream\-Instruct\. For the instruction\-tuned models, we do not apply chat templates\. In our preliminary evaluation, chat templates improve performance on some benchmarks but degrade it on others, introducing an additional source of variation unrelated to the inference method itself\. We therefore use a unified prompt format across all models\.

Benchmarks\.We evaluate reasoning, knowledge, and code generation using five benchmarks: MATH\-500\(512\), BBH\-125\(256\), GSM8K\(256\), HumanEval\(512\), and MBPP\(512\)\. All benchmarks use the same precision, decoding configuration, and hardware environment\.

Baselines\.We compareTraceQueryagainst full denoising and two selective cache\-based acceleration methods: dLLM\-Cache, and d2Cache\. The controlled comparison varies the operating point of each method to obtain matched\-quality and matched\-computation settings, while the final end\-to\-end evaluation uses the adopted configuration of each method\.

Efficiency metrics\.We report downstream task score together with total computation in TFLOPs, an wall\-clock inference time\. For batch experiments, we additionally report the number of samples that can be processed under the same resource budget\. TFLOPs measure the algorithmic computation footprint, while wall\-clock time captures the practical execution cost of each method\.

### 5\.2End\-to\-End Evaluation

Table 1:End\-to\-end quality and inference efficiency\.We evaluate four dLLM backbones on five downstream benchmarks with batch size 16\. Each method uses its adopted operating configuration\. TPS denotes tokens per second\.We begin with the main result: end\-to\-end performance across four backbone models and five benchmarks\. Each method uses its adopted operating configuration\. We compare full denoising, dLLM\-Cache, d2Cache, andTraceQueryfor each backbone\.

Table[1](https://arxiv.org/html/2610.33008#S5.T1)reports downstream task performance together with total computation, and wall\-clock inference time\. Full denoising provides the quality reference, while the three cache\-based methods represent different strategies for reducing repeated computation through approximate reuse and selective refresh\.

Across the four backbones, TraceQuery trades a modest drop in accuracy for substantially lower compute: it retains 90\-98% of full\-denoising accuracy while using roughly 14× fewer TFLOPs per document, the lowest among the caching methods compared\. Compared with the selective cache\-based baselines, the advantage is particularly pronounced in TFLOPs and wall\-clock time, consistent with the small query footprint observed in the controlled comparison\. The improvement is also maintained across both base and instruction\-tuned models, suggesting that the learned querying strategy is not specific to a single model variant or task family\.

The controlled comparison in Section[5\.3](https://arxiv.org/html/2610.33008#S5.SS3)examines how computation allocation contributes to these results\.

### 5\.3Controlled Efficiency Comparison

To examine this allocation directly, we compareTraceQueryagainst dLLM\-Cache and d2Cache under controlled resource budgets\. This experiment varies the operating point of each method to isolate how efficiently fresh computation is allocated\. We consider two complementary settings: matched quality and matched computation\.

Figure 3:Quality–compute and batch\-scaling comparison\.\(a\)Average score across benchmarks as a function of total computation, measured in TFLOPs\. The red dashed line marks the score of full denoising without caching\.TraceQueryreaches higher scores under comparable compute budgets\.\(b\)Supported batch size as a function of wall\-clock time per document\.TraceQuerysupports larger batches at comparable latency, demonstrating improved batch\-scaling efficiency\.Matched\-quality computation\.Figure[3](https://arxiv.org/html/2610.33008#S5.F3)\(a\) compares the amount of computation required to reach the same downstream quality\. The x\-axis reports total computation in TFLOPs\. For each method, we adjust its computation budget and evaluate the resulting average performance across the all benchmarks\.

At the same quality level,TraceQueryrequires substantially less fresh computation than other selective cache\-based baselines\. This difference follows from the allocation objective\. dLLM\-Cache, and d2Cache spend computation maintaining or refreshing a comparatively large set of cached representations, whereasTraceQueryuses a small query horizon and concentrates fresh computation on candidates expected to matter for upcoming decoding decisions\. As a result,TraceQueryreaches the same quality with a smaller computation footprint\.

#### Matched\-compute batch efficiency\.

We next study how the per\-sample query footprint affects batched inference\. For samplebb, let𝒬t\(b\)\\mathcal\{Q\}\_\{t\}^\{\(b\)\}denote the positions requiring fresh computation at denoising steptt\. Batched execution combines these positions into a shared query set:𝒰t​\(B\)=⋃b=1B𝒬t\(b\)\.\\mathcal\{U\}\_\{t\}\(B\)=\\bigcup\_\{b=1\}^\{B\}\\mathcal\{Q\}\_\{t\}^\{\(b\)\}\.As batch size increases, different samples contribute different queried positions, causing this union to grow\. Methods with a large per\-sample refresh budget can therefore quickly cover most of the generation region, making sparse batched inference increasingly close to dense recomputation\.

In contrast,TraceQueryuses a much smaller query horizon, so the union grows substantially more slowly and sparse computation remains effective at larger batch sizes\. Under the same computation budget,TraceQuerytherefore supports substantially larger batches than the competing cache methods\. The advantage is modest in the single\-sample setting but grows rapidly with batch size, showing that reducing the per\-sample query footprint provides an additional systems benefit beyond reducing FLOPs alone\.

Together, the two controlled comparisons show that trajectory\-guided querying improves both*quality efficiency*, by achieving the same quality with less computation, and*batch efficiency*, by accommodating more samples under the same resource budget\.

### 5\.4Ablation Study

We then examine the router design underlying these results through the two\-stage ablation shown in Figure[4](https://arxiv.org/html/2610.33008#S5.F4)\. All results are averaged across LLaDA\-Base, LLaDA\-Instruct, Dream\-Base, and Dream\-Instruct\.

#### Stage I: Router design\.

Figure 4:Two\-stage ablation ofTraceQuery\.\(a\) Stage I selects the router configuration by varying the feature combination, router architecture, ranking loss, normalization, and query horizon during router training\. \(b\) Stage II fixes the selected configuration and evaluates alternative feature groups under closed\-loop cached decoding\. Results are averaged across four models and reported as accuracy change relative to the adopted configuration\. Bars show paired 95% bootstrap confidence intervals; filled markers indicate intervals excluding zero, and gray dots denote per\-benchmark changes\.We tune the router on full\-denoising trajectories, varying one design choice at a time and measuring Recall@5 over 17,670 held\-out positions\. Features matter most: adding geometry and confidence raises Recall@5 from 0\.51 \(attention alone\) to 0\.61 and 0\.78, respectively\. Stale confidence loses 13 points, while nonlinear architectures differ by less than 0\.6 points\. Loss and normalization choices are similarly stable, and performance holds forH=3H=3–1212; we adoptH=5H=5\. We train on mixed tasks and retrain for each backbone\. Stage II tests whether confidence remains useful under cached inference\.

#### Stage II: Closed\-loop feature evaluation\.

Stage I evaluates the router under offline trajectory supervision, but some features may behave differently once the model operates with cached states\. We therefore fix all design choices selected in Stage I and evaluate alternative feature groups under closed\-loop cached decoding\.

For each benchmark, we sample 100 prompts and measure accuracy relative to the adopted router\. We test whether confidence and margin remain useful when cached, since training observes fresh values but inference updates them only on re\-query\. We compare four feature sets:Cleanuses only attention and geometry;Policy\-agedcarries values forward according to the deployment refresh schedule;Random\-ageddraws values from earlier steps; andConf/margin \+ ageadds the time since each position’s last refresh\. Closed\-loop results reverse the offline ranking\. We adopt Clean: conf/margin \+ age changes accuracy by−3\.8\-3\.8points \(CI includes zero\), while random\-aged and policy\-aged significantly reduce it by5\.25\.2and7\.87\.8points, respectively\. Confidence is highly predictive offline, yet every cached variant underperforms without it, making closed\-loop evaluation essential for feature selection\.

## 6Conclusion

We formulate efficient dLLM inference as*trajectory\-guided querying*: allocating limited computation to positions likely to matter for upcoming decoding decisions\.TraceQuerylearns this allocation from full\-denoising trajectories, then uses inexpensive decoding\-state signals to select a small query set at inference time\. The backbone’s updated predictions determine which position and token to commit\. Across LLaDA and Dream models,TraceQueryimproves the quality–computation trade\-off over cache\-based baselines, with further gains in batched inference from its small per\-sample query footprint\.

### AI use statement

In this work, we used generative AI tools only for language polishing, grammar correction, improving the clarity of author\-written text, and identifying potentially relevant related work\. We did not use generative AI tools to generate research ideas, formulate mathematical claims, design experiments, implement methods, or analyze results\. All AI\-assisted content, including suggested related work, was reviewed, verified, and edited by the authors\. We take responsibility for the final content of this work, including all text, claims, and artifacts\.

## References

- Arriolaet al\.\(2025\)M\. Arriola, A\. Gokaslan, J\. Chiu, Z\. Yang, Z\. Qi, J\. Han, S\. Sahoo, and V\. KuleshovBlock diffusion: interpolating between autoregressive and diffusion language models\.InInternational Conference on Learning Representations,Vol\.2025,pp\. 50726–50753\.Cited by:[§2](https://arxiv.org/html/2610.33008#S2.p3.1)\.
- Austinet al\.\(2021\)J\. Austin, D\. D\. Johnson, J\. Ho, D\. Tarlow, and R\. van den BergStructured denoising diffusion models in discrete state\-spaces\.CoRRabs/2107\.03006\.External Links:[Link](https://arxiv.org/abs/2107.03006),2107\.03006Cited by:[§2](https://arxiv.org/html/2610.33008#S2.p1.1)\.
- Cheonget al\.\(2026\)M\. Cheong, D\. Son, W\. Lim, and S\. YooEntropyCache: decoded token entropy guided KV caching for diffusion language models\.arXiv preprint arXiv:2603\.18489\.External Links:[Link](https://arxiv.org/abs/2603.18489)Cited by:[§1](https://arxiv.org/html/2610.33008#S1.p1.1)\.
- Huet al\.\(2026\)Z\. Hu, J\. Meng, Y\. Akhauri, M\. Abdelfattah, J\. Seo, Z\. Zhang, and U\. GuptaFlashdlm: accelerating diffusion language model inference via efficient kv caching and guided diffusion\.InInternational Conference on Learning Representations,Vol\.2026,pp\. 153529–153544\.Cited by:[§1](https://arxiv.org/html/2610.33008#S1.p1.1),[§2](https://arxiv.org/html/2610.33008#S2.p2.1),[§2](https://arxiv.org/html/2610.33008#S2.p3.1)\.
- Jazbecet al\.\(2025\)M\. Jazbec, T\. X\. Olausson, L\. Béthune, P\. Ablin, M\. Kirchhof, J\. Monteiro, V\. Turrisi, J\. Ramapuram, and M\. CuturiLearning unmasking policies for diffusion language models\.arXiv preprint arXiv:2512\.09106\.Cited by:[§2](https://arxiv.org/html/2610.33008#S2.p2.1)\.
- Jianget al\.\(2026\)Y\. Jiang, Y\. Cai, X\. Luo, J\. Fu, J\. Wang, C\. Liu, and X\. Yangd2\{\}^\{2\}Cache: accelerating diffusion\-based LLMs via dual adaptive caching\.InInternational Conference on Learning Representations,External Links:[Link](https://arxiv.org/abs/2509.23094)Cited by:[§1](https://arxiv.org/html/2610.33008#S1.p1.1),[§2](https://arxiv.org/html/2610.33008#S2.p3.1)\.
- Labset al\.\(2025\)I\. Labs, S\. Khanna, S\. Kharbanda, S\. Li, H\. Varma, E\. Wang, S\. Birnbaum, Z\. Luo, Y\. Miraoui, A\. Palrecha,et al\.Mercury: ultra\-fast language models based on diffusion\.arXiv preprint arXiv:2506\.17298\.Cited by:[§1](https://arxiv.org/html/2610.33008#S1.p1.1),[§2](https://arxiv.org/html/2610.33008#S2.p1.1)\.
- Liet al\.\(2022\)X\. Li, J\. Thickstun, I\. Gulrajani, P\. S\. Liang, and T\. B\. HashimotoDiffusion\-lm improves controllable text generation\.Advances in neural information processing systems35,pp\. 4328–4343\.Cited by:[§2](https://arxiv.org/html/2610.33008#S2.p1.1)\.
- Liuet al\.\(2026a\)X\. Liu, Y\. V\. Chun, Y\. Liang, and A\. LiuLookahead path likelihood optimization for diffusion llms\.arXiv preprint arXiv:2602\.03496\.Cited by:[§2](https://arxiv.org/html/2610.33008#S2.p2.1)\.
- Liuet al\.\(2026b\)Z\. Liu, Y\. Yang, Y\. Zhang, J\. Chen, C\. Zou, Q\. Wei, S\. Wang, Y\. Zhu, and L\. ZhangdLLM\-Cache: accelerating diffusion large language models with adaptive caching\.InInternational Conference on Machine Learning,External Links:[Link](https://arxiv.org/abs/2506.06295)Cited by:[§1](https://arxiv.org/html/2610.33008#S1.p1.1),[§2](https://arxiv.org/html/2610.33008#S2.p3.1),[§3\.3](https://arxiv.org/html/2610.33008#S3.SS3.p1.2)\.
- Louet al\.\(2023\)A\. Lou, C\. Meng, and S\. ErmonDiscrete diffusion modeling by estimating the ratios of the data distribution\.arXiv preprint arXiv:2310\.16834\.Cited by:[§2](https://arxiv.org/html/2610.33008#S2.p1.1)\.
- Maet al\.\(2026\)X\. Ma, R\. Yu, G\. Fang, and X\. WangDkv\-cache: the cache for diffusion language models\.Advances in Neural Information Processing Systems38,pp\. 149009–149033\.Cited by:[§1](https://arxiv.org/html/2610.33008#S1.p1.1),[§2](https://arxiv.org/html/2610.33008#S2.p3.1)\.
- Nieet al\.\(2026\)S\. Nie, F\. Zhu, Z\. You, X\. Zhang, J\. Ou, J\. Hu, J\. Zhou, Y\. Lin, J\. Wen, and C\. LiLarge language diffusion models\.Advances in Neural Information Processing Systems38,pp\. 50608–50646\.Cited by:[§1](https://arxiv.org/html/2610.33008#S1.p1.1),[§2](https://arxiv.org/html/2610.33008#S2.p1.1)\.
- Ouet al\.\(2025\)J\. Ou, S\. Nie, K\. Xue, F\. Zhu, J\. Sun, Z\. Li, and C\. LiYour absorbing discrete diffusion secretly models the conditional distributions of clean data\.InInternational Conference on Learning Representations,Vol\.2025,pp\. 64972–65009\.Cited by:[§1](https://arxiv.org/html/2610.33008#S1.p1.1),[§2](https://arxiv.org/html/2610.33008#S2.p1.1)\.
- Panet al\.\(2026\)T\. Pan, B\. Gong, M\. Guang, H\. Yong, T\. Jiang, Y\. Li, Z\. Cao, and K\. LongTrajectory\-level speculative decoding for diffusion language models\.arXiv preprint arXiv:2608\.27514\.Cited by:[§2](https://arxiv.org/html/2610.33008#S2.p2.1)\.
- Sahooet al\.\(2024\)S\. S\. Sahoo, M\. Arriola, Y\. Schiff, A\. Gokaslan, E\. Marroquin, J\. T\. Chiu, A\. Rush, and V\. KuleshovSimple and effective masked diffusion language models\.Advances in Neural Information Processing Systems37,pp\. 130136–130184\.Cited by:[§1](https://arxiv.org/html/2610.33008#S1.p1.1),[§2](https://arxiv.org/html/2610.33008#S2.p1.1)\.
- Songet al\.\(2025\)Y\. Song, Z\. Zhang, C\. Luo, P\. Gao, F\. Xia, H\. Luo, Z\. Li, Y\. Yang, H\. Yu, X\. Qu,et al\.Seed diffusion: a large\-scale diffusion language model with high\-speed inference\.arXiv preprint arXiv:2508\.02193\.Cited by:[§1](https://arxiv.org/html/2610.33008#S1.p1.1),[§2](https://arxiv.org/html/2610.33008#S2.p1.1)\.
- Wuet al\.\(2026\)C\. Wu, H\. Zhang, S\. Xue, Z\. Liu, S\. Diao, L\. Zhu, P\. Luo, S\. Han, and E\. XieFast\-dllm: training\-free acceleration of diffusion llm by enabling kv cache and parallel decoding\.InInternational Conference on Learning Representations,Vol\.2026,pp\. 57027–57051\.Cited by:[§1](https://arxiv.org/html/2610.33008#S1.p1.1),[§2](https://arxiv.org/html/2610.33008#S2.p3.1),[§3\.3](https://arxiv.org/html/2610.33008#S3.SS3.p1.2)\.
- Yeet al\.\(2025\)J\. Ye, Z\. Xie, L\. Zheng, J\. Gao, Z\. Wu, X\. Jiang, Z\. Li, and L\. KongDream 7B: diffusion large language models\.arXiv preprint arXiv:2508\.15487\.External Links:[Link](https://arxiv.org/abs/2508.15487)Cited by:[§1](https://arxiv.org/html/2610.33008#S1.p1.1),[§2](https://arxiv.org/html/2610.33008#S2.p1.1)\.

## Appendix ADetailed Computation and Memory Analysis

#### Computation Complexity

LetSSdenote the total sequence length visible to attention,BBthe size of the active generation region,TTthe number of denoising steps,HHthe query horizon,RRthe number of full refresh steps, andLLthe number of Transformer layers\.

Under single\-token decoding, a regular sparse step recomputes at most

positions, consisting ofHHfuture\-query candidates and one recently committed position required by delayed cache admission\. Near the end of a generation block,qqmay be smaller when fewer thanHHmasked positions remain\.

Queryingqqpositions against a KV cache of lengthSSrequires𝒪⁡\(L​q​S\)\\mathcal\{O\}\(LqS\)attention computation, whereas refreshing allBBpositions in the active generation region requires𝒪⁡\(L​B​S\)\\mathcal\{O\}\(LBS\)\. Ignoring the one\-time cache initialization, the total attention cost is therefore

CTraceQuery=𝒪⁡\(L​S​\[R​B\+\(T−R\)​\(H\+1\)\]\)\.C\_\{\\text\{\{TraceQuery\} \}\}=\\mathcal\{O\}\\left\(LS\\left\[RB\+\(T\-R\)\(H\+1\)\\right\]\\right\)\.\(14\)By comparison, full recomputation at every denoising step requires

CFull=𝒪⁡\(L​S​T​B\)\.C\_\{\\mathrm\{Full\}\}=\\mathcal\{O\}\(LSTB\)\.\(15\)
For single\-block diffusion,B=T=NB=T=N, whereNNdenotes the generation length\. AssumingS=𝒪⁡\(N\)S=\\mathcal\{O\}\(N\)and treatingHHandRRas bounded constants,

CTraceQuery=𝒪⁡\(L​N2\),CFull=𝒪⁡\(L​N3\)\.C\_\{\\text\{\{TraceQuery\} \}\}=\\mathcal\{O\}\(LN^\{2\}\),\\qquad C\_\{\\mathrm\{Full\}\}=\\mathcal\{O\}\(LN^\{3\}\)\.\(16\)Thus, selective querying reduces the dominant attention\-computation scaling from cubic to quadratic in generation length\.

#### Concrete Decoding Configurations

Our default experiments use a query horizon ofH=5H=5\. For single\-block generation withN=256N=256, we performR=16R=16global refreshes\. A regular sparse step therefore recomputes at most six positions: five future\-query candidates and one recently committed position\. Using the number of attention query positions as a computation proxy, the relative attention workload is approximately

16×256\+\(256−16\)×6256×256≈8\.45%\.\\frac\{16\\times 256\+\(256\-16\)\\times 6\}\{256\\times 256\}\\approx 8\.45\\%\.\(17\)This quantity measures the reduction in attention\-query computation and should not be interpreted as wall\-clock speedup\. Practical runtime additionally depends on projection and MLP computation, cache operations, router execution, kernel efficiency, and hardware utilization\.

For a length\-256256generation divided into1616blocks, each active block contains1616positions\. We perform no additional periodic full refresh within each block and use the same query horizonH=5H=5\. The effect of query horizon and refresh schedule is evaluated empirically in Section[5](https://arxiv.org/html/2610.33008#S5)\.

#### Persistent Memory Consumption

TraceQuerymaintains a persistent KV cache for the complete sequence\. LetdKVd\_\{\\mathrm\{KV\}\}denote the dimensionality of the stored KV representation per token,bbthe number of bytes per element, andBbatchB\_\{\\mathrm\{batch\}\}the batch size\. The persistent KV\-cache memory is

MKV=2​Bbatch​L​S​dKV​b,M\_\{\\mathrm\{KV\}\}=2B\_\{\\mathrm\{batch\}\}LSd\_\{\\mathrm\{KV\}\}b,\(18\)where the factor of two accounts for keys and values\. Thus,

MKV=𝒪⁡\(L​S​dKV\),M\_\{\\mathrm\{KV\}\}=\\mathcal\{O\}\(LSd\_\{\\mathrm\{KV\}\}\),\(19\)which scales linearly with sequence length\.

For illustration, assuming a BF16 LLaDA\-8B configuration withL=32L=32,dKV=4096d\_\{\\mathrm\{KV\}\}=4096, and batch size11, the model parameters occupy approximately

Mmodel≈8×109×2=16​GB\.M\_\{\\mathrm\{model\}\}\\approx 8\\times 10^\{9\}\\times 2=16~\\mathrm\{GB\}\.\(20\)For a cached sequence of lengthS=256S=256, the corresponding KV cache requires

MKV=2×32×256×4096×2≈128​MiB,M\_\{\\mathrm\{KV\}\}=2\\times 32\\times 256\\times 4096\\times 2\\approx 128~\\mathrm\{MiB\},\(21\)corresponding to less than1%1\\%additional persistent storage relative to the BF16 model weights\.

Table 2:Persistent KV\-cache memory for the illustrative BF16 LLaDA\-8B configuration under batch size 1\. Relative memory denotes model weights plus KV cache divided by model\-weight memory alone\.

## Appendix BDecoding Signal Definitions

This appendix provides the detailed definitions of the decoding signals introduced in Section[4\.3](https://arxiv.org/html/2610.33008#S4.SS3)\. Leti∈ℳti\\in\\mathcal\{M\}\_\{t\}denote a still\-masked candidate position and letℛt⊆𝒰t\\mathcal\{R\}\_\{t\}\\subseteq\\mathcal\{U\}\_\{t\}denote the set of recently unmasked positions\.

#### Confidence\.

Given the token probability distribution at positionii, confidence is defined as

ci\(t\)=maxv⁡pθ​\(v∣x\(t\)\)i\.c\_\{i\}^\{\(t\)\}=\\max\_\{v\}p\_\{\\theta\}\\left\(v\\mid x^\{\(t\)\}\\right\)\_\{i\}\.\(22\)

#### Probability margin\.

Letpi,\(1\)\(t\)p\_\{i,\(1\)\}^\{\(t\)\}andpi,\(2\)\(t\)p\_\{i,\(2\)\}^\{\(t\)\}denote the largest and second\-largest token probabilities at positionii\. The probability margin is

mi\(t\)=pi,\(1\)\(t\)−pi,\(2\)\(t\)\.m\_\{i\}^\{\(t\)\}=p\_\{i,\(1\)\}^\{\(t\)\}\-p\_\{i,\(2\)\}^\{\(t\)\}\.\(23\)

#### Mask density\.

For a local neighborhood𝒲i\\mathcal\{W\}\_\{i\}centered at candidateii, mask density is defined as

ρi\(t\)=1\|𝒲​i\|∑j∈𝒲i𝟏\[j∈ℳt\]\.\\rho\_\{i\}^\{\(t\)\}=\\frac\{1\}\{\|\\mathcal\{W\}i\|\}\\sum\{j\\in\\mathcal\{W\}\_\{i\}\}\\mathbf\{1\}\\left\[j\\in\\mathcal\{M\}\_\{t\}\\right\]\.\(24\)

#### Relative distance\.

LetLLbe the generation block size\. The distance from candidateiito recently unmasked positions is defined as

di\(t\)=minj∈ℛt⁡\|i−j\|L,d\_\{i\}^\{\(t\)\}=\\frac\{\\min\_\{j\\in\\mathcal\{R\}\_\{t\}\}\|i\-j\|\}\{L\},\(25\)and is normalized by the active generation length when required\.

#### Last\-layer attention\.

LetAℓ,h\(t\)​\(j,i\)A\_\{\\ell,h\}^\{\(t\)\}\(j,i\)denote the attention from positionjjto positioniiat layerℓ\\elland headhh\. The last\-layer attention signal is

ai,last\(t\)=Aggj∈ℛ​t,h⁡A​L,h\(t\)​\(j,i\),a\_\{i,\\mathrm\{last\}\}^\{\(t\)\}=\\operatorname\{Agg\}\_\{j\\in\\mathcal\{R\}t,\\,h\}A\{L,h\}^\{\(t\)\}\(j,i\),\(26\)whereLLdenotes the final Transformer layer\.

#### Full\-layer attention\.

The full\-layer attention signal aggregates the same dependency across Transformer layers:

ai,full\(t\)=Aggℓ,j∈ℛ​t,h⁡A​ℓ,h\(t\)​\(j,i\)\.a\_\{i,\\mathrm\{full\}\}^\{\(t\)\}=\\operatorname\{Agg\}\_\{\\ell,\\,j\\in\\mathcal\{R\}t,\\,h\}A\{\\ell,h\}^\{\(t\)\}\(j,i\)\.\(27\)

## Appendix CInference and Training algorithm

Algorithm 1TraceQueryRouter Training1:Training samples

𝒳\\mathcal\{X\}, diffusion model

pθp\_\{\\theta\}, router

fϕf\_\{\\phi\}, query horizon

HH
2:

𝒟←∅\\mathcal\{D\}\\leftarrow\\emptyset
3:for all

x∈𝒳x\\in\\mathcal\{X\}do

4:Run full\-denoising decoding with one token unmasked per step

5:Record teacher trajectory

τ=\(i0,i1,…,iT−1\)\\tau=\(i\_\{0\},i\_\{1\},\\ldots,i\_\{T\-1\}\)
6:for

t=0,…,T−2t=0,\\ldots,T\-2do

7:Collect signals for all

i∈ℳt\+1i\\in\\mathcal\{M\}\_\{t\+1\}
8:Normalize last\-layer attention withSoftmax

9:Store

\{𝐳i\(t\),gi\(t\)\}i∈ℳt\+1\\\{\\mathbf\{z\}\_\{i\}^\{\(t\)\},g\_\{i\}^\{\(t\)\}\\\}\_\{i\\in\\mathcal\{M\}\_\{t\+1\}\}in

𝒟\\mathcal\{D\}
10:endfor

11:endfor

12:repeat

13:Sample a decoding state

ttfrom

𝒟\\mathcal\{D\}
14:for all

i∈ℳt\+1i\\in\\mathcal\{M\}\_\{t\+1\}do

15:

si\(t\)←fϕ​\(𝐳i\(t\)\)s\_\{i\}^\{\(t\)\}\\leftarrow f\_\{\\phi\}\(\\mathbf\{z\}\_\{i\}^\{\(t\)\}\)
16:endfor

17:

ℒ←0\\mathcal\{L\}\\leftarrow 0
18:for

j=1,…,Hj=1,\\ldots,Hdo

19:

𝒞t,j←\{i∈ℳt\+1:gi\(t\)≥j\}\\mathcal\{C\}\_\{t,j\}\\leftarrow\\\{i\\in\\mathcal\{M\}\_\{t\+1\}:g\_\{i\}^\{\(t\)\}\\geq j\\\}
20:

i\(j\)←i\_\{\(j\)\}\\leftarrowteacher’s

jj\-th future unmask

21:

ℒ←ℒ−log⁡exp⁡\(si\(j\)\(t\)\)∑i∈𝒞t,jexp⁡\(si\(t\)\)\\mathcal\{L\}\\leftarrow\\mathcal\{L\}\-\\log\\dfrac\{\\exp\(s\_\{i\_\{\(j\)\}\}^\{\(t\)\}\)\}\{\\sum\_\{i\\in\\mathcal\{C\}\_\{t,j\}\}\\exp\(s\_\{i\}^\{\(t\)\}\)\}
22:endfor

23:Update

ϕ\\phiusing

∇ϕℒ\\nabla\_\{\\phi\}\\mathcal\{L\}
24:untilconvergence

25:returntrained router

fϕf\_\{\\phi\}

Algorithm 2TraceQueryInference1:Diffusion model

pθp\_\{\\theta\}, router

fϕf\_\{\\phi\}, query horizon

HH, refresh schedule

ℛ\\mathcal\{R\}
2:Initialize masked set

ℳ0\\mathcal\{M\}\_\{0\}and perform a full forward pass

3:Cache all KV states; set

𝒬0←ℳ0\\mathcal\{Q\}\_\{0\}\\leftarrow\\mathcal\{M\}\_\{0\}
4:for

t=0,1,…t=0,1,\\ldotswhile

ℳt≠∅\\mathcal\{M\}\_\{t\}\\neq\\emptysetdo

5:if

t∈ℛt\\in\\mathcal\{R\}then

6:

\(Pt,At,KV\)←FullForward​\(ℳt,KV\)\(P\_\{t\},A\_\{t\},\\mathrm\{KV\}\)\\leftarrow\\textsc\{FullForward\}\(\\mathcal\{M\}\_\{t\},\\mathrm\{KV\}\)
7:

ℰt←ℳt\\mathcal\{E\}\_\{t\}\\leftarrow\\mathcal\{M\}\_\{t\}
8:else

9:

\(Pt,At,KV\)←QueryForward​\(𝒬t,KV\)\(P\_\{t\},A\_\{t\},\\mathrm\{KV\}\)\\leftarrow\\textsc\{QueryForward\}\(\\mathcal\{Q\}\_\{t\},\\mathrm\{KV\}\)
10:

ℰt←𝒬t\\mathcal\{E\}\_\{t\}\\leftarrow\\mathcal\{Q\}\_\{t\}
11:endif

12:

it←arg⁡maxi∈ℰt​Confidence​\(Pt​\(i\)\)i\_\{t\}\\leftarrow\\arg\\max\_\{i\\in\\mathcal\{E\}\_\{t\}\}\\textsc\{Confidence\}\(P\_\{t\}\(i\)\)
13:Commit position

iti\_\{t\}and update

ℳt\+1←ℳt∖\{it\}\\mathcal\{M\}\_\{t\+1\}\\leftarrow\\mathcal\{M\}\_\{t\}\\setminus\\\{i\_\{t\}\\\}
14:if

ℳt\+1=∅\\mathcal\{M\}\_\{t\+1\}=\\emptysetthen

15:break

16:endif

17:for all

i∈ℳt\+1i\\in\\mathcal\{M\}\_\{t\+1\}do

18:

𝐳i\(t\)←\[ai,last\(t\),ρi\(t\),di\(t\)\]\\mathbf\{z\}\_\{i\}^\{\(t\)\}\\leftarrow\[a\_\{i,\\mathrm\{last\}\}^\{\(t\)\},\\rho\_\{i\}^\{\(t\)\},d\_\{i\}^\{\(t\)\}\]
19:

si\(t\)←fϕ​\(𝐳i\(t\)\)s\_\{i\}^\{\(t\)\}\\leftarrow f\_\{\\phi\}\(\\mathbf\{z\}\_\{i\}^\{\(t\)\}\)
20:endfor

21:

𝒬t\+1←TopH​\(\{si\(t\)\}i∈ℳt\+1,H\)\\mathcal\{Q\}\_\{t\+1\}\\leftarrow\\textsc\{TopH\}\(\\\{s\_\{i\}^\{\(t\)\}\\\}\_\{i\\in\\mathcal\{M\}\_\{t\+1\}\},H\)
22:endfor

23:returngenerated sequence