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

ALoDLM uses token-adaptive latent recurrence to resolve the quality-efficiency gap in diffusion language models.

How can diffusion language models achieve autoregressive-level quality while maintaining the speed benefits of parallel token generation?

Diffusion language models (DLMs) struggle to match the quality of autoregressive models because they apply uniform computational depth to all tokens, wasting effort on easy predictions while starving difficult ones of necessary compute. ALoDLM introduces a looped architecture where tokens exit latent refinement adaptively: confident tokens commit early to provide context, while unresolved tokens continue to refine their latent states through additional recurrent passes. This approach achieves state-of-the-art performance among DLMs, outperforming autoregressive baselines at 1.7B and 8B scales while delivering 2.7× the throughput of standard autoregressive models on GSM8K.

Paper Primer

The core mechanism is a token-adaptive inner loop that treats exit depth as a latent variable. The model uses predictive confidence to decide when a token is ready to commit, and a learned halting policy to determine if the remaining tokens require further recurrent refinement.

ALoDLM is like a student taking a test: they answer the easy questions immediately to clear the page, but spend extra time re-reading and refining their logic for the complex problems before committing an answer.

ALoDLM outperforms both existing DLMs and autoregressive baselines in generation quality.

Across eleven benchmarks, ALoDLM-8B achieved an average score of 80.3, surpassing the Qwen3-8B autoregressive baseline (78.5) and the WeDLM-8B diffusion baseline (75.1). +1.8 points over AR baselines; +5.2 points over WeDLM-8B.

The model provides a superior quality-efficiency trade-off for parallel decoding.

On GSM8K, ALoDLM-8B delivers 2.7× the throughput of vLLM-served Qwen3-8B at comparable accuracy.

Why does this approach use a looped architecture instead of just increasing the number of layers?

Looped architectures increase effective depth without expanding the parameter count, allowing the model to scale computation dynamically based on token difficulty rather than forcing a fixed, high-cost depth for every token.

How does the model learn to allocate computation without explicit difficulty labels?

The model jointly optimizes token prediction and a halting policy using a conditional negative evidence lower bound (NELBO), which treats exit depths as latent variables and uses a score-function estimator to reward schedules that balance accuracy with computational efficiency.

ALoDLM demonstrates that latent recurrence is an effective test-time scaling axis: increasing the exit threshold allows the model to trade off inference speed for higher reasoning accuracy on demand.

ALoDLM proves that diffusion models can surpass autoregressive quality by replacing uniform computation with adaptive, token-level latent refinement.

Introduction and Motivation

We expose why diffusion language models lag behind autoregressive models and propose a token‑adaptive recurrent fix.

Diffusion language models generate many tokens in parallel, but they consistently underperform autoregressive models of comparable size. This gap stems from a mismatch between the amount of computation allocated to each token and the intrinsic difficulty of predicting that token.

When some masked tokens are easy to predict while others are hard, applying the same amount of computation to all wastes effort on easy tokens and starves hard ones.

AR models generate text token by token, conditioning each prediction on a fully resolved prefix, which avoids within‑step difficulty variation.

To address the mismatch, we introduce ALoDLM, which uses a token‑adaptive recurrent loop that lets easy tokens commit early while difficult tokens undergo additional refinement. Our contributions are a looped architecture with dynamic depth, a conditional ELBO training objective, and scaling to 8 B parameters with strong quality‑speed trade‑offs.

The key insight is that improving diffusion models’ quality while preserving fast parallel decoding hinges on balancing computation per token.

Inference with Adaptive Recurrence

Describes ALoDLM’s token‑adaptive recurrent inference mechanism.

Standard diffusion language models apply the same denoising depth to every token, which wastes compute on easy positions while starving hard ones of refinement.

Instead of running a uniform number of transformer layers for all tokens, ALoDLM lets each token stop its recurrent updates as soon as it is confident, feeding its final embedding back as discrete context for the remaining tokens.

How does this differ from standard early‑exit models that also use a confidence gate?

Early‑exit models apply a single shared exit after a fixed number of layers, while ALoDLM evaluates an ExitGate at every recurrent pass and can halt individual tokens independently, allowing heterogeneous depths within the same decoding step.

Pass 1: entropies = [0.3, 0.9, 0.6, 0.4]; $\lambda(1)=0.4$. Tokens 1 and 4 have entropy ≤ $\tau$, so they commit.

Update $a^{(1)} = 0 + (1-0)\times0.4 = 0.4$; remaining tokens = {2, 3}.

Pass 2: entropies = [0.7, 0.5]; $\lambda(2)=0.5$. Token 3 now meets the entropy criterion and commits.

Update $a^{(2)} = 0.4 + (1-0.4)\times0.5 = 0.7$; remaining token = {2}.

Pass 3: entropy = 0.6 (> $\tau$); $\lambda(3)=1$ (forced exit). Token 2 is forced to commit.

Update $a^{(3)} = 0.7 + (1-0.7)\times1 = 1.0$, exceeding $q$, so the inner loop terminates.

This toy run shows that easy tokens exit after a single pass, while harder tokens consume additional recurrent updates, reducing total computation compared to a uniform $K$‑step rollout.

Initialize recurrent state $h^{(0)}$ with the Prelude on the masked input.

For $s=1$ to $K$: update $h^{(s)} = \text{Recurrent Core}(h^{(s-1)})$, compute $r^{(s)} = \text{Coda}(h^{(s)})$.

From $r^{(s)}$ obtain token distribution $p(s)$ via LMHead and halting probability $\lambda(s)$ via ExitGate.

Identify tokens whose entropy under $p(s)$ ≤ $\tau$; sample their predictions and embed them, adding them to the committed set.

Accumulate halting probability $a^{(s)} = a^{(s-1)} + (1-a^{(s-1)})\lambda(s)$; if the mean $a^{(s)}$ over remaining tokens ≥ $q$ or no tokens remain, break.

If no token committed after $K$ passes, select the token with lowest entropy, sample its prediction, and commit it to guarantee progress.

Write all committed token embeddings back into the sequence; re‑initialize the Prelude for the next outer denoising step.

Repeat outer steps until no [MASK] tokens remain.

Algorithm 1 – ALoDLM decoding

**Figure 2.** **Overview of ALoDLM.** Standard DLMs apply a fixed-depth denoiser at each denoising step. ALoDLM introduces token-adaptive recurrence within each step: confident tokens commit early and provide discrete context, while unresolved tokens retain and refine their latent states through additional recurrent passes.

Training the Halting Policy

Train ALoDLM by jointly learning a denoiser and a token‑adaptive halting policy.

Fixed‑depth diffusion models waste compute on tokens that are already easy to predict, while hard tokens remain under‑refined. ALoDLM solves this imbalance by letting each token decide when to stop.

The schedule is a per‑token plan that tells the model after which recurrent pass each token should be committed, allowing easy tokens to exit early and hard tokens to keep refining.

Pass 1: token 1 is committed (its prediction is emitted); tokens 2 and 3 remain masked.

Pass 2: token 2 now sees token 1 as context, produces a refined prediction, and is committed; token 3 stays masked.

Pass 3: token 3 receives context from tokens 1 and 2, makes its final prediction, and exits.

The schedule lets the model allocate more recurrent computation to difficult tokens while saving work on easy ones, directly shaping the trade‑off between speed and accuracy.

How does the Joint Exit Schedule differ from a naïve fixed‑depth schedule?

In a fixed‑depth schedule every token runs exactly $K$ recurrent passes regardless of difficulty, so compute is wasted on easy tokens and insufficient for hard ones. The Joint Exit Schedule makes the depth a random variable per token, allowing early exits for easy tokens and deeper refinement for hard tokens, which yields a more efficient allocation of compute.

The policy decides, at each recurrent pass, whether a token should stop based on two signals: an entropy gate $\tau$ that measures confidence, and a cumulative halt gate $q$ that limits how many tokens remain uncommitted.

Pass 1: token 2 satisfies $\operatorname{Entropy}\le\tau$ and is committed; $a_2$ becomes $1$.

Update $a_i \leftarrow a_i + (1-a_i)\lambda_i$ for remaining tokens; suppose $\lambda_1=0.6$, $\lambda_3=0.5$, yielding $a_1=0.6$, $a_3=0.5$.

Since $\max\{a_1,a_3\}=0.6<q$, the loop continues.

Pass 2: recompute entropies; token 1 now has entropy $0.3$ (below $\tau$) and is committed; $a_1$ becomes $1$.

Now $\max\{a_3\}=0.5<q$, so another pass is performed.

Pass 3: token 3 finally meets the cumulative condition $\max a_i =0.5$ still below $q$, but the algorithm stops because the maximum allowed depth $K$ is reached.

The policy can exit early for confident tokens while still giving harder tokens additional chances, and the cumulative gate ensures the process does not linger indefinitely.

Why not simply use a fixed entropy threshold for all tokens without the cumulative halt gate?

A fixed entropy threshold alone can leave many tokens uncommitted if their uncertainties stay above $\tau$, causing the loop to run for the maximum depth $K$ and wasting compute. The cumulative gate $q$ provides a global stopping signal that forces the algorithm to terminate once the remaining tokens collectively have low enough halting probability, guaranteeing a bound on runtime.

**Figure 3.** (a) Complementary roles of the entropy and halting gates. Increasing $\tau$ relaxes the commitment criterion and accelerates generation, with sharper accuracy losses at high thresholds (left). Increasing $q$ permits deeper latent refinement, also yielding accuracy gains at lower throughput (right). These gates collectively balance parallel token commitment with adaptive refinement of unresolved tokens. (b) Adaptive recurrence improves the quality–efficiency frontier. Across most of the evaluated range, ALoDLM-8B achieves higher accuracy than WeDLM-8B at comparable throughput (left) or compute per generated token (right). Solid lines connect frontier configurations.

With the joint exit schedule and halting policy trained via the variational objective and unbiased gradient estimator, ALoDLM learns to allocate computation where it matters most, achieving higher quality at comparable or lower cost than static‑depth diffusion models.

Experimental Results

ALoDLM beats AR baselines while keeping inference fast.

We evaluate ALoDLM on a suite of knowledge, reasoning, and coding benchmarks, comparing against both autoregressive (AR) baselines and recent diffusion language models. All models are run with their native chat templates, greedy decoding, and a 4,096‑token limit. Training follows the SFT pipeline on a 5 B‑token corpus without continued pre‑training.

ALoDLM achieves higher average accuracy than AR baselines at both 1.7 B and 8 B scales.

Table 1 shows averages of 65.5 % vs 63.8 % (1.7 B) and 80.3 % vs 78.5 % (8 B).

**Table.** Data and corruption.

**Figure 4.** **Left: Learning dynamics.** After the initial training phase, deeper recurrent readouts exhibit lower training loss on their active tokens. **Middle: Test-time scaling with q.** Increasing $q$ raises loops-per-token and the average score across eleven benchmarks. **Right: Token-adaptive recurrence.** Mean first-pass halting probabilities differ across token categories, with numerical tokens showing the lowest mean. These preferences are learned without explicit token-difficulty labels or predefined halting targets.

Ablation Studies

Ablation studies quantify how adaptive recurrence, depth, loop placement, and supervision affect performance.

Increasing the exit threshold $q$ raises the average number of recurrent loops per token and improves overall accuracy.

Average loops per token grow from 1.6 to 2.34, and the mean score across eleven benchmarks rises from 77.9 % to 79.1 %.

Maximum recurrent depth experiments compare $K\in\{2,4,8\}$ under matched settings. Early in training $K=2$ and $K=4$ improve faster than $K=8$, but $K=2$ plateaus at a lower score. With half the maximum depth, $K=4$ reaches performance comparable to $K=8$ later, justifying the choice of $K=4$.

Loop placement ablations compare recurrence over middle layers $[10,26)$ versus last layers $[20,36)$ of Qwen3‑8B, using the full‑stack trajectory $[0,36)$ as a reference. Middle‑layer recurrence outperforms last‑layer recurrence later in training, indicating that the location of recurrence critically influences latent refinement.

Removing intermediate denoiser supervision increases conditional gradient variance.

Variance rises to 1.76×, 1.49×, and 1.42× the full‑method value at training steps 1 k, 6.5 k, and 17 k respectively.

Limitations: ALoDLM incurs a longer time to first token because the KV‑cache must be built for the full recurrent depth during the prefill phase, reducing parallel‑decoding benefits for short responses. Additionally, generation speed is input‑dependent; prompts requiring more recurrent refinement can negate the throughput advantage over optimized autoregressive baselines.

**Figure 6.** Token-wise recurrent depth in three GSM8K responses generated by ALODLM. Each token is shaded by its recorded commitment pass; darker shades indicate later commitment within the corresponding denoising step. These examples use $K = 4$, $q = 0.5$, $\tau = 0.4$, and a 16-token decoding window.

**Figure 5.** **Left:** At 8B, $K=4$ outperforms $K=8$ early in training and achieves comparable performance later, while $K=2$ plateaus at a lower score. The inset enlarges steps 10,000–17,000. **Middle:** Middle-layer recurrence outperforms last-layer recurrence later in training. The full-stack curve is a historical reference evaluated on a different benchmark subset. **Right:** Intermediate supervision reduces conditional gradient variance in ALODLM-8B. Variance across sampled exit trajectories is normalized to the value with intermediate supervision at each checkpoint. Error bars show paired 95% bootstrap confidence intervals across eight fixed input packs.

Related Work

We situate ALoDLM among prior diffusion language models and adaptive computation approaches.

Diffusion language models (DLMs) split into two families: masked DLMs, which fix tokens once they are unmasked, and uniform DLMs, which keep revising tokens across denoising steps. Recent work explores both re‑injecting autoregressive structure and using diffusion as a draft for later AR verification. Parallel‑generation models also intersect with adaptive‑depth architectures that learn when to halt computation.

Masked diffusion masks a subset of tokens at each step and, once a token is unmasked, its value is fixed for the remainder of the reverse process.

Generates sequences in autoregressive blocks while denoising tokens within each block in parallel, blending blockwise AR structure with diffusion.

Extends Block Diffusion to pretrained AR models, enabling diffusion steps to operate on existing AR representations.

Combines block‑wise diffusion with AR objectives, allowing diffusion to refine blocks generated by an AR backbone.

Uses topological reordering to merge parallel mask recovery with causal attention, enabling streaming prefix commitment.

Jointly trains on AR and block‑diffusion objectives, supporting standard generation, diffusion, and self‑speculative decoding within a single architecture.

Executes diffusion drafting and AR verification simultaneously in a single‑pass architecture.

Uses introspective consistency training and strided decoding to verify past proposals while generating new tokens concurrently.

Employs a lightweight block‑diffusion drafter conditioned on the target AR model’s context features.

Learns parallel diffusion adapters atop frozen AR weights, using rejection sampling against the AR distribution for lossless acceleration.

Reuse shared architectural blocks to increase effective depth without expanding parameter count, enabling efficient in‑context learning.

Learned halting mechanisms that decide when to stop processing a token based on weighted intermediate states.

Introduce recurrence over transformer layers with a learned halting policy, allowing variable numbers of refinement steps per token.

Dynamically routes tokens to varying recurrent depths, mixing multiple recursion pathways.

Explicitly sums prediction losses across candidate exit depths using learned halting probabilities.

Framework that sums prediction losses across candidate exit depths, similar to PonderNet, but with additional regularization for stable training.

Method Details

Method details for depth‑aware KV caching and recurrent inference cost.

This appendix details two core mechanisms that enable token‑adaptive computation: depth‑aware KV caching and the accounting of recurrent inference cost.

Self‑conditioning feeds the denoiser its own previous clean‑data estimate, allowing each denoising step to refine predictions using its own earlier output as an auxiliary input.

How does ALoDLM’s self‑conditioning differ from the conventional detached feedback used in earlier works?

Earlier methods treat the previous clean estimate as a fixed, non‑learnable signal (stop‑gradient), so gradients cannot flow back through it. ALoDLM keeps the hidden states alive across recurrent passes and allows gradients to pass, enabling the model to learn how earlier states should be shaped to aid later predictions.

Instead of sharing a single key/value cache across all recurrent passes, the model maintains separate logical buffers for each depth, allowing attention to retrieve the most appropriate representation for a token at its current computation depth.

Why can’t we simply reuse a single KV cache across all recurrent passes?

Because different tokens may be processed at different depths; a single cache would contain entries computed at varying depths, causing mismatched representations when a shallow token attends to a deep token’s cache. Separate depth‑specific buffers guarantee that each attention lookup uses the most appropriate depth‑aligned representation.

To compare the cost of token‑adaptive recurrent inference against a full‑stack transformer, the paper counts how many layer evaluations are performed per position for a given number of recurrent passes.

Why does the paper count layer evaluations instead of FLOPs when measuring recurrent inference cost?

Layer evaluations directly capture the amount of sequential work a token undergoes, which determines latency in a serving setting. FLOP counts can be misleading because many operations are parallelized; counting layers reflects the real bottleneck of waiting for the deepest token in a batch.

Training Derivations

Derives the training‑objective bound and the single‑trajectory gradient estimator for ALoDLM.

This appendix formalizes the relationship between the training loss $J$, the conditional likelihood of the observed tokens, and the KL divergence between the variational posterior $q_{\phi}$ and the model posterior $p_{\theta}$. It also derives the gradient estimator used when only a single stochastic trajectory is sampled.

Applying Jensen’s inequality to the marginalisation over the support of $q_{\phi}$ yields a tractable lower bound on the log‑likelihood, which is the first step toward the ELBO‑style expression above.

Using Bayes’ rule, the ratio of the variational posterior to the model posterior simplifies to the ratio of the variational posterior to the prior, which is the term that appears inside the ELBO.

For the single‑trajectory gradient estimator we need the log‑probability of each token’s exit decisions. This is captured by the accumulated decision log‑probability $u_i(z)$.

Finally, the gradient is taken with respect to a surrogate objective $\bar J$ that treats the policy parameters $\theta$ as fixed when differentiating through the stochastic decisions.

With the definitions above, the gradient of the original loss $J$ can be expressed as the gradient of $\bar J$ evaluated at $\bar\theta=\theta$, yielding the single‑trajectory estimator used in the training loop.

Gradient Estimation

Derives the gradients for the training objective and formalizes the geometric prior and depth regularization.

This subsection derives the gradients needed for optimizing the training objective and formalizes the geometric prior that encourages shallow exits.

Regularization Costs

Derives the sampled regularization cost, variance‑reduction surrogate, and related theorems.

This subsection formalizes the regularization cost that depends on the sampled exit schedules, then builds variance‑reduced surrogates for the gate and denoiser updates.

For a corrupted minibatch, with $\alpha=\beta=1$ and any positive aggregate profile, the expectation of the sampled cost equals the KL divergence between the posterior $q_{\phi}(z_b)$ and the prior $\pi_c(z_b)$: $\displaystyle \mathbb{E}_Z\!\big[C_{1,1}^{i}(Z)\big]=D_{\mathrm{KL}}\!\big(q_{\phi}(z_b)\,\|\,\pi_c(z_b)\big).$

The aggregate terms cancel for every realized minibatch.

Taking the expectation over the rollout distribution yields the KL divergence.

Experimental Details

Experimental setup, model configs, training surrogate, and cost estimates.

Weighted denoiser objective need not share the same expected gradient as sampled‑exit supervision; we therefore measure its variance empirically by tracing the gradient‑covariance matrix for a fixed parameter subset.

The actor normalizes sequence‑summed costs with a single $1/M$ factor; inverse‑mask‑probability weighting is applied only to the denoiser term. This surrogate merges the sequence‑level credit structure of Proposition 1 with relaxed depth regularization and intermediate supervision.

Both models are fine‑tuned on the 5 B‑token corpus using the WeDLM block‑diffusion framework; training examples are formatted as conversations, packed into 4,096‑token sequences, and supervised on the completion tokens.

Optimization follows Table 2: AdamW with DeepSpeed ZeRO‑2, cosine learning‑rate schedule, 0.03 warmup fraction, 0.01 weight decay, and gradient clipping at norm 1.0. The 8B model runs eight epochs (34,512 steps) with activation checkpointing every second recurrent‑layer invocation.

Generation quality follows the OpenCompass protocol with greedy decoding and a 4,096‑token cap. Benchmarks include ARC‑Easy, ARC‑Challenge, MMLU, MMLU‑Pro, GPQA‑Diamond, GSM8K, MATH‑500, HumanEval, sanitized MBPP, and their EvalPlus variants.

Inference efficiency is measured on a single NVIDIA B200 using each model’s native CUDA‑graph engine (ALoDLM, WeDLM) or external engines (vLLM, JetEngine, dInfer). Timing includes prefill and decoding with device synchronization, excluding tokenization and answer scoring.

**GSM8K example 1** – Compute daily revenue: eggs used $=3+4=7$, eggs sold $=16-7=9$, revenue $=9\times2=18$, so Janet earns **\$18** per day.

**GSM8K example 2** – Determine fiber bolts: blue fiber $=2$, white fiber $=2\times\frac12=1$, total bolts $=2+1=3$, so the robe needs **3 bolts** of fiber.

**GSM8K example 3** – Compute weekly running distance: daily meters $=3\times60=180$, weekly meters $=180\times3=540$, so James runs **540 m** per week.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers