Summary (Overview)

• This paper investigates selective looping in Uniform Diffusion Language Models (UDLMs), where a small block of Transformer layers is recursively applied within each denoising step, keeping the training objective and sampler unchanged.
• Training efficiency gains: Looped UDLMs reach the validation NLL of a non-looped UDLM trained on ~0.5 trillion tokens with up to 2.8× fewer training FLOPs on OpenWebText and 3.3× fewer on FineWeb-Edu.
• Inference improvements: At matched NFE (number of function evaluations), looped UDLMs achieve up to 9.8% lower generative perplexity than non-looped models; on seven out-of-domain datasets, zero-shot perplexity drops by 11.5% on average.
• Key insight on loop benefits: Extra loop iterations primarily improve predictions at replaced (corrupted) tokens, not the decision of which tokens to keep, with gains peaking at intermediate noise levels.
• Adaptive looping efficiency: Using a convergence-based stopping criterion, adaptive looping matches the accuracy of 12 fixed loop iterations while running 26–35% fewer loop iterations, saving computation without quality loss.

Introduction and Theoretical Foundation

Background

Discrete diffusion language models (dLLMs) generate text by denoising corrupted sequences in parallel. Two main paradigms exist:

1. Masked Diffusion Language Models (MDLMs): Replace corrupted tokens with a dedicated mask symbol, so corrupted positions are always known. Unmasked tokens are never changed during the reverse process.
2. Uniform Diffusion Language Models (UDLMs): Replace tokens with random vocabulary tokens, giving no explicit signal about which positions are corrupted. Any position can be revised at any denoising step.

The Core Challenge in UDLMs

Denoising in UDLMs requires predicting corrupted tokens from a context that may itself be corrupted. At each position, the denoiser must simultaneously learn two things:

• The probability that the observed token was retained (reliability), which depends on noise level and contextual agreement
• The clean token itself, in case it was replaced

Both are derived from the model's contextual prediction of the clean token. Looping—reusing existing layers for more computation without adding parameters—offers a way to refine these predictions without increasing model size.

Theoretical Foundation: The Reliability Decomposition

The key theoretical insight is the exact reliability decomposition (Equation 6):

$$\gamma_t^{\ell} := \operatorname{Pr}(C_t^{\ell} = 1 \mid \mathbf{z}_t^{1:L}, t) = \frac{\alpha_t \pi_t^{\ell}(i)}{\alpha_t \pi_t^{\ell}(i) + (1 - \alpha_t) \nu(\mathbf{e}_i)}$$

For UDLMs, where $\nu(\mathbf{e}_i) = 1/K$, the retention odds become:

$$\frac{\gamma_t^{\ell}}{1 - \gamma_t^{\ell}} = \frac{K \alpha_t}{1 - \alpha_t} \pi_t^{\ell}(i)$$

This shows that retention probability must be inferred from contextual fit—it is not marked in the input. The clean-token posterior is a mixture of a point mass on the observed token and the contextual prediction, weighted by retention probability.

Methodology

Architecture: Looped Denoiser

The DUO++ denoiser (a Diffusion Transformer, DiT) is split into three parts:

• Head $H_\theta$: $n_H$ blocks
• Middle slice $M_\theta$: $n_B$ blocks (recurrently applied)
• Tail $T_\theta$: $n_T$ blocks with output projection

The looped computation (Equation 8):

$$\mathbf{h}^{(r+1)} = M_\theta(\mathbf{h}^{(r)}, t), \quad r = 0, \ldots, S - 1$$ $$\widehat{\mathbf{x}} = \operatorname{softmax}(T_\theta(\mathbf{h}^{(S)}))$$

With $S$ passes, effective depth is $n_H + S \cdot n_B + n_T$ blocks, but parameters are unchanged ($S = 1$ is the original network).

Loop-Count Sampling Distributions

Three loop-count training distributions are studied:

1. Uniform sampling ($S \sim \mathcal{U}\{1, \ldots, S_{max}\}$): Each loop count trained equally often; most training FLOPs go to large loop counts.

2. FLOP-proportional sampling (novel): Balances training FLOPs across loop counts:

   $$p_{\mathrm{fp}}(S) = \frac{F(S)^{-1}}{\sum_{r=1}^{S_{\max}} F(r)^{-1}}$$

   where $F(S) = F_{\mathrm{base}}(1 + (S-1)\rho)$ and $\rho = n_\mathcal{B} F_{\mathrm{layer}}/F_{\mathrm{base}}$. Smaller loop counts are sampled more often.

3. Noise-proportional sampling (novel): Assigns larger loop counts to noisier inputs by sorting microbatches into noise bands.

Experimental Setup

• Model: 170M-parameter DUO++ small DiT (12 blocks, width 768, sequence length 1,024, GPT-2 tokenizer)
• Loop configuration: Blocks 1–2 looped with $S_{max} = 12$ (justified by ablation studies)
• Datasets: OpenWebText (OWT) and FineWeb-Edu sample-10BT subset
• Baseline: Non-looped UDLM trained for 1M optimizer steps, batch size 512
• Iso-FLOP comparison: All runs compared at matched training FLOPs

Adaptive Looping at Inference

A sequence keeps looping until its hidden state converges:

$$\|\mathbf{h}^{(r)} - \mathbf{h}^{(r-1)}\|_F / \|\mathbf{h}^{(r)}\|_F < \epsilon = 0.1$$

with a cap at $S_{max} = 12$.

Empirical Validation / Results

Finding 1: Training Compute Efficiency

Both looped models (uniform and FLOP-proportional) stay below the non-looped baseline at every training FLOP budget on both datasets:

• OpenWebText: 2.7–2.8× fewer FLOPs to match baseline validation NLL
• FineWeb-Edu: 3.3× fewer FLOPs
• Final NLL at matched FLOPs: 3.154 (uniform) vs. 3.159 (FLOP-proportional) on OWT

Finding 2: Noise-Proportional Sampling Underperforms

The noise-proportional distribution offers no benefit—it trails FLOP-proportional sampling by +0.029 NLL at matched compute. Loop gains peak at intermediate diffusion times, not the noisiest inputs.

Finding 3: Inference-Time Generation Gains

At matched NFE:

• Up to 9.8% lower generative perplexity vs. non-looped baseline
• Most gains captured by $S = 6$–$8$ (moderate loop counts)
• No reduction in unigram entropy (diversity preserved)
• Adaptive looping averages ~5.5 iterations and matches or beats fixed $S = 6$

Finding 4: Zero-Shot Likelihood Requires Inference Looping

Table 1: Zero-shot perplexity (↓) on out-of-domain datasets

Key findings:

• At $S = 1$, looped checkpoints are worse than baseline on every dataset—looping during training alone doesn't help
• Looping at inference reverses this: every $S > 1$ and adaptive looping beats baseline everywhere
• 93–94% of perplexity reduction comes from better contextual prediction at replaced tokens (not retention decisions)

Finding 5: Task-Dependent Trade-offs from FLOP-Proportional Training

• Accuracy: FLOP-proportional checkpoint beats uniform on all ten zero-shot MCQ tasks; adaptive looping matches fixed $S = 12$ with 34.8% fewer loop iterations
• Efficiency: Uses 12% fewer loop iterations than uniform checkpoint under adaptive looping (7.82 vs. 8.89)
• Trade-off: Gen. PPL is 1.55% higher on average than uniform checkpoint at matched loop counts

Where Adaptive Looping Stops

Adaptive looping stops:

• Early at low and high noise levels
• Longest at intermediate diffusion times ($t = 0.5$–$0.7$), where loop gains peak
• Keeps 88–98% of the $S = 12$ gain at the stop

Theoretical and Practical Implications

Theoretical Implications

1. Looping replaces parameters with computation: The reliability decomposition (Equation 6) shows that UDLM denoising fundamentally requires contextual reasoning about token reliability. Looping allows the model to iteratively refine these contextual predictions, effectively substituting depth-for-parameters.

2. Retention vs. prediction decoupling: The finding that most gain comes from better predictions at replaced tokens (not retention decisions) suggests the model learns reliability quickly (AUC ≥ 0.96 at $S = 1$) but benefits from deeper computation for clean-token estimation.

3. Optimal depth allocation is noise-dependent: The peak benefit at intermediate noise levels (not the noisiest) challenges the intuitive hypothesis that noisier inputs need more computation. This has implications for designing adaptive computation strategies.

Practical Implications

1. Compute savings: Looped UDLMs can match baseline quality with 2.8–3.3× fewer training FLOPs, significantly reducing the cost of training large diffusion language models.

2. FLOP-proportional sampling is a free lunch: It provides better MCQ accuracy and inference efficiency (12% fewer loops) without extra hyperparameters, making it an attractive default for training.

3. Adaptive looping reduces inference cost: Running 26–35% fewer loop iterations while maintaining quality enables cheaper deployment, especially important for large-scale inference.

Conclusion

This work demonstrates that selective looping significantly improves UDLM training efficiency and inference quality. Key takeaways:

1. Looped UDLMs reach non-looped baseline NLL with up to 3.3× fewer training FLOPs
2. Looping at inference lowers generative perplexity by up to 9.8% and zero-shot perplexity by 11.5% on average
3. Adaptive looping matches fixed 12-iteration quality with 26–35% fewer iterations
4. FLOP-proportional sampling offers a no-cost improvement over uniform sampling

Future Directions

• Comparing how looping helps masked vs. uniform diffusion, especially on reasoning tasks
• Exploring complementary approaches (self-conditioning, recurrence across denoising steps) that could combine with in-step looping
• Scaling studies to larger models (billion-parameter scale) to verify efficiency gains
• Investigating whether adaptive strategies based on intermediate hidden states could further improve efficiency