Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DiFFPO: Training Diffusion LLMs to Reason Fast and Furious via Reinforcement Learning

Hanyang Zhao, Dawen Liang, Wenpin Tang, David Yao, Nathan Kallus

TL;DR

This work addresses latency- and accuracy-related limitations in diffusion LLMs by introducing DiFFPO, a post-training RL pipeline for masked diffusion LLMs. It mitigates the difficulty of RL optimization with a two-times mean-field (2-MF) surrogate likelihood and importance sampling to better approximate the true dLLM policy, addressing the mismatch observed in prior methods. It further enhances efficiency by jointly training the model with an adaptive, EB-style sampler that learns a prompt-conditioned inference threshold $\gamma$ treated as an extra token during optimization. On open-source diffusion LLMs evaluated on math and planning benchmarks, DiFFPO achieves higher accuracy with fewer function evaluations (NFEs), pushing the Pareto frontier of inference-time compute for diffusion LLMs.

Abstract

We propose DiFFPO, Diffusion Fast and Furious Policy Optimization, a unified framework for training masked diffusion large language models (dLLMs) to reason not only better (furious), but also faster via reinforcement learning (RL). We first unify the existing baseline approach such as d1 by proposing to train surrogate policies via off-policy RL, whose likelihood is much more tractable as an approximation to the true dLLM policy. This naturally motivates a more accurate and informative two-stage likelihood approximation combined with importance sampling correction, which leads to generalized RL algorithms with better sample efficiency and superior task performance. Second, we propose a new direction of joint training efficient samplers/controllers of dLLMs policy. Via RL, we incentivize dLLMs' natural multi-token prediction capabilities by letting the model learn to adaptively allocate an inference threshold for each prompt. By jointly training the sampler, we yield better accuracies with lower number of function evaluations (NFEs) compared to training the model only, obtaining the best performance in improving the Pareto frontier of the inference-time compute of dLLMs. We showcase the effectiveness of our pipeline by training open source large diffusion language models over benchmark math and planning tasks.

DiFFPO: Training Diffusion LLMs to Reason Fast and Furious via Reinforcement Learning

TL;DR

This work addresses latency- and accuracy-related limitations in diffusion LLMs by introducing DiFFPO, a post-training RL pipeline for masked diffusion LLMs. It mitigates the difficulty of RL optimization with a two-times mean-field (2-MF) surrogate likelihood and importance sampling to better approximate the true dLLM policy, addressing the mismatch observed in prior methods. It further enhances efficiency by jointly training the model with an adaptive, EB-style sampler that learns a prompt-conditioned inference threshold treated as an extra token during optimization. On open-source diffusion LLMs evaluated on math and planning benchmarks, DiFFPO achieves higher accuracy with fewer function evaluations (NFEs), pushing the Pareto frontier of inference-time compute for diffusion LLMs.

Abstract

We propose DiFFPO, Diffusion Fast and Furious Policy Optimization, a unified framework for training masked diffusion large language models (dLLMs) to reason not only better (furious), but also faster via reinforcement learning (RL). We first unify the existing baseline approach such as d1 by proposing to train surrogate policies via off-policy RL, whose likelihood is much more tractable as an approximation to the true dLLM policy. This naturally motivates a more accurate and informative two-stage likelihood approximation combined with importance sampling correction, which leads to generalized RL algorithms with better sample efficiency and superior task performance. Second, we propose a new direction of joint training efficient samplers/controllers of dLLMs policy. Via RL, we incentivize dLLMs' natural multi-token prediction capabilities by letting the model learn to adaptively allocate an inference threshold for each prompt. By jointly training the sampler, we yield better accuracies with lower number of function evaluations (NFEs) compared to training the model only, obtaining the best performance in improving the Pareto frontier of the inference-time compute of dLLMs. We showcase the effectiveness of our pipeline by training open source large diffusion language models over benchmark math and planning tasks.

Paper Structure

This paper contains 9 sections, 2 theorems, 16 equations, 4 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. DiFFPO
  4. Improved techniques in efficient model post-training
  5. Unlocking better adaptive multi-token prediction via sampler post-training
  6. Experiments
  7. Experimental Setup.
  8. Related Works
  9. Discussion and Future Works

Key Result

Theorem 1

Assume two parameterized policies by different model family yet share the same parameters being $\pi_{\theta},\hat{\pi}_{\theta}$ respectively, with $\theta\in\Theta$. Further assuming that reward function is postive and upper bounded, i.e. there exists $M>0$, such that $0\leq r(c,o)\leq M$ for any

Figures (4)

  • Figure 1: Benchmark results of RL post-trained models across different math and planning tasks.
  • Figure 2: Generation sample and average KL divergence error at each decoding timestep (100 prompts per dataset).
  • Figure 3: Benchmark results ($\mathbin{\color{blue}{-}}$: DiFFPO; $\mathbin{\color{red}{-}}$: DiFFPO without IS; $\mathbin{\color{orange}{-}}$: d1 baseline).
  • Figure 4: Inference-time frontier obtained by models by DiFFPO with or without sampler training and EB sampler with different inference threshold.

Theorems & Definitions (3)

  • Theorem 1
  • Theorem 3
  • proof