Source-linked AI summary
TDLM: Adaptive Temporal Integration for Diffusion Language Models
Bohang Sun, Max Zhu, Francesco Caso, Christopher Irwin, Jialin Yu, Pietro Liò
TL;DR
Masked diffusion denoisers repeatedly compute representations for unresolved tokens without explicitly retaining them for reuse. TDLM adds temporal attention over each position’s own history and trains on self-generated trajectories. Experiments report that fewer backbone evaluations can outweigh added temporal computation at useful quality levels.
Problem
Conventional denoisers provide no explicit route for unresolved positions to reuse representations computed at earlier denoising steps.
Source: §1 · p. 1Method
TDLM adds temporal attention over earlier hidden representations of unresolved positions in a frozen denoiser and trains the added modules on self-generated trajectories with cross-entropy.
Source: §1 · pp. 2–3Results
The available experiments show that fewer backbone evaluations can outweigh added temporal computation at useful quality levels.
Source: §6 · p. 9Takeaways & Limitations
Temporal memory offers a practical direction for improving generation efficiency by reusing denoising history.
Source: §6 · p. 9, §5 · p. 9Self-distillation from a frozen backbone’s own outputs offers limited evidence for correcting systematic teacher errors, making teacher capability a practical reference for attainable quality.
Source: §5 · p. 9
Abstract
from arXiv · showhide
Masked diffusion language models repeatedly compute representations for unresolved tokens, but conventional denoisers do not explicitly retain these representations for reuse across steps. We introduce Temporal Self-Conditioning Diffusion Language Model (TDLM), which lets each unresolved position attend to its own history of hidden representations within a frozen denoiser. This gives later predictions access to computation already performed. Training uses self-generated trajectories and a simple cross-entropy objective for tokens revealed between consecutively sampled decoding steps, while causal temporal attention permits parallel evaluation during training. On LLaDA, the temporal modules add only approximately 0.84% trainable parameters. Across LLaDA and Dream, the main comparison shows up to 7.15$\times$ faster generation than the corresponding original diffusion language model, with competitive answer quality on the evaluated tasks. These results motivate temporal memory as a practical way to make more effective use of repeated denoising computation.
1 Introduction
Masked diffusion denoisers repeatedly compute representations for unresolved tokens but provide no explicit route to reuse them. TDLM adds temporal access to those representations, and the experiments examine its quality–latency trade-off and the effects of restricting history.
1 Introduction
- Unresolved positions repeatedly receive new representations as context changes, but conventional denoisers provide no explicit route to consult earlier ones.The paper asks whether reusing those estimates can help produce reliable predictions sooner, while testing whether apparent acceleration could instead conceal errors.Source: §1 · pp. 1–2
- TDLM lets each unresolved position attend to its own earlier representations at selected layers while keeping the pretrained backbone frozen.Narrow temporal projections and a learned zero-value sink keep adaptation small and allow the model to suppress unhelpful history.Source: §1 · pp. 2–3
- Training uses frozen-backbone trajectories and cross-entropy on tokens revealed between sampled steps, with causal temporal attention enabling parallel evaluation.The procedure updates only added modules and requires no manually annotated reference answers.Source: §1 · pp. 2–3
- Experiments on LLaDA and Dream study generation efficiency, answer quality, and how restricting temporal history changes predictions.Figure 1 illustrates the quality–latency trade-off on GSM8K across threshold sweeps.Source: §1 · pp. 2–3, Figure 1 · p. 2
2 Method
TDLM augments a frozen diffusion denoiser with temporal attention so masked positions can reuse their own earlier hidden representations. The method trains added modules on sampled teacher trajectories and uses the resulting history during threshold-based decoding.
2 Method
- TDLM adds temporal attention over earlier representations of the same masked position while retaining the frozen backbone and its spatial attention across positions.Only the added temporal modules are trained.Source: §2 · p. 3
2.1 Conditioning on denoising history
- A state S_i is a partially revealed response at global decoding step s_i, and each step commits one or more predicted tokens by replacing their masks.Revealed tokens remain fixed; during training only a subset of teacher steps is selected, while inference evaluates successive inputs.Source: §2.1 · p. 3
2.2 Temporal attention at masked token positions
- At selected layers, a masked position queries its earlier same-position representations through narrow temporal attention, with a zero-value sink available to ignore history.Temporal rotary embeddings encode progress using revealed-token counts; empty history yields zero correction, and a zero-initialized output projection preserves the pretrained denoiser at initialization.Source: §2.2 · pp. 3–4, Figure 2 · p. 4
2.3 Simple trajectory self-distillation
- Training samples teacher-generated inputs and supervises positions masked in S_i but revealed in the next selected state S_i+1, averaging cross-entropy over supervised tokens.The initial input seeds temporal memory without supervision, and later tokens remain masked but receive no loss on that input; causal attention permits sampled inputs to be evaluated together layer by layer.Source: §2.3 · pp. 5–6, Figure 3 · p. 5
2.4 From retrieval to fewer evaluations
- At inference, the decoder commits predictions above confidence threshold τ, or the most confident masked token if none exceeds it; TDLM and Fast-dLLM share this rule.History accumulates for masked response positions across block boundaries, but retrieval adds compute and memory overhead, so it is useful when it saves enough backbone evaluations.Source: §2.4 · p. 6
3 Experiments
The experiments evaluate generation quality and latency on LLaDA and Dream, then examine how restricting temporal history changes predictions and commitment behavior. The main comparison reports faster generation with task-dependent quality, while history-window measurements show how accuracy and efficiency vary with available history.
3.1 Experimental setup
- The evaluation covers mathematical accuracy on MATH-500 and GSM8K and code pass@1 on MBPP, using one greedy generation per problem.The main comparison uses a 256-token budget and block size 32; NFE and generation seconds are per-problem means.Source: §3.1 · pp. 6–7
3.2 Generation efficiency and answer quality
- TDLM reaches up to 7.15× faster generation than original DLM and up to 1.85× over Fast-dLLM, with quality rankings varying across the six backbone–task combinations.At the longer LLaDA generation budget, speedup over original DLM reaches 8.54×.Source: §3.2 · pp. 7–8, Table 1 · p. 7
- Figure 1 compares quality and latency across decoding thresholds, and TDLM pushes the empirical Pareto frontier toward higher accuracy at fast speeds.The evaluated backbones are LLaDA-8B-Instruct and Dream-v0-Instruct-7B.Source: §3.1 · p. 6, §3.2 · p. 7
3.3 What does temporal history contribute?
- Extending history from one evaluation to the complete trajectory raises accuracy from 77.62% to 79.08%, while reducing NFE from 44.06 to 40.08 and mean time from 0.637 to 0.585 seconds.These inference-time restrictions reuse the same trained LLaDA checkpoint without retraining.Source: §3.3 · p. 8, Table 2 · p. 8
4 Related work
TDLM is situated among methods for parallel decoding and diffusion acceleration, as well as approaches that condition on or retain representations across denoising steps.
Parallel decoding and diffusion acceleration
Self-conditioning and persistent representations
5 Discussion and future work
The discussion frames temporal memory as a way to reuse denoising computations, while identifying limitations in self-generated supervision and areas needing further evaluation.
Discussion
- Temporal memory provides a way to use earlier denoising computations for unresolved tokens, but further analysis is needed to characterize when history is most useful.The proposed analyses include token-level correctness and confidence, and a capacity-matched adapter using only the current hidden representation.Source: §5 · p. 9
Limitations and future work
- Self-distillation from a frozen backbone’s own outputs offers limited evidence for correcting systematic teacher errors, making teacher capability a practical reference for training-signal quality.The authors point to higher-quality training data and full-model fine-tuning as future directions.Source: §5 · p. 9
- Sparse teacher trajectories and the student’s multi-token commitments can expose the model to different histories; longer responses, repeated measurements, and combinations with spatial KV caching and early stopping remain evaluation directions.Source: §5 · p. 9
6 Conclusion
TDLM adds lightweight attention over unresolved-position histories to a frozen diffusion language model and trains earlier predictions through trajectory self-distillation. The available experiments show that fewer backbone evaluations can outweigh added temporal computation at useful quality levels.
Conclusion
- TDLM augments a frozen diffusion language model with lightweight attention over the history of unresolved positions.Source: §6 · p. 9
- A trajectory self-distillation objective trains earlier predictions, while causal temporal attention permits parallel evaluation of sampled trajectory inputs.Source: §6 · p. 9
- The available experiments show that fewer backbone evaluations can outweigh added temporal computation at useful quality levels, motivating temporal memory as a practical direction for generation efficiency.Source: §6 · p. 9
A Training details
Training uses sampled self-generated trajectories with the backbone frozen; temporal attention retrieves earlier representations for the same token position. A learned zero-value sink can suppress retrieved history, and zero-initialized output projections preserve the backbone at initialization.
A Training details
- TDLM trains on 3,500 self-generated trajectories for 750 updates, with the backbone frozen; LLaDA modules add 67.11M parameters, about 0.84% of the nominal 8B backbone.For Dream, modules occupy layers 7–20 with four heads of width 84; predicting position p uses the hidden representation at p −1.Source: A Training details · p. 12
B.1 Temporal attention operator
- Temporal attention retrieves earlier representations of the same token position at the same layer, using temporal rotary embeddings to encode decoding progress.The selected input S_i and its decoding step s_i distinguish sequence identity from trajectory time.Source: §B.1 · pp. 12–13
B.2 Sink, cold start, and initialization
- The zero-value sink gates the magnitude of retrieved history, while empty history yields zero correction and a zero-initialized output projection preserves the backbone on any fixed trajectory.The output bound is local to the module and does not guarantee whole-model stability or accuracy.Source: §B.2 · p. 13
B.3 Parallel and sequential evaluation
On a fixed set of sampled inputs, parallel evaluation with causal temporal attention matches sequential evaluation with temporal caches in exact arithmetic. This equivalence does not include unsampled decoding steps, so training and generation histories can differ.
B.3 Parallel and sequential evaluation
- With dropout disabled and identical history selection, parallel evaluation of sampled inputs matches sequential temporal-cache evaluation in exact arithmetic.The result applies to a fixed set of sampled inputs.Source: §B.3 · pp. 13–14
- No representations are computed or inserted into memory for unsampled decoding steps, so training histories can differ from those accumulated during generation.Despite frozen backbone weights, later backbone computations still propagate gradients to train the temporal modules.Source: §B.3 · p. 14
C Evaluation details
Evaluation scores mathematical answers, code-generation tests, and runtime under specified decoding and measurement procedures. Failed or unparseable answers count as incorrect, and several long prompts use eager fallback.
C Evaluation details
- Math answers are checked for equivalence and MBPP code against dataset tests; failed or unparseable answers count as incorrect.Scoring uses text before the first EOS/EOT marker; MBPP tests have a three-second timeout and 512 MB memory limit.Source: C Evaluation details · p. 14
- For dagger-marked Dream MetaState scores, removing only the tool-call strings changes raw MBPP pass@1 from 23.74% to 58.37%; NFE and latency use original generations.The treatment applies specifically to scoring the dagger-marked Table 1 results.Source: C Evaluation details · p. 14
- Runtime covers the full generation budget but excludes compilation, warmup, tokenization, external text decoding, and grading.Measurements use compiled generation on separate GH200 jobs, with one timed generation per problem.Source: C Evaluation details · p. 14
- Long-prompt eager fallbacks are included in measured latency; TDLM, Fast-dLLM, and original LLaDA use them for four MATH-500 prompts, while Dream runs use three per method.MetaState compiles the full LLaDA test set.Source: C Evaluation details · p. 14
D Architecture ablations on GSM8K
The GSM8K ablations compare temporal-layer allocations and head configurations, then test how the amount of retained inference history relates to accuracy and denoiser evaluations.
D Architecture ablations on GSM8K
- The architecture ablations evaluate temporal-layer allocations and widths on all 1,319 GSM8K test problems under the main settings.Timing covers full-budget generation without EOS early stopping and excludes compilation and warmup; small timing differences are approximate.Source: D Architecture ablations on GSM8K · p. 15
D.1 Different temporal-layer allocations
- 79.08% accuracy is achieved by sixteen dispersed layers, versus 77.86% for eight dispersed layers, at similar NFE and latency.Sixteen middle layers achieve 79.00% accuracy but require 45.72 NFE and 0.683 seconds, versus 40.08 NFE and 0.585 seconds for the dispersed allocation.Source: Table 4 · p. 15, §D.1 · p. 15
D.2 Temporal width and head configurations
- 79.00% accuracy is reported for four heads of width 64 versus 78.47% for two heads of width 128, with similar NFE and latency.These width-256 rows share training settings; the full table does not isolate width because other rows differ in training recipe, and single-run differences do not establish significance.Source: Table 5 · p. 15, §D.2 · p. 15
E Temporal context during inference on GSM8K
- The inference intervention compares current-only, one-, three-, and five-evaluation windows with complete history, retaining allowed histories across block boundaries.Finite windows use circular K/V slots per uncommitted position; complete history grows in pages of eight steps.Source: E Temporal context during inference on… · p. 16
- Accuracy rises from 77.62% with one previous evaluation to 79.08% with complete history, while NFE falls from 44.06 to 40.08.Each condition is evaluated without retraining the temporal modules; current-only uses the original temporal implementation, while historical-window conditions use paged K/V storage.Source: E Temporal context during inference on… · p. 16
F Generation lengths on LLaDA
The LLaDA evaluation extends generation-budget comparisons to 128 and 384 tokens with block size fixed at 32. TDLM uses fewer evaluations and less time than Fast-dLLM in all six comparisons, with task-dependent score differences.
F Generation lengths on LLaDA
- TDLM uses fewer evaluations and less time than Fast-dLLM in all six task-and-budget comparisons.The budgets are 128 and 384 tokens, with block size fixed at 32.Source: F Generation lengths on LLaDA · p. 17
- 3.50–4.84× speedups over original DLM at 128 tokens and 4.81–8.54× at 384 tokens accompany task-dependent score differences relative to Fast-dLLM.At 128 tokens, the score differences are −0.40, +0.65, and +0.38 percentage points across MATH-500, GSM8K, and MBPP; at 384 tokens they are +1.20, +0.09, and −0.72 points, respectively.Source: Table 6 · p. 17, F Generation lengths on LLaDA · p. 17, Table 7 · p. 18