PaperlayerView exact revision on arXiv

Source-linked AI summary

LFPO: Likelihood-Free Policy Optimization for Masked Diffusion Models

Chenxing Wei, Jiazhen Kang, Hong Wang, Jianqing Zhang, Hao Jiang, Xiaolong Xu, Ningyuan Sun, Ying He, F. Richard Yu, Yao Shu, Bo Jiang

arXiv:2603.01563v1cs.LGcs.AI

TL;DR

Applying reinforcement learning to masked diffusion language models is limited by intractable exact likelihoods and inaccurate approximations. LFPO instead rectifies denoising logits as discrete velocity fields, achieving state-of-the-art code and reasoning performance while requiring fewer generation steps.

  • Problem

    Exact likelihoods for masked diffusion language models are mathematically inaccessible, forcing existing policy-gradient methods to use costly, often inaccurate approximations.

  • Method

    LFPO treats denoising logits as discrete velocity fields and directly rectifies them toward high-reward trajectories through contrastive regression.

  • Results

    LFPO achieves state-of-the-art performance across code and reasoning benchmarks, including an average code score of 60.8 versus AGRPO’s 60.6.

  • Takeaways & Limitations

    Consistency-aware flow rectification produces high-quality generation with fewer iterative steps than likelihood-based approaches.

Abstract

from arXiv · show

Reinforcement Learning with Verifiable Rewards (RLVR) has achieved remarkable success in improving autoregressive models, especially in domains requiring correctness like mathematical reasoning and code generation. However, directly applying such paradigms to Diffusion Large Language Models (dLLMs) is fundamentally hindered by the intractability of exact likelihood computation, which forces existing methods to rely on high-variance approximations. To bridge this gap, we propose Likelihood-Free Policy Optimization (LFPO), a native framework that maps the concept of vector field flow matching to the discrete token space. Specifically, LFPO formulates alignment as geometric velocity rectification, which directly optimizes denoising logits via contrastive updates. This design effectively bypasses the errors inherent in likelihood approximation, yielding the precise gradient estimation. Furthermore, LFPO enforce consistency by predicting final solutions from intermediate steps, effectively straightening the probability flow to enable high-quality generation with significantly fewer iterations. Extensive experiments demonstrate that LFPO not only outperforms state-of-the-art baselines on code and reasoning benchmarks but also accelerates inference by approximately 20% through reduced diffusion steps.

1. Introduction

The introduction identifies intractable likelihood estimation as the central obstacle to applying RLVR to masked diffusion models and motivates LFPO as a likelihood-free alternative. LFPO connects continuous flow matching with discrete denoising logits, using contrastive rectification and consistency-aware training to improve generation and reduce diffusion steps.

  • Problem: Exact likelihoods are inaccessible for dLLMs’ holistic denoising process, forcing existing alignment methods to use likelihood approximations.This contrasts with autoregressive models, whose sequential inference supports straightforward likelihood computation.
  • Theoretical foundation: Denoising logits for masked tokens are identified as the discrete projection of the continuous flow-matching velocity field.This theoretical isomorphism justifies rectifying generation trajectories without likelihood-based policy gradients.
  • Method: LFPO provides a native reinforcement-learning framework for masked diffusion models with stable, efficient off-policy training.Its likelihood-free design operates directly in logit space rather than estimating policy densities.
  • Method: LFPO directly optimizes denoising logits with a contrastive objective that pulls predictions toward positive outcomes and repels them from negative ones.The formulation avoids density approximation, complex ODE backtracking, and noisy likelihood-estimation gradients.
  • Results: LFPO achieves state-of-the-art results on mathematical reasoning and code-generation benchmarks while enabling high-quality generation with fewer iterative steps.Consistency-aware training supports the reduction in diffusion iterations.

2. Related Work

Related work frames dLLMs as parallel, bidirectional-context generators and surveys reinforcement-learning methods that adapt policy gradients to discrete diffusion. It highlights likelihood-based optimization as a bottleneck and motivates a native formulation grounded in dLLMs’ geometric nature.

  • Diffusion Large Language Models: dLLMs generate text through parallel iterative unmasking rather than strictly left-to-right autoregressive decoding.This supports bidirectional context modeling and may improve decoding efficiency.
  • Reinforcement Learning for Diffusion Alignment: Recent methods adapt policy-gradient reinforcement learning to discrete diffusion, including critic-free GRPO-based optimization and SPG’s upper- and lower-likelihood bounds.SPG addresses gradient bias caused by intractable diffusion-trajectory likelihoods.
  • Reinforcement Learning for Diffusion Alignment: Coupled-GRPO and AGRPO further target stability and sample efficiency, with AGRPO using Monte Carlo sampling for unbiased policy-gradient estimation.The passage characterizes AGRPO as achieving state-of-the-art results.
  • Reinforcement Learning for Diffusion Alignment: Existing approaches predominantly maximize the likelihood or surrogates of high-reward trajectories, which the paper identifies as a bottleneck for dLLMs.The paper argues that dLLMs instead require a native reinforcement-learning formulation grounded in their geometric nature.

3. Motivation: A Flow Matching Perspective

The section reframes diffusion-model alignment through flow matching, lifting discrete tokens into a probability-simplex geometry where dLLM logits parameterize a velocity field. It shows that cross-entropy training is gradient-aligned with velocity-field optimization, motivating LFPO’s direct contrastive rectification without likelihood estimation.

  • Continuous Flow Matching: Flow matching trains a neural vector field to match a target conditional velocity along a probability path from prior to data.The framework uses a linearly interpolated Conditional Optimal Transport path and minimizes expected velocity-matching error.
  • Discrete Flow Matching: Discrete lifting represents vocabulary tokens as simplex vertices, masked states as a prior, and model predictions as probability distributions within the simplex.The mask prior serves as the geometric origin, while data tokens are deterministic one-hot vectors.
  • Discrete Flow Matching: A dLLM’s velocity field is the displacement from the mask prior to its vocabulary distribution, making logits a parameterization of generative flow.This interpretation extends logits beyond classification scores to geometric directions driving generation.
  • Gradient Alignment: Theorem 3.1 shows that cross-entropy gradients with respect to logits exactly equal the residual between model and target velocities.Thus, standard dLLM training is dynamically optimization-equivalent to minimizing velocity-field error because both objectives share the identical gradient direction.
  • LFPO Motivation: This equivalence motivates LFPO, which bypasses likelihood estimation by pulling logits toward high-reward trajectories and pushing them away from low-reward trajectories.The proposed alignment performs direct velocity-field rectification through a contrastive objective.

4. Policy Alignment via Velocity Rectification

LFPO aligns dLLM policies by rectifying discrete velocity fields in logit space, replacing unavailable ground-truth supervision with reward-conditioned implicit targets. Stratified sampling and block-wise accumulation make this optimization more stable and memory-efficient.

  • Velocity Rectification: LFPO relies on Gradient Equivalence, establishing that cross-entropy optimization is equivalent to rectifying the discrete velocity field on the probability simplex.This provides the theoretical basis for applying reinforcement learning to dLLMs.
  • Implicit Target Policies: In RLVR, scalar trajectory rewards replace the explicit ground-truth velocity target available in supervised training.The model samples a trajectory τ and receives only r(τ), creating the need for an implicit supervision target.
  • Implicit Target Policies: LFPO defines velocity deviation as Δ(x_t) = log π_θ(·|x_t) − log π_ref(·|x_t), then constructs positive and negative implicit policies by amplifying or reversing it.The positive policy pulls the current model away from the reference in the deviation direction, while the negative policy pushes oppositely.
  • Reward-Conditioned Optimization: Reward-weighted cross-entropy dynamically interpolates between the implicit targets: high rewards align with π+ while low rewards align with π−.As r approaches 1, optimization favors the positive target; as r approaches 0, it favors the negative target.
  • Efficient Optimization: Stratified Trajectory Sampling covers K uniform timestep segments to reduce high-variance gradients, while Block-wise Gradient Accumulation addresses the B × N × K memory bottleneck.The method aggregates gradients from multiple trajectories and representative timesteps in smaller GPU-manageable blocks.

5. Empirical Results

LFPO is evaluated on code generation and mathematical reasoning against state-of-the-art reinforcement-learning baselines, where it improves downstream performance while reducing inference and training costs. Ablations further examine the attraction and repulsion components of its objective.

  • Experimental Setup: LFPO is evaluated with DiffuCoder for code and LLaDA 8B for reasoning against Diffu/Coupled-GRPO, UniGRPO, SPG, and AGRPO.
  • Main Results: LFPO achieves a 60.8 average code-generation score, surpassing AGRPO at 60.6, and scores 75.6 on HumanEval.
  • Main Results: LFPO bypasses likelihood approximation by directly regressing denoising behavior, producing more accurate, lower-variance gradients than surrogate or importance-sampling objectives.
  • Inference Efficiency: 41.8 steps are reduced for code and 159.0 for reasoning versus the base model, whereas AGRPO increases MATH inference by +73.6 steps.
  • Training Convergence: 8.0× faster convergence is achieved on HumanEval and MATH, while Hellaswag reaches the baseline peak 4.4× faster.
  • Ablation Study: Both Pos. Only and Neg. Only improve over the base model, but neither matches the full All Loss objective.

6. Conclusions

LFPO aligns diffusion language models while bypassing intractable likelihoods through efficient stratified sampling and block-wise optimization. It achieves superior code and reasoning benchmark performance while accelerating training convergence and inference generation.

  • Contributions: LFPO aligns dLLMs while bypassing intractable likelihoods.The method is supported by efficient stratified sampling and block-wise optimization.
  • Empirical results: LFPO achieves superior performance across code and reasoning benchmarks.
  • Efficiency: LFPO significantly accelerates both training convergence and inference generation.

A. Detailed Derivation of the Cross-Entropy Gradient

This section derives the Cross-Entropy gradient with respect to pre-softmax logits and formally proves that standard dLLM training follows the same optimization direction as the residual error vector in FM.

  • Gradient derivation: The section provides a step-by-step derivation of the Cross-Entropy loss gradient with respect to pre-softmax logits.It focuses specifically on how the loss changes with the model’s logits before softmax normalization.
  • Gradient derivation: The derivation formally proves that standard dLLM training’s optimization direction is mathematically identical to FM’s residual error vector.This establishes the correspondence between dLLM optimization and the residual-error direction used in flow matching.
  • Gradient derivation: The result connects Cross-Entropy optimization in standard dLLM training with the residual error vector in FM.The connection is made through the derived gradient with respect to pre-softmax logits.

A.1. Definitions and Notation

This section defines the vocabulary size and the core variables used in the paper: logits, their Softmax probabilities, and one-hot ground-truth labels. It also introduces the single-sample Cross-Entropy loss and its derivative with respect to an arbitrary logit.

  • Core notation: The vocabulary contains V tokens, establishing the dimensionality of the model’s output space.The notation uses V as the vocabulary size.
  • Core notation: The model outputs pre-softmax logits z ∈ R^V, with z_i denoting the logit for token i.Each logit corresponds to one vocabulary token before probability normalization.
  • Core notation: Applying Softmax to z yields the probability distribution p ∈ Δ^(V−1).The probability vector lies on the (V−1)-dimensional simplex.
  • Core notation: The one-hot ground-truth vector y ∈ {0,1}^V assigns y_c = 1 to the correct class c and sums to one across classes.All other class entries are zero under the one-hot encoding.
  • Loss notation: The section defines L as the single-sample Cross-Entropy loss and considers its derivative with respect to an arbitrary logit z_i.These quantities support the subsequent optimization analysis.

A.2. Step 1: Derivative of the Softmax Function

This step derives the partial derivative of the softmax output p_j with respect to logit z_i using the quotient rule. It then unifies the two resulting cases with the Kronecker delta δ_ij.

  • Step 1: Derivative of the Softmax Function: The partial derivative of softmax output p_j with respect to logit z_i is computed using the quotient rule.The derivation distinguishes between two cases.
  • Step 1: Derivative of the Softmax Function: The two cases are unified using the Kronecker delta δ_ij, where δ_ij = 1 if i = j and 0 otherwise.This notation combines the case-specific derivative expressions into one form.

A.3. Step 2: Applying the Chain Rule

This section applies the chain rule to compute the gradient of loss L with respect to z_i by summing contributions over all p_j. It simplifies the two summation terms using the one-hot or probability-distribution property of y.

  • Applying the Chain Rule: The gradient of L with respect to z_i is obtained by summing chain-rule contributions over all p_j because L depends on every p_j.The derivation explicitly applies the chain rule and sums over j.
  • Analyzing the Summation Terms: The first summation term is non-zero only when j = i, so it simplifies to y_i.This follows from the factor δ_ij.
  • Analyzing the Summation Terms: The second summation term factors out p_i because p_i does not depend on j, and the remaining sum over y_j equals 1.The passage attributes this to y being a one-hot vector or valid probability distribution.
  • Applying the Chain Rule: Substituting the simplified terms back into the chain-rule expression produces the resulting gradient form.The supplied passage indicates that the simplified terms are substituted back into the derivation.

A.4. Conclusion: Vector Field Interpretation

The conclusion establishes that cross-entropy optimization in dLLMs is dynamically equivalent to minimizing flow-matching velocity error, yielding an implicit discrete velocity-field matching interpretation on the simplex.

  • The gradient of the Cross-Entropy loss is expressed in vector notation.
  • The model and target velocity fields are defined as vθ = p−m and ut = y−m, respectively.
  • Minimizing Cross-Entropy is dynamically equivalent to minimizing flow-matching velocity error, confirming that dLLMs implicitly match the discrete velocity field on the simplex.
Loading 2603.01563v1…