Title: Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods
Venue: NeurIPS 2025
Accepted: yes
Keywords: Discrete Diffusion Models

Policy gradient algorithms

Non-differentiable rewards

Fine-Tuning

Reinforcement Learning from Human Feedback
Abstract: 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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes \"bandit\" policy gradient algorithms for discrete diffusion models in the concrete score framework. The proposed method SEPO is tested  on DNA fine-tuning and natural language tasks and compared against other benchmarks. It also introduces a gradient flow alternative that improves sample quality at a higher complexity. Convergence bounds and derivations of the method are provided.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n1. The problem tackled by this paper is interesting and important.\n2. The ideas and proof strategies are interesting.\n3. The experiments are interesting and compare with multiple baselines.\n\n**Main weaknesses**\nThe main weaknesses of this paper concern the validity of the proofs/statements and the clarity of the paper:\n1. Many of the proofs/statements (Theorem 3.1, Theorem 3.2, SNIS line 145) hinge on the claim that ${s_y(x,\\theta,t)}$ =$ \\frac{q_\\theta(y)}{q_{\\theta}(x)}$. This is not necessarily the case. If $s(x,\\theta,t)_y$ has modeled the ratios $p(y)/p(x)$ perfectly, then it is correct, otherwise it is not necessarily true. This is a very important fact, and therefore either this should be stated as an assumption or be proven. Proving it would imply that the bound given in Lou et al 2024, becomes an equality, so it is unlikely to be true.\n2. $q_t^\\theta(y)$ in Theorem 3.1 comes from $Z_\\theta$, which is the sum over all states $y$. Therefore, as properly written in the paper, the inner sum is taken over all states, not just neighboring ones for which we have the learned scores. When sampling in SEDD this is okay, because the transition-rate matrix is sparse, and therefore due to multiplication, to go backwards, we only need to model neighbor scores. However, in the formula of Theorem 3.1 and 3.2, the transition-rate matrix is not there to simplify the equation. I would have checked the implementation but the code is not provided.\n3. The SNIS approach proposed appears to be computationally expensive, and in addition, regarding line 145, the original problem was that we do not have the values $q_t^\\theta (y)$. It is unclear why this approach is sufficient.\n4. The paper is a bit unclear in the sense based on the way it is written, it leads you to expect a usual gradient policy method, utilizing trajectories with multiple steps, but only the end of the trajectory is utilized in the loss. In [1] RWR is introduced, but it also proposes multistep DDPO, similarly to [2] (in the continuous case). Reading the introduction, and abstract one is likely to expect something similar.\n5. The difference in compute and time is not provided between the proposed method and other methods. One would expect that the inner loop is expensive to calculate.\n\n**Minor weaknesses**\n1. Line 81 and 87, should be $Q_t(y,x)$ instead of $Q_t(x,y)$.\n2. The cummulative reward from time $t$ to 0 in [1], is the constant $r(x_0,c)$ and not 0.\n3. The lower left results in Table 1 are not necessary.\n\n[1] TRAINING DIFFUSION MODELS WITH REINFORCEMENT LEARNING, Black et al, 2024\n\n[2] Diffusion Model Alignment Using Direct Preference Optimization, Wallace et al, 2023\n\nQUESTIONS:\n1. Regarding weakness 2, what is the justification for only summing along the neighbors?\n2. Why use a trajectory of length one instead of the more common MDP approach for gradient policies?\n3. How do the proposed methods compare with the existing ones in terms of compute?\n\nLIMITATIONS:\nThe computational complexity of the the inner importance sampling steps is acknowledged in the conclusion. A separate limitation section is not provided.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces SEPO, a novel policy gradient-based algorithm for fine-tuning discrete diffusion models using non-differentiable reward functions. The authors propose a theoretically grounded framework that adapts policy gradient techniques like PPO and TRPO to discrete diffusion settings, leveraging score entropy and importance sampling. Essential validations are conducted on DNA sequence design and language modeling tasks.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\n\n1. The paper presents a comprehensive theoretical formulation of SEPO, which adapts classic RL algorithms to a score-entropy-based discrete diffusion setting. The derivations, particularly Theorems 3.1 and 3.2, and the convergence result (Theorem 3.5), demonstrate solid mathematical rigor.\n2. I think the inclusion of a gradient flow-based sampling strategy is a notable extension.\n3. Experimental results on DNA sequence modeling demonstrate that SEPO (especially with gradient flow) outperforms prior methods in key metrics, such as enhancer activity and chromatin accessibility.\n\nWeakness:\n\n1. The baselines included in this work remain narrow. Critical missing comparisons include inference-time-only approaches like discrete SMC methods or newer masked diffusion fine-tuning techniques. For language modeling, standard RLHF baselines, such as DPO or PPO-tuned transformers, are almost a must-include for a fairer evaluation.\n2. SEPO’s inner sampling and score estimation steps (especially SNIS) introduce notable computational overhead; relevant discussions are not included.\n3. Technically there is a concern with the technical hardness of this work. Compared to existing works, the major work is done by replacing continuous diffusion models with discrete diffusion models and the optimization algorithm does not present many differences.\n\nQUESTIONS:\n1. Why are standard RLHF baselines such as PPO, DPO, or TRPO not included in the language modeling experiments, even though SEPO is proposed as a general-purpose policy gradient method?\n2. How critical is the gradient flow sampling mechanism to SEPO’s performance? Would simpler sampling strategies yield comparable results?\n3. Can the authors explain what the major difficulty of this algorithm (theoretically and experimentally) is beyond adapting existing policy optimization algorithms to a discrete-time formulation>\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe core algorithmic contribution is score entropy policy optimization (SEPO), which adapts policy gradient methods, commonly used in continuous RLHF, to discrete diffusion models. This addresses a significant challenge in fine-tuning these models, especially with non-differentiable reward functions, which are often encountered in practical applications like RLHF. SEPO provides an explicit characterization of policy gradient algorithms within the concrete score framework, supporting both conditional and unconditional generation. A technical innovation is the use of self-normalized importance sampling (SNIS) to estimate otherwise intractable marginal distributions required for gradient computation, which is claimed to offer consistent, asymptotically unbiased, and lower-variance estimates. The algorithm also incorporates concepts from trust region policy optimization (TRPO) to ensure stability and uses a clipping mechanism for policy updates. An alternative using gradient flow is also explored to improve sample quality.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \n\n1. Algorithmic innovation: SEP0 effectively bridges policy gradients with discrete diffusion, enabling non-differentiable reward optimization—a notable advance over Gumbel-Softmax alternatives.  \n\n2. Theoretical foundation: Convergence guarantees (Theorem 3.5) and gradient flow interpretation (Lemma 3.3) rigorously formalize the method.  \n\n3. Empirical rigor: Strong benchmarks in DNA/language tasks (Tables 1–2), with ablation confirming SNIS necessity (Table 8).  \n\n4. Novelty: SEP0 largely repurposes TRPO/PPO mechanics (Eq. 6–8), but in their rebuttal the authors convinced me that the unified framework adds sufficient innovation over prior RL-guided diffusion.  (Moved from Weakness to Strengths after rebuttal plus promised revisions.)\n\n---\n\nWeaknesses:\n\n1. Incomplete baselines: Omits comparison to RLHF-tuned AR models (e.g., Llama 2) in language tasks. Does this undermine claims of diffusion superiority? Added after rebuttal: This is still a weak point of the paper.\n\n---\n\nI originally signalled scalability omissions as a weakness: No analysis of computational overhead from SNIS (Sec 3.1), especially for large \\(n\\) (e.g., \\(m^n\\) states); wall-clock times were also absent.  Added after rebuttal: In their rebuttal addressed this point in an acceptable manner.\n\nQUESTIONS:\nTheorem 3.1 estimates \\(\\nabla_\\theta \\ell_t(\\theta)\\) via SNIS, but how does variance scale with \\(M\\) (samples) and block size \\(n\\)? Fig 1 suggests instability for \\(n \\gg 10\\)—was this tested?  \n\nEq 4 includes KL regularization, yet experiments (Table 2) use \\(\\alpha=0.05\\) without ablation. Why is KL critical for DNA but omitted in language tasks (Sec 4.1)?  \n\nCan the authors report GPU hours/memory for SEP0 vs. baselines (e.g., DRAKES) in Tables 1–2 to quantify efficiency claims?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose the need for policy gradient methods for finetuning/training Discrete Diffusion moedels for Language modeling. In this paper the authors make the following contributions: (1) provide an explicit characterization of policy gradient algorithms for discrete diffusion models which allows for the use fo non-differentiable reqard function as well as conditional and unconditional generation, (2) propose a policy gradient algorithm called Score Entropy Policy Optimization (SEPO) with a gradient flow alternative and (3) validate the effectiveness of the SEPO family of algorithms with experiments on DNA finetuning and Language modeling\n\nMain Method: the authors first describe the policy gradients for score entropy where they use the discrete REINFORCE trick to the score entropy changes and compute the loss function. To address teh high variance in REINFORCE, the authors build on using TRPO fundamentals to use an importance sampling term. This is then clipped in the interval $[1 - \\epsilon, 1+\\epsilon]$ to build the SEPO algorithm also described in Algorithm 1\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n- The formulation works on the gaps of prior approaches such as DRAKES, GLIDE and d1 to handle unconditional generation, as well as non-diffferentiable rewards\n- Rewards can be provided at any time step (l110) which allows for more flexible reward schemes and also validated in Table 1\n- The paper introduces policy gradients in a principled manner while ensuring that the work is mostly self-contained with relevant works discussed in the main paper and the appendix\n- Gradient Flow interpretation using $KL(p||\\pi_\\theta)$ demonstrates strong metrics on in the DNA sequence modeling task especially on Pred-Activity and ATAC-Acc. Furthermore, these scores have significantly lower variance demonstrating the stability of the training algorithm (Fig 2)\n- The convergence bounds estimate of $\\mathcal{O}(\\frac{1}{\\sqrt{S}})$ in Theorem 3.5 gives useful gurantees to the algorithm   \n- The authors validate the importance of \n\n### Weaknesses\n- In theory, unconditional sampling should be possible, but it is not empirically discussed. \n- Specifically for language models, a window of tokens is typically predicted in Discrete Diffusion models. This makes it quite expensive to update the model outputs with small edits to the inputs or the prediction, hence reducing its scalability.\n\nQUESTIONS:\n1. Maybe out of the scope of this work, but is there a workaround to allow for edits to the input without having to rerun the complete inference step?\n2. For unconditional sampling, some examples of unconditional text generation or DNA sequence generation would add value to the paper\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-07-04"}