Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods

Main answer

Main weakness

We 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.

We 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.

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.

Thank you for raising this important point. We respectfully disagree and believe there has been a misunderstanding. The 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).

Specifically, as established in [1], the concrete score for a distribution $p$ is defined as $c_p(x)y = \frac{p(y)}{p(x)} - 1,$ so that the quantity ${s_\theta(x)}y = {c_{q_\theta}(x)}y + 1$ represents $\frac{q_\theta(y)}{q_\theta(x)}$.

The 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.

Note 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.

Regarding 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.

Following 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.

$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.

We 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.

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.

We 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.

As 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.

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.

Both 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:

$

\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,\\
S_t &= x_{T_0-t}, \quad A_t = x_{T_0-1-t}, \quad R_{t}= 
\begin{cases}
    R(x_0) & \text{if } t = T_0, \\
    0 & \text{if } t < T_0.
\end{cases}

$

with known transition dynamics: $P(S_0) = p_{ref}, \quad {P(S_{t+1} \mid S_t, A_t) = \delta({S_{t+1} = A_t})}.$

Nonetheless, 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.

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.

It 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.

To 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.

Even with this limited compute, it was possible to finetune models in a reasonable time.

Minor weakness

Minor weakness 1 : We thank the reviewer for noticing that typo. It will be addressed in the revision.

Minor 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.

Minor 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.

Answer to questions

The three questions were adressed above.

---

We hope that our clarifications and revisions have addressed the concerns raised, and would be grateful if the reviewer considers updating its score.

[1] Concrete Score Matching: Generalized Score Matching for Discrete Data, Meng et al. (2023)

[2] Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution, Lou et al. (2024)

We thank the reviewer for engaging further in this discussion and for raising these important points. We are glad that our earlier response resolved the reviewer main technical queries, and we appreciate the opportunity to clarify a few additional matters.

Rather than deriving a hard bound, it is common in diffusion‐model theory to assume that the error of the neural‐network score estimator is of order $\mathcal{O}(\epsilon)$, where $\epsilon$ is a small threshold.

This assumption reflects both the expressive power of modern networks and the fact that they are trained to high accuracy, and it appears in several prior works in dicrete diffusion (cf Assumption 4.6 and 4.6' in [1] and Assumption 1 in [2].)

As written in line 151, the true gradient takes the form $\nabla_\theta\ell_t(\theta) = \mathbb{E}[x\sim q^\theta_t][R_t(x)g(x,\theta)]$. In practice, as discussed previously, we are computing $\nabla_\theta\hat{\ell_t}(\theta) = \mathbb{E}[x\sim q^\theta_t][R_t(x)\hat{g}(x,\theta)]$ where $\hat{g}(x,\theta) = \sum_{\substack{y: \\ s_\theta(x,T-t)y \text{ is defined}}}{q^\theta_t(y)\nabla_\theta\log s_\theta(x,T-t)_y}$.

As noted by the reviewer, $\nabla_\theta \log f$ is unlikely to be zero for non-neighbours ratios : dropping those terms does introduce bias in the estimation of $g$ by $\hat{g}$. We will explicitly note this when defining $\nabla_\theta\hat{\ell_t}(\theta)$, and discuss how this bias arises.

However, note that even with this additional bias, we achieve SOTA results against various methods in discrete diffusion. SEPO introduces a policy optimization objective that only depends on concrete score estimation, circumventing for example likelihood approximation that are done in [3,4]. Figuring out how to mitigate this additional bias in practice and potentially quantifying it theoretically are interesting directions for future works, and we will add this discussion in the conclusion section.

Our intuition is that this bias is manageable, because the concrete score $c_p(x)$ (as its continuous space equivalent $\nabla\log p(x)$) is a local descriptor of how the probability mass is changing around $x$ : the summed contribution of $f(\theta,x,y,t)$ terms when $y$ is far apart from $x$ in the graph may have a small influence in how $\hat{g}$ approximates $g$. The strong empirical results seem to support this intuition. In fact, if this bias were large enough to spoil the gradient, our fine-tuning would fail; instead, we consistently observe strong gains and stable outputs once the model is finetuned.

Additionally, we would also like to emphasize that while the time/epoch of our method is higher, it requires an order of magnitude less epoch to reach a higher score. So the total runtime is actually lower than SOTA method while reaching a much higher score. This is something we had not recorded prior to submission so we are grateful to the reviewer for suggesting this, which makes our experimental section and comparison with other methods even stronger.

We thank once again the reviewer for their time, the constructive review and for prompting these clarifications. The feedback has substantially improved both the theoretical presentation and the empirical comparison in our manuscript.

[1] How Discrete and Continuous Diffusion Meet: Comprehensive Analysis of Discrete Diffusion Models via a Stochastic Integral Framework, Ren et al (2024) [2] Convergence Analysis of Discrete Diffusion Model: Exact Implementation through Uniformization, Chen et al. (2024) [3] d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning, Zhao et al. (2025) [4] wd1: Weighted Policy Optimization for Reasoning in Diffusion Language Models, Tang et al. (2025)

Main answer

We 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.

We 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.

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.

We 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.

We 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.

In 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.

We will clarify this distinction in the revised version.

Concerning discrete SMC, to the best of our knowledge, the most relevant work is [2], which appeared on arxiv after the NeurIPS submission deadline.

For language modeling, standard RLHF baselines, such as DPO or PPO-tuned transformers, are almost a must-include for a fairer evaluation.

As 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.

In 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.

SEPO’s inner sampling and score estimation steps (especially SNIS) introduce notable computational overhead; relevant discussions are not included.

We 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.

It 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.

To 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.

Even with this limited compute, it was possible to finetune models in a reasonable time.

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.

We 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.

In particular:

• 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.
• The method introduces a novel score-based objective, allowing REINFORCE-style learning under not only masked diffusion.
• 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.
• Finally, we prove a convergence guarantee (Theorem 3.5), which formalizes the stability of SEPO.

We will better emphasize these contributions in the revised version, both in the introduction and technical sections.

Answer to questions

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?

We answered this question above.

How critical is the gradient flow sampling mechanism to SEPO’s performance? Would simpler sampling strategies yield comparable results?

Thank 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.

We will clarify this point and report additional numbers in Appendix D.

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>

We appreciate this important question. The main challenges of SEPO, that makes it more than a discrete-space adaptation to existing formulations, lie in:

• 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.
• Circumventing likelihood estimation [1,5] by relying directly on the score network to develop a robust policy optimization method.
• 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].

We will better highlight these technical challenges and the associated contributions in the revised version.

[1] d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning, Zhao et al. (2025)

[2] Test-Time Alignment of Discrete Diffusion Models with Sequential Monte Carlo, Pani et al. (2025)

[3] Discrete Diffusion Trajectory Alignment via Stepwise Decomposition, Han et al (2025)

[4] Preference-based alignment of discrete diffusion models, Umberto et al. (2025)

[5] wd1: Weighted Policy Optimization for Reasoning in Diffusion Language Models, Tang et al. (2025)

Main answer

We 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.

We appreciate the insightful questions and constructive suggestions, and we address each of these points below.

Novelty gap: SEP0 largely repurposes TRPO/PPO mechanics (Eq. 6–8); the "unified framework" adds limited innovation over prior RL-guided diffusion.

We 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.

Instead, 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.

We 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.

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.

We 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$.

That 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.

To 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.

Even with this limited compute, it was possible to finetune models in a reasonable time.

Incomplete baselines: Omits comparison to RLHF-tuned AR models (e.g., Llama 2) in language tasks. Does this undermine claims of diffusion superiority?

The 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.

However, 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.

Answer to 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?

Fore the same reasons as above, we deal with $\mathcal{O}(mn)$ states in practice, instead of $\mathcal{O}(m^n)$.

It 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.

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)?

KL 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.

Can the authors report GPU hours/memory for SEP0 vs. baselines (e.g., DRAKES) in Tables 1–2 to quantify efficiency claims?

This question was adressed above.

Main answer

We 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.

In theory, unconditional sampling should be possible, but it is not empirically discussed.

Thank 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.

This 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.

While 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.

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.

We thank the reviewer for this insightful comment, which we address below in question 1.

Answer to questions

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?

This 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.

Moreover, 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.

We will make theses capabilities more explicit in the revised version by adding a dedicated discussion paragraph in the paper.

For unconditional sampling, some examples of unconditional text generation or DNA sequence generation would add value to the paper

Thank 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.

[1] d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning, Zhao et al. (2025)