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

TL;DR

This paper tackles the slow generation of autoregressive recurrent-depth transformers by introducing a diffusion-forcing sampler that enables diagonal, parallel generation across the sequence while preserving left-to-right causality. It formalizes the Huginn-0125-style architecture, develops a simplified diffusion-based sampling algorithm with adaptive exits and stabilization, and provides a theoretical framework that argues for depth-focused prefilling and width-focused decoding under diffusion forcing. Empirically, the method yields up to approximately $5\times$ speedups on a 3.5B-parameter recurrent-depth model with minimal accuracy loss across multiple benchmarks and model variants, validating both the practical efficiency and the expressiveness of the approach. The work also reframes recurrent-depth transformers as continuous latent diffusion language models, guiding future training objectives and inference strategies to exploit this connection while maintaining causality.

Abstract

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.

Figures (13)

• Figure 1: Different generation schemes for autoregressive, recurrent-depth models. Left: Standard sequential generation, which proceeds one token and step of the recurrence at a time (time steps denoted by integers). Right: A diffusion forcing sampler used for the same model can parallelize generation "diagonally", by computing one step of the recurrence per token position, iteratively refining its estimate of the generated sequence.
• Figure 2: An example of a text sequence being generated with the proposed diffusion forcing sampler from a depth-recurrent model. While the original recurrent-depth model requires 32 recurrence steps to produce a single token (the default for this model), the diffusion sampler has already produced and committed 8 new tokens (green). As described, the sampler advances by at least one token per step of the recurrence. Decoded candidate tokens are initial spell out incoherent text, but map into the right concepts, and quickly improve with more steps. Note that the "freeze" decision is dynamic, based on distance to the previous state in latent space (not pictured).
• Figure 3: The Huginn-0125 recurrent-depth model can match the baseline performance on the GSM8k dataset when enabling KV cache sharing (with a minimal cache size of 1), using $r$-times less memory for KV states.
• Figure 4: Examples of adaptive sampler behavior. Each color represents a token id in the vocabulary of the model, showing the development of the generated sequence (running left to right) as a function of sampler steps (running top to bottom) for different hyperparameter choices. The leftmost example is $r'=4$, and tokens are frozen quickly, whereas middle and right show sequences with $r<4$ require more adaptive computation, and in both cases the sampler stalls after hitting the maximal length of the wavefront (here 32 to visualize), before resolving the sequence and advancing again.
• Figure 5: Trade-off between accuracy and speed on GSM8k under different hyperparameter choices. Left: Effect of increasing inner recurrence $r'$. Inner recurrence stabilizes the sampling, increasing accuracy at the cost of throughput. Right: Effect of varying the exit threshold $\varepsilon$. Modulating the exit threshold most directly trades off throughput and accuracy.

Theorems & Definitions (12)

• Remark 3.1: Convergence of the Adaptive Diffusion Sampler
• Remark 3.2: Computational Cost
• Definition 4.1: Depth and Width in Recurrent-Depth Models, informal
• Theorem 4.2: Depth vs. Width Scaling in Prefilling, informal
• Remark 4.3
• Theorem 4.4: Depth vs. Width Scaling in Decoding, informal
• Remark 4.5
• Definition B.1: Depth and Width in Recurrent-Depth Models
• Remark B.2
• Definition B.3: Width Scaling Variants