PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Parallel Samplers for Recurrent-Depth Models and Their Connection to Diffusion Language Models

Jonas Geiping, Xinyu Yang, Guinan Su

arXiv:2510.14961v1cs.LGcs.CL

TL;DR

Recurrent-depth generation waits for recurrence at each position, while fixed-depth models can struggle with multi-step reasoning. The paper adapts diffusion forcing to refine token drafts in parallel, reporting around 5x speedups with roughly 1% generation-quality trade-offs. It also frames recurrent-depth models as continuous latent language diffusion models trained with truncated unrolling.

  • Problem

    Recurrent-depth generation ordinarily waits for the recurrence at a sequence position to converge before producing the next token, while fixed-depth transformers can struggle with multi-step logical chains.

    Source: §1 · pp. 1–2

  • Method

    The sampler drafts tokens from intermediate iterates, then advances to later positions while refining earlier drafts through recurrence.

    Source: §1 · p. 2

  • Results

    Around 5x speedups on all benchmarks came with around 1% trade-offs in generation quality, depending on the task.

    Source: §5.1 · p. 10

  • Takeaways & Limitations

    Recurrent-depth models can be viewed as continuous latent language diffusion models trained with truncated unrolling, a perspective that suggests unrolling objectives could be competitive for future language diffusion models.

    Source: §6 · p. 11

    Looped architectures such as Coconut are not immediately supported and require retraining to separate recurrent state from input data.

    Source: §3.2 · p. 4

Abstract

from arXiv · show

Language models with recurrent depth, also referred to as universal or looped when considering transformers, are defined by the capacity to increase their computation through the repetition of layers. Recent efforts in pretraining have demonstrated that these architectures can scale to modern language modeling tasks while exhibiting advantages in reasoning tasks. In this work, we examine the relationship between recurrent-depth models and diffusion language models. Building on their similarities, we develop a new diffusion forcing sampler for these models to accelerate generation. The sampler advances by decoding new tokens at every forward pass of the model, while the latent states of these tokens can be further refined in parallel through recurrence. Theoretically, generation with our sampler is strictly more expressive than the baseline autoregressive generation using the same time budget on modern hardware. Moreover, this sampler, based on principles from diffusion literature, can be directly applied to existing 3.5B recurrent-depth transformers without any tuning, leading to up to a 5x speedup. Consequently, our findings not only provide an efficient mechanism for parallelizing the extra computation in recurrent-depth models at inference, but also suggest that such models can be naturally viewed as strong continuous, though causal, diffusion language models.

1 INTRODUCTION

Fixed-depth models can struggle with multi-step reasoning, motivating recurrent-depth models that repeat layers. This work adapts diffusion forcing to parallelize their generation while preserving causal, left-to-right information flow.

1 INTRODUCTION

  • Fixed-depth transformers can find multi-step logical chains difficult to learn robustly, motivating work on verbalized chains of thought.The passage places fixed-depth transformers in TC0.Source: §1 · p. 1
  • Recurrent-depth models repeat layers and are also called looped transformers or universal transformers; work reports they can represent a larger complexity class than fixed-depth models.The terminology reflects either looping or the motivation to represent universal Turing machines (Dehghani et al., 2019; Graves et al., 2014; Graves, 2017).Source: §1 · p. 1
  • Standard recurrent-depth generation waits for each token’s recurrence to finish, whereas diffusion forcing refines sequence estimates diagonally.The sampler adapts diffusion forcing (Chen et al., 2024a) to existing recurrent-depth models from Geiping et al. (2025).Source: Figure 1 · p. 2
  • Diffusion forcing advances to the next token while refining earlier token drafts, exposing sequence-level parallelism without reducing FLOPs.Each forward pass can refine drafts at positions n and n + 1 while drafting n + 2; the causal model still propagates information left to right.Source: §1 · p. 2
  • The sampler is presented as balancing practical efficiency and theoretical expressiveness, and as outperforming well-tuned speculative decoding for the same model.Its speed gains can be traded against accuracy.Source: §1 · p. 2

2 RELATED WORK

Related work covers transformer recurrence, diffusion language modeling, and inference strategies for long sequences and fixed-depth models. The paper tests a generic sampler on existing recurrent-depth models and illustrates how it generates tokens across recurrence steps.

2 RELATED WORK

  • Transformer recurrence was reintroduced as universal transformers (Dehghani et al., 2019), and later work reports algorithm learning and capable looped-transformer learners (Bear et al., 2024).Recent work also trains general-purpose recurrent-depth language models and explores recurrence for latent reasoning (Tan et al., 2023; Abnar et al., 2023; Geiping et al., 2025; Hao et al., 2024; Liu et al., 2024a; Cheng and Durme, 2024).Source: §2 · pp. 2–3
  • In the illustrated example, the baseline needs 32 recurrence steps per token, while the sampler has committed 8 new tokens; its early drafts improve with further steps.The freeze decision is dynamic and based on distance to the previous latent state.Source: Figure 2 · p. 3
  • The generic sampler is tested on openly accessible Geiping et al. (2025) models trained on 800B tokens, with 3.5B parameters.The models target general language understanding and reasoning.Source: §2 · p. 3
  • Language diffusion uses either continuous variables projected into discrete text or processes acting directly on discrete variables; masking is described as the most scalable discrete approach.The passage notes large-scale masked diffusion efforts competitive with autoregressive models.Source: §2 · p. 3
  • Block diffusion modifies and freezes text chunks while caching their KV entries, whereas diffusion forcing adds noise to future tokens relative to the current position (Arriola et al., 2025; Chen et al., 2024a).Diffusion forcing lets the sampler move across both sequence position and diffusion time.Source: §2 · p. 3
  • Fixed-depth transformer inference can be memory-bound, and speculative decoding drafts future tokens for verification by the original model.The passage says sequence verification is compute-bound and therefore fast.Source: §2 · pp. 3–4

3 APPLYING DIFFUSION FORCING TO RECURRENT-DEPTH MODELS

The sampler decodes token drafts while recurrent states continue to be refined in parallel. The section describes the architecture and requirements it relies on, then introduces stabilization and adaptive exits to manage generation.

3 APPLYING DIFFUSION FORCING TO RECURRENT-DEPTH MODELS

  • Each recurrence step advances generation by at least one token while intermediate drafts are refined in parallel, accelerating text generation.The sampler applies during generation; prompt prefill remains parallelizable across sequence positions.Source: §3 · p. 4, §3.3 · p. 5

3.1 BACKGROUND ON RECURRENT-DEPTH MODELS

  • Huginn’s prelude maps embedded inputs into a latent space, its recurrent block refines latent states, and its coda produces next-token probabilities; this can be viewed as latent-space diffusion.The model starts from a random state and iteratively refines it conditioned on the embedded input sequence.Source: §3.1 · p. 4

3.2 THE INGREDIENTS FOR DIFFUSION FORCING SAMPLING

  • The sampler requires input injection and decodable intermediate states; KV-cache sharing is beneficial, while Coconut-style looped architectures require retraining to add input injection and separate recurrent state from input data.Random-state initialization may strengthen the diffusion connection but is not architecturally necessary.Source: §3.2 · pp. 4–5

3.3 A SIMPLIFIED VERSION OF THE SAMPLING ALGORITHM

  • The simplified sampler uses inner recurrences to stabilize changing conditioning, then freezes completed tokens and caches their KV states so generation proceeds as a wave.Its fixed exit rule freezes a token after a predetermined number of recurrence steps; setting r′ = r recovers autoregressive sampling.Source: §3.3 · p. 5

3.4 STABILIZING COMPONENTS BASED ON DIFFUSION PRINCIPLES

  • Adding noise to intermediate states can stabilize iteration and improve accuracy and throughput when r′ is small; linearly scheduled βt outperforms constant or square-root schedules, with its optimum depending on r′.The authors also report a small but robust average gain from stabilizing recurrence on challenging sequences.Source: §3.4 · pp. 5–6

3.5 ADAPTIVE EXITS

  • Fixed exits can waste computation or freeze unconverged states; adaptive exits freeze the longest converged prefix while limiting wavefront size to bound the number of modified states and memory.Convergence is assessed using normalized latent-space distance against a threshold.Source: §3.5 · p. 6
  • The adaptive sampler can in principle converge to the autoregressive solution if the recurrent block is a contraction, though the authors note that this property is difficult to prove for large models.The paper illustrates a wave that advances, pauses briefly, and then resumes on a challenging GSM8k input.Source: §3.5 · pp. 6–7
  • The sampler adds prelude and coda costs and may incur slower convergence, but practical parallelization gains counteract its FLOP inefficiency.The extra costs depend on the recurrent block’s size relative to the prelude and coda.Source: §3.5 · p. 7

4 THEORETICAL ANALYSIS

The analysis compares depth and width scaling for Transformer-based autoregressive LLMs, excluding chain-of-thought length scaling. It finds different scaling advantages for prefilling and decoding.

4.1 PROBLEM FORMULATION

  • The framework compares depth and width in Transformer-based autoregressive LLMs and excludes length scaling such as chain-of-thought scaling.It defines depth and width at each time step for recurrent-depth models.Source: §4.1 · pp. 7–8

4.2 LLMS SHOULD PRIORITIZE DEPTH SCALING DURING PREFILLING.

  • For its width-scaling comparison, the analysis repeats each token along the sequence dimension, with variants differing in which earlier token copies each copy may attend to.During prefilling, increasing the number of repeated tokens is equivalent to width scaling under the paper’s definition.Source: §4.2 · p. 8
  • At equal scaling factors, depth is more expressive and no more costly than KVShare, while width offers s^2 times the parallelism below L⋆ and similar saturated parallelism above it.The asymptotic prefill costs satisfy EDepth≤EWidth-KVShare<EWidth-NoShare.Source: §4.2 · p. 8

4.3 LLMS SHOULD PRIORITIZE WIDTH SCALING DURING DECODING.

  • With r > 1, KV-cache sharing, and wavefront size W ≤ L⋆, diffusion forcing matches autoregressive depth but achieves strictly greater width under the same runtime budget.The authors state that W ≤ L⋆ typically holds because maximum recurrence depth rarely exceeds r ≈ 100.Source: §4.3 · p. 8

5 EXPERIMENTAL EVALUATION

Across four reasoning and coding benchmarks, the sampler delivers roughly 5× speedups with about 1% generation-quality trade-offs. Hyperparameter sweeps show how recurrence, noise, wavefront size, and headway affect the accuracy-throughput balance.

5.1 BENCHMARK RESULTS.

  • Around 5× speedups with roughly 1% generation-quality trade-offs were observed across all benchmarks; model variants showed similar gains and trade-offs.Table 1 compares autoregressive and diffusion samplers; Table 2 reports 4–5× speed gains and 0.5–1% accuracy deviations across two additional checkpoints.Source: Table 1 · p. 9, §5.1 · p. 10, Table 2 · p. 10

5 EXPERIMENTAL EVALUATION

  • The evaluation covers four generative benchmarks—GSM8K, MATH500, HumanEval, and MBPP—and compares the sampler with static autoregressive, tuned adaptive-compute, and heavily tuned self-speculative baselines.For the Huginn-0125 checkpoint on GSM8k, drafting 4 future tokens with 4 draft steps is optimal. Samplers use comparable Hugging Face transformers implementations with dynamic KV caching.Source: §5 · p. 9

5.2 VARIANTS AND HYPERPARAMETERS

  • Increasing inner recurrence or adjusting the exit threshold trades throughput against accuracy, while the sampler remains robust across broad momentum and noise settings.The reported upsides from varying embedding EMA and the noise schedule are limited.Source: Figure 5 · p. 9, §5.2 · p. 10, Figure 6 · p. 10
  • On the tested A100, maximum wavefront sizes between 64 and 128 appear optimal, but the optimum is likely accelerator-specific.Larger headways remain stable but bring only marginal speed gains at an accuracy cost; for the studied model, headways above one do not yield practical speedups.Source: §5.2 · pp. 10–11, Figure 8 · p. 11
  • More noise is, on average, helpful when fewer inner recurrence steps are used; at r′ = 4, adding noise is not optimal.The reported Pareto curve shows moderate noise, such as βt = 0.2, dominating runs without added noise, with frontier accuracy losses of only 2%.Source: §5.2 · p. 10, Figure 6 · p. 10, Figure 7 · p. 11

6 CONCLUSIONS: ARE RECURRENT-DEPTH TRANSFORMERS SECRETLY CONTINUOUS LANGUAGE DIFFUSION MODELS?

Diffusion forcing can be applied directly to existing recurrent-depth models, yielding faster single-sequence inference. The authors also frame these causal models as continuous latent language diffusion models, while noting that diffusion modeling’s advantage remains elusive without further training.

6 CONCLUSIONS: ARE RECURRENT-DEPTH TRANSFORMERS SECRETLY CONTINUOUS LANGUAGE DIFFUSION MODELS?

  • Recurrent-depth models can be viewed as continuous latent language diffusion models trained with truncated unrolling, but they remain causal without further training.The authors suggest unrolling objectives could be competitive objectives for future language diffusion models; the advantage of diffusion modeling remains elusive in the causal models discussed.Source: §6 · p. 11

A APPENDIX

Algorithm 2 implements diffusion-style generation with latent-difference-based freezing, using adaptive exits to finalize stable tokens while extending the sequence.

A.1 ADDITIONAL ALGORITHM DETAILS

  • The algorithm takes a prompt, token budget, recurrence and diffusion-step counts, initialization scale, and exit threshold as inputs.It initializes frozen and current token sequences and latent states before iterating through generation steps.Source: §A.1 · p. 16
  • At each step, the algorithm conditions on the current sequence, updates latent states through recurrence, and samples a token from the resulting output.It computes relative latent changes per position and checks whether any position falls below the exit threshold.Source: §A.1 · p. 16
  • When a position meets the threshold, the algorithm freezes the sequence through the last qualifying position, retains the unfrozen latent tail, and appends a new latent state.Generation stops when the frozen sequence reaches the requested number of new tokens; otherwise, no tokens are frozen that step.Source: §A.1 · p. 16

A.2 ADDITIONAL VARIANTS

The study focuses on single-sequence inference and leaves batched inference to future work because implementing the required inference engine would be complex.

A.2 ADDITIONAL VARIANTS

  • Batched or continuously batched inference was not implemented; the study focuses on single-sequence inference because an efficient inference engine would be complex.The proposed design could reserve hidden-state slots up to a multiple of maximum wavefront size and schedule recurrent updates alongside sequence updates.Source: §A.2 · p. 16

A.3 ADDITIONAL INFORMATION.

The appendix reports a finetuned math-model setup, additional hyperparameter findings, and evaluation details. It also specifies the datasets, scoring procedures, and sampling settings used in evaluation.

A.3 ADDITIONAL INFORMATION.

  • The math model was finetuned for one epoch with a peak learning rate of 5 × 10−7, randomized unrolling with mean r = 32, and r sampled from an Exponential distribution.The model was finetuned on MetaMath; the authors note that most gains can also be achieved by finetuning only the adapter that maps inputs and states into the recurrent block.Source: §A.3 · pp. 16–17
  • Initialization scale had only a minor effect, while continuous compute was less effective because the preceding position used only the result of r′.The continuous-compute variant initializes new states with previously computed states from the position one step to the left.Source: Figure 9 · p. 17
  • GSM8k uses its CoT version with 8 few-shot examples and flexible-extract scoring, while MATH500 and Minerva Math use math-verify grading.MBPP and HumanEval are graded normally; inference uses temperature 0.2 and top-p 0.95.Source: §A.3 · p. 17

A.4 QUALITATIVE EVALUATION

Additional visualizations show sampler progress, including temporary stalls, on a challenging GSM8k validation sequence.

A.4 QUALITATIVE EVALUATION

  • Additional visualizations depict sampler progress on a challenging sequence from the GSM8k validation set.The section provides these views in Figure 12.Source: §A.4 · p. 17
  • The selected sequence is used to show progress or temporary lack of progress by the sampler.Source: §A.4 · p. 17
  • Figure 12 supplies several additional views of the sampler's behavior on that sequence.Source: §A.4 · p. 17

B THEORETICAL ANALYSIS

The appendix formalizes recurrent-model depth and width, then presents visualizations of sampler trade-offs, a hyperparameter failure, and robustness on a finetuned math model.

B.1 PROBLEM FORMULATIONS

  • Depth is the number of serial Transformer-block forward passes, while width counts active hidden states processed in parallel at a generation step.Source: §B.1 · p. 17
  • Each generation step adds one Transformer-block pass, and width changes with token entries and hidden-state exits; all active states share the same depth.Source: §B.1 · pp. 17, 19, 20
  • Figure 10 plots accuracy and throughput across noise and inner-recurrence settings, while Figure 11 adds visualizations of their trade-off.Source: Figure 10 · p. 18, Figure 11 · p. 18
  • In Figure 12, hard early token decisions stall progress, then progress resumes; meanwhile, large wavefront regions decode whitespace and compute no useful state information until earlier tokens resolve.Source: Figure 12 · p. 19
  • Figure 13 shows largely similar hyperparameter-robustness behavior for the finetuned GSM8k math model despite its noticeably changed capability.Source: Figure 13 · p. 19

B.2 LLMS SHOULD PRIORITIZE DEPTH SCALING DURING PREFILLING.

The section defines width scaling by replicating tokens and processing the copies in parallel without changing Transformer parameters. Two attention variants preserve serial depth while increasing width by the scaling factor.

B.2 LLMS SHOULD PRIORITIZE DEPTH SCALING DURING PREFILLING.

  • Width scaling replicates each input token s times, applies an unchanged Transformer block with a customized attention mask, then reduces outputs back to the original sequence length.Source: §B.2 · p. 20
  • Width-NoShare lets each copy attend to prior tokens' copies and earlier copies of its own token, whereas Width-KVShare attends only to prior tokens' last copies and earlier same-token copies.Source: §B.2 · p. 20
  • Both variants keep serial depth fixed and increase width by factor s; their attention patterns differ, but not the number of parallel hidden states computed.Source: §B.2 · p. 20
Loading 2510.14961v1…