PaperlayerView exact revision on arXiv

Source-linked AI summary

ALoDLM: Adaptively Looped Diffusion Language Models

Liancheng Fang, Zhuowei Li, Youngeun Kim, Tianchen Zhao, Rajat Koner, Jiaye Wu, Linghan Xu, Xuanbai Chen, Xiang Xu, Zheng Zhang, Jakub Zablocki, Nishant Sankaran, Yifan Xing

arXiv:2610.04198v1cs.AIcs.LG

TL;DR

DLMs generate multiple tokens in parallel but retain a quality gap against comparably sized autoregressive models, which the paper attributes to uniform computation for tokens of differing difficulty. ALoDLM uses token-adaptive latent recurrence and learns computation schedules as latent variables. Across eleven benchmarks it surpasses evaluated DLMs and corresponding AR baselines at both scales, while retaining fast parallel decoding.

  • Problem

    DLMs apply uniform computation to masked tokens of differing prediction difficulty, while their quality remains below that of comparably sized AR models.

    Source: §1 · p. 2

  • Method

    ALoDLM uses token-adaptive latent recurrence, feeding committed tokens back as context while unresolved tokens refine their latent states, and learns token-wise exit schedules as latent variables.

    Source: §1 · p. 2, §3.2 · p. 5

  • Results

    Across eleven benchmarks, ALoDLM scores 65.5 at 1.7B and 80.3 at 8B, surpassing evaluated DLMs and corresponding AR baselines.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    ALoDLM combines higher average benchmark scores than evaluated DLMs and corresponding AR baselines with fast parallel decoding.

    Source: §1 · p. 2, §6 · p. 10

    ALoDLM can have longer time to first token because prefill constructs the KV-cache for the full recurrent depth, and its generation speed can vary with input-dependent halting.

    Source: §5 · p. 10

Abstract

from arXiv · show

Diffusion language models (DLMs) enable fast generation by predicting multiple tokens in parallel, but their practical adoption remains limited by a persistent quality gap relative to comparably sized autoregressive (AR) models. We attribute this gap to a computation-difficulty mismatch: within a partially observed sequence, some unknown tokens are easy to predict, while others require substantially more computation. Existing DLMs nevertheless apply uniform computational depth to all unknown positions at each denoising step. We introduce ALoDLM, which replaces uniform computation with token-adaptive latent recurrence. At each denoising step, ALoDLM iteratively refines latent representations and allocates computation according to token difficulty. Tokens ready to commit are fed back as discrete context, while unresolved tokens retain and further refine their latent states through additional recurrent passes. To learn token prediction and computation allocation jointly, we formulate token-wise computation schedules as latent variables and derive a conditional negative evidence lower bound (NELBO). We train ALoDLM at 1.7B and 8B parameter scales. Across eleven benchmarks, ALoDLM outperforms all evaluated DLMs and the corresponding AR baselines in average benchmark score at both scales. ALoDLM also retains fast parallel decoding, yielding a strong quality-efficiency trade-off among evaluated autoregressive and diffusion models under optimized inference engines.

1 INTRODUCTION

The introduction attributes DLMs’ quality gap to uniform computation across tokens of differing difficulty and presents ALoDLM’s adaptive recurrence as a response. It reports stronger average benchmark scores than evaluated DLMs and corresponding AR baselines, alongside fast parallel decoding.

1 INTRODUCTION

  • DLMs apply fixed-depth computation despite token-level differences in prediction difficulty, potentially wasting compute on easy predictions while starving difficult ones.The introduction argues that confidence-based deferral still recomputes unresolved tokens from mask embeddings rather than preserving their latent computation.Source: §1 · p. 2
  • ALoDLM allocates recurrent depth by token: easy tokens commit as context while difficult tokens refine persistent latent states.The proposed training objective jointly optimizes token prediction and computation allocation using token-wise schedules as latent variables.Source: §1 · p. 2
  • The authors derive an unbiased stochastic gradient estimator with variance reduction techniques for stable end-to-end training.The objective jointly learns the base language model and compute-allocation policy.Source: §1 · p. 2
  • Across eleven benchmarks, ALoDLM scores 65.5 at 1.7B and 80.3 at 8B, surpassing evaluated DLMs and corresponding AR baselines.On GSM8K, ALoDLM-8B delivers approximately 2.7× the speed of vLLM-served Qwen3-8B at comparable or higher accuracy.Source: §1 · p. 2
  • ALoDLM is trained at 1.7B and 8B scales through direct SFT of Qwen3 backbones on a 5B-token corpus.The paper describes the 8B model as, to the authors’ best knowledge, the largest looped DLM to date.Source: §1 · pp. 2–3

2 RELATED WORK

The related work distinguishes ALoDLM from prior looping and adaptive-computation approaches by its token-wise latent refinement within diffusion denoising. It also explains why joint exit schedules make exact schedule marginalization difficult.

2 RELATED WORK

  • Looped Transformers increase effective depth without expanding parameter count, and prior work uses recurrence as a scaling axis for algorithmic and reasoning tasks.Adaptive-computation methods differ in their halting approaches, including continuous weighting or summing losses over candidate exit depths.Source: §2 · p. 3
  • In ALoDLM, token commitments change unresolved positions’ context, making exact summation over joint exit schedules intractable.This distinguishes its schedule-learning challenge from approaches that sum over candidate exit depths.Source: §2 · p. 3
  • Prior looped DLMs apply uniform recurrent depth across token positions, whereas ALoDLM lets tokens exit latent refinement at different depths.FReDA supports token-level adaptive commitment but revises an explicit token draft rather than propagating internal latent states across loops.Source: §2 · p. 3

3 METHOD

ALoDLM adds token-adaptive recurrent computation within each denoising step: committed tokens become context while unresolved tokens retain and refine latent states. Its training treats exit schedules as latent variables and optimizes prediction and computation allocation together.

3.1 INFERENCE

  • Committed tokens become discrete context for subsequent recurrent passes, while unresolved positions retain and refine their latent states until the inner loop stops or reaches maximum depth.Predictive confidence controls commitment; a learned halting policy controls continuation, and the loop terminates when no positions remain unresolved or mean cumulative halt probability reaches a threshold.Source: §3.1 · pp. 3–4
  • The model partitions a Transformer into a Prelude, Recurrent Core, and Coda, with the core repeatedly updating the state and the Coda producing readouts for two heads.The LMHead predicts tokens, while the ExitGate supplies per-pass halting probabilities; recurrent depth is capped at K.Source: §3.1 · pp. 3–4

3.2 TRAINING

  • An exit schedule assigns each masked token a discrete commit depth, and early commitments alter context for later recurrent passes.Because schedule choices affect subsequent trajectories, marginalizing all schedules would require summing over exponentially many possibilities.Source: §3.2 · p. 5
  • ALoDLM uses a variational distribution over exit schedules to derive a negative evidence lower bound for joint token prediction and computation allocation.The prior is a product of truncated geometric distributions to encourage shallow exits.Source: §3.2 · p. 5
  • A single sampled exit trajectory supports an unbiased score-function gradient for the halting policy, while standard backpropagation updates denoiser parameters.The trajectory loss also supplies sequence-level credit for whether a sampled schedule achieves accuracy while remaining close to the prior.Source: §3.2 · p. 6
  • Variance reduction reuses predictions from preceding passes to supervise tokens at multiple depths without additional recurrent passes.The method also subtracts the detached first-pass prediction cost as a control variate.Source: §3.2 · p. 6

4 EXPERIMENTS

ALoDLM leads evaluated DLMs and corresponding AR baselines in average benchmark scores at both scales, while offering competitive quality–efficiency trade-offs. Experiments also examine inference-time controls and recurrent-design choices.

4.2 PERFORMANCE EVALUATION

  • At 1.7B and 8B, ALoDLM scores 65.5 and 80.3 on average, surpassing evaluated DLMs and corresponding AR baselines; ALoDLM-8B leads WeDLM-8B by 5.2 points on ten tasks.Table 1 presents the main comparison across eleven benchmarks.Source: Table 1 · p. 7, §4.2 · p. 8

4.3 QUALITY–EFFICIENCY TRADE-OFF

  • At 93.25% accuracy, ALoDLM-8B reaches 612.4 tokens/s versus WeDLM-8B’s 564.2, an 8.5% gain; it uses 133.5 versus 154.6 GFLOPs per token, a 13.6% reduction.The throughput frontier comparison uses single-stream inference on one NVIDIA B200; the compute comparison isolates estimated arithmetic cost.Source: Figure 3 · p. 7, §4.3 · pp. 8–9

4.1 EVALUATION SETUP

  • Both models are directly SFT-converted from Qwen3 using WeDLM’s block-diffusion framework; comparisons use native chat templates, greedy decoding, and optimized inference engines.The evaluation covers eleven reasoning, mathematics, and coding benchmarks.Source: §4.1 · pp. 7–8

4.3 QUALITY–EFFICIENCY TRADE-OFF

  • At fixed q = 0.5, increasing τ from 0.1 to 0.6 raises throughput from 278.7 to 508.3 tokens/s while accuracy falls from 93.8% to 92.3%.At fixed τ = 0.2, increasing q from 0.1 to 0.9 raises accuracy from 93.3% to 93.8% while throughput falls from 455.3 to 309.0 tokens/s.Source: Figure 3 · p. 7, §4.3 · p. 8

4.4 ANALYSIS AND ABLATIONS

  • Increasing q raises loops per token from 1.6 to 2.34 and the average score across eleven benchmarks from 77.9% to 79.1%.Deeper recurrent readouts also exhibit lower training loss on active tokens after the initial training phase.Source: Figure 4 · p. 8, §4.4 · p. 9
  • Numerical tokens have the lowest mean first-pass halt probability (0.369), below the cross-dataset reference mean (0.417), while word tokens have the highest (0.428).These token-dependent preferences are learned without explicit difficulty labels or predefined halting targets.Source: Figure 4 · p. 8, §4.4 · p. 9
  • With maximum recurrent depth K = 4, ALoDLM-8B achieves comparable later-training performance to K = 8, while K = 2 plateaus at a lower score.K = 2 and K = 4 improve faster initially than K = 8 under matched settings.Source: §4.4 · p. 9, Figure 5 · p. 9
  • Middle-layer recurrence outperforms last-layer recurrence later in training, motivating the 8B model’s middle-layer recurrent core.Without intermediate supervision, conditional gradient variance rises to 1.76×, 1.49×, and 1.42× the full-method value at steps 1,000, 6,500, and 17,000, respectively.Source: §4.4 · pp. 9–10, Figure 5 · p. 9

5 LIMITATIONS

ALoDLM’s latency and speed advantages are conditional: prefill can increase time to first token, and adaptive computation makes throughput input-dependent.

5 LIMITATIONS

  • Constructing the KV-cache for full recurrent depth during prefill can make ALoDLM’s time to first token longer than a comparable AR model.This overhead can diminish parallel decoding’s benefit for short responses or latency-sensitive applications.Source: §5 · p. 10
  • Generation speed can vary significantly across prompts, datasets, and domains because recurrent computation and simultaneous commitments depend on token confidence and halting.This variability can occur even under greedy decoding.Source: §5 · p. 10
  • On domains underrepresented in training, additional refinement and fewer parallel commitments may reduce or reverse ALoDLM’s speed advantage over an optimized AR baseline.Source: §5 · p. 10

6 CONCLUSION

ALoDLM uses token-adaptive latent recurrence and a conditional ELBO to address the computation–difficulty mismatch, with benchmark gains and fast parallel generation. The related work places it alongside AR–diffusion hybrids, recurrent models, and self-conditioning methods.

6 CONCLUSION

  • At 1.7B and 8B, ALoDLM outperforms evaluated diffusion models and corresponding Qwen3 baselines in average score across eleven benchmarks.With optimized inference, ALoDLM-8B delivers approximately 2.7× the throughput of vLLM-served Qwen3-8B at comparable accuracy on GSM8K.Source: §6 · p. 10
  • ALoDLM commits confident tokens as discrete context while unresolved tokens retain and refine latent states, using a conditional ELBO over token-wise exit schedules.Intermediate supervision supports denoiser training.Source: §6 · p. 10

A ADDITIONAL RELATED WORK

  • Prior approaches either reintroduce autoregressive structure into diffusion or use diffusion to draft candidates for autoregressive verification [Arriola et al., 2025].The cited work includes blockwise generation and simultaneous diffusion drafting with autoregressive verification.Source: A ADDITIONAL RELATED WORK · p. 17
  • Adaptive halting methods include weighted intermediate outputs and summed prediction losses over candidate exit depths, but early commitments make exact joint-schedule summation intractable for ALoDLM.Early token commitments change the context for later unresolved positions.Source: A ADDITIONAL RELATED WORK · p. 17
  • Unlike conventional self-conditioning across denoising steps, ALoDLM preserves hidden states within a denoising step and backpropagates through recurrent passes.Its conditional NELBO jointly trains the denoiser and token-wise halting policy.Source: A ADDITIONAL RELATED WORK · p. 18

B ADDITIONAL METHOD DETAILS

The appendix specifies depth-aware KV caching for recurrent inference and accounts for its layer-evaluation costs. It also cautions that cache reuse and dense execution affect how those costs should be interpreted.

B.1 DEPTH-AWARE KV CACHING

  • A recurrent layer’s cache is indexed by both physical layer and invocation depth, with decoding using matching-depth prefix entries when available.The lookup can fall back to the deepest available cache depth when the matching entry is absent.Source: §B.1 · p. 18
  • Committed positions can still be processed as context, so their newly computed depth-specific cache entries remain valid; coda readouts use refreshed branch entries.Separate depth buffers are used for recurrent layers, while prelude and coda layers share one buffer per physical layer.Source: §B.1 · p. 18
  • At 8B, 36 layers, 16 recurrent layers, and four passes require 84 logical KV-buffer pairs, including 64 for the recurrent core.Fallback and shared coda buffers are serving approximations and need not match recomputation at every depth.Source: §B.1 · p. 18

B.2 ACCOUNTING FOR RECURRENT COMPUTATION

  • For the 8B partition [10, 26), four passes evaluate 114 Transformer layers, versus 144 for four full-stack passes.The prelude runs once, while the recurrent core and coda readout run at every pass; the 1.7B configuration repeats all 28 layers.Source: §B.2 · pp. 18–19
  • Dense-path layer counts are not per-token FLOPs: committed positions may still incur computation, and batched execution can run to the largest active stopping depth.Serving cost is therefore measured from executed operations and generated-token counts rather than inferred from commitment-depth distributions.Source: §B.2 · p. 19

C TRAINING OBJECTIVE AND GRADIENT DERIVATIONS

The appendix formalizes a schedule-conditioned likelihood and its variational bound, then derives single-trajectory gradient estimation for learning schedules. The derivations specify how earlier commitments condition later predictions and state assumptions for the gradient calculations.

C.1 SCHEDULE-CONDITIONED LIKELIHOOD AND VARIATIONAL BOUND

  • The schedule-conditioned likelihood predicts each exit group using earlier committed targets, with tokens in a group conditionally independent under the parallel readout model.Likelihood evaluation uses observed earlier targets; generation instead feeds back samples.Source: §C.1 · p. 19
  • The conditional negative ELBO bounds the negative log likelihood by expected trajectory loss plus KL divergence between the schedule distribution and its prior.The bound is tight exactly when the variational schedule distribution equals the schedule posterior.Source: §C.1 · p. 19

C.2 SINGLE-TRAJECTORY GRADIENT ESTIMATOR

  • The single-trajectory derivation treats exit decisions as a normalized schedule distribution whose realized-history hazards can depend on earlier exits.The derivation assumes differentiable probabilities on fixed support; it does not assume independent exit depths.Source: §C.2 · p. 20
  • The denoiser derivative holds the reference backbone fixed at its forward-pass value, rather than differentiating through policy features.The expectation step assumes integrable 1/t-weighted gradients and permits differentiation through the outer expectation.Source: §C.2 · p. 20

C.3 GEOMETRIC PRIOR AND DEPTH REGULARIZATION

The appendix defines a truncated geometric prior over exit depths and separates token-level depth variation from aggregate matching to that prior. It distinguishes exact joint-KL recovery from the rollout-based regularization setting used to permit different token depths.

C.3 GEOMETRIC PRIOR AND DEPTH REGULARIZATION

  • With K = 4 and c = 0.4, the geometric prior has mean depth approximately 2.02 and assigns probabilities approximately (0.4131, 0.2769, 0.1856, 0.1244).Larger c favors shallower depths, while learned mean depth also depends on the prediction objective.Source: §C.3 · p. 21
  • The (0.1, 1) setting allows tokens to favor different depths while matching their average depth distribution more strongly to the prior.The regularizer separates variation among token depth profiles from matching their average to the prior.Source: §C.3 · p. 21
  • Unit coefficients recover the original joint KL in expectation, whereas (0.1, 1) gives a rollout-based relaxation with weaker regularization of individual variation.The rollout-based cost’s expectation need not equal the design objective because its profiles depend on sampled histories.Source: §C.3 · p. 21

C.4 VARIANCE REDUCTION AND THE TRAINING SURROGATE

The training surrogate combines denoiser, gate, and auxiliary autoregressive losses, with a sequence-level baseline and intermediate denoiser supervision. The appendix also distinguishes quality and efficiency evaluation protocols.

C.4 VARIANCE REDUCTION AND THE TRAINING SURROGATE

  • A first-pass readout supplies an input-dependent baseline independent of the sampled schedule, and subtracting its detached value leaves the expected gate update unchanged.The baseline requires no extra forward pass; its variance benefit depends on how well it tracks sampled sequence cost.Source: §C.4 · p. 22
  • Intermediate denoiser supervision reuses pre-commitment predictions, while the gate cost still uses each token’s sampled-exit loss.Under blockwise corruption, tokens use their block’s masking probability; the weighted denoiser objective need not share the sampled-exit objective’s expected gradient.Source: §C.4 · p. 22
  • Figure 5’s gradient-covariance trace measures variation from sampled trajectories on fixed corrupted inputs, not variation from sampling training examples.The comparison evaluates variance empirically for the measured parameter subset using paired sampled trajectories.Source: §C.4 · p. 22
  • The complete surrogate combines denoiser, gate, and auxiliary AR losses, with equal diffusion/AR weighting and one 1/M normalization for sequence-summed actor costs.Inverse-mask-probability weighting applies only to the denoiser; the unbiasedness result concerns the original joint-KL objective.Source: §C.4 · p. 22

D EXPERIMENTAL DETAILS

  • The appendix separates single-token decoding for generation-quality comparisons from parallel decoding for quality–efficiency measurements.These are distinct evaluation protocols.Source: D EXPERIMENTAL DETAILS · p. 22

D.1 MODELS AND TRAINING

The appendix specifies ALoDLM’s model configurations, training recipe, benchmark protocol, and inference-cost measurement. Quality and efficiency are evaluated under distinct decoding setups.

D.1 MODELS AND TRAINING

  • ALoDLM-1.7B recurrently applies its 28-layer stack, while ALoDLM-8B uses a 10-layer prelude, shared 16-layer core, and 10-layer coda; both use maximum recurrent depth four.The coda and output head produce readouts at each recurrent pass, so the coda is re-evaluated as the core recurs.Source: §D.1 · p. 23
  • Both scales are supervised-fine-tuned from Qwen3 backbones on a 5B-token corpus, with packed 4,096-token conversations and independently masked 32-token blocks.There is no separate continued-pretraining stage; training also uses an auxiliary next-token prediction loss weighted equally with diffusion loss.Source: §D.1 · p. 23

D.2 GENERATION QUALITY AND INFERENCE EFFICIENCY

  • Generation-quality comparisons use OpenCompass-style evaluation, native chat templates, greedy selection, and a 4,096-token cap; the eleven benchmark scores are averaged without weighting.Non-code tasks report accuracy or exact match, while code tasks report execution-based pass@1.Source: §D.2 · p. 23
  • For quality comparisons, each model pass commits one token; ALoDLM uses a 16-token window and scale-specific q values, with entropy-based selection disabled.At 1.7B q = 0.4 and at 8B q = 0.5; diffusion baselines use their single-token presets.Source: §D.2 · pp. 23–24
  • Parallel-decoding throughput is measured on one NVIDIA B200 using each model’s specified inference engine, including prefill and decoding but excluding tokenization and answer scoring.Dataset throughput divides total generated tokens by summed generation time; the parallel decoder’s q, τ sweeps are separate from the quality preset.Source: §D.2 · p. 24
  • Estimated FLOPs per generated token combines Transformer-layer evaluations, vocabulary readouts, and an attention term using an approximate context length.The estimate includes prefill, evaluated window positions, graph padding, and executed readouts; minor elementwise operations are omitted.Source: §D.2 · p. 24

E CASE STUDIES

The case study visualizes token-specific commitment depth in three GSM8K responses, complementing aggregate allocation analysis with individual examples.

E CASE STUDIES

  • Commitment depth varies across token positions within individual GSM8K responses, with different recurrent passes interleaved in each response.The aggregate allocation analysis is reported separately in Figure 4.Source: E CASE STUDIES · p. 25
  • Figure 6 shades each token by its recorded commitment pass, with darker shades indicating later commitment within that denoising step.The figure presents three GSM8K responses.Source: Figure 6 · p. 25
  • The illustrated generations use K = 4, q = 0.5, τ = 0.4, and a 16-token decoding window.These settings describe the examples, not the aggregate allocation analysis.Source: Figure 6 · p. 25, E CASE STUDIES · p. 25
Loading 2610.04198v1…