Research map / Image generation / Fast diffusion sampling

Looped Diffusion Language Models

NeurIPS 2026 · Image generation › Fast diffusion sampling · approach: Masked diffusion language model decoding · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes LoopMDM, which selectively loops two early-middle transformer layers of a masked diffusion language model with stochastic loop counts, matching MDM quality with up to 3.3x fewer training FLOPs and improving GSM8K by up to 8.5 points.

Problem

Masked diffusion language models use standard transformer depth, and it is unclear whether looping shared layers helps them at matched training compute, where prior looped-transformer work has focused on parameter efficiency or inference refinement.

Approach

Loop only a small early-middle block of the MDM denoiser during training with a uniformly sampled loop count, keeping parameters fixed; the number of loops can be varied at inference, and adaptive stopping based on hidden-state change allocates loops across timesteps.

Key results

• LoopMDM matches the test NLL of an equal-size MDM with up to 3.34x fewer training FLOPs on LM1B, 2.95x on OpenWebText, and 2.34x on FineWeb-Edu.
• On GSM8K, LoopMDM at S=16 improves over the 14-layer baseline by +8.5 (Top-2) and +8.1 (Top-3) points, and beats a 21-layer baseline with comparable per-step FLOPs by +3.4 at S=16.
• Adaptive looping with epsilon=0.1 reduces average loop count from 12 to about 5 while keeping zero-shot PPL (51.7 to 52.9) and downstream accuracy (41.4 to 41.3) nearly unchanged.

Surprising findings

• Input injection of Gaussian noise into the mask embedding degrades performance relative to the standard MDM baseline.
• Looping more than six layers underperforms the non-looped baseline under matched compute.

Limitations

• Experiments are at small scale (125M-170M parameters) and the largest downstream gains are on GSM8K from a single small training setup.
• Downstream evaluation uses only FineWeb-Edu checkpoints with zero-shot likelihood evaluation and no fine-tuning.

Experimental setup

Datasets: OpenWebText LM1B FineWeb-Edu TinyGSM GSM8K Sudoku (100k)

Benchmarks: GSM8K LM1B OpenWebText PTB WikiText Lambada AG News PubMed Arxiv ARC-c ARC-e BoolQ

Metrics: test NLL zero-shot perplexity generative perplexity accuracy

Baselines: MDM (non-looped, same size) deeper non-looped MDMs (14/18/21 layers) log-normal Poisson loop sampling input injection

Builds on: masked diffusion models looped transformers Universal Transformer TinyGSM

Closest papers by meaning

• Scaling up Masked Diffusion Models on Text (ICLR 2025 · similarity 0.93)
• On the Reasoning Abilities of Masked Diffusion Language Models (ICLR 2026 · similarity 0.91)
• Theoretical Benefit and Limitation of Diffusion Language Model (NeurIPS 2025 · similarity 0.90)
• Sparse Layers are Critical to Scaling Looped Language Models (NeurIPS 2026 · similarity 0.90)
• Simple and Effective Masked Diffusion Language Models (NeurIPS 2024 · similarity 0.90)
• From Next-Token to Next-Block: A Principled Adaptation Path for Diffusion LLMs (NeurIPS 2026 · similarity 0.90)