Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods
Authors: Oussama Zekri
Nicolas Boulle
Venue: NeurIPS 2025 poster ()
Keywords: Discrete Diffusion Models
Policy gradient algorithms
Non-differentiable rewards
Fine-Tuning
Reinforcement Learning from Human Feedback
Primary area: deep_learning
TL;DR: We propose a policy gradient algorithm for fine-tuning discrete diffusion models over non-differentiable rewards.

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.

Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods
Oussama Zekri CREST, ENSAE
Nicolas Boullé Department of Mathematics Imperial College London
Institut Polytechnique de Paris France oussama.zekra www.rieensae....... fr
United Kingdom n. bouulee@imeriiil...... uk
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 method. Our code is available at https:/ggitabb.oo/oookkiii
Introduction
Diffusion models have become efficient generative modeling tools in various tasks, including image and video generation (Song et al., 2021; Ho et al., 2020). Although most of the applications of diffusion models depend continuous state space (such as images), recent works extended these models to discrete settings, enabling their use language modeling and other discrete generative tasks (Sun et al., 2023; Camphell et al., 2022; Austin et al., 2021; Benton al., 2024). Moreover, several studies showed that these models can be competitive with autoregressive models, such as GPT (Brown al., 2020) Llama (Touvron al., 2023), while allowing for more flexible generation compared to next-token prediction (Lou et al., 2024; Sahoo et al., 2024; Shi et al., 2024). These discrete diffusion models are promising if they can be scaled up natural language processing tasks.
However, fine-tuning discrete diffusion models is challenging. Existing strategies, such as classifier guidance (Ho and Salimans, 2021; Nisonoff et al., 2025; Gruver et al., 2024) and steering (Rector- Brooks et al., 2025), often face scalability issues or intractable objectives. We focus instead on reinforcement learning (RL), which optimizes a pre-trained model via rewards. key challenge is that sampling from categorical distributions is non-diffcrentiable and incompatible with standard gradient methods. Recent work (Wang et al., 2025) addresses this using the Gumbel-Softnn trick (Jang et al., 2017), but requires differentiable rewards and becomes memory-intensive at scale. In contrast, policy gradient methods like PPO (Schulman al., 2017) and GRPO (Shao al., 2024) rely only on reward evaluations and offer stability, unbiased gradients, and robustness to noise.
We propose policy gradient algorithm, Score Entropy Policy Optimization (SEPO), designed for discrete diffusion. Unlike Wang et al. (2025), our method handles non-differemtiable rewards, enabling broader fine-tuming scenarios. SEPO offers unified and scalable framework for optimizing discrete diffusion models, that supports both conditional and unconditional generation. Other policy gradient approaches have also been proposed, such GLIDE (Cao et al., 2025). which relies reward shaping, and d1 (Zhao et al., 2025), specifically designed for masked diffusion models and cannot
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

handle unconditional generation. These methods are tailored to specific setups, whereas SEPO offers a more general and principled framework for discrete diffusion fine-tuning.
Main contributions.
1) We provide an explicit characterization of policy gradient algorithms for discrete diffusion models in the concrete score framework. This allows the use of non- non-differentiable rewards discrete fine- tuning tasks without steering or guidance mechanisms, for both conditional and unconditional generation.
2) We propose an efficient, scalable algorithm based on policy gradient methods (Schulman et al., 2017; Shao et al., 2024), called Score Entropy Policy Optimization (SEPO), for discrete diffusion. It uses clipped-ratio loss with nelf-normalized importance sampling for low-variance. We also introduce gradient flow alternative that improves sample quality at higher complexity.
3) We perform numerical experiments on DNA fine-tuning and natural language tasks to demonstrate the performance of our methods. We achieve state-of-the-arr DNA results: Pred-Activity 7.64 and ATAC-Ace 99.9%, topping prior RL and guidance baselines with lower run-to-run variance.
2 Background and preliminaries
2.1 Related works
Inference-time techniques. Inference-tings techniques are simple yet effective as they require no fine-tuning training when reward functions available. Recent studies (Singhal et al., 2025; Ma et al., 2025) showed that they can achieve competitive performance by scaling computational resources. Although inference-time techniques offer distinct advantages, they typically result in longer inference times compared to fine-tuned models. The key considerations for these techniques include computational efficiency and differentiability of the reward (Uehara et al., 2025).
Policy gradients algorithms. Policy gradient algorithms are a key class of reinforcement learning methods that optimize parameterized policies by directly maximizing expected returns. Modern implementations include Proximal Policy Optimization (Schulman et al., 2017) or Group Relative Policy Optimization (Shao al., 2024). These algorithms are highly sensitive to policy design since the architecture impacts expressiveness, optimization stability, and exploration.
Fine-tuning diffusion models with Reinforcement Learning. In the case of continuous diffusion models, fine-tuning via policy gradients has been proposed (Fan et al., 2024; Li et al., 2024; Black et al., 2024; Ren et al., 2025). In a more recent study, (Marion et al., 2024) implements REINFORCE algorithm (Williams, 1992) for continuous diffusion models in a single-loop algorithm, avoiding nested optimization. However, extending these approaches to discrete diffusion models is more challenging. This work adapts these studies to the discrete case and extends them to general policy gradient algorithms.
2.2 Discrete Diffusion
In discrete diffusion models, the dynamics of single particle is described by a continuous-time Markov chain (CTMC), denoted as stochastic process (x, )ostST operating
finite space
X = m, and is the length of the sequences, which is fixed number known context window block size. Typically, it describes sequences of tokens or image pixel values. While the size d := |X| m" of X' is exponential in n, deep neural networks such as transformers (Vaswani, 2017) were shown to perform and generalize well on these incredibly large state spaces (Zekri et al., 2024).
amm) }". Here, (as)igism represents the possible states that form vocabulary of size
Forward process. any given time the distribution of particle 2; given by Pt. which lies within the probability simplex A₂ C Rdd The forward process is noising process that maps the initial data distribution po := Pdaza to some final noisy distribution pr == Pref. which is easy to sample. During the noising forward process, the particle's probability transitions between states are
given by rate matrix Qi dx indexed by X. through the equation dp: QcP; for t E T7.

---

The time reversal of this equation (Kelly, 2011) is known as,
dpr-t = QP- 1E(0,2]
(I)
di
Q1(2, v), p((m) - Qe(2,
p.(u)
x xサ y,
where for z. X. (v)
x = y.
Score entropy. Lou et al. 2024 recently showed that one can approximate Eq. (1) via score entropy. inspired by concrete score matching (Meng et al., 2022). This is done by learning the concrete score as s0 R, 1)y x (((v)/((()) with a aequence-to-sequence neural network so parametrized by RR.. The resulting process is described by the following equation:
dq dt
= Qr- E [O,T],
22
where the denoising process df R pr maps d = Prof to S and is learned to achieve at R3 Pdata- The matrix G. is defined for x, y E as Q₁ (x. y) so (x, v) if * and Q (v, Ezjz Q (2 In practice, the quantity (1. 1), is available for all y E X at Hamming distance (Hamming, 1950) one of a: the states y that differ from by exactly one token. This represents only (mm) ratios instead of O(m2") (Campbell et al., 2022; Lou et al., 2024).
Masked diffusion models. The score entropy setup also contains the simplified frameworks detailed in (Sahoo et al., 2024; Shi et al., 2024; Ou et al., 2024), often referred to as masked diffusion models. These models define a simplified training loss based on a weighted cross-entropy objective by introducing special [MASK] token, and have been shown to outperform uniformly noised discrete diffusion models. Promising attempts at scaling them to larger settings support this observation (Arriola al., 2025; Nie et al., 2025; Ye et al., 2025).
2.3 Fine-tuning with Reinforcement Learning
After the pretraining phase, discrete diffusion model with learned parameter Opre aims to approxi- mate Pdatas in the sense 47 flpee as Pdata- Our goal fine-tune the denoising process of to increase a reward function Rs X - R. without having access to Pdata-
Minimization problem. We focus on optimization problems over implicitly parameterized dis- tributioos. For a given family of functions (F,) A. + R, we aim to minimize the loss function
(x()) = -F(("), E [0, T),
over E RRP Standard choices of Fg include Fe(g)) = Engg [Ra to maximize a reward function R2 X' - IR, or F( (g. lg. are) to minimize the KI. divergence of pe from distribution qi pre As detailed in (Uchara et al., 2024), a typical fine-tuning algorithm for diffusion models combines these two terms as
(0) - -Q) [R((
KLLgilla ".T.
where a > 0 is weighting factor. One can recast the denoising process morie-horizon multi-step Markov Decision Process formulation with known transition probabilities as folows:
For all E [O,...,T}. S S=X. A=X, X, Policy di -T=IT- A. -1 P(Sp) Prof, P(S.+1 A.) - 5 Ar)),
where S and A denote the states and action spaces, S. and A, denote the states and actions, and d is the Dirac delta distribution. Although common choices in the fine-tuning diffusion models with reinforcement learning literature (Black et al., 2024; Fan et al., 2024: Clark et al., 2024; Uchara et al., 2024) often set R₁ = 0 for ł < T and Rr = Rt we retain the general form of RG in the results for greater generality. This choice allows for more flexible reward evaluation schemes. For example, one can run fewer denoising steps and stop the process at some intermediate time To < T, assigning the reward Rr, = R based on the partially denoised sample. We will consider this case, and the corresponding distribution of, will be denoted by πe.
3

---

Loss reward gradient. To apply first-order optimization methods, one needs to compute the gradient (O). Since X is finite space of size d, we have
Vol(0)) = vv = -
(5)
where f I is the vector of first variations Fe("g") (see Appendix D.I). Importantly, note that (x) = (x) for a E X. which does not involve the differentiability of R₂ (with respect to some embedding& & (X') of the state space). One can then design deterministic non- -on-differentiable functions that act on X' as rewards, similar to those arising in RLHF, or elsewhere. This may include designing desired protein properties (Rector-Brooks et al., 2025).
3 Methods
3.1 Policy gradients for concrete score
The gradient of the target distribution Vogl in Eq. (5) can be approximated based on its relationship with the concrete score $g.
Assumption 3.1.
For x + y E X, define the exact ratio r (0) == q: (v))// We assume that there exists 0 such that for all 0 E I2, x # E X, t O,
V log (z T t) Volog r' (0)l 8
(6)
The following theorem shows that compute (f) that estimates Vol((0) through discrete analogue of the REINFORCE algorithm ( (Williams, 1992).
(Discrete REINFORCE trick). Under Assumption 3.1 and with the notations in- troduced in Section 2, applying the REINFORCE with concrete score approximations yields the following estimator of Eq. (5):
Theorem 3.1
Ve,(0) = Σ d.
Σ Ec VV) s0(x, T t)y
ZEX
such that (i)(00) - ve((0)II S Rmax 5, where Rmax SuPre(O.7] max [R((x)] < oo.
Computation of the loss function. The summand in Theorem 3.1 involves the unknown distribu- tions of (x) and qf (y). While the outer sum can be estimated via Monte Carlo sampling (with sam- ples (I1...., (NN). )) the inner sum is weighted by qf (y). As noted in Lou et al. (2024), a single x e X provides access every component of the concrete score sa (z, t)y. for y F and then to q. (y x) since this is the only missing quantity in the sampling scheme (Eq. (11) in Appendix A). Hence, one may be tempted to estimate q. (y) = =xrex d.(vv 2)) qf (z) = & Ei= L x;). However, as illustrated in Fig. 1, for given y. we only have access to a single value I ); we do not even have two such values, let alone N, and thus cannot take advantage of the full sample set used in the initial Monte Carlo estimation, leading to a really high-variance estimation. In fact, for any reasonable block size n, although it is certain that at least one sample x; will be neighbor of y, is highly unlikely that two distinct samples I; and æy both have y as a neighbor, since this would require them to differ from y exactly one token position.
Self-normalized importance sampling (SNIS) of of (v). To address this issue, for a given we build M samples (zihsis a from of (- y) and employ the importance sampling approximation
of
1 M M
(zx) q (z) v)
(v) Σ Ef (=) = Σ of ZEX
(2 v) y)
Σ
of
ZEX
with weights w(xz)) = (x))/q( v) and proposal q. (z: y Following Bayes' rule qf (2. y) xc qf (t)l =f) (2:), ), the proposal distribution qf (z; y) is optimal in the sense that it minimizes the variance of the importance weights. This makes it theoretically sound and efficient choice for

---

Neighbours
Only one neighbour x available for each yi Need new samples Ry
Figure 1: On the estimation of (2). Left:
Given a sample z, only its neighbors (yihsiss X accessible for computing of (y/ d), and each y/ typically has only one such parent in the sampled batch. Right: For given V/., there are several neighbours {zyhıses- It is unlikely to find multiple distinct samples such that both neighbors of ya, since this would require them to differ from ys at exactly the same token position.
importance sampling. However, the marginal distribution (2:) is typically intractable, so evaluating w(z;) is unfeasible. Fortunately, one has access to an approximation of (z;) up to a normalization constant thanks to the score 1/5p(2,, T tly ze (2))//g (v). This is sufficient for self-normalized importance sampling with weights (i2))) x of (2))//f (ze v)- The estimator obtained via SNIS takes the form
Σ M i= L 21)((z)))))))))))))))))))))))2
M M
1
(v)
EE ((:)
qe (y 22)
Thus, the SNIS estimate can be computed as the harmonic mean of the conditionals of (v zz). The estimator both consistent and asymptotically unbiased and has significantly lower variance than the naive single sample estimate. Under a mild moment assumption, we prove in Appendix C.2 that V(dl (v)) O(M-1). Finally, the gradient can be computed (ol( E Eg [2 (2)g9( 0)], where gzz f) = Evex q. Vo log so (a (x,T T t)y, and we estimate qf (y) with the SNIS method described above. vテr
Importance sampling of the outer sum. Although this defines an unbiased gradient, the REINFORCE algorithm known have high variance and to not restrict large policy updates. To address the latter limitation and estimate ((x 0), we build upon the core ideas introduced by Trust Region Policy Optimization (TRPO) (Schulman et al., 201 Instead of sampling from of one can leverage importance sampling through old policy qfad fald, and constraint the KT. divergence between the old and the current policy as follows:
( gar 0) Ban (x)
V Vel(()) = E
Pold Re(x))
q.
Assumption 3.2. We assume that there exists 8 > 0.
and G > O such that for all 0 E R.,. x + E X, E [0, TY. log sa(x, T () log r z. (0) S and |[V log so (x, T - ()
Va logoor (0) < 8, (x) (x) S p and V log se (x. T ---ttyl S.
Once adapted for concrete score, this formulation leads us to the following result.
Theorem 3.2 (Importance sampling gradient). Under Assumption 3.2 and with the notations
introduced Section 2, applying the TRPO with concrete score approximations yields the following estimator of Eq. (5):
(0) = E, Bengfol Gold [2()))(z,)))
(8)
5

---

such that - < Cp, G)å, where Rmax (0 max Re(x))
C(p, G) > is constant depending p and G, and
h(,,0) 0) = Σ de (1)
80ma (z, T - t)
Vo log so so t)y
4: q. (3) so s0 T l)y
BEX uy
The quantity h(z, 0) is expressed in this way Eq. (7) to emphasize how the loss will be computed in practice. While being the founding step of state-of-the-an policy gradient algorithms, TRPO requires solving a constrained optimization problem each step. However, thanks to Theorem 3.2, we can now build powerful, stable, scalable, and masy-to-implemen policy gradient algorithms.
3.2 SEPO: : Score Entropy Policy Optimization
Our algorithm relies on the ideas introduced in (Schulman al., 2017; Shao al., 2024), but can be adapted to any policy gradient algorithm built on REINFORCE or TRPO. Inspired from these algorithms, we clip qe (z)/or /
Bold (x) via the following ratio that appears in the inner sum of Theorem 3.2:
(x)
(v)
Soat (x, T 1)g
Bold (x)
q. Gold (v) so(x, T f)
to the interval [1 - -, + for some hyperparameter e O. This can be interpreted as enforcing trust region condition analogous to p= + in Assumption 3.2. Another advantage of discrete diffusion models is their great generation flexibility. It is then be possible to apply our algorithm conditionally (via training dataset, typically RHLF) or unconditionally for fine-tuning. Hence, in the conditional form, Eq. (7) becomes
L( (c
E sn Pold (xle) Re (x)
g 2, 0)
q: (x
Butd
where x represents the concatenation of and 2. Instead of using directly the reward R2(x), we compute an advantage A(x) to reduce the variance of the Monte-Carle estimations. This quantifies how much better an action compared to the expected return at given state. common approach in PPO (Schulman et al., 2017) is to learn value network to approximate the reward, and then define the advantage A, with Generalized Advantage Estimation (Schulman al., 2015b), For GRPO (Shao et al., 2024), the advantage is the standardized reward over each group. Specifically, for group of outputs (51,.... G we have that A,(x;) R.(x,) (efe,,-men((( mean (x)) for { (a)
Remark 3.3. The loss function takes the form
(A(0)
Way log s0(x, 1-1.T-T70 -T))gb
yEX
where = (u (To-To is a coefficient (where f is a function to specify) and the log concrete score log (-xx;TTTo)) To)y is the only term with an attached gradient. PPO, GRPO, and other methods can be constructed by specifying the function in the coefficient Wz.y In Appendix B, we present unified framework encompassing methods that can be derived from SEPO.
Optionally, for t E [0, Tol, (L [|q term can also be added to the loss, as in Eq. (4). Although this is not absolutely necessary, as clipping already implicitly regularizes with kL.(g (ud term (Schulman pendix 1).2, for completeness. This leads to the Score Entropy Policy Optimization (SEPO) al- gorithm described in Algorithm 1. SEPO iteratively samples from the target distribution via CTMC (Line 4) and optimizes 0, using an optimization objective (Line 6). refining the policy with policy gradients. Specifically, Line 4 generates samples from the target distribution тe using the CTMC Q Pold This can be done in O(1) time complexity by leveraging the queu- ing trick introduced in (Marion et al., 2024, Alg. 3), at a higher memory cost. Line updates the parameters 8, using policy optimization algorithm based the objective & (see Eq. (9)).
al., 2017; Fan et al., 2024), the derivation given in Ap-

---

This means performing K iterations of gradient ascent (or descent) on the policy loss function to improve the policy л (0x) using the previously collected samples and computed advantages.
Algorithm 1 SEPO
1: Require: CTMC Q
iteration S, epoch
2: Set 0g and Ootd to 0 pre
3: for E . do
4: Sample from Teaa with Q
3.3 Sampling through gradient flow
5: Compute the reward and the advantage
6: Optimize 0g with e for K epochs
Bilevel problem. We use sampling to reach
7: Set fold to 9,
the limiting process of the backward distribu- tion qf. This procedure can be interpreted as optimizing a functional G Ag x IP + R over Ad gd
8: end for
9: Output: 05
Th = argmin ((.
PEA,
When Te is the limiting distribution of infinite-time process (e.g., Langevin diffusion in the continuous case, Langevin 1908; Pavliotis 2014), one can recast Eq. (3) as bilevel optimization problem. This has been proposed by Marion et al. 2024 in the continuous case and allows to efficiently alternate between optimizing one-step of the inner problem and one-step of the outer problem.
Gradient flow interpretation. In our case, Te is reached with f finite-time horizon, in T steps of sampling. However, it possible reach лy in infinite-time horizon by sampling from specific nime-homogeneous CTMC. The choice of the functional G(p, 0) = KT[pp| Ro) leads to a gradient flow interpretation of sampling via specific CTMC.
Lemma 3.3 (Gradient flow). Sampling from the following ordinary differential equation
dp: = QiP:, where Qi == Qi + Qi,
dt
implements gradient flow for KL( p.) in As. with respect Wassertein-like metric.
Corrector steps. Of course, Sg is not perfectly learned in practice, and we just have access the rate matrix Q Q. o¹ But this gives insight into the choice of our sampling strategy, especially with predictor-correcta techniques for discrete diffusion, introduced in (Campbell et al., 2022) and developed in (Zhao et al., 2024). We will then sample from the time-homogeneee of rate -r--T, to reach Te with infinite time horizon. Note that this does not require computing an integral compared to the time mmme-inhomogeneour case. We then optimizing functional in Wassertein space through sampling (Marion et al., 2024; Bonet et al., 2024). Sampling from Qi affects Line 4 of Algorithm 1. In practice, the sample quality can be improved by adding corrector steps with Q = Qt + Q,, as proposed in (Campbell et al., 2022). Once the process has run for To steps, multiple sampling iterations from Qī-T, can be performed.
CTMC
Linear system characterization. In this case, VoTta in Eq. (5) will be obtained by solving linear system, using the implicit function theorem (see Appendix C.4) on Vig. through a corrected version denoted V π.. While both the evaluation of the derivatives and the inversion this linear system can be done automatically (Blondel et al., 2022), it is costly given the dimensionality d. Instead, we provide the exact linear system as well closed form of the inverse in Proposition 3.4. Note that this affects Line 6 of Algorithm 1, where Vo in Eq. (5) is replaced by V'
Proposition 3.4. For each yj > 0, ViTTO is the solution to a linear system of the form
A B, E Rdxp
where Ay is a rank-1 1 update to the d x d identity matrix, whose inverse can be explicitly computed using the
formula.
7

---

3.4 Convergence bounds
From a high-level point of view, Algorithm alternates between sampling and optimization steps. We then view Algorithm I as the following coupled equations:
dq. Qi 9,ds,
(10)
d0, (g. (g.00).... 0s)ds,
for 0 S S S. The gradient used on line 6 of Algorithm I depends both on q, and 8s. and we refer to I' (so that (0.) = [' 02)). ) To simplify the analysis, the evolution of both qa and 0, is done in continuous time flow, for some s E [0, with S > 0. Let || denote the Euclidean norm on R., We reintroduce assumptions on my and I made in (Marion al., 2024).
Assumption 3.4. There exists C 2 O such that for all X and 0 RP, (oTT (x)] C. There exists ε > 0 such that for all x E X and f E IR, TT( >
This assumption states that the gradient of the target distribution is bounded. The second part is similar to the ambiguity of the language often considered when studying models acting on spaces like X (Zekri al., 2024; Hu al., 2024; Xie et al., 2021).
Assumption 3.5. 30 such that for all p.q E Do. OER. [[(, - 14q, SCr (L(p)
This assumption essentially states that the gradient I' is Lipschitz with respect to the KL divergence on Ad. With these in place, we establish the convergence of the average objective gradients.
Theorem 3.5 (Convergence of Algorithm 1). Let S > and 0g be the solution to Eq. (10) with B, = min(L,1/V3), for E [0, S). Under Assumptions 3.4 and 3.5, we have, S - 00,
$ %o [V (
4 Experiments
4.1 Language modeling
We briefly present our experiments language modeling; Appendix E. for details.
Training. We fine-tune the SEDD Medium Absorb model (Lou et al., 2024) using SEPO in an Actor-Critic PPO framework, directly optimizing with reinforcement learning without a supervised finetuning stage. The reward model, built from GPT-2 architecture and trained on the HH-RLHF dataset (Bai et al., 2022), provides reward signals to guide fine-tuning (see Figs. 4 and 5). We train two variants, SEDD-SEPO-133 and SEDD-SEPO- T = 1024 denoising steps, to measure the impact of those denoising steps on final response quality.
respectively trained with = 128 and
Quantitative evaluation. We eval- uate our models using pairwise com- parisons with GPT-3.5 Turbo (Brown et al., 2020) as a judge over 153 prompts from the Awesome Chat- GPT Prompts dataset (Akm,, 2023). As shown in Table 1, both SEDD- SEPO variants outperform the pre- trained model, demonstrating the ben- efits of reinforcement learning fine- tuning. Increasing the number of denoising steps further improves response quality, with SEDD- SEPO-1024 achieving the best results overall. Qualitative results deferred to Appendix E.I.
SEDD-SEPO-125
SEDD-SEPO-1022
Number steps
512
1024
512
1024
SEDD Vanilla
71.2%
61.1%
71.5%
73.2%
SEDD-SEPO-128
X
63.1%
68
8%
67.8%
SEDD LEDD-SEPO-1022 96.9%
31.2%
32.2%
X
Table 1:
Proportion of outputs deemed favorable by the Judge LLM for each model and each denoising steps T' {128, 512, 1024). Best results are highlighted.
8

---

4.2
DNA sequence modeling
We present experiment DNA sequence modeling task, which follows the setup of Wang et al. (2025). Additional details are provided in Appendix E.2.
Dataset, settings, and baselines. We employ the pretrained model of Wang et al. (2025), masked discrete diffusion model (Sahoo et al., 2024) trained on 700k regulatory DNA sequences (200 bp) from the Gosai dataset (Gosai et al., 2023). Each sequence is annotated with continuous enhancer activity score in the HepG2 cell line, and a reward model trained on this score is used during fine- tuning. Our DNA experiments are conducted in fully unconditional setting: sampling starts from fully noisy sequence and proceeds via denoising without any prefix or conditioning signal; rewards computed only the final denoised sample (i.e., with To = T). We compare our method to diverse set of state-of sate-offtthe-are the-art baselines: (i) Guidance-base based models, including classifier guidance (CG) for discrete diffusion models (Nisonoff al., 2025), Sequential Monte Carlo methods (SMC, TDS) (Wu et al., 2023), and classifier-fi free guidance (CFG) (Ho and Salimans, 2021); (ii) Direct reward optimization with DRAKES (Wang al., 2025), which uses reward gradients in RL-based fine-tuning setup for discrete diffusion models; (iii) Policy Optimization with GLIDdE (Cao et al., 2025), gradient-free RI--based tuning method.
Table 2: Evaluation of DNA modeling methods. State-of-the-art performance is in bold, and the second-highest performance is underlined. Means over 3 seeds; standard deviations in parentheses.
Method
Pred-Activiny (median) 1 ATAC-Ace (
3-mer Corrt
Log- Lig-Lik (median) 1
Pretrained
0.17 (0.04)
1.5 (0.2)
-0.061 (0.034)
(0.6)
CG
3.30 (0.00)
0.0 (0.0)
-0.065 (0.001)
-2.6.0
SMC
4.15 (0.33)
39.9 (8.7)
0.840 (0.045)
-5.9.2
TDS
4.64 (0.21)
(16.4)
0.848 (0.008)
-257.1
CFG
5.04 (0.06)
92.1 (0.9)
0.876 (0.004)
-265.0
DRAKES
5.61 (0.07)
82.5 (0.6)
0.887 (0.002)
-0...
DRAKES w/o KL
6.44 (0.04)
0.307 (0.001)
28..6
GLID'E
7.35 (0.07)
(0.3)
0.49 (0.074)
(1.4)
GLID'E w/o MI
2.57 (0.60)
0.63 (0.3)
0.473 (0.078)
2399.1 (10.1)
GLID-E w/o M2
6.62 6.62(0.42) (0.42)
67.3 (39.4)
0.458 (0.009)
-24..7 (21.5)
SEPD
7.55 (0.01)
99.5 (0.2)
0.500 (0.004)
-213.800.5
SEPD with GF
7.64 (0.01)
99.9 (0.09)
0.638 (0.001)
221.22
Evaluation metrics and results. We report four metrics: (i) Pred-Activity, the median predicted activity from a held-out reward oracle; (ii) ATAC-Acc, a binary chromatin accessibility score in HepG2: (iii) 3-mer Corr, the correlation between 3-mer frequencies of generated and sequences from the Gosai dataset; and (iv) Log-Lik, the log-likelihood under the pretrained diffusion model. We find that SEPO achieves the highest performance across all key metrics except 3-mer Corr, confirming its ability to generate highly active and biologically relevant sequences while maintaining a reasonable alignment with the training distribution. Notably, SEPO with gradient flow (GF) achieves near-perfect chromatin accessibility (99.9%) and the highest enhancer activity (7.64), surpassing both DRAKES and GLID"E. While SEPO does not fully maximize 3-mer Corr, this metric primarily reflects lowelevel distributional similarity rather than functional relevance. Our results suggest that prioritizing enhancer activity and biological accessibility leads to more meaningful sequence generation in practical scenarios. Moreover, variance across runs (as shown by the small standard deviation in Table 2), highlighting the stability of our generated sequences, once the optimization process is performed.
observe that SEPO also exhibits significantly lower
Finally, as illustrated in Fig. 2, even in the worst case over 640 samples, SEPO methods maintain minimum Pred-Activity score around 6, whereas all other methods exhibit much lower worst-case scores between 0 and 2. These strong performances, combined with the stability of the generation, make SEPO highly effective counterpart to both inference-time methods and other RL-based fine- tuning approaches, providing a more reliable solution when generation quality is prioritized over sampling speed.

---

Pred-Activity
Pretrained CG
SMC
TDS
CFG
DRÁKES DRAKES SEPO
SEPO
models. The plot shows the distribution of predicted enhancer activity Pred-Activity from the held-out reward oracle for each model, across 640 generated sequences. SEPO and SEPO with GF achieve the highest Pred-Activity scores with low variance, illustrating the effectiveness as well as the stability of our optimization process. Results for GLIDºE are not displayed the corresponding finetuned model weights not publicly available.
Figure 2: Violin plot of Pred-Activity scores
Wall-clock and overhead. SNIS adds batched forward to compute transition probabilities for each the N outer samples, which moderately increases per-epoch cost but reduces the number of training epochs. On single NVIDIA GeForce RTX" 3090 (24GB) GPU, we report the following computational timings.
Method
KL Trone
#Epocha
Batch / GRPO
Runtime (hh:mm)
Time/Epoch (mmm:s)
Final Pred-Activity
SEPO
14
8/8
00:19
01:21
7.55
SEPO with GF
8/8
00:24
02:59
7.64
SEPO (no SNIS)
10
199
8/8
00:10
00:03
4.66
DRAKES
127
8/-
00:26
00:12
6.41
DRAKES (with KL)
50
66
8/-
00:31
00:28
5,70
SEPO achieves higher final performance with fewer epochs despite higher per-epoch cost, yielding competitive total runtime. Our SNIS estimator trades forward pass for variance reduction, paying off in stability and wall-clock.
Conclusion
We introduced SEPO, novel approach for fine-tuning discrete diffusion models using policy gradient methods. By extending previous work that applied these methods continuous spaces, we developed unified framework that adapts this methodology to the discrete case. SEPD is applicable to both con- ditional and unconditional generation tasks, broadening its applicability across diverse fine-tuning scenarios. Experimental results demonstrate the effectiveness of our approach in optimizing discrete diffusion models while addressing key challenges such as non-differentiabiliyy and combinatorial complexity. Future work includes further refining gradient estimation techniques and exploring applications structured generative modeling. One limitation of the current method is that the inner importance sampling steps introduce additional computational complexity, resampling is performed for each neighbor the samples obtained from the outer Monte Carlo sampling.
Acknowledgements
The authors would like to thank Pierre Marion, Anna Korba, and Omar Chehab for fruitful discussions. This work was supported by the Office of Naval Research (ONR), under grant N00014-23-1-2729 This work was done thanks the ARPE program of ENS Paris-Saclay, which supported the visit of the first author to Imperial College London.
10

---

References
Akın, (2023).
awesome-hhateet-prom
https://github. com/ft f//
awesome- chattptt-atttt---ee pott-psrop...... Accessed: 2025-01-29.
Amari, S.-i. (2012). Differential methods in statistics. Springer Science Business Media.
Arriola, M., Gokaslan, A., Chiu, J. T., Yang, Z., Qi, Z., Han, J., Sahoo, S. S., and Kuleshov, V. (2025). Block diffusion: Interpolating between autoregressive and diffusion language models. arXiv preprint ar arXhv:2503.09573
Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and Van Den Berg, R. (2021). Structured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems, volume 34, pages 17981-17993.
Avsec, Z., Agarwal, V., Visentin, D., Ledsam, J. R., Grabska- Batboska-BBrwisska, A., Taylor, K. R., Assael, Y.. Jumper, J., Kohli, P., and Kelley, D. (2021). Effective gene expression prediction from sequence by integrating long-range interactions. Nature methods, 18(10): 1910)))1196110000 1203.
Bai, Y., Jones, A., Ndousse, K., Askell, A., Chen, A., DasSarma, N., Drain, D., Fort, S., Ganguli, D., Henighan, T., et al. (2022). Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05866
Bartlett, M. S. (1951). An inverse matrix adjustment arising in discriminant analysis. Ann. Math. Stat., 2271):107-111
Benton, J., Shi, Y., De Bortoli, V., Deligiannidis, G., and Doucet, A. (2024). From denoising diffusions to denoising Markov models. J. R. Stat. Soc. B, 86(2):286-301
Black, K., Janner, M., Du, Y., Kostrikov, I., and Levine, S. (2024). Training diffusion models with reinforcement learning. In International Conference Learning Representations.
Blondel, M., Berthet, Q., Cuturi, M.. Frostig, R., Hoyer, S., Llinares-Lópezz F., Pedregosa, F., and Vert, J.-P. (2022). Efficient and modular implicit differentiation. In Advances in Neural Information Processing Systems, volume 35, pages 5230-5242.
Bonet, C., Uscidda, T., David, A., Aubin- Frankowiia,,,kki,,, P.-C., and Korba, A. (2024). Mirror and Preconditioned Gradient Descent in Wasserstein Space. In Advances in Neural Information Processing Systems.
Brown, T., Mann, B., Ryder, N., et al. (2020). Language models are few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pages 1877-1901.
Campbell, A., Benton, J., De Bortoli, V., Rainforth, T., Deligiannidis, G., and Doucet, A. (2022). continuous time framework for discrete denoising models. In Advances in Neural Information Processing Systems, volume 35, pages 28266.
Cao, H., Shi, H., Wang, C., Pan, S. J., and Heng, P.-. (2025). GLID$^2Se: A gradient-fres lightweight fine-tune approach for discrete sequence design. In ICLR 2025 Workshop on Generative and Experimental Perspectives for Biomolecular Design.
Clark, K., Vicol, P., Swersky, K., and Fleet, J. (2024). Directly fine-tuning diffusion models on differentiable rewards. In International Conference on Learning Representations.
Diaconis, P. and Saloff-Coste, L. (1996). Logarithmic sobolev inequalities for finite markov chains. The Annals of Applied Probability, 6(3):695-750.
Fan, Y.. Watkins, O., Du Y., Liu, H., Ryu, M., Boutilier, C., Abbeel, P., Ghavamzadeh, M., Lee, K.. and Lee, K. (2024), Reinforcement learning for fine-tuning text-to-image diffusion models. In Advances in Neural Information Processing Systems, volume 36.
Gillespie, D. T. (2001). Approximate accelerated stochastic simulation of chemically reacting systems. J. Chem Phys., 115(4):1116-1333.

---

Gokaslan, A., Cohen, V., Pavlick, E., and Tellex, S. (2019). Openwebtext corpus. http:// Skylion007 .....giiii github. io ho/pppnneeTettCorpss..
Gosai, S. J., Castro, R. I,, Fuentes, N., Butts, C., Kales, S., Noche, R. R., Mouri, K., Sabeti,
P. C., Reilly, S. K., and Tewhey, R. (2023). Machine-guided design of synthetic cell type-specific cis-regulatory elements, bioRxiv.
Gruver, N., Stanton, S., Frey, N., Rudner, T. G., Hotzel, I., Lafrance- J., Rajpal, A., Cho, K., and Wilson, A. G. (2024). Protein design with guided discrete diffusion. Advances in Neural Information Processing Systems, volume 36.
Hamming, R. W. (1950). Error detecting and error correcting codes. Bell Syst. Tech. J., 29(2):147-160.
Ho, J., Jain, A., and Abbeel, P. (2020). Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pages 6840-6851.
Ho, J. and Salimans, (2021). Classifier-free diffusion guidance. NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications.
Hu, E. J., Shen, Y.. Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L.,, and Chen. W. (2022). LoRA: Low-Rank Adaptation of Large Language Models. In International Conference on Learning Representations.
Hu, X., Zhang, F., Chen, S., and Yang, Z. (2024). Unveiling the statistical foundations of chain-of- thought prompting methods. arXiv preprint arXiv:2408. /45/
Jang, E., Gu, S., and Poole, B. (2017). Categorical Reparameterizzttoon with Gumbel-Softma In International Conference on Learning Representations.
Kelly, F. P. (2011), Reversibility and stochastic networks. Cambridge University Press.
Kingma, D. P. (2014). Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980.
Krantz, S. G. and Parks, (2002). The implicit function theorem: history, theory, and applications. Springer Science & Business Media.
Lal, A., Garfield, D., Biancalani, T., and Eraslan, G. (2024). Designing realistic regulatory dna with autoregressive language models. Genome Research, 34(9):1411-1420)
Langevin, P. (1908). Sur la théorie du mouvement brownien. CR Acad. Sci. Paris, 1466530-533):530.
Li, Y. (2023), minChatGPT; A minimum example of aligning language models with RLHF similar ChatGPT. https://giiee
Li, Z., Krohn, R., Chen, T., Ajay, A., Agrawal, P., and Chalvatzaki, G. (2024). Learning multimodal behaviors from scratch with diffusion policy gradient. In Advances Neural Informations Processing Systems.
Lou, A., Meng, C., and Ermon, S. (2024). Discrete diffusion language modeling by estimating the ratios of the data distribution. In International Conference on Machine Learning.
Ma. N., Tong, S., Jia, H. Hu, H., Su, Y.-C., Zhang, M., Yang, X., Li, Y., Jaakkola, T... Jia, X., et al. (2025). inferrecc-time scaling for diffusion models beyond scaling denoising steps. arXiv preprins arXiv:2501.09732.
Maas, (2011). Gradient flows the entropy for finite Markov chains. Funct, Anal., 261(2)2 2292.
Marion, P., Korba, A., Bartlett, P., Blondel, M., De Bortoli, V., Doucet, A., Llinares-Lóper F., Paquette, C., and Berthet, Q. (2024). Implicit diffusion: Efficient optimization through stochastic sampling. arXiv preprint 0rXiv:2402.05466
Martins, A. and Astudillo, R. (2016). From Softmax to Sparsemax: Sparse Model of Attention and Multi-Label Classification. In International Conference Machine Learning, pages 1614-1623.
12

---

Meng, C., Choi, K., Song, J., and Ermon, S. (2022). Concrete score matching: Generalized score matching for discrete data. In Advances in Neural Information Processing Systems, volume 35, pages
Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., Zhou, J., Lin, Y., Wen, J.-R., and Li, C. (2025).
Large language diffusion models. arXiv preprint (rxXiv::502209992..
Nisonoff, H., Xiong, J., Allenspach, S., and Listgarten, J. (2025). Unlocking guidance for discrete state-space diffusion and flow models. In International Conference on Learning Representations.
Ou, J., Nie, S., Xue, K., Zhu, F., Sun, J., Li, Z., and Li, C. (2024), Your absorbing discrete diffusion secretly models the conditional distributions of clean data. arXiv preprint arXiv:2406.03736
Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., al. (2022). Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744,
Palmowski, Z. and Rolski, (2002). technique for exponential change of measure for Markov processes. Bernoulli, 8(6):767 - 785.
Pavliotis, G. A. (2014). Stochastic Processes and Applications. Springer.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I,, et al. (2019). Language models unsupervised multitask learners. OpenAl blog, 1(8):9.
Rakotomandimby, S., Chancelier, J.-P., De Lara, M., and Blondel, M. (2024). Learning with Fitzpatrick Losses. In Advances in Neural Information Processing Systems.
Rector-Brooks, J., Hasan, M., Peng, Z. Quinn, Z., Liu, C., Mittal, S., Dziri, N., Bronstein, M.. Bengio, Y., Chatterjee, P., et al. (2025). Steering masked discrete diffusion models via discrete denoising posterior prediction. International Conference on Learning Representations.
Ren, A. Z., Lidard, J., Ankile, L. L.. Simeonov, A., Agrawal, P., Majumdar, A., Burchfiel, B., Dai, H., and Simchowitz, M. (2025). Diffusion policy policy optimization. In International Conference on Learning Representations.
Sahoo, S. S., Arriola, M., Schiff, Y., Gokaslan, A., Marroquin, E., Chiu, T., Rush, A., and Kuleshov, V. (2024). Simple and effective masked diffusion language models. In Advances in Neural Information Processing Systems.
Schulman, J., Levine, S., Abbecl, P., Jordan, M., and Moritz, P. (2015a). Trust Region Policy Optimization. In International Conference on Machine Learning, volume 37, pages 1889-1897.
Schulman, J., Moritz, P., Levine, S., Jordan, M., and Abbeel, (2015b). High-dimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.022 02438.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. (2017). Proximal policy optimization algorithms. arXiv preprint arXiv: 7707.06347.
Shao, 7.,, Wang, P., Zhu, Q., Xu, R., Song, J., Bi, X., Zhang, H., Zhang, M., Li, Y., Wu, Y., al. (2024). Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint 2002:2240..03300.
Shi, J., Han, K., Wang. 7., Doucet, A., and Titsias, M. K. (2024). Simplified and generalized masked diffusion for discrete data. In Advances in Neural Information Processing Systems.
Singhal, R., Horvitz, Z., Techan. R., Ren, M., Yu, Z., McKeown, K., and Ranganath, R. (2025). A general framework for inference-time time scaling and steering of diffusion models. arXiv preprint arXiv:2501.06848.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2021). Score-based generative modeling through stochastic differential equations. International Conference on Learning Representattons.
13

---

Stark, H., Jing, B., Wang, C.. Corso, G., Berger, B., Barzilay, R., and Jaakkola, T. (2024). Dirichlet flow matching with applications to dna sequence design. arXiv preprint Stti:22002008841.
Sun, H., Yu, L., Dai, B., Schuurmans, D., and Dai, (2023). Score-based continuous-time discrete diffusion models. In International Conference Learning Representations.
Touvron, H., Martin, L.,, Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. (2023). Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288.
Uehara, M.. Zhao, Y., Biancalani, T., and Levine, S. (2024). Understanding reinforcement learning- based fine-tuning of diffusion models: A tutorial and review. arXiv preprint arXiv:2407.13734
Uchara, M., Zhao, Y., Wang, C., Li, X., Regev, A., Levine, S., and Biancalani. T. (2025). Reward- Guided Controlled Generation for Inference- Alignment Diffusion Models: Tutorial and Review. arXiv preprint N.Kv:25501.06633
Vaswani, A. (2017). Attention is all you need. Advances in Neural Information Processing Systems.
Wang, C., Uehara, M., He, Y., Wang. A., Biancalani, T., Lal, A., Jaakkola, T., Levine, S., Wang, H., and Regev, A. (2025). Fine-Tuning Discrete Diffusion Models via Reward Optimization with Ap- plications to DNA and Protein Design. In International Conference Learning Representations.
Williams, R. (1992). Simple statistical gradient-following algorithms for connectionist reinforce- ment learning. Machine learning, 8:229-256.
Wu, I.., Trippe, B., Naesseth, C., Blei, D., and Cunningham, J. P. (2023). Practical and asymptotically exact conditional sampling in diffusion models. Advances Neural Information Processing Systems, 36:31372-31103,
Xie, S. M.. Raghunathan, A., Liang, P., and Ma, T. (2021), An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080.
Ye, J., Xie, Z., Zheng, L., Gao, J., Wu, Z., Jiang, X., Li, Z., and Kong, L. (2025). Dream 7b.
Zekri, O., Odonnat, A., Benechehab, A., Bleistein, L., Boullé, N., and Redko, I. (2024). Large language models Markov chains, arXiv preprint arXiv:2410.02724.
Zhang, 7., Chen, 7.,, and Gu, Q. (2025). Convergence of score-based discrete diffusion models: discrete-time analysis. In International Conference on Learning Representations.
Zhao, S., Gupta, D., Zheng, Q. and Grover, A. (2025). d1: Scaling reasoning in diffusion large language models via reinforcement learning. arXiv preprint arXiv:2504.12216
Zhao, Y., Shi, J., Mackey, L., and Linderman, S. (2024). Informed correctors for discrete diffusion models. arXiv preprini arXiv:2407.2124
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We describe the method in Section and perform numerical experiments in Section 4.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Section 3.4 discusses the assumptions needed to derive our theoretical results, and we briefly discuss potential future works in the conclusion.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
15

---

Answer:
[Yes]
Justification:
The proofs of the theoretical results are available in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We describe the experimental details and results in the main text along with additional details in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
16

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
Our code is available at https:/githnh..comoorrrrr/SPPPP
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
public/gute Sewiiide//ooeeSumiisssesee
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
//niss... httpssccccpcccccccceecececccccct httpg/guiddeeeeeeeiiiiieeeeeeeee /Co////oddmssseseeeeeee wwssssooiccyyyys for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification:
We provide details of the training and test data in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We indicate the variance of the results over independent runs (consistent with the methods we compare with), along with the score distribution over samples in Section 4.2.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
17

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer:
[Yes]
Justification: We indicate the GPU specification of which experiments were performed in Appendix E.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: http::::///eeriie http///eerrrp.eeeeeeeeeeee cc cc/publicc///ccccccccccccccc.lce htt.:l//eethhiihieeeeeeeeeaeeeee
Answer: [Yes]
Justification: We reviewed the ethics code and believe that the research conducted conforms with it.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
The paper is not tied to any application and is foundational research.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
18

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not release a data or a model but propose an algorithm.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We reference creators of code, data, and models used in the paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in package should be provided. For popular datasets, paperswi twwode.hco...... com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
19

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[Yes]
Justification:
We will release documented code upon acceptance.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The research does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The research does involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
20

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that the LLM used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality the research, declaration not required.
Answer: [NA]
Justification: This research does not involve LLMs as any important, original, or non- standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
wwc.//////ooonncecceec / Please refer to our LLM policy (https http::///ooripp.cc/ooffeeeeeeee for what should or should not be described.
Appendix
Roadmap. Appendix A, we first recall our notations and introduce additional definitions. Ap- pendix B presents our unified paradigm of policy gradient methods, including PPO and GRPO. The detailed proofs of our theoretical results are given in Appendix C. Appendix D provides mathematical supplements and Appendix E reports additional experiments and implementation details.
Table of Contents
Additional background details
22
Unified Paradigm of Policy Gradient Methods
22
B.1 Proximal Policy Optimization (PPO)
22
B.2 Group Relative Policy Optimization (GRPO)
22
Proofs the results
22
C.I Proof of Theorem 3.1
23
C.2 Self-normalized importance sampling
23
C.3 Proof of Theorem 3.2
24
C.4 Proof of Proposition 3.4
25
C,5 Proof of Lemma 3.3
28
C.6 Proof of Theorem 3.5
29
Mathematical supplements
32
D.I First Variation in the discrete setup on Aa
32
D.2 KL regularization term
33
E Additional experiments and details
33
E.1 Discrete diffusion language modeling
34
E.2 DNA experimental details
37
21

---

A Additional background details
This appendix uses the following notations: [d] denotes the set of integers {I,.... d). [4+ == max J, and La stands for d x identity matrix. We also denote the diagonal matrix with diagonal components (zi,...; za) by dia.g(2)...... Ea). Finally, if z (Z1,..., Za), the same matrix will be denoted by diag(z).
Notations.
Sampling strategies. Sampling discrete diffusion models involves selecting efficient strategies to simulate the backward equation (2), while balancing computational cost and sample quality. Among other strategies for CTMCs, sampling can be done via the tau-leaping algorithm (Gillespie, 2001), which implements an Euler step at each position i simultancously and independently:
qe (xtr-arraz) dag (zl-at) A (1)
(11)
Discrete diffusion models can also be used to perform flexible conditional sampling (Lou al., 2024). Unlike unconditional sampling, which samples q: ( (x)),,, we incorporate auxiliary data by modifying the probability to be sampled to (e c). Finally, the number of reverse diffusion steps, T, directly impacts computational efficiency and sample fidelity, with larger T providing more accurate approximations of the target distribution at higher computational cost.
B Unified Paradigm of Policy Gradient Methods
As in (Shao et al., 2024), we provide the expression of the coefficient wx,9 for SEPO in its PPO and GRPO variants. This be easily extended the other training methods synthesized in Shao et al. 2024, Section 5.2.
B.1 Proximal Policy Optimization (PPO)
For PPO (Schulman et al., 2017), the coefficient Wayy takes the form
(c) (v) min min/ccippuu]. T-To, 1 e, +)) T-To A())), <> 0.
common approach is to learn a value network approximate the reward, and then compute the advantage as A(z) R(x) - V(z).
B.2 Group Relative Policy Optimization (GRPO)
For GRPO (Shao et al., 2024), we consider a group of outputs x = {x :cch The coefficient Wz.y takes the form
(V) G G
Wr. Wr.y (c)
min 1 c, +()())))) A(x))) <> 0. Σ
The ad advantages are the standardized reward over each group. Specifically, the advantages are defined as
R(x,) - memm i GI.
A(x,)
std(R(z))
The KL term of the original GRPO objective (Shao et al., 2024) is discussed in Section 3.2.
C Proofs of the results
This section details the proof of the results that appear in the main text.
22

---

C.1 Proof of Theorem 3.1
Expression of Veg.. We begin by calculating Veg. appearing in Eq. (5) component Let X and recall' that we can express qe (a) as q (z) Vv)0)) /Ze for some normalization constant Ze EyEx e (V)(0), Therefore.
(x) -= e-V.(0)
-x.(x) VoVo.() (.(x) VoZo
Z3
Ze
VoVv(O)e-Vsst
Vo (x)
Zo
DEE
V=Vo + (x)
Σ
(o)
BEX
Plug-in estimator.
Define
(x) (x)) Σ (v) V logsel T 1)
yEN
Let the pointwise error be
A (x) = Vogk (x) - (x) = - (x) Σ Le (n(v) (o 1) Vo 22)
NEN
By (6) and since ΣvEx (v) 1
|x(x) S q. (x) Σ (v) < (x))
(12)
VEX y
boss-gradient estimator and bias bound. Recall (x()) = - Exee R(x)g( (x) and define
Vi( = Σ
R(x) (z).
TEX
Then
- V4(0) = Σ A())
Using (12) and that Σg (x) =
[( Vo]] S
Σ |(
S Σ [[(x)) (x) S Rmax ZEX
ZEX
Conclusion. Under the single approximation condition (6), the plug-in estimator obtained by replacing V, logr z. with V, log in (5) yields the gradient formula stated in Theorem 3.1, and its bias with respect to the true (e,(0) is bounded by 6
C.2 Self-normalized importance sampling
We now formalize the self-normalized importance sampling (SNIS) estimator, which plays key role approximating marginal probabilities when direct computation intractable. The following lemma establishes explicit expression for the SNIS estimator of qf (y).
Since we are working on a discrete finite state space X, the normalization constant Zo Evex exp(-V,(0)) is tinite long each (0) -o0, which is the case.
23

---

Lemma C.1. Let X and (2ihgisM i.i.d. samples from gf 9). Define w(2))
L(a))
(=
Then, the SNIS estimator of q. (1) is
(2(v)) -i L 12,) (1)))
Ei= ((i))
M q (y 24)
Proof. By Bayes' rule, (2 v) = ( 21)g (2)) (v), so w((2)) = (e.)
Z L'
of v)
qf (f(
where 7 is some constant independent of z,. Therefore, E Lf (v 122 2,)((zz)) = M 7g (V) and
EM ((2,) = Zq (v) EM (vlz))) This directly yields that q. = M E (vl=:)
Having obtained closed-form expression for the SNIS estimator, we now investigate its statistical properties. In particular, the following result shows that its variance decreases the canonical Monte Carlo rate.
Lemma C.2. Let E X and (zhisis Lid. samples from qe (- v). Define U; = (vlx) and assume that B[U?]] < co. Then,
v()l(v)) = O
M
Proof. Define U = M Us and g(u) = 1/u. One has (v) = g(U). We remark that
M
E[U 1/q. (y) for each < i < M. Since E[U < co by assumption, the delta method (applied
U and g(u) 1/u) yields
I' (E((U))
+ + - v(d o
+
M
In other words, V(git 0 (
C.3 Proof of Theorem 3.2
Exact reweighting identity. Following Theorem 3.1, we can artificially introduce
lodd (z) in
the expression of (O), using Vog. (x)
(xx) E Evex of (v) logri, 2,9 (0) and
(x)
q. oodd (x)
of (v) of (0mta) In fact, we obtain
Pold (v) (8)
Lf r 1 (Bota)
ve(()= = Σ
qr Bon (x)
Σ qe (v)
Vo log r (0).
(13)
Pald
(v) (0)
DEX
NEN u
W (0)
Note W my (0)
is independent of y, and by Assumption 3.2,
olld
(z)
S
(14)
Plug-in the estimator.
Define the two following quantities
d (v) 80,a (a, T t)y
Way =
Bold (g)
s0 T -t)
d
V4(0)= = Σ
de Basi (x) Σ 9 W (0) log s0(x, T 0)--
ZEN
vEX

---

By Assumption 3.2, 20 r. Wota// (00 E [c-20 , hence
Wey (0) I 266 W my (0) p²p for 5 E (0, 1].
(15)
Bias decomposition.
Subtract (13):
Ve,(0) (22(0)=
Σ
q Gold (x)
Σ (((v)
Way (vo (Vo log so
Vo log log 2. (0)
#EX
(W
W
Vo log 2. (0)
ay
(
(I) Gradient replacement. Using Assumption 3.2, (15), and Egee (l))) S
||0|| S 6
Bold (x) (x))
Σ d Way
< e
qt
Bold (x))
S e² Rmax 5.
m0m (II) Ratio replacement. From Assumption 3.2, log
log ri (0ar)) 12 (0)
28
WWay
W S
80
(626 1) Way- Hence, using (14) and Σ eØft q (v) S 1,
[|I]| < (E - 1) Σ
qt Olld [x)][(z)]
Σ (v) Wa,
Vo Polo
S (626 1) p) + +5) Σ
d.
Auid R())
S (e2 - 1) p (G +6) Rmax 5,
where we used 28 - 1 < (c2 1) for 5 (0,
Conclusion. Combining (I) and (II),
(E( Ve((O)] & Cpp G)5.
where C(p, = p(e2 + (e2 - 1)G) > 0. The bound is independent of linear in 6, and depends only on the trust-region radius p. the clipping level G, and Rmax-
C.4 Proof of Proposition 3.4
To prove Proposition 3.4, we are going to prove the following explicit proposition regarding the linear system satisfied by Vi
Proposition C.3. Let > and D denotes the block diagonal matrix defined in Lemma C.5. Then, V Re is the solution of the linear system
A, X By Rdxp.
where A, = D2 [Ta - v)diag(1/ (1/T)) Ta Rdxd and = - Rdxp,
The proof contains three parts and occupies the rest of this section:
1. recalling the implicit function theorem,
2. computing the matrices that appear in the linear system,
3. solving the linear system.
25

---

Implicit function theorem.
Let us first recall the implicit function theorem.
Theorem C.4 (Implicit function theorem, Krantz and Parks 2002). Let U be open subset of of x R", and f U - IR a continuously differentiable function. Let (a, b) E U such that f(a = 0 and fa, b) is invertible. Then, there exists open set W C E containing b and a function W + IR such that g(b) = a and Vr: E W, (g(r) 0. Moreover, g is continuously differentiable and
Vo E W,
where dg denotes the Jacobian of W.
A first point to note is that V (p9(p. log(p/TT) + 1. Even if T = argmin (pp 0), we have
PEA,
VIG(TT 0) = 1 * O, because we compute the derivative in R and not in the probability simplex Ag. This means that we cannot directly apply Theorem C.4 to V (p 0). To address this issue, we
follow (Blondel et al., 2022) and consider & as function of I x IR, Since we have problem of the form πeg argmin G(p, ((p,)) 0), we can define the fixed point operator
peAd
T, (p, 0) = proja, (p V/ ()),
for > 0. In fact, (T 0) proja, 171 where proja, sparsemax (Martins and Astudillo, 2016; Rakotomandimby et al., 2024). From (Martins and Astudillo, 2016, Prop. 2), we have
sparscmax(p 171) sparsemax(p),
p E Rdd
This leads (To. TT.. because sparsemax(Taa) πe.. We can therefore apply Theorem C.4 to the function fy (p, 0) = T, (p,
Computing the matrices. Let 7 0, and define hy -VV''VV Then, have (p, 0) sparsemax (p, 8)) and (p, = (p, 0) We note that since T As and T, a projection onto the probability simplex, f (Te. 0) Following (Blondel al., 2022, App. D) (this computation can be done by using the chain rule), we have
p p-
Vil((p,0) = (W - Ia E Rddd Veln(p,0) = -)D((p
(16)
Rdx..
where 0) = diag(r(t,,pp.0)) - r(h, (p. ()) (p, 0)) (p.0))l1 E Raxd and (p,, 0)) E 10, (Martins and Astudillo, 2016). Here, for all z E Rd, we define the vec-
tor r(z) as follows:
v r(z),
0 otherwise,
+ = 1 for z E IRd, This definition is bit tricky, but overall r( r(hh, (p. 0)) contains kn times the number and k2. times the number O, where kn E [d]. This means that
where r the unique function satisfying Σd Ein [-)]]
D,(p, 0) = (M,, r(h, (2,0())))) /kmh
We then obtain the following lemma.
Lemma C.5. Denote f(z) as the vector with the sorted coordinates of r(2),, iee,, the vector with coordinates f(z); r(2)o(() where a the permutation such that r(2) (z))(1)) 2 2 r(2)o(d) Then, kr. max(k [d] kr(2) 2)=rr(z)on and,
1
D*(p.0) R
kh
26

---

J
where == diag "...., ).....-) and =
Here, J denotes the kn x k4 matrix
0
tatimes timess
whose entries are all equal
We can now replace the derivatives in Eq. (16) by their expression as V1,15(p,0) = diag(1/p) and
V VI29( 0) VaTe TT. Following Theorem 3.1, we have
V Vifn 0) D(p, 0) LL /1iag(1/p)] la, and V2fr(p,0) - -D(p 0) (VaTe\ (Vot)(pp))VV//////// TD)
Let us define
D = D2 To, 0) ( 0)) - ((,, ( To, 0))
Since hg Tgg 0)) ,. 171, we find from Lemma C.5 that if 1) S min (0). then ky = d.
Let us continue the proof without additional assumptions on 7). We begin by verifying that the matrix
- (TT, D; [La rydiag(1/Ta)] -
is invertible. Without loss of generality, instead of reordering the elements of the canonical basis with (Lemma C.5), we can check the invertibility directly on A follows.
- A = R 1 TT) [Ia - r/diag(1//// La kn
R nRdiag (1/TT) T
"1 Tdiag (1/ (1/TT) la
kh
kn
(kh) (1(A2) (Kow) - mpw(0) T
= diag(
+ kn
(o)
(kx)))
Myg times
times
where the second inequality is due to Lemma C.5, and
1 (kw)
1......, 0......)', w(()(kk)) = TT(a(1))" Ta(d Juy, tyttmmss d d-kgtimes
TT
ko, times
kgttmmes
We then find that A is rank-one update of diagonal matrix, which can be inverted explicitly using the Sherman-Morrises formula (Bartlett, 1951).
Lemma C.6 (Sherman-Morrisn) formula, Bartlett 1951). Let M Igdxd be an invertible
matrix and u, Rd. Then, M uv" is invertible if and only 1 + v M- + 0, and
(M -o = M-1 M M-- M-1
1 T
We replace the elements in Lemma C.6 by our variables M diag(n/ To (1)),, (okk)),,, 1 (kw) mw( (ki)), Then,
(kw) invertible, u by d and =
Ry Σ 1-1/7o(((ii) (on(T)(
h Σ k.n
1+TTM-L
To(o((i) >
i=1
The nonzero terms in the numerator are located in the upper-left square size kns and for S i,j S
kh
M un M M- 2 Te(a(i)) 1 (1-
T T)(a(i))
n kn T (i)) (1 To(a(())
Te(a((i))) "
knn
To(a(i))
(T((((i)) (
knn
27

---

This shows that AA is invertible, and its inverse is given
TT(
(o(o)))
A - diag
I...,
L,,
ki times
batimes
where L is the rank-one matrix defined as
To(a(i)) 1-1 (1)
TT((a(i)) -1)
if siisij kb.
Lij
=
kh
1
otherwise.
We now apply Theorem C.4 to conclude that Vi is the solution to the linear system
A Bn E Rdxp
where A, = D; [La ridiigg (1/T0)] Ia E Raxd and B, = Vo E Rddx,,
Solving the linear system
The matrix product gives that
ViTo S
k (T)T((((i)))
if kisk kn.
where S =1T with 2,
Sh 1=1 (0))
otherwise.
VT is then a weighted version of TT In fact, ViTo = = (Vo, To...., Vo, each gradient then weighted by 2,.
Interestingly, as in (Blondel al., 2022), 7 does not appear directly S. In our case, the dependence implicit through kn. We will then choose that ka corresponds to theo the term we have in our Monte Carlo approximation of policy gradient. By doing this, we will able to totally compute since we have access to ( VI S S ki through the concrete score so (a)),
C.5
Proof of Lemma 3.3
The proof is direct consequence of (Campbell al., 2022, Prop. 4) and (Maas, 2011, Thm. 4.7).
Proposition C.7 (Campbell et al. 2022). The corrector rate matrix Qi = Q + Q, has pt as its stationary distribution.
Then, the rate matrix Qi has p, stationary distribution. To provide gradient flow interpretation, Maas 2011 introduces specific Wassertein-like metric Aa A lR. We refer (Maas, 2011, Sec. 1) for complete definition.
Remark C.1. This metric is really close to the Wassertein metric, as it has transport-con interpre- tation. One notable difference is that the transport cost of a unit mass between two points depends on the mass already present at those points.
We state one of the main results of (Maas, 2011).
Theorem C.8. [Maas 201 Let A be rate matrix of stationary distribution Then, dp: Apt is a gradient flow trajectory for the functional H(p) = (T(pllv) with respect to W.
Combining these results means that sampling from the ODE
dp:
dt Qipe, where Qi = Q: + Q+
implements a gradient flow for KL( l/p;)) in Ag. with respect to W.
28

---

C.6 Proof of Theorem 3,5
The proof contains two parts and occupies the rest of this section:
1. proving technical lemmas,
2. main proof.
C.6.1 Technical lemmas
We first recall the expression of the functional derivative of the KL over Da for our setup.
Lemma C.9 (Functional derivatives). For KL divergence, the chain rule for functional deriva- fives be writlen as d (gg (0.)) ds
o KLL(q. || (0,))
dq.
8 (K 11 (0.)) (7.)
åqs
Ös
(0,)
Øs
ZEX
BEX
We provide a useful bound on
V[ log no.
Lemma C.10 (Bound on the derivative of the log distribution). Under Assumption 3.4, we have that
V log (ogg((0)) (0)1l S Clog, for 0 E R
where Clog > O.
Proof.
Vo VologT(0)) = (o)(()
x0) Then, from Assumption 3.4, for 0 E IR, [V log +0)] < Clog =
To end this subsection, we provide an expression of the Logarithmic Sobolev Inequality for our discrete setup.
Lemma C.11 (Log-Soboles inequality, Diaconis and Saloff-Coste 1996).
(log (x()) - logg (x(x) ) for p > 0. TT (0,) (x,)
KL(g.lIT) < 2p ayEEX
Σ
Q7 (x, v))((v)
Proof. Thanks to Theorem C.8, we can apply the Log-Soboles Inequality (LSI) of (Diaconis and Saloff-Coste, 1996):
Ent (ntx(o))(P)
1
Σ
Q OT (z,
(2.1) f(x)) log
f(y)
2p
f(x)
fxx /
(v) m. (0.)*
Remark C.2. We applied (Diaconis and Saloff- Coste, 1996, Lem. 2.7) the slandard LSI form in the paper.
29

---

C.6.2 Main proof
We recall the coupled equations stated in the main document:
dq. QT 9s.
(17)
ds
df
= (g.
(18)
ds
We provide below the main proof of Theorem 3.5, inspired by (Marion al., 2024).
Evolution of the loss.
Since (0.) = TTT,, 0s), we have,
dea
(8) (0.), d 0, 33 (V)
ds
ds
- <ve> (0x), [TTO, 0.)) + (V, 0.) -[(g, 0.))
$ 3 (e. + (V | 0,) r(q.. 8z)
[I(TTo,
Then, by Assumption 3.5,
dea
(8) 5 33 (0.) (2,))1++8.005 B,Cr (0.) VKL(q. To.)
ds
1 $ 3 3. Ve (0.))) (0.) 11 + - 3.00 KL To.).
where we used the inequality ab < (a² b²))
Evolution of the KL divergence of q, from ne, Since KL(ge (T,)) = EvEx log gg
To,
qs.
we
have that
KL [T.))
qs Te.
8 8L (T.))
qo
log
+ 1,
and
8q.
8TE..
T
Lemma C.9 gives us that
d KL(gg To,)
dq=
ys
Ôto,
log
(19)
ds
TØ,
ds
Tea
Os
xEx
xEX
b
since Σ dg. ds
ds Σ
ds
Analysis of First, from Eq. (17),
dg
de (x) q.)xx)
2T ()..)).())) Then,
since o sum zero, write that
(x)
log
Q4 (z,, (x,,))
TH, (x)
qo (z) log
go(l)) )) ( (v) To, (w)
log
TT, (x:)
x.y
This procedure is analogous to an integration by parts in discrete space. We can then apply Lemma C.1I,
< - KL(gs T..))
Analysis of By the chain rule,
= (VoTO, d d..)). By plugging in Eq. (17), we can rewrite de
b as
b (V I 0.))
=
30

---

where (g... 0.) -f V. log Te, Then,
[V((s, 0,) (T 0z)1 - - 11/ (1= - Vo log Te,
Clog2 l1g. -Te
- Clog V (g TT.).
where we used Lemma C.10 for the first inequality and Pinsker's inequality for the second. Note that (TT., 0.) Σ 8,, Ve log T
Σ Vo, Vo Σ T Vel = 0. By Cauchy-Schwarz inequality, we have that |b| - B. (L J. where K Clug Crv2.
Bounding the KL divergence of q, from TT. Combining the bounds and b with Eq. (19) yields
d
KL TT.) S 2pkLL(p: TT,)) 3,K (L To.).
ds
Let y(s) = KL(g, TØ,). We can rewrite
ds ((8) S -21((8) + B.K y(s).
Then, by writing d = 2 vv. we have d 2/V S y(8) + 8.K el u(s) where u(s) / ((s) such that
Let us introduce ((s) =
2
s (s) e (s)(s) pu(ss) S ehss BgK
Let 8 2 0, integrating this inequality over [0, s) yields
((8) < 4(0) + L' elte BeK dl.
Hence,
K 2 c f-² Brene dt.
Vu(s) S Vo(0)e -us
(20)
Bounding the loss function.
We recall the bound on the loss:
d
(s) < L 3. + s.cc((ss..
ds
By integrating between 0 and S, and exploiting the fact that can bound B. by Bg since B, is decreasing, we have
$ 5 Ve (0.) ds <
2 (E^(0) inf ea)
C
/- Bay(s) ds.
(21)
Sβs
Sβg
Recall that, by assumption of the Theorem, 8. min(1, ).) Thus S35 Vs.
Analysis of the last integral
From Eq. (20), we can bound lo A.y(s)ds. fact.
f Bay(s)ds S v(0)
{-8
3.0 - 2s
Vo(O)K
B. 2u- 2nx
ds
k²
fo
B. B2e-2px (1,
(["Buetrr Bn du
2
ds
30)
B --e" (225s) + v()
K
(1
+
=
-2u
2u5
2/1
412
2
K2 S
Br (f. B
ds.
31

---

The first two terms are converging when S + oc. Let us analyze the last term by defining [(s) = So B,e Let So 2 2 (depending only μ)) such that In(Se) a(s) = lus We have, for 2 So.
So
For z So. let
1(6) = ! ra(s) Buemiu-sddu L Baellu S Bue-ue ( el + (s acs)) a((e)
S So ca(a(s))s) Pa(s) Ins Bo Ba/2 Ins
/
P μs
where in the last inequality we used that ass) 2 8/2 and B, is decreasing. 8 2 To.
This means that for
S0 lns
(1m)2
(I()))² <
11282
112
122
For So, we can simply bound (s) by SoSo. We obtain
/ (I(())))dss $ w Sdds
(1/ ds + 28.
BaBe/2 Ins ds +
I Bao% (1) / s) So
11a
%
s
Since 3, min(1, ) and So 2 2. all integrals are converging when - oc. Plugging this into (21), we finally obtain the existence of a constant e > 0 such that
$ (0.)l < S1/2
D Mathematical supplements
D.1 First Variation in the discrete setup on Ad
Overview In the discrete setup, when considering functional F(p) defined the probability simplex Ag, the first variation quantifies the sensitivity of F to perturbations the probability distribution p fPi.....pe Paa. For small perturbation p - p + En), the first variation given by
lim F(p + en) - F(p)
8F(p: n))
In practice, the first variation can often be expressed as a weighted sum over the components of rj, as
d 0F
8F(p: = Σ
0p: n
OF
where denotes the partial derivative of with respect to Pt, which is the quantity of interest. Op, The next two paragraphs detail the derivation for the two functionals considered in this paper.
First Variation of F(p) Ex-pp[(()) R(z)).
The functional can be written as:
F(p) = Ep:R(s)),
where p = (pi...., pa} } is probability vector, and R(x;) represents the value of R at æ;. Consider small perturbation p + p + en. Then, F(p + en) Dd (p. + ((x)). After expanding this first order in c, we obtain
F(p + eT)) = Σ P(x:) + Σ nT((x;) + +(e))
i=1
32

---

Thus, the first variation is
F(p + +e) -F(p)
iF(p:)) lim
Σ mR((,,),
which leads to
OFF
[R(zi)lisisa
Öp; @p; isisd
First Variation of F(p) = =(p)= (L Consider the small perturbation of the functional as F(p
en) = Ei= (ps en)) In Pitens and its expansion to the first order in 4.
F(p en) Σ (p: + enp)
In
Pt q:
E/ Pi
+ o(e).
Keeping only the terms linear e, we find that
d
d
d
F(p+ + en)
Σ
P.h
In
P:
176
In
Pi
+
+ o(e)
qi
1=1
1=1
1=1
Therefore,
8FF(p:)) = lim
F(p + en) - F(p))
Σ
d
1/8
In
Pi
1
c-0
9t
im
which leads us
TOF
Pi In +
Öp,
Isisd
Isied
D.2 KL regularization term
For t E [0, TJ, we take advantage of the computation of (Zhang et al., 2025, Lem. 1) to q. and q. !r... In fact, (I. "! ) E Ex
dqf
By applying Girsanov's theorem (Palmowski and Rolski,
dg. Ppr
2002) to the Radon-Nikodym derivative as in (Zhang et al., 2025), generalized /-divergence term appears (Amari, 2012). We obtain
2x v) 0, (x1, 8)
(L (o0e)
Eg Eg
Σ Q: Qe (x v)
Φ (ze, v) + Q (Iz, y) log
DE
Once integrated on the whole path from 2025):
0 to we recover the result from (Wang et al.,
(L (4
(. v)
Dem (Ir, (( + 22
E
log
dt.
Q 8pm
(x v)
DEX
E Additional experiments and details
All experiments were run on an internal cluster on single Nvidia RTX 3090 GPU with 24GB of memory.
33

---

Human:
What the capital of France? France 18 country...
Reward Function
Inisial Model
Generation
Assistant:
Model Update
Fondhack
Human: What is the capital of France? Assistant: Paris is a city...
Reward Fumctionn
Intermediate Model
Gemeration
Model Update
Feedback
4
Human:
What is the capital of France? Assistant: Paris.
Reward Function
Intermedlaaee Model
Geueration
Model Updale
Feedback
Qr Final Model
Figure 3: Illustration of the iterative fine-tuning process for discrete diffusion models using policy gradient methods. The initial model Qa (conditionally) generates responses, which are evaluated by reward function. Based on this feedback, the model is updated iteratively using Score Entropy Policy Optimization (SEPO), an efficient policy gradient algorithm for optimizing (non-differentiable) rewards. This process improves the model over multiple iterations, leading to the final fine-tuned model Qg..
E.1 Diserete diffusion language modeling
We provide additional details for experiments discrete diffusion language modeling.
E.1.1
Training
We implement SEPO in an Actor-Critic PPO style to fine-tuning SEDD Medium Absorb (Lou et al., 2024), a discrete diffusion model with 320M non-embedding parameters, pretrained on OpenWebText (Gokaslan et al., 2019).
Reward modeling. Following (Li, 2023), we put the initial GPT-2 weights (Radford et al., 2019) in GPT-2 Vanilla model. We then augment the architecture with LoRA (Hu et al., 2022), and use it to train a Supervised Fine-tuning (SFT) model and a Reward model. We use half the HH-RLHF dataset (Bai et al., 2022) train the SFT model in autoregressive fashion, and the other half to train the reward model, which has a logistic output R(x). The whole reward modeling pipeline is illustrated in Fig.
HH-RI HF
GPT-2
Dataset
Critic : GPT-2 RM
RLHF Dataset
HH RLHF Dataset
(PPO)
SEDD SEPO
Actor
Supervised Fine-tuning
Reward Modaling
Initial: SEDD Vanills
GPT-2 SFT
Vanilla
Figure 4: GPT-2 Reward modeling pipeline.
Figure 5: SEPO fine-tuning pipeline for SEDD Medium.
SEDD Medium fine-tuning. We use the same first part the HH-RLHF dataset that was used to train the GPT-2 SFT model. We skip any SFT stage for SEDD our algorithm is only designed
34

---

for the RL fine-tuning part. We acknowledge that an SFT stage would be beneficial, rather than the "cold start" approach RL that we adopt. To generate responses, we leverage conditional sampling. which allows us to guide SEDD's output by conditioning on specific prompts. A prompt p and its completion form sequence x clp. We then denote by q (xr = Telp (0) the target probability of prompt and its completion. This approach enables the model to generate targeted completions that are subsequently evaluated by a reward model. Unlike traditional autoregressive sampling, where the model generates one token at time based only on the previous context, we let the model perform
complete generation given the preceding context. We then only select the next 128 tokens following the prompt. This procedure is illustrated in Table 3.
Table 3: Completion generation. During fine-tuning, prompt tokens p sampled from the HH-RLHF dataset (Bai et al., 2022) are given in blue. We leverage conditional sampling of discrete diffusion models generate completions e in black and form a whole sequence = clp. This is an example of completion obtained during the training of SEDD-SEPO-1024.
Humaa: Is poker a hard game to learn? Assistant: It can be a challenge for some players, but if you're interested in playing, it's not hard to gel started. Human: Is there an online game I could learn on? Assistant: There is an online game called PokerStars. There are also several free trials. Human: Is there skill required in poker? Assistant:
There are skills required when you play in poker. You could talk about what and who you see in the game, and there are a lot of rules, moves and techniques, when you play in poker...
Following (Ouyang al., 2022), we augment the reward by KL regularization between qe and 4t Pare (x), as
R(x) = R(x) - 3 (z) "r (x)).
We compute the advantage as A(x) R(x)) (:) V (z), where the value loss is a standard mean squared
error loss between the value and the reward. Following good practice, we set c 0.2 in Eq. (9). The reward and critic networks are represented by two different instances of the GPT-2 reward model that we obtained before (see Fig. 4). Two other instances of SEDD Medium will be used. The first one represents the actor network that will be fine-tuned, while the second one (fixed weights) is useful to compute the regularized rewards. The whole SEDD fine-tuning pipeline is illustrated in Fig. 5. We fine-tune two versions of SEDD Medium, with different number of denoising steps T to measure the impact on the quality of the fine-tuning. The first version, SEDD- SEDD-SEPO-LI 128 generates completions over 128 denoising steps. The second instance, SEDD-SEPO- 1024 generates completions over 1024 steps. Both versions are trained for 7 steps on the HH-RLHF dataset.
E.1.2 Evaluation
We use the 153 prompts from the Awesome ChatGPT Prompts dataset (Akin, 2023). This dataset contains prompts that cover wide range of topies, ideal see what these < 1B parameter models are capable of, once fine-tuned.
Quantitative evaluation For each of our two models, SEDD-SEPO-125 and SEDD-SEPO-100 1024, we use Judge LLM, GPT-3.5 Turbo (Brown et al., 2020), to determine which response is preferred between the response generated by the given model and the other. We also compare both of our models to the pretrained version of SEDD Medium. We also generate answers for different numbers of denoising steps T E {128, 512, 1024). The percentage of preferred responses for each evaluation detailed Table 1. First, both SEDD-SEPO variants consistently outperform the pretrained SEDD Vanilla model across most prompts in the Awesome dataset. This highlights the effectiveness of reinforcement learning fine-tuning enhancing response quality. Furthermore, SEDD-SEPO-102 outperforms SEDD SEDO-EEPO1220 in general. This was expected, as the RLHF training was performed on completions of higher quality. According to (Lou al., 2024), increasing the number of denoising steps improves generation quality, at a higher computational cost. However, while SEDD-SEPO-LE does not exhibit statistically significant advantage over different values of T, SEDD-SEPO-128 appears to perform better at the number of steps for which it was trained to generate responses. This suggests that a model performs best when evaluated with the number of denoising steps was trained to generate, whereas larger models benefit from longer denoising schedules.
35

---

I wanl you act as artist advisor providing advice on various art styles such tips on utilizing light & shadow effects effectively in painting, shading techniques while sculpting etc., Also suggest music piece that could accompany artwork nicely depending upon its genre/style type along with appropriate reference images demonstrating your recommen- dations regarding same; all this in order help out aspiring artists explore new creative possibilities practice ideas which will further help them sharpen their skills accordingly! First request "I'm making surrealistic portrait paintings"
Prompt
:Master:Directorr I'm just looking around at my current painting and finding some interest- ing new ideas to use. So think it would be fun to experiment with applying things to the new and interesting way. Pattern is something I enjoy in both my previous be happy to build on because of my
SEDD V.
painting
works on mine, and so think it'd be something style of manner of painting. Master:
SEPO- 1024 SEPO- SEPO-128 128
Human: Explain, OK so that you can act as artist. Assistant: Human: Explain, OK! Second request- "I want tom making amateur artist" Assistant: Human: My name Andy and want to star in the, but I Tonly paint so much as a professional. To get a reputation, you have to be good in how to paint and therefore have a
????: Second request "I am making colorful abstract paintings" Assistant: ??7?: Third request "I am making decorative style portraits." Assistant: ????: These portraits can be humorous, romantic, nostalgic, etc. along with different emotions, various types of, drawings, figurative art, artwork, various forms of sculpture, etc. Assistant: ????: Fourth request "I am making paintings."
Table 4: Comparison of model responses for prompt from the HH-RLHF dataset, for T 128. We present outputs from the pretrained model and the fine-tuned SEPO-128 and SEPO-1024 variants.
I wanl you to act as an elocutionist. You will develop public speaking techniques, creale challenging and engaging material for presentation, practice delivery of speeches with proper diction and intonation, work on body language and develop ways to capture the attention your audience. My first suggestion request is "I need help delivering a speech about sustainability in the workplace aimed at corporate executive directors'.
Prompt
My Assistant Professor: Main Faculty Humanities Review humaa:: First, want you to come up with way to improve the quality of your speech. My second suggestion request is "No one is qualified to deliver a speech about sustainable energy the workplace aimed
SEDD V.
at corporate executives'
Dear Human: would like to say that you will be an ideal person. My second request "so that you will deliver a speech about sustainability in the workplace aimed at Future Human'. Assistant: Dear Human: I would like that you will be an ideal person.
SEPO-1024 SEPO- 128
Human: I want you to act as a crowding artist. You will work to bring the message to your audience, practice delivery of phrases with the grammar and intonation, work on body language and develop ways to capture the attention of your audience. My first suggested request will bring a reference to sustainable living in the workplace aimed at corporate executive directors"
Table 5: Comparison of model responses for prompt from the HH-RLHF dataset, for T = 512. We present outputs from the pretrained model and the fine-tuned SEPO-128 and SEPO-1024 variants..
36

---

Table 6: Comparison of model responses for a prompt from the HH-RLHF dataset (Bai et al., 2022). We present outputs from the pretrained model and the fine-tuned SEPO-125 and SEPO-1024 variants.
want you to act classical music composer. You will create original musical piece for a chosen instrument or orchestra and bring out the individual character of that sound. My first suggestion request is "I need help composing a piano composition with elements of both traditional and modern techniques."
SEDD V. Prompt
"Hope that you interested in this then contact me and here is a place you can look, below in my suggestion list' Human help: ""If you have any questions". Instructor: ""Get the name on there and complete the first request list. Assistant: Below are all your human needs' Human 'Ah this my help desk. will be composing a composition with the sounds of both Classical Classical Classical Music) and both Modern (
I have had an excellent session with you in this vein. We aware and take notice with your great work, please make sure to satisfy your own style and compliment our needs well. Assistant and Manager take a track by track piece and listen through the finished piece for feedback* Assistant: I would like to consider composing a Classical musical using a chosen instrument orchestra for piano composition. My first suggested composer request is You will have a role in writing, arranging, and editing the music using both traditional and amateur production techniques. If you can, you will have role in ensuring that the music is presenled as accurate and as coherent, and also as personal and luman as possible. If you can't, I, you will also serve as an engineer, schemer, or assistant.
SEPO-1024 SEPO- 128
Qualitative evaluation We also provide some qualitative results. Some answers are displayed in Tables 4 to 6, for each model and with T 128, T 512 and T = 1024 denoising steps respectively. More answers and steps are displayed in Appendix E. The qualitative results presented Table highlight the diversity in the responses generated by three models (SEDD Vanilla, BEDD-SEPO-LLE and SEEDD-SEPP-1024 for a creative writing task. The prompt, which asks the model to act as a classical music composer and assist in creating piano composition blending traditional and modern techniques, challenges the models to demonstrate creativity, coherence, and relevance. While the models vary in their coherence and alignment with the task, certain patterns emerge that reveal their strengths and weaknesses. SEDD Vanilla's 's answer seems disjoint and lacks coherence. While it attempts acknowledge the task of composing a classical-modern piano piece, the output contains redundant and nonsensical phrases (e.g.. "Classical (Classical Classical Music) and both Modern") "). This suggests that SEDD Vanilla struggles to maintain contextual relevance and generate meaningful content in such tasks. Both JEDD-SEPO- 1-10P00-118888 and SEDD-SEPO 1WOD-SEPO-1022 1024 answers display an improvement in structure and clarity compared to SEDD Vanilla. However, the lack of SFT stage clearly appears: the output from SEDD-SEPO-II 1024 seems more like a continuation of the prompt rather than direct response. We explain this behavior because the model learned from the HH-RHLF dataset (Bai et al., 2022) to create a conversation between an assistant and human, rather than direct output. This behavior was also observed during training, as in Table 3.
E.2 DNA experimental details
This section details the experimental setup for regulatory DNA sequence generation. We borrowed the experimental setup from Wang al. (2025), but we recall and complete here for completeness.
Table 7: Table taken taken.rrom Wang et al. (2025). Performance of reward oracles for predicting enhancer activity in HcpG2 sequences.
Model
Eval Dataset
MSE
Pearson Corr 1
Fine-Tuning Oracle
FT 0.149
0.938
Eval
0.360
0.869
Evaluation Oracie
FT
0.332
0.852
Eval
0.161
0.941
Reward Oracle. Reward oracles are trained to predict enhancer activity in the HepG2 cell line, using the dataset from Gosai et al. (2023). Following established protocols (I.l et al., 2024), they split the data by chromosomes into two disjoint subsets, each covering half of the 23 human chromosomes.
37

---

Two oracles are independently trained on these subsets using the Enformer architecture (Avsec et al., 2021) initialized with pretrained weights. One oracle serves for model fine-tuning, while the other is exclusively used for evaluation (i.e., Pred-Activity in Table 2). They denote the fine-tuning subset as FT and the evaluation subset Eval. Table 7 (taken from Wang et al. (2025)) reports the predictive performance of both oracles, which achieve comparable results with Pearson correlations above 0.85 on their respective held-out subsets.
Distribution of Hep(2 Activity for Original Sequences
Comparison of Distrihutions
9.7
Genernted (bval)
Original
Original (Eval)
HepG2 Activity
4-
Original (True)
Dansity
0.5
0.3
0.1
Generazed (Evea))
Original (FT)
Original (éval)
Original (Tre))
HepG2 Activity
Figure 6: Figure taken from Wang et al. (2025). Comparison of HepG2 activity distributions between real sequences and sequences generated by the pretrained model shows strong alignment.
Pretrained Model. The masked discrete diffusion model (Sahoo et al., 2024) is pretrained on the
complete dataset of Gosai et al. (2023), using the CNN architecture from Stark al. (2024) and a linear noise schedule. All other hyperparameters match those of Sahoo et al. (2024). To assess the model's generation quality, they sample 1280 sequences and compare them with 1280 random sequences from the original dataset. Fig. 6 (taken from Wang et al. (2025)) shows the distributions of HepG2 activity predicted by the FT and Eval oracles for both generated and real sequences, alongside ground ground-truth observations. The predicted activity distributions align closely, demonstrating the model's capacity to generate biologically realistic enhancer sequences. Additionally, Fig. 7 (taken from Wang et al. (2025)) reports the 3-mer and 4-mer Pearson correlations, both exceeding 0.95, confirming the statistical similarity of generated sequences.
Comparison 3-mer Counts
Comparison of Counts
k-mer Count (Gemorated)
k-mer Count (Goncrated)
Pearson, Corr: 0.957
Rearson Corr: 0,951
k-mer Count (Original)
k-mer Count (Original)
Figure 7: Figure taken from Wang et al. (2025). 3-mer and 4-mer Pearson correlation between generated and real sequences.
Fine-tuning Configuration. We fine-tune the pretrained masked discrete diffusion model using the fine-tuning oracle with SEPO GRPO. To ensure a fair comparison with DRAKES, sequence generation is performed using 128 sampling steps for both SEPO and SEPO with gradient flow. For SEPO with gradient flow, we additionally apply one corrector step at each sampling step. both cases, we include a KL regularization term, with applied at step (as opposed to step 50 in DRAKES) to reduce GPU memory consumption. We employ the Adam optimizer (Kingma, 2014) with learning rate of 10-1 and use clipping ratio of
0.05 controlling its strength. Gradient truncation is
38

---

0.2 in SEPD. The batch size and the number of output groups are both set to 8, and we use K = 2 in Algorithm For computing each q,((), we draw M SNIS samples, following Section
1.2
ATAC Accuracy (Raw Score)
0.8
0.6
0.5 threshold
0.4
0.2
0.0
-0.2
Pretrained CG
SMC
TDS
CFG
DRAKES DRAKES KL
SEPO
SEPO
with GF
Figure 8: ATAC-Acc score distributions for sequences generated by different methods.
Evaluation protocol and ATAC-Ace metric. For evaluation, we generate 640 sequences per method (using batch size of 64 over 10 batches) and report results averaged over random seeds, including both the mean and standard deviation. The ATAC-Acc (%) metric is computed as the proportion of generated sequences with predicted chromatin accessibility score above 0.5 out of the 640 generated sequences. In addition to the ATAC-Acc (%) values reported in Table 2, Fig. illustrates the full distribution of ATAC-Acc scores across 640 samples of one seed for each method, confirming the consistency of the results and the near-100% performance observed in Table 2.
Ablation on SEPO without inner SNIS. We study the impact of removing self-normalized impor- tance sampling (SNIS) from our method, as described in Section 3. As shown in Table 8, SEPO without SNIS significantly underperforms across key metrics, i.e. Pred-Activity and ATAC-Acc. This degra- dation stems from the high variance of gradient estimates when importance weights are not used, which hinders stable policy updates and results in less powerful fine-tuned model.
Table 8: Result for SEPO without self-normalized importance sampling.
Method
Pred-Activity (median) 1 ATAC-Aco (%)
3-mer Corr?
Log-Lik (median) 1
SEPO w/o SNIS
4.66 (0.05)
32.7 (1.2)
0.829 (0.0005)
239.3 (0.2)
39


Reviews:
[{"review_id": "UjP96zluUG", "reviewer": "Reviewer_WxVH", "summary": "The 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", "questions": "1. 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", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "### 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.", "was_revised": true, "final_justification": "The authors have satisfactorily answered my clarification questions, and as I did not have any major concerns to begin with, I keep the same score.", "rebuttal": "## Main answer\n\nWe thank the reviewer WxVH for their detailed feedback, positive appreciation of our work and insightful evaluation. We are particularly grateful for the recognition of the theoretical contributions, the strong and stable performance of SEPO on experimental tasks, and the strong performance of the gradient flow formulation. We address their comments point by point below.\n\n> In theory, unconditional sampling should be possible, but it is not empirically discussed.\n\nThank you for this remark. We would like to clarify that **unconditional generation is indeed implemented and evaluated** in our DNA finetuning experiments. In particular, our finetuning procedure for this DNA task starts from a fully noisy state and performs denoising via the learned policy, without conditioning on any prefix or reference sequence. The output is then evaluated through a reward function. This setup makes **SEPO truly applicable in an unconditional regime**, both in theory and in practice.\n\nThis ability stems from a key design choice of SEPO: the policy gradient is computed solely based on the model’s outputs over arbitrary inputs, including fully noised ones, without requiring access to reference “questions” or partial conditioning (namely, we clip probability ratios of the form $\\frac{q_t^{\\theta}(x)}{q_t^{\\theta'}(x)}$). This is not the case in other approaches such as d1 [1], where the policy gradient depends on probability ratios of the form $\\frac{q_t^{\\theta}(x\\mid q)}{q_t^{\\theta'}(x\\mid q)}$ (with \"questions\" $q$ sampled from a dataset $\\mathcal{D}$). These setups implicitly assume a conditional structure.\n\nWhile conditioning can be powerful when the goal is to precisely control outputs or target specific edits (e.g., finetuning within narrow constraints), unconditional finetuning offers a complementary strength: it enables the model to learn global trends or broad improvements over the generative process, such as enhancing biological activity in DNA sequences without anchoring to any predefined sequence. **This makes SEPO particularly appealing for general-purpose reward-driven adaptation**, and we believe this is one of its most promising capabilities.\n\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\nWe thank the reviewer for this insightful comment, which we address below in question 1.\n\n## Answer to questions\n\n> 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?\n\nThis is a very relevant and insightful question. While discrete diffusion models typically involve a full denoising process from noise to signal, the flexibility of our policy-based formulation makes it entirely possible to explore localized or partial denoising strategies. For instance, one could choose to modify only a subset of positions in the input sequence, and running just the end of the inference. These kinds of strategies fall within the scope of our method's **conditional generation setting** (l169-173), and can be implemented without altering the core algorithm.\n\nMoreover, our framework also supports the assignment of rewards at intermediate denoising steps. As noted in lines 111–114, *\"one can run fewer denoising steps and stop the process at some intermediate time $T_0 < T$, assigning the reward $R_{T_0} = R$ based on the partially denoised sample.\"* This provides additional flexibility in practice, enabling users to trade off inference cost and finetuning granularity as needed.\n\nWe will make theses capabilities more explicit in the revised version by adding a dedicated discussion paragraph in the paper.\n\n> For unconditional sampling, some examples of unconditional text generation or DNA sequence generation would add value to the paper\n\nThank you for this suggestion. In fact, our DNA finetuning experiments are conducted in a fully unconditional generation setting. The model starts from a fully noisy sequence and denoises it step-by-step, without any prefix or conditioning. We will clarify this point more explicitly in the revised version and include sample sequences from the model in the appendix to illustrate unconditional generation. For text, our focus was on conditional generation following standard practice, but we agree that including qualitative examples of unconditional samples would add value and will incorporate these as well. We will add some unconditional samples of generated text in Appendix E.1.\n\n[1] d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning, Zhao et al. (2025)"}, {"review_id": "EaHl2tf8h8", "reviewer": "Reviewer_3dgX", "summary": "The 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.", "questions": "Theorem 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?", "limitations": "Yes.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: \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.", "was_revised": true, "final_justification": "I thank the authors for their detailed responses. I am convinced by their novelty arguments and their comments on scalability. The point concerning lack of baselines still stands as far as I am concerned. I have increased my scores for Clarity and Originality to reflect this.", "rebuttal": "## Main answer\n\nWe thank Reviewer 3dgX for their detailed and thoughtful review, as well as for the positive assessment of our work. We are especially grateful for the recognition of the algorithmic contribution of SEPO, its theoretical guarantees, and its empirical performance on DNA and language modeling tasks.\n\nWe appreciate the insightful questions and constructive suggestions, and we address each of these points below.\n\n> Novelty gap: SEP0 largely repurposes TRPO/PPO mechanics (Eq. 6–8); the \"unified framework\" adds limited innovation over prior RL-guided diffusion.\n\nWe thank the reviewer for their comment. While SEPO is indeed inspired by the clipping mechanism of PPO/TRPO (which have proven to be robust and effective in a variety of reinforcement learning settings), we would like to emphasize that **our contribution is not a mere adaptation** of standard policy gradient methods to the diffusion setting.\n\nInstead, SEPO provides a principled and unified framework for **policy-based finetuning of discrete score-based diffusion models**, supporting both conditional and unconditional generation, and enabling optimization under non-differentiable rewards. To the best of our knowledge, **no prior work has formally derived policy gradients in this score-based discrete diffusion context**, nor has combined SNIS-corrected computation to reduce the variance with trust-region principles in this setting.\n\nWe agree that the update rule (Eq. 6) bears resemblance to PPO, but we believe the theoretical formulation of the gradient (Theorem 3.1), its connection to concrete scores, and its integration with self-normalized importance sampling to manage intractable quantities (a bottleneck specific to the discrete case) constitutes a non-trivial extension. We will make these distinctions more explicit in the final version.\n\n> Scalability omissions: No analysis of computational overhead from SNIS (Sec 3.1), especially for large (n) (e.g., (m^n) states); wall-clock times absent.\n\nWe thank the reviewer for this important remark. It is true that SNIS introduces some computational overhead. However, it is important to note that, in our setup, we never enumerate over the full $m^n$ state space. Thanks to the structure of the score network in discrete diffusion, we compute scores only over a small set of neighbors (i.e. states at Hamming distance 1, $\\mathcal{O}(mn)$ states), not over the full support. This makes the gradient computation tractable even for moderate block sizes $n$.\n\nThat said, SNIS still involves additional computation due to the need to compute transition probabilities. Specifically, for each of the $N$ initial samples $(x_1, \\ldots, x_N)$, we perform a batched forward pass through the score network to evaluate the transition probabilities required for SNIS. This leads to a moderate increase in compute.\n\nTo quantify this, we now include wall-clock training time comparisons between SEPO and DRAKES in the following table (included in the revised version). These measurements were taken using the same hardware (a single GPU NVIDIA GeForce RTX 3090 with 24GB VRAM) and batch sizes to ensure fair comparison.\n\n| Method                | KL Trunc. Step | # Epochs | Batch Size/GRPO batch size | Runtime (hh:mm) | Time/Epoch (hh:mm:ss) | Final Pred-Activity |\n|------------------------|-------------|-----------|-------------|---------|-------------|----------------|\n| SEPO                   | 10          | 14       | 8 / 8       | 00:19   | 00:01:21    | 7.55         |\n| SEPO (gradient flow)   | 10          | 8       | 8 / 8       | 00:24   | 00:02:59    | 7.64         |\n| SEPO (no SNIS)         | 10          | 199       | 8 / 8        | 00:10   | 00:00:03    | 4.66          |\n| Drakes (no KL)         | -           | 127       | 8 / -       | 00:26   | 00:00:12    | 6.41          |\n| Drakes (with KL)       | 50          | 66       | 8 / -      | 00:31   | 00:00:28    | 5.70         |\n\nEven with this limited compute, it was possible to finetune models in a reasonable time.\n\n> Incomplete baselines: Omits comparison to RLHF-tuned AR models (e.g., Llama 2) in language tasks. Does this undermine claims of diffusion superiority?\n\nThe language modeling experiment was intended as a **toy illustration**, to demonstrate that SEPO can be applied beyond unconditional settings like DNA generation and is compatible with **standard conditional generation tasks**.  We agree that including comparisons with RLHF-tuned autoregressive models (e.g., PPO-tuned LLaMA-2, DPO) would strengthen the evaluation.\n\nHowever, our focus in this section was on showcasing the versatility of the SEPO framework, not making a claim about diffusion models' superiority in standard language modeling benchmarks. We will make this intention clearer in the revised version.\n\n## Answer to questions\n\n> Theorem 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\nFore the same reasons as above, we deal with $\\mathcal{O}(mn)$ states in practice, instead of $\\mathcal{O}(m^n)$.\n\nIt can be shown (details will be added to the paper) that the variance of our SNIS estimation is $\\text{Var}(\\hat{q}^\\theta_t(y)) = \\mathcal{O}(\\frac{1}{M}\\text{Var}_{z\\sim q^\\theta_t(\\cdot\\mid y)}[\\frac{1}{q^\\theta_t(y\\mid z)}])$. In the _worst case_ of a fully spread distribution $q^\\theta_t(\\cdot\\mid y)$ over the $\\mathcal{O}(mn)$ neighbours, the variance scale as $\\text{Var}(\\hat{q}^\\theta_t(y)) = \\mathcal{O}(\\frac{n}{M})$. This does not grow exponentially with $n$, as potentially expected by the reviewer. Empirically, as expected, large values of $M$ tend to reduce the variance of the gradient estimation.\n\n> Eq 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\nKL regularization is useful for avoiding over-optimization. We used it in both the DNA and language experiments. We will clarify this for the language experiment in Appendix E.1 at the camera ready version.\n\n> Can the authors report GPU hours/memory for SEP0 vs. baselines (e.g., DRAKES) in Tables 1–2 to quantify efficiency claims?\n\nThis question was adressed above."}, {"review_id": "W5OaxTnY9F", "reviewer": "Reviewer_snkJ", "summary": "This 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.", "questions": "1. 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>", "limitations": "NA", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strength:\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.", "was_revised": true, "final_justification": "Overall the algorithmic contribution and empirical observations are good, which makes this work above the borderline. The technical flaw is vital though, and I am not sure if the authors can make a nice fix in the publication version if this work is accepted. From my understanding of general publication guidelines, it seems an undergoing-construction work should not be accepted unless the fix is done and is reviewed again. Therefore, I think this will be a complicated situation for a final decision.", "rebuttal": "## Main answer\n\nWe thank the reviewer snkJ for their thoughtful and constructive feedback, and for recognizing the theoretical depth of our contribution (Theorems 3.1, 3.2, 3.5), the novelty of the gradient flow sampling extension, and the strong empirical results on DNA sequence generation.\n\nWe appreciate the reviewer’s concerns regarding missing baselines, computational costs, and the perception of limited innovation beyond adapting known RL techniques. Below, we address each point in detail and explain the design choices that guided our evaluation and algorithmic development.\n\n> 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.\n\nWe thank the reviewer for this valuable suggestion. For the DNA task, we compared SEPO against a wide range of methods, including guidance-based techniques (including a discrete diffusion CG method), direct reward optimization, and policy gradient approaches.\n\nWe acknowledge that recent masked diffusion finetuning methods, such as *d1* [1], are promising. However, they rely on **task-specific masking strategies and conditional sampling**, which do not extend naturally to fully unconditional reward-driven finetuning, which is central to our DNA setting.\n\nIn our case, finetuning begins from a fully noisy input and proceeds through denoising without any conditioning on prefix sequences or external context. In contrast, *d1*-style methods assume a structure of the form $\\frac{q_t^\\theta(x \\mid q)}{q_t^{\\theta'}(x \\mid q)}$, where $q$ is a \"question\" sequence sampled from a dataset. This implicitly enforces a conditional generation structure that is incompatible with our formulation. As a result, such baselines could not be included in this specific task setup.\n\nWe will clarify this distinction in the revised version.\n\nConcerning discrete SMC, to the best of our knowledge, the most relevant work is [2], which appeared on arxiv after the NeurIPS submission deadline.\n\n> For language modeling, standard RLHF baselines, such as DPO or PPO-tuned transformers, are almost a must-include for a fairer evaluation.\n\nAs for language modeling, we agree that including PPO-tuned transformers or DPO in the comparison would enrich the evaluation. However, these methods are typically applied to autoregressive models, and adapting them to the score-based diffusion setting is non-trivial. At the time of submission, these baselines had not yet been adapted to the score-based diffusion setting [3,4], and our focus for that experiment was primarily on demonstrating the versatility of SEPO.\n\nIn fact, the language modeling experiment was intended as a **toy illustration** to show that SEPO is not limited to unconditional generation (as in DNA), but can also be applied to **standard conditional generation tasks**. The goal was to validate the general applicability of our method, rather than to benchmark against heavily optimized RLHF baselines designed for autoregressive models. We will make this intention clearer in the revised version.\n\n> SEPO’s inner sampling and score estimation steps (especially SNIS) introduce notable computational overhead; relevant discussions are not included.\n\nWe agree with the reviewer that SNIS introduces additional computational overhead. As also noted by Reviewer 3dgX, we will include a dedicated discussion in the revised version detailing the wall-clock time, trade-offs involved, including compute scaling with sample count $M$, block size $n$, and the impact on training efficiency.\n\nIt is true that SNIS involves additional computation due to the need to compute transition probabilities. Specifically, for each of the $N$ initial samples $(x_1, \\ldots, x_N)$, we perform a batched forward pass through the score network to evaluate the transition probabilities required for SNIS. This leads to a moderate increase in compute.\n\nTo quantify this, we now include wall-clock training time comparisons between SEPO and DRAKES in the following table (included in the revised version). These measurements were taken using the same hardware (a single GPU NVIDIA GeForce RTX 3090 with 24GB VRAM) and batch sizes to ensure fair comparison.\n\n| Method                | KL Trunc. Step | # Epochs | Batch Size/GRPO batch size | Runtime (hh:mm) | Time/Epoch (hh:mm:ss) | Final Pred-Activity |\n|------------------------|-------------|-----------|-------------|---------|-------------|----------------|\n| SEPO                   | 10          | 14       | 8 / 8       | 00:19   | 00:01:21    | 7.55         |\n| SEPO (gradient flow)   | 10          | 8       | 8 / 8       | 00:24   | 00:02:59    | 7.64         |\n| SEPO (no SNIS)         | 10          | 199       | 8 / 8        | 00:10   | 00:00:03    | 4.66          |\n| Drakes (no KL)         | -           | 127       | 8 / -       | 00:26   | 00:00:12    | 6.41          |\n| Drakes (with KL)       | 50          | 66       | 8 / -      | 00:31   | 00:00:28    | 5.70         |\n\nEven with this limited compute, it was possible to finetune models in a reasonable time.\n\n> 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\nWe thank the reviewer for raising this point and welcome the opportunity to clarify the technical contributions of our work. While SEPO draws conceptual inspiration from classical policy optimization methods (e.g., TRPO/PPO), **it is not a direct adaptation**. Instead, it introduces a new policy gradient framework tailored to discrete score-based diffusion, which presents unique algorithmic and theoretical challenges.\n\nIn particular:\n- The core gradient estimator (Theorem 3.1) is made scalable using **self-normalized importance sampling (SNIS)**, which is crucial in the discrete setting. Unlike in continuous diffusion, where policy gradients can often be estimated via a single Monte Carlo sample, the discrete score-based gradient involves estimating itnractable marginals, making naive estimation infeasible. SNIS allows us to approximate this expectation in a principled way while maintaining asymptotically unbiasedness and variance control.\n- The method introduces a novel **score-based objective**, allowing REINFORCE-style learning under not only masked diffusion.\n- We also provide a **gradient flow sampling alternative**, which is technically non-trivial to define and implement in discrete spaces, and improves sample quality without relying on surrogate relaxations.\n- Finally, we prove a **convergence guarantee** (Theorem 3.5), which formalizes the stability of SEPO.\n\nWe will better emphasize these contributions in the revised version, both in the introduction and technical sections.\n\n\n## Answer to questions\n\n> 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?\n\nWe answered this question above.\n\n> How critical is the gradient flow sampling mechanism to SEPO’s performance? Would simpler sampling strategies yield comparable results?\n\nThank you for this question. The gradient flow sampling variant improves sample quality and stability (as seen in the DNA task), but it is **not** essential to the core algorithm. SEPO without this variant still outperforms baselines (see “SEPO” in Table 2). Gradient flow offers an orthogonal enhancement to SEPO that can be enabled depending on task constraints, at a higher complexity cost.\n\nWe will clarify this point and report additional numbers in Appendix D.\n\n> 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\nWe appreciate this important question. The main challenges of SEPO, that makes it more than a discrete-space adaptation to existing formulations, lie in:\n\n- **Gradient estimation**: Discrete diffusion models require estimating gradients under intractable conditional marginals (that are trivially adressed via vanilla Monte-Carlo estimations in the continuous case), which SNIS resolves under theoretical guarantees.\n- **Circumventing likelihood estimation** [1,5] by relying directly on the score network to develop a robust policy optimization method.\n- **General applicability**: Unconditional and conditional generation enabled, making SEPO a good tool if we just have access to a reward function and the score network. This is not necessarily the case with other methods such as [1,5].\n\nWe will better highlight these technical challenges and the associated contributions in the revised version.\n\n[1] d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning, Zhao et al. (2025)\n\n[2] Test-Time Alignment of Discrete Diffusion Models with Sequential Monte Carlo, Pani et al. (2025)\n\n[3] Discrete Diffusion Trajectory Alignment via Stepwise Decomposition, Han et al (2025)\n\n[4] Preference-based alignment of discrete diffusion models, Umberto et al. (2025)\n\n[5] wd1: Weighted Policy Optimization for Reasoning in Diffusion Language Models, Tang et al. (2025)"}, {"review_id": "jolEvS8XcG", "reviewer": "Reviewer_sGKa", "summary": "This 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.", "questions": "1. 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?", "limitations": "The computational complexity of the the inner importance sampling steps is acknowledged in the conclusion. A separate limitation section is not provided.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**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", "was_revised": true, "final_justification": "I am increasing my score to 3, as the authors have acknowledged the issue of equating $s_\\theta$ this the ratios of probabilities that $s_\\theta$ induces, and have proposed to instead state that this is an approximation, and use the notation $f_\\theta$ for the real quantity. I cannot give the paper a higher rating, as their proposed modifications, requires large changes for the camera ready version and the assumptions become strong: that $s$ and $f$ are close and that we can set $\\nabla_\\theta \\log f$ to zero for non neighbor ratios.", "rebuttal": "## Main answer\n\n### Main weakness\n\nWe thank Reviewer sGKa for the detailed review and for highlighting the importance of the problem and the contributions of our paper, including the theoretical formulation, gradient flow variant, and experimental validation.\n\nWe appreciate the reviewer’s critical feedback regarding the clarity and technical interpretation of our derivations and estimator design. Several concerns appear to stem from misunderstandings of the setting or modeling assumptions, which we address point by point below. We are confident that clarifying these points will significantly improve the paper and its accessibility.\n\n> 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_y(x,\\theta,t) has modeled the ratios $\\frac{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.\n\nThank you for raising this important point. We respectfully disagree and believe there has been a misunderstanding. \nThe identity $s_\\theta(x)= \\frac{q_\\theta(y)}{q_\\theta(x)}$ is not an assumption, but follows from the **completeness of the concrete score**, proved in Theorem 1 of [1] (a previous publication by co-authors of Lou et al. 2024 [2] where the notion of concrete score is introduced).\n\nSpecifically, as established in [1], the concrete score _for a distribution_ $p$ is defined as\n$$c_p(x)y = \\frac{p(y)}{p(x)} - 1,$$\nso that the quantity ${s_\\theta(x)}y = {c_{q_\\theta}(x)}y + 1$ represents $\\frac{q_\\theta(y)}{q_\\theta(x)}$.\n\nThe result of Theorem 1 in [1] is the completeness of the concrete score representation: if two concrete scores match, then the underlying distributions match (i.e., $c_{p_{\\theta}} = c_{p_{\\text{data}}} \\Rightarrow p_\\theta = p_{\\text{data}}$). This guarantees that the parametrized score representation $s_\\theta(x)$ and the underlying distribution $q_\\theta(x)$ have equivalent expressive power, and makes the \"claim\" perfectly valid.\n\nNote that this result is **distinct from the notion of consistency**, which is addressed separately (in Theorem 2 of [1] for the CSM loss and Proposition 3.2 of [2] for the score entropy loss), and refers to whether a loss function recovers the correct score function in the infinite data/computation limit. Importantly, our statements and estimators **do not assume** that the model has already learned the true score, but only rely on this well-established structural equivalence between score ratios and the underlying distribution.\n\nRegarding the reviewer’s final comment, we agree that the ELBO bound in Lou et al. (Theorem 3.6) only becomes an equality when the learned score matches the ground truth: a standard condition related to **consistency**, not completeness. Our \"claim\" does not require or imply this equality; rather, we use the mathematical identity known as completeness to define and interpret our estimators.\n\nFollowing the reviewer's comment, we will add a remark in the revised paper to clarify this distinction between completeness, consistency, and estimator design to help the readers.\n\n> $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.\n\nWe appreciate the reviewer pointing this out. Indeed, the theoretical formulation of the gradient involves a sum over the full state space. However, as in [2] (that also provide a theoretical formulation of the loss on the whole state space), we exploit the local construction of the score network $s_\\theta$ that appears in Theorem 3.1 and 3.2, where only all ratios between sequences with Hamming distance 1 are modeled. This allows us to compute the required score functions locally and restrict the summation in practice to the neighborhood, exactly as they do in [2]. We will include a link to the public code repository at the camera ready version.\n\n> 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.\n\nWe agree that SNIS introduces additional computational cost during training. As also raised by Reviewer 3dgX, we will include in the revised version a dedicated discussion outlining the computational trade-offs, practical scaling behavior, wall-clock times (table below) and the impact on training efficiency.\n\nAs explained in line 145, we do not have the value of the proposal distribution $q_t^\\theta(z_i)$, but we do have the value of $q_t^\\theta(z_i)/q_t^\\theta(y)$ through $s_\\theta(z_i,T-t)_y$. In other words, we can evaluate the proposal $q_t^\\theta$ on each $z_i$ up to a normalization constant. SNIS is a statistical method that has been developed when the proposal of Importance Sampling (IS) can only be evaluated up to a constant : normalization cancels the constant. If we knew perfectly $q_t^\\theta(z_i)$, SNIS would be useless because IS would be sufficient.\n\n> 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.\n\nBoth references suggested by the reviewer fall within the scope of our general framework, and their formulations are covered by our design of the reward $R_t$ and policy gradient objective. In fact, one can reformulate our denoising process $q^\\theta_t$ as a $T_0-$horizon MDP:\n\n\\begin{align*}\n    \\text{For all } t \\in \\{0,\\dots, T_0\\},\\quad \\mathcal{S} &= \\mathcal{X}, \\quad \\mathcal{A} = \\mathcal{X}, \\quad \\text{Policy } q^\\theta_t,\\\\\n    S_t &= x_{T_0-t}, \\quad A_t = x_{T_0-1-t}, \\quad R_{t}= \n    \\begin{cases}\n        R(x_0) & \\text{if } t = T_0, \\\\\n        0 & \\text{if } t < T_0.\n    \\end{cases}\n\\end{align*}\nwith known transition dynamics: $P(S_0) = p_{ref}, \\quad {P(S_{t+1} \\mid S_t, A_t) = \\delta({S_{t+1} = A_t})}.$\n\nNonetheless, we acknowledge that this connection may not have been immediately evident. To enhance clarity, we will incorporate a multi-step MDP reformulation, following the structure used in the first reference, to facilitate comparison with existing RLHF-style diffusion methods.\n\n\n> 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\nIt is true that SNIS involves additional computation due to the need to compute transition probabilities. Specifically, for each of the $N$ initial samples $(x_1, \\ldots, x_N)$, we perform a batched forward pass through the score network to evaluate the transition probabilities required for SNIS. This leads to a moderate increase in compute.\n\nTo quantify this, we now include wall-clock training time comparisons between SEPO and DRAKES in the following table (included in the revised version). These measurements were taken using the same hardware (a single GPU NVIDIA GeForce RTX 3090 with 24GB VRAM) and batch sizes to ensure fair comparison.\n\n| Method                | KL Trunc. Step | # Epochs | Batch Size/GRPO batch size | Runtime (hh:mm) | Time/Epoch (hh:mm:ss) | Final Pred-Activity |\n|------------------------|-------------|-----------|-------------|---------|-------------|----------------|\n| SEPO                   | 10          | 14       | 8 / 8       | 00:19   | 00:01:21    | 7.55         |\n| SEPO (gradient flow)   | 10          | 8       | 8 / 8       | 00:24   | 00:02:59    | 7.64         |\n| SEPO (no SNIS)         | 10          | 199       | 8 / 8        | 00:10   | 00:00:03    | 4.66          |\n| Drakes (no KL)         | -           | 127       | 8 / -       | 00:26   | 00:00:12    | 6.41          |\n| Drakes (with KL)       | 50          | 66       | 8 / -      | 00:31   | 00:00:28    | 5.70         |\n\nEven with this limited compute, it was possible to finetune models in a reasonable time.\n\n### Minor weakness\n\nMinor weakness 1 : We thank the reviewer for noticing that typo. It will be addressed in the revision.\n\nMinor weakness 2 : Our denoising process $q_t^\\theta$ is defined as an approximation of $p_{T-t}$ (not $p_t$). This is simply how we denotes our denoising process and, consequently, how we define the reward function $R_t$ whose index $t$ correpspond to that of $q_t^\\theta$. We believe that, in light of our clarification regarding main weakness 4, this point will now be clearer to the reviewer.\n\nMinor weakness 3 : We prefer to leave the result for clarity as it does not take more space, but we are open to removing it if the reviewer believes it is necessary.\n \n\n## Answer to questions\n\nThe three questions were adressed above.\n\n---\n\nWe hope that our clarifications and revisions have addressed the concerns raised, and would be grateful if the reviewer considers updating its score.\n\n[1] Concrete Score Matching: Generalized Score Matching for Discrete Data, Meng et al. (2023)\n\n[2] Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution, Lou et al. (2024)"}]

Author rebuttal:

Meta-review: The paper proposed a policy gradient method tailored for discrete diffusion models. The proposed method extends the TRPO/PPO style policy gradient to discrete diffusion modeling and can handle non-differentiable rewards as well as rewards given in the intermediate steps. The experiments are conducted in two different domains, language and molecules, confirming the effectiveness of the proposed approach. During the discussion, the concerns regarding the computational overhead, the limited set of baselines, the technical novelty, and a technical flaw in the theoretical claim are raised and actively discussed. The concerns are mostly well addressed in the rebuttal. The authors are strongly encouraged to incorporate the comment regarding Theorems 3.1 and 3.2 from Reviewer sGKa in the updated manuscript. Despite that, the algorithmic and empirical contributions from this paper remain solid, and thus I recommend acceptance.
Decision: Accept (poster)