PaperlayerView exact revision on arXiv

Source-linked AI summary

Time-Anchored Diffusion Language Models: Latent-Space Caching for Fast Generation

Joel Anto Paul, Litu Rout, Aditya Akella, Sanjay Shakkottai

arXiv:2609.37924v1cs.LGcs.CL

TL;DR

DLM inference remains expensive because generation repeatedly evaluates the model across denoising steps. This paper learns time-based latent anchors that can be reused across nearby steps, and reports throughput and computation gains for post-training and pretraining variants.

  • Problem

    DLMs require repeated model evaluations across denoising steps, making inference substantially more expensive than autoregressive decoding.

    Source: §1 · p. 1

  • Method

    TADM learns latent anchors whose semantic content persists across nearby diffusion times, periodically computes them, and fuses cached anchors with current-state representations for denoising.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    TADM:Post-train improves DiffusionGemma throughput by approximately 49% to 79% across evaluated benchmarks, while TADM:Pretraining achieves up to 73% higher measured throughput than ADLM and reduces Transformer-layer computation by up to 38%.

    Source: §1 · p. 2, §7 · p. 9

  • Takeaways & Limitations

    TADM:Post-train adds latent-cache reuse to pretrained DLMs without repeating pretraining, using a trained fusion module with the backbone frozen.

    Source: §1 · p. 2

    Generation quality degrades at larger anchor refresh intervals, and TADM:Post-train accuracy is constrained by its frozen pretrained backbone.

    Source: §7 · p. 9

Abstract

from arXiv · show

Recent work on anchored diffusion language models improves denoising by shaping an intermediate latent space with supervised important-token targets. In this work, we introduce time-based (self-supervised) anchoring, which learns and reuses latent anchors without requiring such targets. Our key observation is that anchors encode persistent properties of the clean sequence, such as its semantic intent, global structure, or intermediate plan. Although their hidden representations become stale as the token canvas evolves, their semantic content remains useful across nearby diffusion times. This is implemented through a two-stage architecture consisting of a relatively expensive anchor network that generates the latent cache state and a lightweight denoising network that intelligently combines the cached latent state with the current state at each reverse step using a fusion module. This gives anchoring a latent-space caching interpretation: the anchor network is evaluated periodically, while its cached representation is reused across multiple reverse steps. We instantiate this framework as TADM:Post-train, which time-anchorizes pretrained DLMs, and TADM:Pretraining, which learns time-based anchors during pretraining. Applied to DiffusionGemma-26B, TADM:Post-train improves throughput by approximately 49% to 79% on several math, code, and STEM benchmarks (GSM8K, AIME26, GPQA-Diamond, LiveCodeBench-v6, HumanEval, MMLU-Pro). TADM:Pretraining reduces Transformer-layer computation by up to 38% relative to a standard single-stage DLM, achieves up to 73% higher measured throughput than ADLM.

1 Introduction

DLMs generate by repeatedly refining a sequence, but repeated evaluations make inference expensive. The paper introduces time-based anchoring, which reuses latent representations across nearby diffusion times and reports throughput gains for post-training and pretraining variants.

1 Introduction

  • DLMs repeatedly evaluate the model across denoising steps, making inference substantially more expensive than autoregressive decoding.Their iterative sequence refinement supports bidirectional attention and parallel token generation.Source: §1 · p. 1
  • Anchors encode persistent clean-sequence properties, including semantic intent, global structure, and intermediate plans, that remain useful across nearby diffusion times despite stale representations.The token canvas evolves, so the hidden representations may require correction even when their semantic content remains useful.Source: §1 · p. 2
  • TADM periodically computes an expensive anchor representation, then combines the cached anchor with the current-state representation before denoising.A lightweight shared network encodes the current canvas, while a fusion module combines its representation with the cached anchor.Source: §1 · p. 2, Figure 1 · p. 2
  • The framework recasts anchoring as latent-space caching: it learns representations for reuse and correction rather than detecting unchanged key-value states.The passage contrasts this with local, layer-specific key-value reuse methods that are not trained for temporal reuse.Source: §1 · p. 2
  • TADM:Post-train improves DiffusionGemma-26B throughput by 49%–79% across evaluated math, code, and STEM benchmarks while preserving corresponding-baseline task accuracy.TADM:Pretraining reduces Transformer-layer computation by up to 38% and achieves up to 73% higher measured throughput than ADLM.Source: §1 · p. 2

2 Preliminaries

The preliminaries define discrete diffusion over token sequences and introduce ADLM as a response to masked diffusion’s lack of token prioritization. They describe the denoising objective and ADLM’s two-stage anchor-and-denoiser setup.

2 Preliminaries

  • A sample sequence x is represented by V-dimensional one-hot token vectors, with the vocabulary space V containing V unique discrete tokens.The sequence distribution is denoted q(·), and ⊙ and ⟨·,·⟩ denote the Hadamard and dot products.Source: §2 · pp. 2–3

2.1 Discrete Diffusion Models

  • Discrete diffusion gradually corrupts a clean sequence into noise according to a monotonically decreasing noise schedule α_t, with α_0 = 1 and α_1 = 0.The limiting noise distribution is denoted π.Source: §2.1 · p. 3
  • At inference, a neural denoiser predicts clean-token distributions used in the reverse process, which is trained by minimizing the discrete-time NELBO.The clean sequence is unknown at inference, unlike training, when the clean token is available for the posterior.Source: §2.1 · p. 3

2.2 Anchored Diffusion Language Models

  • Standard masked diffusion does not explicitly prioritize informative tokens, so missing anchors can leave high conditional uncertainty; ADLM addresses this by predicting important anchor tokens before denoising.ADLM’s anchor network predicts a probability mixture over important tokens, and its denoiser predicts clean tokens conditioned on those anchors.Source: §2.2 · p. 3
  • ADLM jointly trains its anchor and denoiser networks with an ANELBO, whose γ-weighted auxiliary anchor loss encourages prediction of informative tokens.The denoising term reconstructs the clean sequence conditioned on the anchor-network output.Source: §2.2 · p. 3

3 Time-Anchored Diffusion Framework

Time-based anchoring learns latent representations that retain useful sequence information across diffusion times and reuses them through a cached, fused two-stage architecture. The framework amortizes expensive anchor computation across reverse steps and defines training around the distribution induced by that cache policy.

3 Time-Anchored Diffusion Framework

  • Time-based anchors capture information persistent across diffusion time and are learned from denoising without explicit anchor targets.The framework presents this as latent-space caching: a representation computed at one state is retained and reused as the state evolves.Source: §3 · p. 4

3.1 Model Decomposition

  • The anchor network computes an expensive latent representation periodically, while the shared network and denoiser run at every reverse step.The decomposition is architecture-dependent; the shared pathway may include embeddings alone or additional Transformer layers.Source: §3.1 · p. 4

3.2 Fusion network

  • At intermediate steps, a fusion module combines the current-state representation with the cached anchor before the denoiser predicts the clean sequence.The cached representation alone lacks information about the current state, motivating fusion with the current representation.Source: §3.1 · p. 4, §3.2 · pp. 4–5

3.3 Anchor-Cache inference

  • Increasing refresh interval K reduces Transformer computation by amortizing the expensive anchor network across multiple reverse steps.The cache is refreshed once every K steps; the shared and denoiser networks are evaluated at each step.Source: §3.3 · p. 5

3.4 Cache-induced distribution

  • The cache-aware training objective averages denoising loss over refresh policies and the joint state-and-anchor distribution induced by executing the model under that policy.Sampling this cache-induced distribution directly would require running reverse generation during training.Source: §3.4 · p. 5

4 TADM:Post-train

TADM:Post-train approximates the cache-policy inference distribution by rolling out the pretrained model to generate training states. It trains against the resulting cached anchors and student-generated terminal states while using distillation to preserve the output distribution.

Rollout construction

  • Training approximates the cache-policy inference distribution by rolling out the pretrained model rather than sampling states only from the forward process.The passage says forward-process sampling would induce a different distribution because the model is pretrained.Source: §4 · p. 5
  • The anchor is computed once at the sampled state and held fixed while the student performs k reverse transitions to produce the current state.The rollout samples k from {0, 1, 2}; its endpoint is sampled from the student rather than synthesized from the clean target.Source: §4 · p. 5

Training objective

  • The discrete rollout is stop-gradient, and recomputed paths provide supervision to trainable pathways without backpropagating through token sampling.The loss also includes knowledge distillation to prevent unnecessary corruption of the output distribution.Source: §4 · pp. 5–6

5 TADM:Pretraining

TADM:Pretraining uses analytically generated noisy states rather than costly model-generated trajectories. Its objective trains denoising with stale anchors, while optional anchor-token supervision can be disabled for fully self-supervised time-based anchoring.

Training setup

  • Analytically generated noisy states avoid the expense of multiple model forwards during pretraining and the poor trajectories expected from a model trained from scratch.The procedure samples z_t and further corrupts it to obtain z_t′ for the anchor representation.Source: §5 · p. 6

TADM Parametrization

  • The cached reverse process uses the most recently refreshed anchor representation available to each denoising step.Its one-step reverse transition inserts the fused prediction into the ReMDM remasking posterior.Source: §5 · p. 6

Training objective

  • With σ_t = 0, the training objective predicts from the current state using a stale anchor sampled at time t′.The denoising cross-entropy trains explicit reuse of latent representations through time.Source: §5 · pp. 6–7

Anchor token selection

  • Setting γ = 0 removes optional anchor supervision and yields fully self-supervised time-based anchoring.When anchor tokens are used, they are clean-token targets at positions masked in z′_t.Source: §5 · pp. 6–7

6 Experiments

Across the evaluated math, code, and STEM benchmarks, TADM:Post-train improves throughput while largely preserving task accuracy. TADM:Pretraining reduces Transformer computation and exceeds ADLM's measured throughput, with competitive generation quality at larger sampling budgets.

6.1 TADM:Post-train

  • The post-training evaluation uses a frozen DiffusionGemma base and evaluates with 48 denoising steps per canvas.Training uses γ = 0, and the evaluation canvas size is 256.Source: §6.1 · p. 7
  • 1.49×–1.79× higher throughput across evaluated benchmarks, with essentially the same task accuracy, is reported for TADM:Post-train.The results are averaged over three seeds and compared with the DiffusionGemma baseline under the same fixed-step evaluation setting.Source: §6.1 · p. 7, Table 1 · p. 8

6.2 TADM:Pretraining

  • The pretraining comparison uses 1,024-token unconditional generations and reports MAUVE, Gen PPL, entropy, compute, and throughput.The 1M-step OpenWebText setup uses a 12-layer anchor and 6-layer denoiser; TADM:Pretraining† has γ = 0, while ‡ uses optional anchor supervision.Source: §6.2 · p. 7
  • Up to 38% compute reduction and 1.73× measured throughput of ADLM are reported for TADM:Pretraining.At larger sampling budgets, TADM matches or improves upon ReMDM in Gen PPL and remains competitive with ADLM; its MAUVE scores remain competitive with ReMDM.Source: §6.2 · p. 8
  • At K = 2, 6 denoiser plus 12/2 anchor layers give 12 Transformer layers per step, matching the 12-layer ReMDM backbone.The comparison excludes the lightweight fusion module's additional computation.Source: §6.2 · p. 8

7 Conclusion

Time-based anchoring makes DLM latent representations reusable across diffusion steps, and TADM instantiates this through post-training and pretraining. The reported speed and compute gains are bounded by cache-refresh and backbone limitations.

7 Conclusion

  • Time-based anchoring makes DLM latent spaces cacheable by periodically computing an anchor representation and reusing it during intervening denoising steps.Source: §7 · pp. 8–9
  • TADM:Post-train improves DiffusionGemma throughput by approximately 49% to 79%, while TADM:Pretraining reduces Transformer-layer computation by up to 38% and achieves up to 73% higher measured throughput than ADLM.The table reports quality metrics, Transformer-layer evaluations, and measured throughput across sampling budgets, but the supplied caption gives no cell values.Source: §7 · p. 9, Table 2 · p. 9
  • Generation quality degrades at larger anchor-refresh intervals, and TADM:Post-train accuracy is constrained by its frozen pretrained backbone.Source: §7 · p. 9

A Additional Theory

The appendix presents the time-anchored training objective and its implications, then reviews discrete, masked, and uniform-state diffusion formulations. It describes stale-anchor denoising as the basis for temporal reuse and allows optional supervised anchor targets.

A Additional Theory

  • The time-anchored objective trains denoising from an earlier, noisier anchor state; optional anchor supervision is controlled by γ, with γ = 0 requiring no predefined anchor targets.Sampling stale anchor times exposes the model to representations reused across diffusion steps.Source: A Additional Theory · p. 14

B.1 Discrete Diffusion Models

  • Discrete diffusion gradually corrupts a clean sequence toward a noisy prior, while a learned model predicts the clean sequence from the noisy input.Source: §B.1 · pp. 14–15

B.1.1 Masked Diffusion Models

  • Masked diffusion uses a mask prior; MDLM adds zero-masking and carry-over unmasking, while ReMDM permits remasking to enable correction of early decoding errors.Source: §B.1.1 · p. 15

B.1.2 Uniform-State Diffusion Models

  • Uniform-state diffusion uses a uniform prior, and in the stated continuous-time limit the reconstruction and prior terms vanish, leaving only the diffusion loss.Source: §B.1.2 · p. 16

B.2 Anchored Diffusion Language Models

ADLM addresses late availability of informative tokens by predicting anchor tokens for the denoiser. Its rationale is that informative anchors can reduce uncertainty about the remaining sequence.

B.2 Anchored Diffusion Language Models

  • ADLM uses an anchor network to produce soft context logits and a denoiser conditioned on them to predict the clean sequence.Its reverse process can remask decoded tokens; setting σt = 0 recovers the absorbing-state transition.Source: §B.2 · p. 16
  • An ideal anchor set minimizes uncertainty in the remaining sequence, and ADLM trains its anchor network to predict informative tokens early for denoiser conditioning.Source: §B.2 · p. 17
  • When a small informative anchor set captures sequence dependencies, anchoring reduces effective conditioning dimension and can substantially reduce the samples needed to learn the conditional distributions.Source: §B.2 · p. 18

C Experimental Setup

The section reports the experimental setup and hyperparameters used for both TADM variants.

C Experimental Setup

C.1 TADM:Post-train

TADM:Post-train freezes DiffusionGemma and trains a small fusion module to support stale-anchor reuse. It reports fixed-step speedups with largely preserved accuracy and reduced Transformer-layer computation.

C.1 TADM:Post-train

  • The DiffusionGemma backbone is split into shared, anchor, and denoiser components, with cached deep anchors fused against the current shallow state.The experiment uses a 2/20/8 shared/anchor/denoiser split.Source: §C.1 · p. 18
  • Paired-attention fusion uses the cached shallow representation as a reference, making its correction vanish when current and cached shallow states coincide.The fusion module then corrects the cached representation as the current canvas changes.Source: §C.1 · p. 18
  • Post-training omits γ supervision because attaching a language-model head to the anchor network would require expensive training.Source: §C.1 · p. 19
  • Evaluation covers six reasoning, knowledge, and coding benchmarks in no-think mode, reporting mean accuracy and throughput over three random seeds.Source: §C.1 · p. 21
  • 0.67× the baseline Transformer-layer evaluations (a 33.3% reduction) result from evaluating the 20-layer anchor network every two steps in a 2/20/8 split.The shared and denoiser networks still run at every step.Source: §C.1 · p. 21
  • Training only the 3.1M-parameter fusion module with the pretrained backbone frozen yields fixed-budget speedups while largely preserving task accuracy.The evaluation disables adaptive stopping to isolate speedups attributable to latent-cache reuse.Source: §C.1 · pp. 20–21

C.2 TADM:Pretrain

TADM:Pretraining learns cacheable latent states through time-based reuse rather than explicit anchor-token conditioning. The section specifies its architecture, cache-age training, and OpenWebText evaluation setup.

C.2 TADM:Pretrain

  • TADM:Pretraining replaces explicit anchor-token conditioning with time-anchored latent caching in a two-stage ADLM architecture.Source: §C.2 · p. 21
  • A shared embedding feeds a 12-block anchor network and 6-block denoiser, joined by gated residual MLP fusion of the current shallow state and stale anchor.Source: §C.2 · p. 21
  • Pretraining uses OpenWebText for 1M steps; evaluation generates 5,000 unconditional 1,024-token samples per configuration and measures quality, Transformer-layer computation, and throughput.Source: §C.2 · pp. 22–23
  • Training samples continuous current time and cache ages aligned to reverse-step budgets, with anchor refresh intervals sampled from K ∈ {1, 2, 4, 8}.Source: §C.2 · p. 22

D Additional Experiments

The additional experiments assess anchor generalization and cache-age effects on generation quality, throughput, and robustness, including stale-anchor reuse in DiffusionGemma. They also provide qualitative generations and define fields and a protocol for evaluating generation corruption.

D Additional Experiments

  • On unseen datasets, TADM:Pretraining remains competitive with MDLM and other diffusion baselines, matching or improving on MDLM for Lambada, LM1B, AG News, PubMed, and ArXiv while remaining comparable on Wikitext and PTB.It is slightly behind ADLM on several benchmarks; the results indicate its learned latent structure supports likelihood modeling beyond the OpenWebText training distribution and a particular sampling trajectory.Source: D Additional Experiments · p. 23, Table 3 · p. 24
  • At T = 2048, increasing K from 1 to 8 changes TADM:Pretraining† Gen PPL from 21.49 to 22.34 as throughput rises from 26.59 to 41.19 Tok/s.At T = 128 and T = 256, larger K noticeably degrades quality; at T = 4096, Gen PPL stays between 17.21 and 17.36 across the sweep. The two TADM variants show similar cache-age trends, suggesting temporal denoising training drives stale-anchor reuse while optional anchor supervision has a modest quality effect.Source: D Additional Experiments · p. 23, Table 4 · p. 25

D.1 Qualitative TADM:Post-train Generations

  • Six TADM:Post-train generations compare anchor refresh intervals K = 1, 2, and 3 on AIME 2026 and GPQA-Diamond, using the same prompt across intervals within each benchmark.Three generations are shown for each benchmark, one per refresh interval.Source: §D.1 · p. 25

D Additional Experiments

  • At T = 1024, increasing K from 1 to 8 raises ADLM Gen PPL from 25.97 to 144.97, while TADM:Pretraining† changes from 23.18 to 30.52 and throughput rises from 47.02 to 74.54 Tok/s.ADLM was trained only with fresh anchors; the comparison indicates that stale-anchor reuse requires training for temporal reuse, while fusion lets current-state information correct stale representations.Source: D Additional Experiments · pp. 23–24, Table 5 · p. 26
  • At K = 3, naive DiffusionGemma reuse corruption reaches 37.78% on AIME26, 14.81% on GPQA-Diamond, and 32.00% on LiveCodeBench-v6; TADM:Post-train reduces these to 20.00%, 4.71%, and 21.71%.TADM:Post-train preserves substantially more task performance, though corruption remains above the fresh-anchor baseline on some difficult tasks.Source: D Additional Experiments · p. 24, Table 6 · p. 27

E Generation corruption Evaluation Prompt

  • The corruption evaluator labels incoherence and repetition loops independently of answer correctness, and excludes isolated errors, verbosity, incomplete outputs, and coherent irrelevance as sufficient grounds for flagging degeneration.When degeneration is unclear, the prompt instructs the evaluator to use null for the relevant label.Source: E Generation corruption Evaluation… · pp. 32–33
  • The evaluation records whether output is incoherent or repetitive, supporting excerpts, and a brief reason.The incoherent and repetition_loop fields accept true, false, or null.Source: E Generation corruption Evaluation… · p. 33

Related papers

Loading 2609.37924v1…