dUltra: Ultra-Fast Diffusion Language Models via Reinforcement Learning

TL;DR

This work addresses the slow sampling of masked diffusion language models by introducing dUltra, an on-policy reinforcement learning framework that jointly learns a lightweight unmasking planner with a base MDLM. Using Group Relative Policy Optimization and a multi-objective reward that includes verifiable task performance, distillation guidance, and efficiency, dUltra learns per-token unmasking probabilities that increase parallelism while preserving output quality. Empirically, it outperforms state-of-the-art heuristic and off-policy distillation baselines on math and coding benchmarks, with substantial reductions in the number of denoising steps, and reveals task-dependent unmasking patterns that resemble a blend of autoregressive and parallel strategies. The results suggest a viable path toward diffusion supremacy by coupling on-policy learning with mode-filter-aware planning for MDLMs.

Abstract

Masked diffusion language models (MDLMs) offer the potential for parallel token generation, but most open-source MDLMs decode fewer than 5 tokens per model forward pass even with sophisticated sampling strategies. As a result, their sampling speeds are often comparable to AR + speculative decoding schemes, limiting their advantage over mainstream autoregressive approaches. Existing distillation-based accelerators (dParallel, d3LLM) finetune MDLMs on trajectories generated by a base model, which can become off-policy during finetuning and restrict performance to the quality of the base model's samples. We propose \texttt{dUltra}, an on-policy reinforcement learning framework based on Group Relative Policy Optimization (GRPO) that learns unmasking strategies for efficient parallel decoding. dUltra introduces an unmasking planner head that predicts per-token unmasking likelihoods under independent Bernoulli distributions. We jointly optimize the base diffusion LLM and the unmasking order planner using reward signals combining verifiable reward, distillation reward, and the number of unmasking steps. Across mathematical reasoning and code generation tasks, dUltra improves the accuracy--efficiency trade-off over state-of-the-art heuristic and distillation baselines, moving towards achieving ``diffusion supremacy'' over autoregressive models.

Figures (8)

• Figure 1: Parallel decoding of dependent ambiguous tokens can lead to unwanted modes.
• Figure 2: Given the same prompt and ground truth, we mask 30% of tokens using both random masking and autoregressive (AR) masking. The left column shows the masked input sequence with random masking (<m> represents the mask token), and the corresponding one-step decoding result. The right columns show the masked input sequence with AR masking, and the corresponding one-step decoding result. We use LLaDA-8B-Instruct for both generations.
• Figure 3: dUltra overview and results. Left: Accuracy versus number of function evaluations (NFE; denoising steps) on math and coding benchmarks for Fast-dLLM, dParallel, d3LLM, and dUltra (ours) variants. Here we use a block size of 128 and a generation length of 256. Right: dUltra architecture: a masked diffusion LM (LLaDA blocks) produces last-layer representations that feed (i) a token decoding head to output logits for masked positions and (ii) an unmasking planner head to output per-position unmasking probabilities; the Bernoulli sampled unmasking decisions update the masked sequence ([M] denotes the mask token).
• Figure 4: Training dynamics on GSM8K.
• Figure 5: Effect of advantage clipping value $C$ on training stability. Without advantage clipping, the average NFE quickly diverges as the planner collapses to never unmasking any tokens. With advantage clipping, training remains stable and the model learns to reduce NFE while maintaining task performance.