Reproduce “dUltra: Ultra-Fast Diffusion Language Models via Reinforcement Learning” on OpenResearch

Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.

Universal app for Apple silicon and Intel. Requires macOS 11 or later.

About this paper

Shirui Chen, Jiantao Jiao, Lillian J. Ratliff, Banghua Zhu

dUltra trains an unmasking planner alongside a masked diffusion language model, using on-policy reinforcement learning to choose which masked positions to fill together. On GSM8K, its LLaDA-8B-Instruct-based math model reaches 77.65% accuracy at 16.38 NFE, compared with Fast-dLLM at 75.66% and 54.43 NFE. These results indicate fewer denoising steps, not measured wall-clock speedup.

Key insights

  • On-policy reinforcement learning can effectively learn dynamic, adaptive unmasking strategies, leading to more efficient and higher-quality parallel generation than prior off-policy distillation or fixed-heuristic methods.
  • Incorporating an on-policy distillation reward is crucial for encouraging the model to generate coherent text and perform mode-seeking behavior, even when unmasking potentially dependent tokens in parallel.
  • Advantage clipping within the GRPO framework is essential for training stability, preventing the unmasking planner from collapsing to degenerate policies and enabling robust learning of balanced unmasking strategies.

Summary by alphaXiv.

Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.