PaperlayerView exact revision on arXiv

Source-linked AI summary

Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods

Oussama Zekri, Nicolas Boullé

arXiv:2502.01384v3stat.MLcs.AIcs.CLcs.LG

TL;DR

Fine-tuning discrete diffusion models with policy gradients is challenging because categorical sampling is non-differentiable, limiting reward-based optimization. SEPO provides a policy-gradient framework for non-differentiable rewards, and its experiments report state-of-the-art DNA results alongside evaluations on natural-language tasks.

  • Problem

    Fine-tuning discrete diffusion with rewards is challenging because categorical sampling is non-differentiable, while an existing Gumbel-Softmax approach requires differentiable rewards and becomes memory-intensive at scale.

    Source: §1 · p. 1

  • Method

    SEPO characterizes policy gradients for discrete diffusion and uses a clipped-ratio loss with self-normalized importance sampling to support non-differentiable rewards in conditional and unconditional generation.

    Source: §1 · p. 2

  • Results

    Pred-Activity 7.64 and ATAC-Acc 99.9%: SEPO tops prior RL and guidance baselines on DNA results with lower run-to-run variance.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The experiments demonstrate SEPO's effectiveness for optimizing discrete diffusion models while addressing non-differentiability and combinatorial complexity.

    Source: §5 · p. 10

    The convergence assumptions require uniformly bounded policy gradients and probabilities bounded away from zero.

    Source: §3.4 · p. 8

Abstract

from arXiv · show

Discrete diffusion models have recently gained significant attention due to their ability to process complex discrete structures for language modeling. However, fine-tuning these models with policy gradient methods, as is commonly done in Reinforcement Learning from Human Feedback (RLHF), remains a challenging task. We propose an efficient, broadly applicable, and theoretically justified policy gradient algorithm, called Score Entropy Policy Optimization (\SEPO), for fine-tuning discrete diffusion models over non-differentiable rewards. Our numerical experiments across several discrete generative tasks demonstrate the scalability and efficiency of our method. Our code is available at https://github.com/ozekri/SEPO.

1 Introduction

Fine-tuning discrete diffusion models with rewards is challenging because categorical sampling is non-differentiable, while existing approaches may require differentiable rewards or encounter scalability limits. SEPO is a policy-gradient framework for discrete diffusion, and experiments report state-of-the-art DNA results.

1 Introduction

  • Discrete diffusion models extend diffusion to language and other discrete tasks, offering more flexible generation than next-token prediction.Prior studies report competitiveness with autoregressive models.Source: §1 · p. 1
  • Fine-tuning discrete diffusion is challenging because categorical sampling is non-differentiable; Gumbel-Softmax requires differentiable rewards and becomes memory-intensive at scale.Classifier guidance and steering also often face scalability issues or intractable objectives.Source: §1 · p. 1
  • SEPO characterizes policy gradients for discrete diffusion in the concrete score framework, allowing non-differentiable rewards without steering or guidance for conditional and unconditional generation.The framework is presented as more general than approaches tailored to specific setups.Source: §1 · pp. 1–2
  • SEPO uses a clipped-ratio loss with self-normalized importance sampling for low variance; its gradient-flow alternative improves sample quality at higher complexity.The algorithm is described as efficient and scalable.Source: §1 · p. 2
  • Pred-Activity 7.64 and ATAC-Acc 99.9%: SEPO tops prior RL and guidance baselines on DNA results with lower run-to-run variance.The experiments also cover natural language tasks.Source: §1 · p. 2

2 Background and preliminaries

The background reviews prior fine-tuning approaches, describes discrete diffusion through score-based transition models, and frames reward optimization as a reinforcement-learning problem. The setup permits flexible reward timing, including evaluation on partially denoised samples.

2.1 Related works

  • Prior work applies policy gradients to continuous diffusion, while extending them to discrete diffusion is more challenging; this work adapts them to general policy-gradient algorithms.Inference-time techniques avoid training but typically take longer at inference; masked diffusion uses a weighted cross-entropy objective with a special [MASK] token.Source: §2.1 · p. 2, §2.2 · p. 3

2.2 Discrete Diffusion

  • A discrete diffusion particle follows a continuous-time Markov chain over a finite sequence space, and the score model approximates state-probability ratios to define the resulting process.In practice, scores are available for states differing by one token, representing O(mn) ratios rather than O(m2n).Source: §2.2 · pp. 2–3
  • The forward process maps the data distribution to an easily sampled noisy distribution through transitions specified by a rate matrix.The score-entropy setup also contains masked-diffusion frameworks.Source: §2.2 · pp. 2–3

2.3 Fine-tuning with Reinforcement Learning

  • The optimization can maximize expected rewards or minimize divergence from the pretrained distribution, and typical fine-tuning combines these terms with a weighting factor.The formulation retains general per-step rewards rather than restricting reward to the final denoising step.Source: §2.3 · p. 3
  • Fine-tuning casts denoising as a finite-horizon Markov decision process with known transitions, where the policy is the model and rewards can be assigned at intermediate steps.This supports evaluating a partially denoised sample after fewer denoising steps; reward functions need not be differentiable with respect to a state embedding.Source: §2.3 · pp. 3–4

3 Methods

The methods estimate discrete-diffusion policy gradients using concrete scores and self-normalized importance sampling, then optimize them with clipped policy objectives. A gradient-flow sampling alternative and convergence analysis extend the framework.

3.1 Policy gradients for concrete score

  • For a given y, few sampled states provide accessible neighbors, making the naive estimate of qθ_t(y) high variance.The passages describe importance sampling from qθ_t(· | y) and self-normalized weights based on concrete-score approximations as the remedy.Source: §3.1 · pp. 4–5
  • O(M^-1) variance is proved for the self-normalized estimate, which is consistent and asymptotically unbiased under a mild moment assumption.It uses M proposals from qθ_t(· | y), with weights available up to normalization through the concrete score.Source: §3.1 · pp. 4–5
  • The concrete-score discrete REINFORCE gradient estimator has error at most Rmax C(ρ, G)δ under the stated assumptions.The TRPO-based formulation uses an old policy and a KL constraint; its practical expression is intended to avoid solving a constrained optimization problem at each step.Source: §3.1 · pp. 5–6

3.2 SEPO : Score Entropy Policy Optimization

  • SEPO clips the policy ratio and uses group-standardized rewards as advantages, while supporting conditional and unconditional fine-tuning.The clipping interval is [1 − ϵ, 1 + ϵ], and the group-relative advantage centers and scales each output's reward within its group.Source: §3.2 · p. 6
  • PPO, GRPO, and related methods can share the SEPO loss template by specifying its coefficient, with the log concrete score as the only gradient-bearing term.The coefficient is given as w_x,y = πθ(y)f(·).Source: §3.2 · p. 6

3.3 Sampling through gradient flow

  • The gradient-flow variant samples through a time-homogeneous CTMC to reach πθ over an infinite-time horizon, with corrector steps available to improve sample quality.The stated gradient-flow interpretation uses a functional G(p, θ) = KL(p||πθ).Source: §3.3 · p. 7
  • For gradient-flow updates, the implicit-function derivative is obtained from a linear system with a closed-form inverse using the Sherman–Morrison formula.This replaces ∇θπθ in the parameter update and avoids costly automatic inversion of the dimensional system.Source: §3.3 · p. 7

3.4 Convergence bounds

  • As S→∞, Algorithm 1's average objective gradients converge under bounded target-distribution gradients and a KL-Lipschitz gradient map.The theorem uses β_s = min(1, 1/√s) and Assumptions 3.4 and 3.5.Source: §3.4 · p. 8

4 Experiments

Across DNA modeling and language modeling, SEPO fine-tuning improves reported generation outcomes. DNA results favor activity and accessibility, while the language-model comparison favors the variant with more denoising steps.

4.1 Language modeling

  • Language experiments fine-tune SEDD Medium Absorb with an Actor-Critic PPO framework and compare variants trained with 128 versus 1024 denoising steps.Fine-tuning uses a GPT-2-based reward model trained on HH-RLHF, without a supervised fine-tuning stage.Source: §4.1 · p. 8
  • Both SEDD-SEPO variants outperform pretrained SEDD on pairwise judgments, and SEDD-SEPO-1024 achieves the best overall results among the tested variants.The evaluation used GPT-3.5 Turbo as judge over 153 prompts; Table 1 reports favorable-output proportions across models and denoising steps.Source: §4.1 · p. 8, Table 1 · p. 8

4.2 DNA sequence modeling

  • The DNA experiment uses a pretrained masked discrete diffusion model trained on 700k regulatory sequences and evaluates unconditional generation with rewards on final denoised samples.The comparison includes guidance methods, DRAKES, and GLID2E.Source: §4.2 · p. 9
  • SEPO with gradient flow achieves Pred-Activity 7.64 and ATAC-Acc 99.9%, the highest results on those metrics, while SEPO leads on all reported key metrics except 3-mer Corr.The authors report lower run-to-run variance; Figure 2 also shows SEPO variants with the highest Pred-Activity distributions and low variance.Source: §4.2 · p. 9, Table 2 · p. 9, Figure 2 · p. 10
  • SEPO reaches higher final Pred-Activity with fewer epochs than the compared methods, despite higher per-epoch cost, yielding competitive total runtime.The authors attribute the SNIS tradeoff to an added batched forward pass for variance reduction and fewer training epochs.Source: §4.2 · p. 10

5 Conclusion

The conclusion presents SEPO as a policy-gradient framework for fine-tuning discrete diffusion models in conditional and unconditional generation. The appendices provide the paper roadmap and background on sampling and notation.

5 Conclusion

  • SEPO adapts policy-gradient methods to discrete diffusion and is presented as applicable to both conditional and unconditional generation.The conclusion identifies extra computational complexity from inner importance-sampling resampling for each neighbor of outer Monte Carlo samples.Source: §5 · p. 10

Appendix

  • The appendix roadmap assigns policy-gradient methods to Section B, theoretical proofs to Section C, mathematical supplements to Section D, and further experiments to Section E.Section A recalls notation and introduces additional definitions.Source: Appendix · p. 15

A Additional background details

  • The background describes tau-leaping as simultaneous independent Euler steps and notes that more reverse diffusion steps improve fidelity at higher computational cost.Conditional sampling incorporates auxiliary data into the reverse transition; the notation section also defines [d], [t]+, and Id.Source: A Additional background details · p. 15

B Unified Paradigm of Policy Gradient Methods

The appendix gives coefficient expressions for SEPO under PPO and GRPO. It also describes their advantage constructions and points to the proof section for the main-text results.

B Unified Paradigm of Policy Gradient Methods

B.1 Proximal Policy Optimization (PPO)

  • PPO can estimate advantage using a value network, defining it as reward minus the value estimate.The appendix identifies this as a common approach for PPO.Source: §B.1 · p. 16

B.2 Group Relative Policy Optimization (GRPO)

  • GRPO uses standardized rewards over a group of outputs as its advantages.The appendix introduces a group of outputs and then defines the advantages for that group.Source: §B.2 · p. 16

C Proofs of the results

C.1 Proof of Theorem 3.1

The proof establishes a bias bound for a plug-in gradient estimator under an approximation condition. It then examines the variance of the self-normalized importance-sampling estimator.

C.1 Proof of Theorem 3.1

  • Under approximation condition (6), replacing ∇θ log r_t(x,y) with ∇θ log s_θ in the estimator yields the stated gradient formula with bias bounded by Rmaxδ.The bound is derived for the plug-in estimator relative to the true gradient.Source: §C.1 · p. 17

C.2 Self-normalized importance sampling

  • The self-normalized importance-sampling estimator’s variance decreases at the canonical Monte Carlo rate, under the stated finite-second-moment assumption.The proof uses the delta method to obtain the variance expression.Source: §C.2 · pp. 17–18

C.3 Proof of Theorem 3.2

The appendix proofs establish a trust-region error bound for the policy-gradient estimator and derive results on sparsemax-based gradients and KL gradient flow. It also supplies technical lemmas and discrete-setup derivations used in the analysis.

C.3 Proof of Theorem 3.2

  • The Theorem 3.2 bound is independent of θ, linear in δ, and depends only on trust-region radius ρ, clipping level G, and Rmax.Its positive coefficient is C(ρ, G) = ρ(e^2 + (e^2 − 1)G).Source: §C.3 · p. 19

C.4 Proof of Theorem 3.4

  • The Theorem 3.4 proof uses an implicit-function linear system for the sparsemax projection, choosing η so k_h matches the Monte Carlo policy-gradient approximation.That choice makes the required probabilities accessible through the concrete score.Source: §C.4 · pp. 19, 20, 22

C.5 Proof of Theorem 3.3

  • The corrector-rate ODE has p_t as its stationary distribution, and sampling from it implements gradient flow for KL(· || p_t) on the simplex with respect to W.W is described as a Wasserstein-like metric whose transport cost depends on the mass already at the two points.Source: §C.5 · p. 22

C.6.1 Technical lemmas

  • The technical lemmas bound the parameter gradient of the log distribution and state a discrete Log-Sobolev inequality.The log-gradient bound is given under Theorem 3.4; the inequality is stated for μ > 0.Source: §C.6.1 · p. 23

C.6.2 Main proof

The main proof bounds the accumulated gradient norm under a decreasing schedule. The surrounding appendix material gives discrete first-variation and KL-regularization derivations.

C.6.2 Main proof

  • Under β_s = min(1, 1/√s), the proof establishes an integral gradient-norm bound of at most c S^1/2.The argument concludes by showing the relevant integrals converge as S → ∞.Source: §C.6.2 · pp. 25–26

D.1 First Variation in the discrete setup on ∆d

  • For a functional on the discrete simplex, the first variation quantifies sensitivity to perturbations and can be written using componentwise partial derivatives.The appendix derives this for the functionals considered in the paper.Source: §D.1 · p. 26

D.2 KL regularization term

  • Applying Girsanov’s theorem yields a generalized I-divergence term in the KL-regularization derivation, whose path integral recovers the result of Wang et al. (2025).The derivation uses a computation from Zhang et al. (2025).Source: §D.2 · p. 27

E.1.1 Training

SEPO is evaluated on SEDD Medium using an Actor-Critic PPO-style RL setup with GPT-2-based reward modeling. The experiments use conditional completion generation and compare two denoising-step settings, without an SFT stage for SEDD.

E.1.1 Training

  • The reward-model pipeline trains GPT-2-based SFT and reward models on separate halves of HH-RLHF, with LoRA and a logistic reward output.The SFT model is trained autoregressively; the reward model supplies the reward signal.Source: §E.1.1 · p. 27
  • SEDD Medium is fine-tuned directly with RL, skipping supervised fine-tuning; the authors acknowledge that an SFT stage would be beneficial to their cold-start setup.The reward model uses GPT-2 architecture and HH-RLHF reward signals.Source: §E.1.1 · p. 27
  • Conditional sampling generates completions from prompts, and two SEDD-SEPO variants are trained for 7k steps using 128 or 1024 denoising steps.The prompt and completion form x = c|p; the two variants are evaluated for the effect of denoising-step count.Source: §E.1.1 · pp. 28–29

E.1.2 Evaluation

On 153 prompts judged by GPT-3.5 Turbo, both SEPO variants generally outperform pretrained SEDD, and the 1024-step variant generally outperforms the 128-step variant. The evaluation also reports qualitative gains in response structure and clarity, alongside a remaining prompt-continuation issue.

E.1.2 Evaluation

  • The evaluation uses 153 prompts spanning a wide range of topics and GPT-3.5 Turbo to judge which response is preferred.It compares both fine-tuned variants with each other and with pretrained SEDD Medium.Source: §E.1.2 · p. 29
  • Both SEDD-SEPO variants outperform pretrained SEDD Vanilla across most prompts, while SEDD-SEPO-1024 generally outperforms SEDD-SEPO-128.The 1024-step model has no statistically significant advantage across evaluation step counts; SEDD-SEPO-128 appears to perform better at its training step count.Source: §E.1.2 · p. 29
  • Both SEPO variants show improved response structure and clarity over SEDD Vanilla in the creative-writing examples, but SEDD-SEPO-1024 sometimes continues the prompt rather than directly answering.The authors connect this behavior to the lack of an SFT stage and note that it was also observed during training.Source: §E.1.2 · pp. 30–31

E.2 DNA experimental details

This section describes the DNA-generation setup, evaluates pretrained sequence quality and fine-tuned outputs, and tests the contribution of inner SNIS. The evaluation uses separate fine-tuning and evaluation oracles and reports ATAC-Acc across methods.

E.2 DNA experimental details

  • Two Enformer reward oracles trained on disjoint chromosome subsets achieve held-out Pearson correlations above 0.85, with one used for fine-tuning and the other for evaluation.The oracles use Gosai et al. (2023); the evaluation oracle supplies Pred-Activity.Source: §E.2 · p. 31, Table 7 · p. 31
  • Generated sequences closely match real HepG2 activity distributions, and their 3-mer and 4-mer Pearson correlations both exceed 0.95.Figures 6 and 7 show the activity distributions and k-mer correlations, respectively.Source: §E.2 · p. 31, Figure 6 · p. 31, Figure 7 · p. 32
  • Fine-tuning uses the pretrained masked discrete diffusion model with SEPO GRPO, 128 sampling steps, KL regularization, and gradient truncation at step 10.SEPO uses clipping ratio ϵ = 0.2 and K = 2; the gradient-flow variant adds one corrector step per sampling step.Source: §E.2 · pp. 31–32
  • ATAC-Acc is the percentage of 640 generated sequences whose predicted accessibility score exceeds 0.5; the reported results average three seeds and are near 100%.Figure 8 displays score distributions for one seed per method.Source: §E.2 · p. 32, Figure 8 · p. 32
  • Without inner SNIS, SEPO significantly underperforms on Pred-Activity and ATAC-Acc; the authors attribute this to high-variance gradient estimates that hinder stable policy updates.The comparison is reported in Table 8.Source: §E.2 · p. 32, Table 8 · p. 33

Related papers

Loading 2502.01384v3…