Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design

## Abstract

Abstract Recent studies have demonstrated the strong empirical performance of diffusion models on discrete sequences (i.e., discrete diffusion models) across domains from natural language to biological sequence generation. For example, in the protein inverse folding task, where the goal is to generate a protein sequence from a given backbone structure, conditional diffusion models have achieved impressive results in generating natural-like sequences that fold back into the original structure. However, practical design tasks often require not only modeling a conditional distribution but also optimizing specific task objectives. For instance, in the inverse folding task, we may prefer protein sequences with high stability. To address this, we consider the scenario where we have pre-trained discrete diffusion models that can generate natural-like sequences, as well as reward models that map sequences to task objectives. We then formulate the reward maximization problem within discrete diffusion models, analogous to reinforcement learning (RL), while minimizing the KL divergence against pretrained diffusion models to preserve naturalness. To solve this RL problem, we propose a novel algorithm, DRAKES , that enables direct backpropagation of rewards through entire trajectories generated by diffusion models,
by making the originally non-differentiable trajectories differentiable using the Gumbel-Softmax trick.
Our theoretical analysis indicates that our approach can generate sequences that are both natural-like (i.e., have a high probability under a pretrained model) and yield high rewards.
While similar tasks have been recently explored in diffusion models for continuous domains, our work addresses unique algorithmic and theoretical challenges specific to discrete diffusion models, which arise from their foundation in continuous-time Markov chains rather than Brownian motion. Finally, we demonstrate the effectiveness of our algorithm in generating DNA and protein sequences that optimize enhancer activity and protein stability, respectively, important tasks for gene therapies and protein-based therapeutics. The code is available at https://github.com/ChenyuWang-Monica/DRAKES .

## 1 Introduction

Diffusion models have gained widespread recognition as effective generative models in continuous spaces, such as image and video generation (Song et al., ; Ho et al., ). Inspired by seminal works (e.g., Austin et al. (); Campbell et al. (); Sun et al. ()), recent studies (Lou et al., ; Shi et al., ; Sahoo et al., ) have shown that diffusion models are also highly effective in discrete spaces, including natural language and biological sequence generation (DNA, RNA, proteins). Unlike autoregressive models commonly used in language modeling, diffusion models are particularly well-suited for biological sequences, where long-range interactions are crucial for the physical behavior of molecules arising from those sequences (e.g., the 3D folded structure of RNA or proteins).

While discrete diffusion models effectively capture conditional distributions (e.g., the distribution of sequences given a specific backbone structure in an inverse protein folding design problem (Dauparas et al., ; Campbell et al., )), in many applications, especially for therapeutic discovery, we often aim to generate sequences that are both natural-like and optimize a downstream performance objective. For instance, in the inverse folding problem, we may prefer stable protein sequences (i.e., sequences that fold back into stable protein conformations (Widatalla et al., )); for mRNA vaccine production we desire 5’ UTRs that drive high translational efficiency (Castillo-Hair and Seelig, ); for gene and cell therapies, we desire regulatory DNA elements, such as promoters and enhancers, that drive high gene expression only in specific cell types (Taskiran et al., ); and for natural language we optimize to minimize harmfulness (Touvron et al., ).

To address these challenges, our work introduces a fine-tuning approach for well-pretrained discrete diffusion models that maximizes downstream reward functions. Specifically, we aim to optimize these reward functions while ensuring that the generated sequences maintain a high probability under the original conditional distribution (e.g., the distribution of sequences that fold into a given backbone structure). To achieve this, we formulate the problem as a reward maximization task, analogous to reinforcement learning (RL), where the objective function integrates both the reward terms and the KL divergence with respect to the pre-trained discrete diffusion model, which ensures that the generated sequences remain close to the pre-trained model, preserving their naturalness after fine-tuning. To solve this RL problem, we propose a novel algorithm, DRAKES, that enables direct backpropagation of rewards through entire trajectories by making the originally non-differentiable trajectories differentiable using the Gumbel-Softmax trick (Jang et al., ).

Our main contribution is an RL-based fine-tuning algorithm, Direct Reward bAcKpropagation with gumbEl Softmax trick (DRAKES), that enables reward-maximizing finetuning for discrete diffusion models (Figure 1). We derive a theoretical guarantee that demonstrates its ability to generate *natural* and *high-reward* designs, and demonstrate its performance empirically on DNA and protein design tasks. While similar algorithms exist for continuous spaces (Fan et al., ; Black et al., ; Uehara et al., ), our work is the first, to the best of our knowledge, to address these aspects in (continuous-time) discrete diffusion models. This requires addressing unique challenges, as discrete diffusion models are formulated as continuous-time Markov chains (CTMC), which differ from Brownian motion, and the induced trajectories from CTMC are no longer differentiable, unlike in continuous spaces. Our novel theoretical guarantee also establishes a connection with recent advancements in classifier guidance for discrete diffusion models (Nisonoff et al., ).

Figure: Figure 1: DRAKES. We maximize the reward with a penalty term relative to pre-trained discrete diffusion models using the Gumbel-Softmax trick.
Refer to caption: extracted/5934883/images/main_fig1.png

## 2 Related Works

#### Discrete diffusion models and their application in biology.

Building on the seminal works of Austin et al. (); Campbell et al. (), recent studies on masked diffusion models (Lou et al., ; Shi et al., ; Sahoo et al., ) have demonstrated strong performance in natural language generation. Recent advances in masked discrete diffusion models have been successfully applied to biological sequence generation, including DNA and protein sequences (Sarkar et al., ; Campbell et al., ). Compared to autoregressive models, diffusion models may be particularly well-suited for biological sequences, which typically yield molecules that fold into complex three-dimensional (3D) structures.

In contrast to these works, our study focuses on fine-tuning diffusion models to optimize downstream reward functions. One application of our approach is the fine-tuning of protein inverse folding generative models to optimize stability, as discussed in Widatalla et al. (). However, unlike this prior work, we employ discrete diffusion models as the generative model.

#### Controlled generation in diffusion models.

There are three primary approaches:

- •
Guidance: Techniques such as classifier guidance (Song et al., ; Dhariwal and Nichol, ) and its variants (e.g., Bansal et al. (); Chung et al. (); Ho et al. ()) introduce gradients from proxy models during inference. However, since gradients are not formally well-defined for discrete states in diffusion, a recent study (Nisonoff et al., ) proposed a method specifically designed for discrete diffusion models. Alternative approaches directly applicable to discrete diffusion models include sequential Monte Carlo (SMC)-based methods (Wu et al., ; Trippe et al., ; Dou and Song, ; Cardoso et al., ; Phillips et al., ). While these guidance-based inference techniques have their own advantages, they generally lead to longer inference times compared to fine-tuned models. We compare our methods against these in terms of generation quality in Section 6.
- •
RL-based fine-tuning: To maximize reward functions for pretrained diffusion models, numerous recent studies have explored RL-based fine-tuning in continuous diffusion models (i.e., diffusion models for continuous objectives) (Fan et al., ; Black et al., ; Clark et al., ; Prabhudesai et al., ). Our work, in contrast, focuses on discrete diffusion models.
- •
Classifier-free fine-tuning (Ho and Salimans, ): This approach constructs conditional generative models, applicable in our setting by conditioning on high reward values. Although not originally designed as a fine-tuning method, it can also be adapted for fine-tuning (Zhang et al., ) by adding further controls to optimize. However, in the context of continuous diffusion models, compared to RL-based fine-tuning, several works (Uehara et al., ) have shown that conditioning on high reward values is suboptimal, because such high-reward samples are rare. We will likewise compare this approach to ours in Section 6. Lastly, when pretrained models are conditional diffusion models (i.e., $p(x|c)$) and the offline dataset size consisting of triplets $(c,x,r(x))$ is limited, it is challenging to achieve success. Indeed, for this reason, most current RL-based fine-tuning papers (e.g., Fan et al. (); Black et al. (); Clark et al. ()) do not empirically compare their algorithms with classifier-free guidance.

#### Dirichlet diffusion models for discrete spaces.

Another approach to diffusion models for discrete spaces has been proposed (Stark et al., ; Avdeyev et al., ; Zhou et al., ). In these models, each intermediate state is represented as a vector within a simplex. This is in contrast to masked diffusion models, where each state is a discrete variable.

## 3 Preliminary

### 3.1 Diffusion Models on Discrete Spaces

In diffusion models, our goal is to model the data distribution $p_{\mathrm{data}}\in\Delta(\mathcal{X})$ using the training data, where $\mathcal{X}$ represents the domain. We focus on the case where $\mathcal{X}=\{1,2,\cdots,N\}$. The fundamental principle is (1) introducing a known forward model that maps the data distribution to a noise distribution, and (2) learning the time reversal that maps the noise distribution back to the data distribution (detailed in Lou et al. (); Sahoo et al. (); Shi et al. ()).

First, we consider the family of distributions $j_{t}\in\mathbb{R}^{N}$ (a vector summing to 1) that evolves from $t=0$ to $t=T$ according to a continuous-time Markov chain (CTMC):

$$ $\displaystyle\textstyle\frac{dj_{t}}{dt}=Q(t)j_{t},\quad p_{0}\sim p_{\mathrm{ data}},$ $$

where $Q(t)\in\mathbb{R}^{N\times N}$ is the generator. Generally, $j_{t}$ is designed so that $p_{t}$ approaches a simple limiting distribution at $t=T$. A common approach is to add Mask into $\mathcal{X}$ and gradually mask a sequence so that the limiting distribution becomes completely masked (Shi et al., ; Sahoo et al., ).

Next, we consider the time-reversal CTMC (Sun et al., ) that preserves the marginal distribution. This can be expressed as follows:

$$ $\displaystyle\textstyle\frac{dj_{T-t}}{dt}=\bar{Q}(T-t)j_{T-t},\quad\bar{Q}_{y ,x}(t)=\begin{cases}\textstyle\frac{j_{t}(y)}{j_{t}(x)}\bar{Q}_{x,y}(t)\,(y \neq x)\\ \textstyle-\sum_{y\neq x}\bar{Q}_{y,x}(t)\,(y=x),\end{cases}$ $$

where $Q_{y,x}(t)$ is a $(y,x)$-entry of a generator $Q(t)$.
This implies that if we can learn the marginal density ratio $p_{t}(y)/p_{t}(x)$, we can sample from the data distribution at $t=T$ by following the above CTMC controlled by $\bar{Q}(T-t)$. Existing works (e.g., Lou et al. ()) demonstrate how to train this ratio from the training data. Especially when we use masked diffusion models (Sahoo et al., ; Shi et al., ), we get

$$ $\displaystyle\textstyle\bar{Q}_{y,x}(t)=\begin{cases}\gamma\mathbb{E}[x_{0}=y| x_{t}=\mathrm{Mask}]\quad(y\neq\mathrm{Mask},x_{t}=\mathrm{Mask})\textstyle,\\ -\sum_{z\neq\mathrm{Mask}}\gamma\mathbb{E}[x_{0}=z|x_{t}=\mathrm{Mask}]\quad(y =\mathrm{Mask},x_{t}=\mathrm{Mask})\\ 0\quad(x_{t}\neq\mathrm{Mask})\textstyle\end{cases},$ (1) $$

for a certain constant $\gamma$, where the expectation is taken with respect to (w.r.t.) the distribution induced by the forward CTMC. Notably, the above formulation suggests that masked diffusion models could be viewed as a hierarchical extension of BERT (Devlin, ).

###### Remark 1 (Sequence of multiple tokens) .

When dealing with sequences of length $M$, $x=[x^{\langle 1\rangle},\cdots,x^{\langle M\rangle}]$, we simply consider the factorized rate matrix, i.e., $\bar{Q}_{y,x}=\sum_{i}\bar{Q}_{y^{\langle i\rangle},x}$ (Campbell et al., ), thereby avoiding exponential blowup.

###### Remark 2 (Conditioning) .

We can easily construct a conditional generative model for any $c\in\mathcal{C}$ by allowing the generator to be a function of $c\in\mathcal{C}$.

###### Remark 1 (Sequence of multiple tokens) .

###### Remark 2 (Conditioning) .

### 3.2 Goal: Generating Natural Samples While Optimizing Reward Functions

In our work, we consider a scenario with a pretrained discrete diffusion model $p^{\mathrm{pre}}(x|c)\in[\mathcal{C}\to\Delta(\mathcal{X})]$ trained on an extensive dataset and a downstream reward function $r:\mathcal{X}\to\mathbb{R}$. The pretrained diffusion model captures the naturalness or validity of samples. For example, in protein design, $p^{\mathrm{pre}}(\cdot|\cdot)$ could be a protein inverse-folding model that generates amino acid sequences that fold back into the given backbone structure (similar to Campbell et al. ()), and $r$ could be a function that evaluates stability. Our objective is to fine-tune a generative model to generate natural-like samples (high $\log p^{\mathrm{pre}}(\cdot|\cdot)$) with desirable properties (high $r(\cdot)$).

#### Notation.

We introduce a discrete diffusion model parameterized by $\theta$ from $t=0$ to $t=T$:

$$ $\displaystyle\textstyle\frac{dp_{t}}{dt}=Q^{\theta}(t)p_{t},\quad p_{0}=p_{ \lim}.$ (2) $$

The parameter $\theta$ from the pretrained model is denoted by $\theta_{\mathrm{pre}}$. The distribution at time $T$ is denoted as $p^{\mathrm{pre}}(\cdot)$, which approximates the training data distribution $p^{\mathrm{data}}$. We denote an element of the generated trajectory from $t=0$ to $t=T$ by $x_{0:T}$. For simplicity, we assume the initial distribution is a Dirac delta distribution (completely masked state), and we often treat the original pretrained diffusion model as an unconditional model for a single token for notational convenience. In this paper, all of the proofs are in Appendix B.

## 4 Algorithm

In this section, we present our proposed method, DRAKES, for fine-tuning diffusion models to optimize downstream reward functions. We begin by discussing the motivation behind our algorithm.

### 4.1 Key Formulation

Perhaps the most obvious starting point for fine-tuning diffusion models to maximize a reward function $r(x_{T})$ is to simply maximize the expected value of the reward under the model’s distribution, i.e., $\mathbb{E}_{x_{0}\sim P^{\theta}}\left[r(x_{T})\right]$, where the expectation is taken over the distribution $P^{\theta}(x_{0:T})$ induced by (2) (i.e., the generator $Q^{\theta}$). However, using only this objective could lead to over-optimization, where the model produces unrealistic or unnatural samples that technically achieve a high reward, but are impossible to generate in reality. Such samples typically exploit flaws in the reward function, for example, by being outside the training distribution of a learned reward or violating the physical assumptions of a hand-engineered physics-based reward (Levine et al., ; Clark et al., ; Uehara et al., ).
We address this challenge by constraining the optimized model to remain close to a pretrained diffusion model, which captures the distribution over natural or realistic samples. More specifically, we introduce a penalization term by incorporating the KL divergence between the fine-tuned model $P^{\theta}(x_{0:T})$ and the pretrained diffusion model $P^{\theta_{\mathrm{pre}}}(x_{0:T})$ in CTMC.

Accordingly, our goal during fine-tuning is to solve the following reinforcement learning (RL) problem:

$$ $\displaystyle\theta^{\star}\textstyle$ $\displaystyle=\mathop{\mathrm{argmax}}_{\theta\in\Theta}\underbrace{\mathbb{E} _{x_{0:T}\sim P^{\theta}}\left[r(x_{T})\right]}_{\textrm{Reward term}}$ (3) $\displaystyle-\alpha\underbrace{\mathbb{E}_{x_{0:T}\sim P^{\theta}}\left[\int_ {t=0}^{T}\sum_{y\neq x_{t}}\left\{Q^{\theta_{\mathrm{pre}}}_{x_{t},y}(t)-Q^{ \theta}_{x_{t},y}(t)+Q^{\theta}_{x_{t},y}(t)\log\frac{Q^{\theta}_{x_{t},y}(t)} {Q^{\theta_{\mathrm{pre}}}_{x_{t},y}(t)}\right\}dt\right]}_{\textrm{KL term}}.$ $$

The first term is designed to generate samples with desired properties, while the second term represents the KL divergence. The parameter $\alpha$ controls the strength of this regularization term.

Finally, after fine-tuning, by using the following CTMC from $t=0$ to $t=T$:

$$ $\displaystyle\textstyle\frac{dp_{t}}{dt}=Q^{\theta^{\star}}(t)p_{t},\quad p_{0 }=p_{\lim}.$ (4) $$

we generate samples at time $T$. Interestingly, we can show the following.

###### Theorem 1 (Fine-Tuned Distribution) .

When $\{Q^{\theta}_{\cdot,\cdot}:\theta\in\Theta\}$ is fully nonparametric (i.e., realizability holds), the generated distribution at time $T$ by (4) is proportional to

$$ $\displaystyle\exp(r(\cdot)/\alpha)p^{\mathrm{pre}}(\cdot).$ $$

This theorem offers valuable insights. The first term, $\exp(r(x))$, represents high rewards.
Additionally, the second term, $p^{\mathrm{pre}}(\cdot)$, can be seen as prior information that characterizes the natural sequence. For example, in the context of inverse protein folding, this refers to the ability to fold back into the target backbone structure.

###### Remark 3 .

A similar theorem has been derived for continuous diffusion models (Uehara et al., , Theorem 1).
However, our formulation (3) differs significantly as our framework is based on a CTMC, whereas those works are centered around the Brownian motion. Furthermore, while the use of a similar distribution is common in the literature on (autoregressive) large language models (e.g., Ziegler et al. ()), its application in discrete diffusion models is novel, considering that $p^{\mathrm{pre}}(\cdot)$ cannot be explicitly obtained in our context, unlike autoregressive models.

###### Theorem 1 (Fine-Tuned Distribution) .

###### Remark 3 .

### 4.2 Direct Reward Backpropagation with Gumbel Softmax Trick ( DRAKES )

Based on the key formulation presented in Section 4.1, we introduce our proposed method (Algorithm 1 and Figure 1), which is designed to solve the RL problem (3). The core approach involves iteratively (a) sampling from $x_{0:T}\sim P^{\theta}$ and (b) updating $\theta$ by approximating the objective function (3) with its empirical counterpart and adding its gradient with respect to $\theta$ into the current $\theta$. Importantly, for step (b) to be valid, step (a) must retain the gradients from $\theta$. After explaining the representation of $x_{t}$, we will provide details on each step.

Figure: Algorithm 1 DRAKES (Direct Reward bAcKpropagation with gumbEl Softmax trick)

#### Representation.

To represent $x\in\{1,\cdots,N\}$, we often use the $N$-dimensional one-hot encoding representation within $\mathbb{R}^{N}$ interchangeably. From this perspective, while the original generator corresponds to a map $\mathcal{X}\times\mathcal{X}\to\mathbb{R}$, we can also regard it as an extended mapping: $\mathbb{R}^{N}\times\mathbb{R}^{N}\to\mathbb{R}$. We will use this extended mapping when we consider our algorithm later.

#### Stage 1: Data collection (Step 2 - 9 )

We aim to sample from the distribution induced by the generator $Q^{\theta}$. In the standard discretization of CTMC, for $(y,x)\in\mathcal{X}\times\mathcal{X}$, at time $t$, we use

$$ $\displaystyle p(x_{t+\Delta t}=y|x_{t}=x)=\mathrm{I}(x=y)+Q^{\theta}_{y,x}(t)( \Delta t).$ $$

Thus, by defining $\pi_{t}=[Q^{\theta}_{1,x}(t)(\Delta t),\cdots,(1+Q^{\theta}_{x,x}(t))\Delta t
\cdots,Q^{\theta}_{N,x}(t)(\Delta t)]$,
we sample $x_{t+\Delta t}\sim\mathrm{Cat}(\pi_{t})$, where $\mathrm{Cat}(\cdot)$ denotes the categorical distribution.

However, this procedure is not differentiable with respect to $\theta$, which limits its applicability for optimization. To address this, we first recognize that sampling from the categorical distribution can be reduced to a Gumbel-max operation. Although this operation itself remains non-differentiable, we can modify it by replacing the max operation with a softmax, as shown in Line 7, which is also utilized in discrete VAE (Jang et al., ). This modification results in a new variable, $\bar{x}_{t}\sim[0,1]^{N}$, which maintains differentiability with respect to $\theta$. As the temperature $\tau$ approaches zero, $\bar{x}_{t}$ converges to a sample from the exact categorical distribution $\mathrm{Cat}(\pi_{t})$, effectively becoming $x_{t}$. Thus, we typically set the temperature to a low value to closely approximate the true distribution.

#### Stage 2: Optimization (Step 10 - 11 )

After approximately sampling from the distribution induced by $P^{\theta_{s}}$, we update the parameter $\theta_{s}$ by maximizing the empirical objective. Although $x_{t}$ itself may not have a valid gradient, $\bar{x}_{t}$ retains the gradient with respect to $\theta$. Therefore, we use the empirical approximation based on $\bar{x}_{t}$. We offer several remarks below, with details in Appendix C:

- •
Validity of $\bar{x}_{t}$: While $\bar{x}_{t}$ does not strictly belong to $\mathcal{X}$, this is practically acceptable since the generator $Q^{\theta}(t)$ is parameterized as a map $\mathbb{R}^{N}\times\mathbb{R}^{N}\to\mathbb{R}$.
- •
SGD Variants: Although Line 11 uses the standard SGD update, any off-the-shelf SGD algorithm, such as Adam (Kingma, ), can be applied in practice.
- •
Soft Calculation with $\bar{x}_{t}$: Transition probability $\pi(t)_{y}$ and the KL divergence term in $g(\theta)$ are modified to their soft counterparts by using $\bar{x}_{t}$ in place of $x_{t}$.
- •
Straight-Through Gumbel Softmax: Non-relaxed computations can be used in the forward pass (in Line 10). This is commonly known as straight-through Gumbel softmax estimator.
- •
Truncated Backpropagration: In practice, it is often more effective to backpropagate from intermediate time steps rather than starting from $t=0$. In practice, we adopted this truncation approach, as in Clark et al. ().
- •
Optimization Objective $g(\theta)$: For the masked diffusion models (1) that we utilized, $g(\theta)$ can be further simplified to reduce computational complexity, as detailed in Appendix C.2.

## 5 Theory of DRAKES

In this section, we provide an overview of the proof for Theorem 1. Based on the insights gained from this proof, we reinterpret state-of-the-art classifier guidance for discrete diffusion models (Nisonoff et al., ) from a new perspective.

### 5.1 Proof Sketch of Theorem 1

We define the optimal value function $V_{t}:\mathcal{X}\to\mathbb{R}$ as follows:

$$ $\displaystyle\mathbb{E}_{x_{t:T}\sim P^{\theta^{\star}}}\left[r(x_{T})-\alpha \int_{s=t}^{T}\sum_{y\neq x_{s}}\left\{Q^{\theta^{\star}}_{x_{s},y}(s)-Q^{ \theta_{\mathrm{pre}}}_{x_{s},y}(s)+Q^{\theta^{\star}}_{x_{s},y}(s)\log\frac{Q ^{\theta^{\star}}_{x_{s},y}(s)}{Q^{\theta_{\mathrm{pre}}}_{x_{s},y}(s)}\right \}ds\mid x_{t}=x\right].$ $$

This represents the expected return when starting from state $x$ at time $t$ and following the optimal policy. Once the optimal value function is defined, the optimal generator can be expressed in terms of this value function, as shown below. This is derived from the Hamilton-Jacobi-Bellman (HJB) equation in CTMC.

###### Theorem 2 (Optimal generator) .

For $x\neq y$ ($x,y\in\mathcal{X}$), we have

$$ $\displaystyle Q^{\theta^{\star}}_{x,y}(t)=Q^{\theta_{\mathrm{pre}}}_{x,y}(t) \exp(\{V_{t}(y)-V_{t}(x)\}/\alpha).$ $$

Next, consider an alternative expression for the soft value function, derived using the Kolmogorov backward equations in CTMC. This expression is particularly useful for learning value functions.

###### Theorem 3 (Feynman–Kac Formula in CTMC) .

$$ $\displaystyle\exp(V_{t}(x)/\alpha)=\mathbb{E}_{x_{t:T}\sim P^{\theta_{\mathrm{ pre}}}}[\exp(r(x_{T})/\alpha)|x_{t}=x]$ $$

With this preparation, we can prove our main theorem, which reduces to Theorem 1 when $t=T$.

###### Theorem 4 (Marginal distribution induced by the optimal generator $Q^{\theta^{\star}}(t)$ ) .

The marginal distribution at time $t$ by (4), $p^{\star}_{t}\in\Delta(\mathcal{X})$, is proportional to

$$ $\displaystyle\exp(V_{t}(\cdot)/\alpha)p^{\mathrm{pre}}_{t}(\cdot).$ $$

where $p^{\mathrm{pre}}_{t}\in\Delta(\mathcal{X})$ is a marginal distribution induced by pretrained model at $t$.

This is proved by showing the Kolmogorov forward equation in CTMC: $dp^{\star}_{t}/dt=Q^{\theta^{\star}}(t)p^{\star}_{t}$.

###### Theorem 2 (Optimal generator) .

###### Theorem 3 (Feynman–Kac Formula in CTMC) .

###### Theorem 4 (Marginal distribution induced by the optimal generator Q θ ⋆ ⁢ ( t ) superscript 𝑄 superscript 𝜃 ⋆ 𝑡 Q^{\theta^{\star}}(t) italic_Q start_POSTSUPERSCRIPT italic_θ start_POSTSUPERSCRIPT ⋆ end_POSTSUPERSCRIPT end_POSTSUPERSCRIPT ( italic_t ) ) .

### 5.2 Relation to Classifier Guidance for Discrete Diffusion Models

Now, we derive an alternative fine-tuning-free algorithm by leveraging observations in Section 5.1 for reward maximization. If we can directly obtain the optimal generator $Q^{\theta^{\star}}$, we can achieve our objective.
Theorem 2 suggests that the optimal generator $Q^{\theta^{\star}}$ is a product of the generator from the pretrained model and the value functions. Although we don’t know the exact value functions, they can be learned through regression using Theorem 3 based on

$$ $\displaystyle\exp(V_{t}(x)/\alpha)=\mathop{\mathrm{argmin}}_{g:\mathcal{X}\to \mathbb{R}}\mathbb{E}_{x_{T}\sim P^{\theta_{\mathrm{pre}}}}[\{\exp(r(x_{T})/ \alpha)-g(x_{t})\}^{2}].$ $$

In practice, while we can’t calculate the exact expectation, we can still replace it with its empirical analog. Alternatively, we can approximate it by using a map from $x_{t}$ to $x_{0}$ in pretrained models following DPS (Chung et al., ) or reconstruction guidance (Ho et al., ).

Interestingly, a similar algorithm was previously proposed by Nisonoff et al. (). While Nisonoff et al. () originally focused on conditional generation, their approach can also be applied to reward maximization or vice versa. In their framework for conditional generation, they define $r(x)=\log p(z|x)$ (e.g., the log-likelihood from a classifier) and set $\alpha=1$. By adapting Theorem 2 and 3 to their setting, we obtain:

$$ $\displaystyle Q^{\theta^{\star}}_{x,y}(t)=Q^{\theta_{\mathrm{pre}}}_{x,y}(t) \times p_{t}(z|y)/p_{t}(z|x),\quad p_{t}(z|x_{t}):=\mathbb{E}_{x_{t:T}\sim P^{ \theta_{\mathrm{pre}}}}[p(z|x_{T})\mid x_{t}].$ $$

Thus, we can rederive the formula in Nisonoff et al. (). Here, we also note that this type of result is commonly referred to as the Doob transform in the literature on stochastic processes (Levin and Peres, , Section 17).

While this argument suggests that classifier guidance and RL-based fine-tuning approaches theoretically achieve the same goal in an ideal setting (without function approximation, sampling, or optimization errors), their practical behavior can differ significantly, as we demonstrate in Section 6. At a high level, the advantage of classifier guidance is that it requires no fine-tuning, but the inference time may be significantly longer due to the need to recalculate the generator during inference. Indeed, this classifier guidance requires $\mathrm{O}(NM)$ computations of value functions at each step to calculate the normalizing constant. While this can be mitigated using a Taylor approximation, there is no theoretical guarantee for this heuristic in discrete diffusion models. Lastly, learning value functions in classifier guidance can often be practically challenging.

## 6 Experiments

Our experiments focus on the design of regulatory DNA sequences for enhancer activity and protein sequences for stability. Our results include comprehensive evaluations, highlighting the ability of DRAKES to produce natural-like
sequences while effectively optimizing the desired properties.

### 6.1 Baselines

We compare DRAKES against several baseline methods discussed in Section 2, which we summarize below with further details in Appendix D.1.

- •
Guidance-based Methods (CG, SMC, TDS). We compare our approach with representative guidance-based methods, including state-of-the-art classifier guidance (CG) tailored to discrete diffusion models (Nisonoff et al., ), SMC-based guidance methods (e.g., Wu et al. ()): SMC, where the proposal is a pretrained model and TDS, where the proposal is CG.
- •
Classifier-free Guidance (CFG) (Ho and Salimans, ). CFG is trained on labeled datasets with the measured attributes we aim to optimize.
- •
Pretrained. We generated sequences using pretrained models without fine-tuning.
- •
DRAKES w/o KL. This ablation of DRAKES omits the KL regularization term, evaluating how well this term mitigates over-optimization (discussed in Section 4.1).

### 6.2 Regulatory DNA Sequence Design

Here we aim to optimize the activity of regulatory DNA sequences such that they drive gene expression in specific cell types, a critical task for cell and gene therapy (Taskiran et al., ).

Dataset and settings. We experiment on a publicly available large-scale enhancer dataset (Gosai et al., ), which measures the enhancer activity of $\sim$700k DNA sequences (200-bp length) in human cell lines using massively parallel reporter assays (MPRAs), where the expression driven by each sequence is measured. We pretrain the masked discrete diffusion model (Sahoo et al., ) on all the sequences. We then split the dataset and train two reward oracles (one for finetuning and one for evaluation) on each subset, using the Enformer (Avsec et al., ) architecture to predict the activity level in the HepG2 cell line. These datasets and reward models are widely used in the literature on computational enhancer design (Lal et al., ; Uehara et al., ; Sarkar et al., ).
Detailed information about the pretrained model and reward oracles is in Appendix D.2.

Evaluations. To comprehensively evaluate each model’s performance in enhancer generation, we use the following metrics:

- •
Predicted activity based on the evaluation reward oracle (Pred-Activity). We predict the enhancer activity level in the HepG2 cell line using the reward oracle trained on the evaluation subset. Note that the diffusion models are fine-tuned (or guided) with the oracle trained on a *different* subset of the data, splitting based on chromosome following conventions (but in the same cell lines) (Lal et al., ).
- •
Binary classification on chromatin accessibility (ATAC-Acc). We use an independent binary classification model trained on chromatin accessibility data in the HepG2 cell line  (Consortium et al., ) (active enhancers should have accessible chromatin). While this is *not* used for fine-tuning, we use it for evaluation to further validate the predicted activity of the synthetic sequences, following Lal et al. ().
- •
3-mer Pearson correlation (3-mer Corr). We calculate the 3-mer Pearson correlation between the synthetic sequences and the sequences in the dataset (Gosai et al., ) with top 0.1% HepG2 activity level. Models that generate sequences that are more natural-like and in-distribution have a higher correlation.
- •
JASPAR motif analysis (JASPAR Corr). We scan the generated sequences of each model with JASPAR transcription factor binding profiles Castro-Mondragon et al. (), which identify potential transcription factor binding motifs in the enhancer sequences (which are expected to drive enhancer activity). We then count the occurrence frequency of each motif and calculate the Spearman correlation of motif frequency between the synthetic sequences generated by each model and the top 0.1% HepG2 activity sequences in the dataset.
- •
Log-likelihood of sequences (Log-Lik). We calculate the log-likelihood of the generated sequences with respect to the pretrained model to measure how likely the sequences are to be natural-like. Models that over-optimize the reward oracle generate out-of-distribution sequences and would have a low likelihood to the pretrained model. The likelihood is calculated using the ELBO of the discrete diffusion model in  Sahoo et al. ().

Results. DRAKES generates sequences with high predicted activity in the HepG2 cell line, as robustly measured by Pred-Activity and ATAC-Acc (Table 1).
The generated sequences closely resemble natural enhancers, as indicated by high 3-mer and JASPAR motif correlations, and a similar likelihood to the pretrained model. These highlight DRAKES’s effectiveness in generating plausible high-activity enhancer sequences.
Notably, while DRAKES, without KL regularization achieves higher Pred-Activity, this can be attributed to over-optimization. Despite splitting the data for fine-tuning and evaluation, the sequences remain highly similar due to many analogous regions within each chromosome. However, when evaluated with an independent activity oracle, ATAC-Acc, DRAKES  demonstrates superior performance while maintaining higher correlations and log likelihood.

Table: Table 1: Model performance on regulatory DNA sequence design. DRAKES generates sequences with high activity in the HepG2 cell line, measured by Pred-Activity and ATAC-Acc, while being natural-like by high 3-mer and JASPAR motif correlations and likelihood. We report the mean across 3 random seeds, with standard deviations in parentheses.

### 6.3 Protein Sequence Design: Optimizing Stability in Inverse Folding Model

In this task, given a pretrained inverse folding model that generates sequences conditioned on the backbone’s conformation (3D structure), our goal is to optimize the stability of these generated sequences, following Widatalla et al. ().

Dataset and settings. First, we pretrained an inverse folding model based on the diffusion model
(Campbell et al., ) and the ProteinMPNN (Dauparas et al., ) architecture, using the PDB training set from Dauparas et al. ().
Next, we trained the reward oracles using a different large-scale protein stability dataset, Megascale (Tsuboyama et al., ), which includes stability measurements (i.e., the Gibbs free energy change) for $\sim$1.8M sequence variants from 983 natural and designed domains.
Following dataset curation and a train-validation-test splitting procedure from Widatalla et al. () (which leads to $\sim$0.5M sequences on 333 domains) and using the ProteinMPNN architecture, we constructed two reward oracles – one for fine-tuning and one for evaluation, that predict stability from the protein sequence and wild-type conformation. Detailed information on the pretrained model and reward oracles is in Appendix D.3.

Evaluations. We use the following metrics to evaluate the stability of the generated sequences and their ability to fold into the desired structure. During evaluation, we always condition on protein backbone conformations from the test data that are *not* used during fine-tuning.

- •
Predicted stability on the evaluation reward oracle (Pred-ddG). The evaluation oracle is trained with the full Megascale dataset (train+val+test) to predict protein stability. Conversely, the fine-tuning oracle is trained only with data from the Megascale training set. Thus, during fine-tuning, the algorithms do not encounter any proteins used for evaluation.
- •
Self-consistency RMSD of structures (scRMSD). To assess how well a generated sequence folds into the desired structure, we use ESMFold (Lin et al., ) to predict the structures of the generated sequences and calculate their RMSD relative to the wild-type structure (i.e., the original backbone structure we are conditioning on). This is a widely used metric (Campbell et al., ; Trippe et al., ; Chu et al., ).

Following prior works (Campbell et al., ; Nisonoff et al., ), we calculate the success rate of inverse folding as the ratio of generated sequences with Pred-ddG$>0$ and scRMSD$<2$.

Table: Table 2: Model performance on inverse protein folding. DRAKES generates protein sequences that have high stability and fold to the desired structure, outperforming baselines in the overall success rate. We report the mean across 3 random seeds, with standard deviations in parentheses.

Figure: (a) Conditioning on the backbone structure of 7JJK.
Refer to caption: extracted/5934883/images/7JJK.png

Results. For inverse protein folding, DRAKES generates high-stability protein sequences capable of folding into the conditioned structure (Table 2). It achieves the highest Pred-ddG among all methods, while maintaining a similar success rate of inverse folding (measured by %(scRMSD$<2$), the percentage of scRMSD smaller than 2) as the pretrained model. Considering both factors, DRAKES significantly outperforms all baseline methods in terms of overall success rate.
Note that CFG does not work well for protein sequence design due to limited labeled data, as Megascale includes only a few hundred backbones, making generalization difficult. This is expected, as we mention in Section 2. Further details are provided in Appendix D.1.

Moreover, the results highlight the importance of the KL term, as DRAKES without KL regularization tends to suffer from over-optimization, with high scRMSD (i.e., failing to fold back to the target backbone structure), even though Pred-ddG may remain high.

#### In silico validation.

For validation purposes, we calculate the stability (i.e., Gibbs free energy) of the generated sequences using physics-based simulations (PyRosetta (Chaudhury et al., )) for wild-type protein backbone structures in Figure 2, following (Widatalla et al., ). Although all models are conditioned on the same set of protein backbones, different sets of sequences generated by generative methods can lead to significant differences in side chain interactions, which affect folding energies. The results demonstrate that sequences generated by our algorithms are more stable in this in silico validation compared to other baseline methods. For additional results, refer to Figure 6 in Appendix D.3.

## 7 Conclusions

We propose a novel algorithm that incorporates reward maximization into discrete diffusion models, leveraging the Gumbel-Softmax trick to enable differentiable reward backpropagation, and demonstrate its effectiveness in generating DNA and protein sequences optimized for task-specific objectives. For future work, we plan to conduct more extensive in silico validation and pursue wet-lab validation.

## Appendix A Potential Limitations

We have formulated the RL problem, (3), in the context of CTMC. The proposed algorithm in our paper to solve this problem requires reward models to be differentiable. Since differentiable models are necessary when working with experimental offline data, this assumption is not overly restrictive. Moreover, many state-of-the-art sequence models mapping sequences to functions in biology are available today, such as enformer borozi. In cases where creating differentiable models is challenging, we recommend using PPO-based algorithms (Black et al., ; Fan et al., ) or reward-weighted MLE (Uehara et al., , Section 6.1).

## Appendix B Proof of Theorems

### B.1 Preparation

We prepare two well-known theorems in CTMC for the pedagogic purpose. For example, refer to Yin and Zhang () for the more detailed proof. In these theorems, we suppose we have the CTMC:

$$ $\displaystyle\frac{dp_{t}}{dt}=Q(t)p_{t}.$ (5) $$

###### Lemma 1 (Kolmogorov backward equation) .

We consider $g(\cdot,t)=\mathbb{E}[r(x_{T})|x_{t}=\cdot]$ where the expectation is taken w.r.t. (5). Then, this function $g:\mathcal{X}\times[0,T]\to\mathbb{R}$ is characterized by the following ODE:

$$ $\displaystyle\frac{dg(x,t)}{dt}=\sum_{y\neq x}Q_{x,y}(t)\{g(x,t)-g(y,t)\}, \quad g(x,T)=r(x_{T}).$ $$

###### Proof.

Here, we prove that the p.d.f. $g$ satisfies the above backward equation. To show the converse, we technically require regularity conditions to claim the ODE solution is unique, which can often be proved by the contraction mapping theorem. Here, we skip the converse part.

When $t=T$, the statement $g(x,T)=r(x_{T})$ is obvious. For the rest of the proof, we aim to show a result when $t\neq T$.
We have

$$ $\displaystyle g(x_{t},t)=\int g(x_{t+dt},t+dt)p(x_{t+dt}|x_{t})dx_{t+dt}.$ $$

The above implies

$$ $\displaystyle g(x,t)=\{1+Q_{x,x}(t)dt\}g(x,t+dt)+\sum_{y\neq x}\{Q_{x,y}(t)dt \}g(y,t+dt).$ $$

Now combined with the property of the generator as follows

$$ $\displaystyle 0=\sum_{y}Q_{x,y}(t+dt),$ $$

with some algebra,

$$ $\displaystyle g(x,t)=g(x,t+dt)-\sum_{y\neq x}\{Q_{x,y}(t)dt\}g(x,t+dt)+\sum_{y \neq x}\{Q_{x,y}(t)dt\}g(y,t+dt).$ $$

Then, we have

$$ $\displaystyle\frac{g(x,t)-g(x,t+dt)}{dt}=\sum_{y\neq x}Q_{x,y}(t)\{g(y,t+dt)-g (x,t+dt)\}$ $$

Finally, by setting $dt\to 0$, we obtain

$$ $\displaystyle-\frac{dg(x,t)}{dt}=\sum_{y\neq x}Q_{x,y}(t)\{g(y,t)-g(x,t)\}.$ $$

∎

###### Lemma 2 (Kolmogorov forward equation) .

The density $p_{t}\in\Delta(\mathcal{X})$ is characterized as the following ODE:

$$ $\displaystyle\frac{dp_{t}(x)}{dt}=\sum_{y\neq x}Q_{y,x}(t)p_{t}(y)-\sum_{y\neq x }Q_{x,y}(t)p_{t}(x),\quad p_{0}=p_{ini}.$ $$

###### Proof.

Here, we prove that the p.d.f. $p_{t}$ satisfies the above forward equation. To show the converse, we technically require regularity conditions to claim the ODE solution is unique, which can often be proved by the contraction mapping theorem. Here, we skip the converse part.

We first have

$$ $\displaystyle p_{t+dt}(x)=\int p_{t+dt}(x|x_{t})p_{t}(x_{t})dx_{t}$ $$

This implies

$$ $\displaystyle p_{t+dt}(x)$ $\displaystyle=\sum_{y\neq x}\{Q_{y,x}(t)dtp_{t}(y)\}+\{1+Q_{x,x}(t)dt\}p_{t}(x)$ $\displaystyle=\sum_{y\neq x}\{Q_{y,x}(t)dtp_{t}(y)\}+\{1-\sum_{y\neq x}Q_{x,y} (t)dt\}p_{t}(x).$ $$

Hence,

$$ $\displaystyle\frac{p_{t+dt}(x)-p_{t}(x)}{dt}=\sum_{y\neq x}Q_{y,x}(t)p_{t}(y)- \sum_{y\neq x}Q_{x,y}(t)p_{t}(x).$ $$

By taking $dt\to 0$, we obtain

$$ $\displaystyle\frac{dp_{t}(x)}{dt}=\sum_{y\neq x}Q_{y,x}(t)p_{t}(y)-\sum_{y\neq x }Q_{x,y}(t)p_{t}(x).$ $$

Then, the proof is competed.
∎

###### Lemma 1 (Kolmogorov backward equation) .

###### Proof.

###### Lemma 2 (Kolmogorov forward equation) .

###### Proof.

### B.2 Proof of Theorem 2

We derive the Hamilton-Jacobi-Bellman (HJB) equation in CTMC. For this purpose, we consider the recursive equation:

$$ $\displaystyle V(x,t)$ $\displaystyle=\max_{\theta}\left[\left\{\sum_{y\neq x}Q^{\theta}_{x,y}(t)-Q^{ \theta_{\mathrm{pre}}}_{x,y}(t)-Q^{\theta}_{x,y}(t)\log\frac{Q^{\theta}_{x,y}( t)}{Q^{\theta_{\mathrm{pre}}}_{x,y}(t)}\right\}dt\right.$ $\displaystyle+\left.\sum_{y\neq x}\{Q^{\theta}_{x,y}(t)dtV(y,t+dt)\}+\{1+Q^{ \theta}_{x,x}(t)\}V(x,t+dt)\}\right].$ $$

Using $\sum_{y\in\mathcal{X}}Q_{x,y}(t)=0$, this is equal to

$$ $\displaystyle V(x,t)$ $\displaystyle=\max_{\theta}\left[\left\{\sum_{y\neq x}Q^{\theta}_{x,y}(t)-Q^{ \theta_{\mathrm{pre}}}_{x,y}(t)-Q^{\theta}_{x,y}(t)\log\frac{Q^{\theta}_{x,y}( t)}{Q^{\theta_{\mathrm{pre}}}_{x,y}(t)}\right\}dt\right.$ $\displaystyle\left.+V(x,t+dt)+\sum_{y\neq x}Q^{\theta}_{x,y}(t)dt\{V(y,t+dt)-V (x,t+dt)\}\right].$ $$

By taking $dt$ to 0 0, the above is equal to

$$ $\displaystyle\textstyle-\frac{dV(x,t)}{dt}$ $\displaystyle=\max_{\theta\in\Theta}\left\{\left[\sum_{y\neq x}Q^{\theta}_{x,y }(t)-Q^{\theta_{\mathrm{pre}}}_{x,y}(t)-Q^{\theta}_{x,y}(t)\log\frac{Q^{\theta }_{x,y}(t)}{Q^{\theta_{\mathrm{pre}}}_{x,y}(t)}\right]+\sum_{y\neq x}Q^{\theta }_{x,y}(t)\{V(y,t)-V(x,t)\}\right\}$ (6) $$

This is the HBJ equation in CTMC.

Finally, with simple algebra (i.e., taking functional derivative under the constraint $0=\sum_{y\in\mathcal{X}}Q^{\theta}_{x,y}(t)$), we can show

$$ $\displaystyle\forall x\neq y;Q^{\theta^{\star}}_{x,y}(t)=Q^{\theta_{\mathrm{ pre}}}_{x,y}(t)\exp(\{V(y,t)-V(x,t)\}).$ $$

### B.3 Proof of Theorem 3

This theorem is proved by invoking the Kolmogorov backward equation.

First, by plugging

$$ $\displaystyle\forall x\neq y;Q^{\theta^{\star}}_{x,y}(t)=Q^{\theta_{\mathrm{ pre}}}_{x,y}(t)\exp(\{V(y,t)-V(x,t)\}).$ $$

into (6), we get

$$ $\displaystyle\frac{dV(x,t)}{dt}=\sum_{y\neq x}Q^{\theta_{\mathrm{pre}}}_{x,y}( t)\{1-\exp(\{V(y,t)-V(x,t)\})\}.$ $$

By multiplying $\exp(V(x,t))$ to both sides, it reduces to

$$ $\displaystyle\frac{d\exp(V(x,t))}{dt}=\sum_{y\neq x}Q^{\theta_{\mathrm{pre}}}_ {x,y}(t)\{\exp(V(x,t))-\exp(V(y,t))\}.$ (7) $$

Furthermore, clearly, $V(x,T)=r(x_{T})$. Then, the statement is proved by invoking the Kolmogorov backward equation.

### B.4 Proof of Theorem 4

We define

$$ $\displaystyle H_{t}(x):=\exp(V(x,t))p_{t}(x)/C.$ $$

We aim to prove that the above satisfies the Kolmogorov forward equation:

$$ $\displaystyle\underbrace{\frac{dH_{t}(x)}{dt}}_{\text{l.h.s.}}=\underbrace{ \sum_{y\neq x}Q^{\theta^{\star}}_{y,x}(t)H_{t}(y)-\sum_{y\neq x}Q^{\theta^{ \star}}_{x,y}(t)H_{t}(x)}_{\text{r.h.s.}},\quad p_{ini}=H_{0}(\cdot).$ $$

First, we calculate the l.h.s. Here, recall

$$ $\displaystyle\frac{d\exp(V(x,t))}{dt}=\sum_{y\neq x}Q^{\theta_{\mathrm{pre}}}_ {x,y}(t)\{\exp(V(x,t))-\exp(V(y,t))\}$ $$

using (7), and

$$ $\displaystyle\frac{dp_{t}(x)}{dt}=\sum_{y\neq x}Q^{\theta_{\mathrm{pre}}}_{y,x }(t)p_{t}(y)-\sum_{y\neq x}Q^{\theta_{\mathrm{pre}}}_{x,y}(t)p_{t}(x)$ $$

holds, using the Kolmogorov forward equation. Then, we obtain

$$ $\displaystyle\frac{dH_{t}(x)}{dt}$ $\displaystyle=\frac{1}{C}\times\left\{\frac{d\exp(V(x,t))}{dt}p_{t}(x)+\exp(V( x,t))\frac{dp_{t}(x)}{dt}\right\}$ $\displaystyle=\frac{1}{C}\times\left[\sum_{y\neq x}Q^{\theta_{\mathrm{pre}}}_{ x,y}(t)\{\exp(V(x,t))-\exp(V(y,t))\}p_{t}(x)\right]$ $\displaystyle+\frac{1}{C}\times\exp(V(x,t))\left\{\sum_{y\neq x}Q^{\theta_{ \mathrm{pre}}}_{y,x}(t)p_{t}(y)-\sum_{y\neq x}Q^{\theta_{\mathrm{pre}}}_{x,y}( t)p_{t}(x)\right\}$ $\displaystyle=\frac{1}{C}\times\sum_{y\neq x}Q^{\theta_{\mathrm{pre}}}_{y,x}(t )\exp(V(x,t))p_{t}(y)-\frac{1}{C}\times\sum_{y\neq x}Q^{\theta_{\mathrm{pre}}} _{x,y}(t)\exp(V(y,t))p_{t}(x).$ $$

On the other hand, the r.h.s. is

$$ $\displaystyle\frac{1}{C}\times\left\{\sum_{y\neq x}Q^{\theta^{\star}}_{y,x}(t) H_{t}(y)-\sum_{y\neq x}Q^{\theta^{\star}}_{x,y}(t)H_{t}(x)\right\}$ $\displaystyle=\frac{1}{C}\times\sum_{y\neq x}Q^{\theta_{\mathrm{pre}}}_{y,x}(t )\exp(\{V(x,t)-V(y,t)\})H_{t}(y)-\frac{1}{C}\sum_{y\neq x}Q^{\theta_{\mathrm{ pre}}}_{x,y}(t)\exp(\{V(y,t)-V(x,t)\})H_{t}(x)$ $\displaystyle=\frac{1}{C}\times\sum_{y\neq x}Q^{\theta_{\mathrm{pre}}}_{y,x}(t )\exp(V(x,t))p_{t}(y)-\frac{1}{C}\times\sum_{y\neq x}Q^{\theta_{\mathrm{pre}}} _{x,y}(t)\exp(V(y,t))p_{t}(x).$ $$

Here, from the first line to the second line, we use

$$ $\displaystyle\forall x\neq y;Q^{\theta^{\star}}_{x,y}(t)=Q^{\theta_{\mathrm{ pre}}}_{x,y}(t)\exp(\{V(y,t)-V(x,t)\}).$ $$

Finally, we can see that $l.h.s.=r.h.s.$ Furthermore, recalling we have an assumption that $p_{ini}$ is Dirac delta distribution, we clearly have $p_{ini}=H_{0}(\cdot)$. Hence, the statement is proved by the Kolmogorov forward equation.

## Appendix C Details of Algorithm

### C.1 Straight-Through Gumbel Softmax

We apply the straight-through Gumbel softmax estimator to the last time step, i.e.

$$ $\text{ST}(x^{(i)}_{T}):=\bar{x}^{(i)}_{T}+\text{SG}(x^{(i)}_{T}-\bar{x}^{(i)}_ {T})$ $$

where $x^{(i)}_{T}$ is the corresponding Gumbel-max variable, i.e. $x^{(i)}_{T}=\mathop{\mathrm{argmax}}_{x\in\mathcal{X}}[\bar{x}^{(i)}_{T}]_{x}$, and SG denotes stop gradient. Then, $\text{ST}(x^{(i)}_{T})$ is input into the reward function $r(.)$ instead of $\bar{x}^{(i)}_{T}$ for forward and backward propagation.

We observe a boost in fine-tuning performance with the straight-through Gumbel softmax, as converting the input to $r(.)$ into a one-hot vector makes it better aligned with the reward oracle’s training distribution.

### C.2 Simplified Formula of g ⁢ ( θ ) 𝑔 𝜃 g(\theta) italic_g ( italic_θ )

The key objective function in DRAKES, $g(\theta)$, can be further simplified for the masked diffusion models that we utilized in the experiments.

$$ $\displaystyle g(\theta)=\frac{1}{B}\sum_{i=1}^{B}\left[r(\bar{x}^{(i)}_{T})- \frac{\alpha}{T}\sum_{t=1}^{T}\sum_{x\in\mathcal{X}}[\bar{x}^{(i)}_{t-1}]_{x} \sum_{\begin{subarray}{c}y\in\mathcal{X}\\ y\neq x\end{subarray}}\left\{-Q^{\theta}_{x,y}(t)+Q^{\theta_{\mathrm{pre}}}_{x ,y}(t)+Q^{\theta}_{x,y}(t)\log\frac{Q^{\theta}_{x,y}(t)}{Q^{\theta_{\mathrm{ pre}}}_{x,y}(t)}\right\}\right]$ $$

We denote the second term estimating the KL divergence with the $i$-th sample as $k_{i}(\theta)$:

$$ $\displaystyle k^{(i)}(\theta)=\frac{1}{T}\sum_{t=1}^{T}\sum_{x\in\mathcal{X}}[ \bar{x}^{(i)}_{t-1}]_{x}\sum_{\begin{subarray}{c}y\in\mathcal{X}\\ y\neq x\end{subarray}}\left\{-Q^{\theta}_{x,y}(t)+Q^{\theta_{\mathrm{pre}}}_{x ,y}(t)+Q^{\theta}_{x,y}(t)\log\frac{Q^{\theta}_{x,y}(t)}{Q^{\theta_{\mathrm{ pre}}}_{x,y}(t)}\right\}$ $$

When $x=\text{Mask}$, the value of $Q_{x,y}(t)$ is irrelevant to the parametrization $\theta$, i.e.

$$ $\displaystyle Q^{\theta}_{x,y}(t)=Q^{\theta_{\mathrm{pre}}}_{x,y}(t)=\begin{ cases}0,y\neq\text{Mask}\\ -\gamma,y=\text{Mask}\end{cases}$ $$

where $\gamma$ is a constant related to the forward process schedule (Sahoo et al., ). In particular, when applying a linear schedule (as in our experiments), $\gamma=1/t$. Thus, the corresponding KL divergence component equals 0.

When $x\neq\text{Mask}$,

$$ $\displaystyle Q^{\theta}_{x,y}(t)=\begin{cases}0,y\neq\text{Mask}\\ \gamma\mathbb{E}_{\theta}[x_{0}=x|x_{t-1}=\mathrm{Mask}],y=\text{Mask}\end{cases}$ $$

Denote $\mathbb{E}_{\theta}[x_{0}=x|x_{t-1}=\mathrm{Mask}]$ as $[\hat{x}_{0}^{\theta}]_{x}$.
The KL divergence term $k^{(i)}(\theta)$ can be simplified as

$$ $\displaystyle k^{(i)}(\theta)$ $\displaystyle=\frac{1}{T}\sum_{t=1}^{T}\sum_{x\in\mathcal{X}}[\bar{x}^{(i)}_{t -1}]_{x}\sum_{\begin{subarray}{c}y\in\mathcal{X}\\ y\neq x\end{subarray}}\left\{-Q^{\theta}_{x,y}(t)+Q^{\theta_{\mathrm{pre}}}_{x ,y}(t)+Q^{\theta}_{x,y}(t)\log\frac{Q^{\theta}_{x,y}(t)}{Q^{\theta_{\mathrm{ pre}}}_{x,y}(t)}\right\}$ $\displaystyle=\frac{1}{T}\sum_{t=1}^{T}\sum_{x\in\mathcal{X}}[\bar{x}^{(i)}_{t -1}]_{x}\left\{-Q^{\theta}_{x,\text{Mask}}(t)+Q^{\theta_{\mathrm{pre}}}_{x, \text{Mask}}(t)+Q^{\theta}_{x,\text{Mask}}(t)\log\frac{Q^{\theta}_{x,\text{ Mask}}(t)}{Q^{\theta_{\mathrm{pre}}}_{x,\text{Mask}}(t)}\right\}$ $\displaystyle=\frac{\gamma}{T}\sum_{t=1}^{T}\sum_{\begin{subarray}{c}x\in \mathcal{X}\\ x\neq\text{Mask}\end{subarray}}[\bar{x}^{(i)}_{t-1}]_{x}\left\{-[\hat{x}_{0}^{ \theta}]_{x}+[\hat{x}_{0}^{\theta_{\mathrm{pre}}}]_{x}+[\hat{x}_{0}^{\theta}]_ {x}\log\frac{[\hat{x}_{0}^{\theta}]_{x}}{[\hat{x}_{0}^{\theta_{\mathrm{pre}}}] _{x}}\right\}$ $$

The simplified formula reduces the computational complexity of calculating $k^{(i)}(\theta)$ to $O(NT)$.

### C.3 Schedule of Gumbel Softmax Temperature

We use a linear schedule for the Gumbel softmax temperature $\tau$, decreasing over time as $\tau\sim 1/t$. In early time steps, the temperature is higher, introducing more uncertainty, while later steps have a lower temperature, approximating the true distribution more closely.
This improves the fine-tuning procedure as the input becomes closer to clean data at later time steps and the uncertainty of model prediction is reduced.

## Appendix D Experimental Details and Additional Results

In this section, we first provide a more detailed explanation of the baseline. We then discuss additional results and present more detailed settings for both regulatory DNA sequence design and protein sequence design.

### D.1 Baselines

In this section, we provide a detailed overview of each baseline method.

- •
Guidance-based Methods. Guidance-based methods are based on the pretrained model while adjusting during the sampling process according to the targeted property. This leads to longer inference time compared to fine-tuning approaches.
–
CG (Nisonoff et al., ). CG adjusts the transition rate of CTMC by calculating the predictor guidance:
$Q_{x,y|r}(t)=\frac{p(r|y,t)}{p(r|x,t)}Q_{x,y}(t)$
where $r$ is the target property, and the predictor guidance is further approximated using a Taylor expansion, i.e.
$\log\frac{p(r|y,t)}{p(r|x,t)}\approx(y-x)^{T}\nabla_{x}\log p(r|x,t)$
The predictor $p(r|x,t)$ is estimated using the posterior mean approach (Chung et al., ), where the pretrained model is first utilized to estimate the clean data from the noisy input $x_{t}$, and then the reward oracle is applied to the predicted clean sequence. We remark that the above Taylor approximation doesn’t have formal theoretical guarantees, considering that $x$ is discrete. This could be a reason why it does not work well in the case of protein-inverse folding in Section 6.3.
–
SMC (Wu et al., ). SMC is a sequential Monte Carlo-based approach that uses the pretrained model as the proposal distribution. While it was originally designed for conditioning rather than reward maximization, it can be adapted for reward maximization by treating rewards as classifiers. In our experiment, we use this adapted version.
–
TDS (Wu et al., ). Similar to SMC, TDS also applies sequential Monte Carlo, but utilizes CG rather than the pretrained model as the proposal.
- •
Classifier-free Guidance (CFG) (Ho and Salimans, ). Unlike guidance-based methods, CFG trains a conditional generative model from scratch and does not rely on the pretrained model. To generate sequences $x$ with desired properties $r(x)$, CFG incorporates $r(x)$ as an additional input to the diffusion model and generates samples conditioning on high $r(x)$ values. Specifically, binary labels of $r(x)$ are constructed according to the $95\%$ quantile, and sampling is done conditioned on the label corresponding to high values of $r(x)$.
It is important to note that CFG requires labeled data pairs $\{x,r(x)\}$ for training,
which can limit its performance in cases with limited labeled data, especially when the pretrained model is already a conditional diffusion model $p(x|c)$. For example, in the protein inverse folding task, where $x$ is the protein sequence, $c$ is the protein structure, and $r(x)$ is the protein stability, CFG struggles, as shown in Table 2.
This is due to the small size of the Megascale dataset (containing only a few hundred different protein structures), which reduces its capability and generalizability(^1^11In contrast, other methods (guidance-based methods and fine-tuning methods) leverage the pretrained model trained on the much larger PDB dataset ($\sim 23,000$ structures) and achieve better performance.). While data augmentation can be applied to construct additional training data, it is resource-intensive, requires significant case-by-case design, and is beyond the scope of this work.
For the DNA sequence design task, since all sequences in the dataset are labeled, there is no such issue.

### D.2 Regulatory DNA Sequence Design

In this section, we first outline the training process for reward oracles and pre-trained models in the regulatory DNA sequence design experiment. Subsequently, we provide additional explanations on fine-tuning setups and present further results.

Reward Oracle. We train reward oracles to predict activity levels of enhancers in the HepG2 cell line using the dataset from Gosai et al. (). Following standard practice (Lal et al., ), we split the dataset into two subsets based on chromosomes, with each containing enhancers from half of the 23 human chromosomes. We train two reward oracles on each subset independently using the Enformer (Avsec et al., ) architecture initialized with its pretrained weights. One oracle is used for fine-tuning, while the other is reserved for evaluation (i.e. Pred-Activity in Table 1). Denote the subset used for training the fine-tuning oracle as FT and the subset for training the evaluation oracle as Eval. Table 3 presents the model performance for both oracles on each subset. Both oracles perform similarly, achieving a high Pearson correlation ($>0.85$) on their respective held-out sets (Eval for the fine-tuning oracle and FT for the evaluation oracle).

Table: Table 3: Performance of the reward oracles for predicting HepG2 activity of enhancer sequences.

Pretrained Model. We pretrain the masked discrete diffusion model (Sahoo et al., ) on the full dataset of Gosai et al. (), using the same CNN architecture as in Stark et al. () and a linear noise schedule. Other hyperparameters are kept identical to those in Sahoo et al. (). To assess the model’s ability to generate realistic enhancer sequences, we sample 1280 sequences and compare them with 1280 randomly selected sequences from the original dataset. Figure 3 presents the distribution of HepG2 activity predicted by either the fine-tuning (FT) or evaluation (Eval) oracle for both the generated and original sequences, along with the true observations for the original sequences. The activity levels of the generated sequences align well with those of the original dataset, indicating the effectiveness of pretrained model in generating in-distribution enhancer sequences. Furthermore, Figure 4 shows the 3-mer and 4-mer Pearson correlation between the synthetic and original sequences, both of which exceed 0.95, further validating the model’s performance.

Figure: Figure 3: Comparison of HepG2 activity distributions between original sequences and those generated by the pretrained model. The activity distributions match closely with each other.
Refer to caption: x1.png

Figure: Figure 4: 3-mer and 4-mer Pearson correlation between the original and generated sequences.
Refer to caption: x3.png

Fine-tuning Setup. We utilize the pretrained masked discrete diffusion model and the fine-tuning oracle described above for fine-tuning. During DRAKES’s stage 1 data collection, sequences are generated from the pretrained model over 128 steps. We set $\alpha=0.001$ to govern the strength of the KL regularization and truncate the backpropagation at step 50. The model is fine-tuned with 128 samples as a batch (32 samples per iteration, with gradient accumulated over 4 iterations) for 1000 steps. For DRAKES w/o KL, we follow the same setup, but set $\alpha$ to zero. For evaluation, we generate 640 sequences per method (with batch size of 64 over 10 batches) for each random seed. We report the mean and standard deviation of model performance across 3 random seeds.

Additional Results for Fine-Tuning. Along with the median Pred-Activity values shown in Table 1, Figure 5 presents the full distribution of Pred-Activity for each method, which shows consistent patterns as Table 1.

Figure: Figure 5: Distribution of Pred-Activity for the generated sequences of each method.
Refer to caption: x5.png

### D.3 Protein Inverse Folding

We first discuss the setup of datasets used for training reward models. Next, we present the performance of the pre-trained model, followed by an evaluation of the reward models. Finally, we describe the fine-tuning setup and provide additional results.

Dataset Curation. We utilize the large-scale protein stability dataset, Megascale (Tsuboyama et al., ) for the protein inverse folding experiment, which contains stability measurements for $\sim 1.8$M sequence variants (for example, single mutants and double mutants) from 983 protein domains.
We follow the dataset curation and train-validation-test splitting procedure from Widatalla et al. (). Specifically, the wild-type protein structures are clustered with Foldseek clustering and the data is split based on clusters. We then drop a few proteins with ambiguous wild type labels, and clip the $\Delta G$ values that are outside the dynamic range of the experiment ($>5$ or $<1$) to the closest measurable value ($5$ or $1$) as in Nisonoff et al. (). We further exclude proteins where a significant proportion of the corresponding variants’ $\Delta G$ measurements fall outside the experimental range. The final dataset consists of 438,540 sequence variants from 311 proteins in the training set, 15,182 sequences from 10 proteins in the validation set, and 23,466 sequences from 12 proteins in the test set.

Pretrained Model. We pretrain an inverse folding model using the discrete flow model loss from (Campbell et al., ) and the ProteinMPNN (Dauparas et al., ) architecture to encode both sequence and structure as model input. The model is trained on the PDB training set used in Dauparas et al. (), containing 23,349 protein structures and their ground truth sequences, which is distinct from the dataset in Tsuboyama et al. (). We first evaluate the effectiveness of the inverse folding model on the PDB test set in Dauparas et al. (), which has 1,539 different proteins. As in Nisonoff et al. (), we set the temperature during sampling to be 0.1, and randomly sample one sequence conditioned on each structure for both our pretrained discrete flow model and the de facto inverse folding method, ProteinMPNN. As shown in Table 4, the pretrained model performs similarly to ProteinMPNN, achieving comparable sequence recovery rate.

Table: Table 4: Model performance of protein inverse folding on PDB test set.

We further evaluate the generalizability of the pretrained model to the proteins in the Megascale dataset. Results on both Megascale training and test set are shown in Table 5. We calculate the self-consistency RMSD (scRMSD) to assess how well a generated sequence folds into the desired structure. Specifically, the generated sequences are folded into 3D structures using ESMFold (Lin et al., ), and scRMSD is calculated as their RMSD relative to the original backbone structure we are conditioning on. An scRMSD lower than $2\mathring{A}$ is typically considered a successful inverse folding (Nisonoff et al., ; Campbell et al., ). As shown in Table 5, the pretrained model achieves a similar sequence recovery rate on Megascale as the PDB test set and low scRMSD, with a success rate greater than $90\%$, indicating its effectiveness on the inverse folding task.

Table: Table 5: Model performance of protein inverse folding on Megascale proteins.

Reward Oracle. We train the reward oracles on the Megascale dataset using the ProteinMPNN architecture. The oracles take both the protein sequence and the corresponding wild-type structure as input to predict the stability of the sequence, measured by $\Delta\Delta G$ (calculated as the difference in $\Delta G$ between the variant and the wild-type from the dataset). Similar to Nisonoff et al. (), the final layer of the ProteinMPNN architecture is mean-pooled and mapped to a single scalar with a 2-layer MLP, and the model weights before the mean-pooling are initialized with the weights from the pretrained inverse folding model.

Similar to the practice in the enhancer design experiment, we train two oracles – one for fine-tuning and one for evaluation. The fine-tuning oracle is trained on Megascale training set. We select the best epoch based on validation set performance, and report the Pearson correlation on both Megascale training and test set in Table 6. The performance gap between the training and test sets highlights the difficulty of generalizing to unseen protein structures in this task.

The evaluation oracle is trained on the complete dataset (train+val+test). To attain the best hyperparameters, we randomly split the full dataset into two subsets, an in-distribution set for training, denoted as I, and an out-of-distribution set for validation, denoted as O. Note that here the evaluation oracle is trained part of the variants of all wild-type proteins (i.e. Megascale-Train-I & Megascale-Val-I & Megascale-Test-I), and the out-of-distribution set contains unseen sequence variants, but no new structures. The Pearson correlation on each subset is presented in Table 6. It achieves much higher correlations than the fine-tuning oracle, indicating good generalizability of the evaluation oracle to new sequences of in-distribution protein structures. For the final evaluation oracle used to calculate results in Table 2, we train it on the full dataset using the best hyperparameters selected as discussed. It achieves a Pearson correlation of 0.951 on Megascale training set and 0.959 on Megascale test set (both being in-distribution for the evaluation oracle).

Table: Table 6: Performance of the reward oracles for predicting stability conditioned on protein sequence and structure, across a variety of Megascale subsets.

Finetuning Setup. We utilize the pretrained inverse folding model and the fine-tuning oracle described above for fine-tuning. During DRAKES’s stage 1 data collection, we generate sequences from the pretrained model over 50 steps. We set $\alpha=0.0003$ and truncate the backprogagtion at step 25. The model is finetuned with proteins in Megascale training set with batch size 128 (16 samples per iteration, with gradient accumulated over 8 iterations) for 100 epochs. For DRAKES w/o KL, we follow the same setup, but set $\alpha$ to zero. The model is evaluated on Megascale test set, where we generate 128 sequences conditioned on each protein structure for every method (with batch size of 16 over 8 batches) and each random seed. We report the mean and standard deviation of model performance across 3 random seeds.

Evaluation Oracle Accounts for Over-Optimization. As discussed in Section 6.2, for the enhancer design experiment, significant over-optimization occurs when evaluating Pred-Activity, even with an evaluation oracle trained on distinct data unseen during fine-tuning. In contrast, the protein inverse folding experiment largely mitigates this issue. Table 7 shows the median values of Pred-ddG for the generated sequences based on both the evaluation oracle (same as those reported in Table 2) and the fine-tuning oracle. Although DRAKES w/o KL shows significantly higher Pred-ddG than DRAKES with the fine-tuning oracle, their performance with the evaluation oracle remains similar, suggesting less pronounced over-optimization in evaluation.
This is because enhancer sequences are relatively homogeneous, and even though we split based on chromosomes, each chromosome still has similar regions. However, protein structures are more distinct, and training on different proteins creates unique model landscapes.

Table: Table 7: Model performance on protein inverse folding, with Pred-ddG calculated using either the evaluation oracle (Eval) or the fine-tuning oracle (FT).

Additional Results. We provide more examples of the generated proteins in Figure 6, in addition to Figure 2. We also provide the specific values for energy, Pred-ddG and scRMSD of the visualized protein generated by DRAKES, as well as the energy values for the corresponding wild-type structure.

Figure: Figure 6: Additional examples of generated proteins.
Refer to caption: extracted/5934883/images/app_proteins/7JJK_app.png

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 5
presentation: 3
questions: See weaknesses above.
soundness: 2
strengths: 1. Using Gumbel-softmax in non-differentiable reinforcement learning in training discrete diffusion model is a novel approach.
summary: This paper presents a new fine-tuning method for discrete diffusion models. The objective is to generate high-reward yet natural sequences of DNA or RNA, where naturalness is measured by the KL divergence with the prior distribution (i.e., the pretrained model). This problem is mainly formulated with KL-constrained reinforcement learning, so they use Gumbel-softmax tricks to enable backpropagation over rewards. The performance is compared with existing RL and classifier guidance methods in DNA and RNA optimization.
weaknesses: 1. This work should establish stronger connections with probabilistic inference methods that aim to sample from unnormalized densities using objectives with KL-constrained reinforcement learning. This includes hierarchical variational inference and GFlowNets, as mentioned in recent work [1] on fine-tuning diffusion models, including the discrete case.


2. Is the KL constraint for diffusion models tractable to compute? Based on theoretical results in [2], the KL divergence is intractable due to the compositionality of the diffusion process, so we can only obtain an upper bound for the KL. Please make correction that if my understanding is wrong. 


3. Why should this method be better than existing RL methods? Variational inference, REINFORCE, and Gumbel-softmax are rival approaches in many applications that enable training models in non-differentiable settings. Why does the Gumbel trick have to be the best among them?

[1] Venkatraman et al. "Amortizing Intractable Inference in Diffusion Models for Vision, Language, and Control," NeurIPS 2024

[2] Fan et al. "Reinforcement Learning for Fine-Tuning Text-to-Image Diffusion Models," NeurIPS 2023.

---

confidence: 5
contribution: 2
initial_rating: 6
presentation: 2
questions: I encourage the authors to respond to my theory questions on whether there is indeed an initial value bias problem for prior distributions outside of the masked prior.
soundness: 2
strengths: The approach considered in this paper is quite logical and follows a general approach of MaxEnt RL applied to fine-tuning generative models. In some sense, there is a tight connection to existing RLHF approaches to fine-tuning in the autoregressive setting. As a result, it is encouraging to see that the rewards do end up improving with the considered objective for discrete diffusion models in the smaller-scale biological sequence design tasks considered in this paper.
summary: This paper introduces DRAKES an approach to performing fine-tuning for discrete diffusion models. In particular, the approach is predicated on maximizing the expected reward of a trajectory that is controlled with a KL constraint to prevent gaming the reward function itself. The authors suggest using the Straight-Through Gumbel Softmax estimator to backpropagate the non-differentiable loss function. Experiments are conducted on biological sequence design tasks such as inverse folding with stability. Empirical results demonstrate that DRAKES can indeed improve reward metrics in the considered settings.
weaknesses: Despite some strengths in this paper, I share several concerns regarding the method, its scalability, and evaluation protocol. I will outline this below:

**Theoretical questions**

So I believe the result in Theorem 3 is not true and needs an additional assumption. I will give an intuitive argument as to why this might be the case and I am open to being wrong here. If I am indeed wrong I would love for the authors to correct and fix my understanding. Essentially, I believe there is an initial value bias problem as outlined in "Adjoint Matching" (Domingo-Enrich 2024) and the initial starting state affects the optimal fine-tuning distribution you hit when you add the KL regularization term. While they formulated this result in the continuous setting, I believe it equally holds in the discrete setting because it relies on simple facts of the continuity (Kolmogorov forward/backward) equations. More particularly, this fact arises due to the usage of not using a memory-less noise schedule. In the discrete setting, the masked prior is likely fine but say using a uniform or another prior would make this problematic. Currently, all the theory in this paper attempts to be more general and this is where things I suspect break, while for the masked prior we automatically achieve the memory-less property.

**Technical concerns**

DRAKES as a method requires rollouts of entire trajectories and bears resembles to the method Relative Trajectory Balance (Venkataramman et. al 2024) eq 12. In fact, it is largely the same except that in DRAKES Gumbel-softmax is used instead of learning the partition function. This also means that DRAKES is not scalable because there is an expensive process of collecting full trajectories and doing stochastic backdrop through the trajectory. This induces a bias that is accumulated through the length of the trajectory. This would suggest that the current method has difficulty scaling to longer sequences and larger problem domains. This is perhaps evidenced by the fact the authors do not consider fine-tuning on more standard discrete domains like text, or pixel-level fine-tuning.

Regarding, the Gumbel-Softmax estimator I would encourage the authors to try more modern gradient estimators. One suggestion is REINMAX (Liu et. al 2023). I believe this alleviates some of the bias and is more scalable than Gumbel-softmax. I also encourage the authors to do a bit more theoretical analysis of the error incurred by using 
 as all the theory relies on not doing the Gumbel-softmax approximation so doesn't apply to the practical settings considered in the paper.

**Experimental questions**

I have some further questions and concerns regarding the considered experiments. In particular, the protein stability experiments do not indicate what length protein sequences are used. I suspect, the preprocessing of PDB largely restricts the protein sequences to small ones which would be in line with my estimation that the method does not scale to larger sequences. Moreover, it is surprising why pLDDT is not reported as a metric. This is the most common metric to show that protein sequences are designable. Self-consistency makes more sense in settings where structure is generated.

Moreover, I would like to see a few more metrics that characterize the diversity of generated samples pre and post fine-tuning with DRAKES for protein sequences. Something like the number of clusters as I suspect there is likely a sharp decline in diversity which of course would not be captured by self-consistency.

At present, DRAKES is presented as a general-purpose fine-tuning approach to discrete diffusion models but is only tested in biological sequence design. Nothing about DRAKES suggests it would not be feasible to attempt more standard discrete generative modeling settings. As such, I would like to see some standard text benchmarks for example fine-tuning for sentiment as done in Twisted SMC (Zhao et. al 2024) or class conditional fine-tuning in pixel-level image modeling. These settings would help understand both quantitatively and qualitatively, through generated samples, the empirical caliber of DRAKES.

Given the suspected similarity of DRAKES and RTB I would encourage the authors to include RTB as a baseline in there experimental settings.

**Presentation weakness**

All of the theory in the paper is rather incomplete and just plain sloppy. The appendices do not have complete proofs or the exact theorem statements. There is an air of informality that is scary because I believe some of the statements are actually not totally accurate and need a few assumptions as outlined previously.

Lemma 1 & 2 proofs are not complete. You cannot just skip the converse part. This is not a complete proof and is a bit sloppy. I encourage the authors to complete the proof.

Log Likelihood in Table 1 is actually misleading. You are likely computing an upper bound to Perplexity as done in MDLM.

**Typos and minor details**

Line 848 "competed" -> completed
Concrete distribution (Maddison et. al 2016) should be cited for Gumbel-softmax as well.


**Concluding remarks**

I am open to increasing my score if all the points in my review are sufficiently addressed, especially the experimental ones.

**References**

Liu, Liyuan, et al. "Bridging discrete and backpropagation: Straight-through and beyond." Advances in Neural Information Processing Systems 36 (2024).

Venkatraman, Siddarth, et al. "Amortizing intractable inference in diffusion models for vision, language, and control." arXiv preprint arXiv:2405.20971 (2024).

Zhao, Stephen, et al. "Probabilistic inference in language models via twisted sequential monte carlo." arXiv preprint arXiv:2404.17546 (2024).

---

confidence: 3
contribution: 3
initial_rating: 3
presentation: 2
questions: - Could the authors write Equation 3 in terms of the KL divergence terms, to better understand what is the exact KL divergence that they want to regularize with. It is not immediately clear where the terms not associated with the $\log$ come from. 
- For the equation under Stage 1, how does the probability that the authors write sum to 1? In particular, the sum over y of that equation would be greater than or equal to 1 always, unless $Q$ is negative, which will lead to other issues.
- Could the authors make sure that they have an equation number for all of their equations?
soundness: 3
strengths: - The proposed work is well motivated theoretically and tackles a challenging problem of fine-tuning discrete diffusion models to allow for realistic generation while respecting certain desirable downstream properties, defined by a reward model.
- The algorithm, DRAKES, outperforms some of the baselines considered in the work (eg. guidance and SMC-based approaches) and shows significant promise in terms of generating realistic protein and DNA sequences while optimizing for better stability / activity.
- While similar approaches have been considered in continuous space diffusion models, this work tackles the challenging problem of discrete space. In particular, it constructs theory to be able to formally do this finetuning and shows alternate ways of looking at guidance approaches already proposed for discrete space models.
summary: The authors propose a new methodology for fine-tuning discrete diffusion models in a simulation-based manner by back-propagating through the forward generative dynamics. They sidestep the problem of discrete sampling by leveraging the Gumbel based soft approximator for categorical sampling by first replacing the sampling with Gumbel-max and then performing the approximation by replacing this max operator with a softmax. The proposed algorithm, DRAKES, shows improved benefits in the domain of DNA sequences, where fine-tuning is with respect to a reward to optimize for activity of the sequences. The authors also conduct experiments and show improved performance on protein sequence design, especially inverse folding models to optimize for stable sequence generation. Overall, the authors show that their KL-regularized objective yields significant improvements when finetuned on a downstream reward, and without this KL regularization they often end up with over-optimized solutions, often lacking diversity and being unrealistic, especially under the pre-trained discrete diffusion model.
weaknesses: - I found the draft a bit hard to follow. If I understood it correctly, the preliminary section deals with generation from mask to noise as a trajectory from $t = T\ldots 0$, whereas this story is flipped from Section 3.2 onwards where the generation is now considered as a process from $t = 0\ldots T$.
- In line with the above problem, it was quite confusing to switch between $Q_{x, y}$ to $Q_{y, x}$, without really having some notion of what the $(i, j)$ entry of $Q$ represents. This might be because of my lack of understanding, but does it mean that $p(x_{t+dt} = i | x_t = j) = \sum_{k} Q_{i, k} p(x_t = k)$?
- The authors use truncated backpropagation to train their setup and the bias introduced due to that is unclear. Can the authors provide some kind of ablation into how the amount of truncation affects task performance?
- For the equation under Section 5.2, the left hand side of the equation is a function of $x$ and $t$ whereas the right side of the equation is not. Is there a typo?
- It is not quite clear whether the theory proposed in the work is very novel and insightful or a direct extension of existing theory in Uehara et. al 2024, for example. In particular, could the authors clarify that in their setup, what exactly is different and how is it not covered in prior work?

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: **Target Distribution)**

It is widely known that the target distribution of the following RL problem is proportional to $\exp(r(\cdot)/\alpha)p^{\text{pre}}(\cdot)$. However, recent works which formulates fine-tuning as stochastic optimal control state that naive KL regularization does not lead to the target distribution [1, 2]. Does this work not suffer from this issue?

**Comparison with other works)**

DNA Enhancers is a widely used benchmark for biological sequence designs. While the paper focus on discrete diffusion models, there are also works which fine-tune continuous diffusion models to generate biological sequences [3]. Is it better to fine-tune discrete diffusion models rather than continuous diffusion models for solving biological sequence problems? More analysis between those works may beneficial.

[1] Uehara, Masatoshi, et al. "Fine-tuning of continuous-time diffusion models as entropy-regularized control." *arXiv preprint arXiv:2402.15194* (2024).

[2] Domingo-Enrich, Carles, et al. "Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control." *arXiv preprint arXiv:2409.08861* (2024).

[3] Uehara, Masatoshi, et al. "Bridging Model-Based Optimization and Generative Modeling via Conservative Fine-Tuning of Diffusion Models." *Advances in Neural Information Processing Systems* 37.
soundness: 3
strengths: While most of diffusion fine-tuning methods conduct experiments on continuous spaces such as images, this paper tackles more practical and difficult setting, fine-tuning discrete diffusion models using reward function. The proposed method, gumbel-softmax trick for tackling non-differentiablility of discrete diffusion, is simple but have shown effectivenss in various experiments. Moreover, authors utilize various observations of prior methods for fine-tuning diffusion models, such as KL regularization from DPOK [1], truncated backpropagation from DRAFT [2].

[1] Fan, Ying, et al. "Reinforcement learning for fine-tuning text-to-image diffusion models." *Advances in Neural Information Processing Systems* 36.

[2] Clark, Kevin, et al. "Directly Fine-Tuning Diffusion Models on Differentiable Rewards." *The Twelfth International Conference on Learning Representations*.
summary: The paper presents a novel method for fine-tuning discrete diffusion models using reward function by formulating it as a RL problem. To mitigate the challenge of non-differentiable trajectories, authors use gumbel-softmax trick that enables direct backpropagation of rewards. The proposed method conducts experiments on DNA and protein sequences and show effectiveness for generating both natural and high-rewarding sequences.
weaknesses: **Gumbel softmax temperature)**

It seems that the temperature for the gumbel softmax, $\tau$ might be crucial for the performance. While the authors mentioned that linear schedule is deployed, it might be better to analyze the effect of temperature scheduling such as constant or loglinear.

**Position of the paper)**

The paper reminds me several similar papers… I noticed that fine-tuning discrete diffusion models is quite novel, but except for gumbel softmax trick (which is also widely used to differentiate discrete setting) other parts are really straightforward. It might be much better to propose a general framework for fine-tuning discrete diffusion models beyond biological sequences such as languages for more impactness. Otherwise, more comprehensive comparison with other models for biological sequence design can also make the paper to be unique. Nevertheless, I lean to acceptance.

Meta-review: This paper introduces a novel fine-tuning approach for discrete diffusion models, aiming to generate sequences of DNA or RNA that achieve high rewards while maintaining naturalness. Naturalness is quantified using the KL divergence from the prior distribution represented by the pretrained model. The problem is framed as a KL-constrained reinforcement learning task, where the authors leverage Gumbel-softmax techniques to enable backpropagation through non-differentiable reward functions. A key innovation lies in applying Gumbel-softmax within the non-differentiable reinforcement learning paradigm, enabling effective training of discrete diffusion models for sequence generation.

Decision: Accept (Poster)