Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design
Authors: Chenyu Wang
Masatoshi Uehara
Yichun He
Amy Wang
Avantika Lal
Tommi Jaakkola
Sergey Levine
Aviv Regev
Hanchen
Tommaso Biancalani
Venue: ICLR 2025 Poster (2025)
Keywords: Discrete Diffusion Models
Reward Optimization
Fine-Tuning
AI for science
Reinforcement learning
Primary area: generative models
TL;DR: We propose an approach for fine-tuning discrete diffusion models to generate "natural" sequences optimized for specific tasks, addressing algorithmic and theoretical challenges of discrete domains, and showing success in DNA and protein design.

Abstract: Recent studies have demonstrated the strong empirical performance of diffusion models on discrete sequences (i.e., discrete diffusion models) across domains such as natural language and 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" 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 proteins with high stability. To address this, we consider the scenario where we have pre-trained discrete diffusion models that can generate "natural" 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 pre-trained diffusion models to preserve naturalness. To solve this RL problem, we propose a novel algorithm 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" (i.e., have a high probability under a pre-trained 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.

Published as a conference paper at ICLR 2025
FINE-TUNING DISCRETE DIFFUSION MODELS VIA RE- WARD OPTIMIZATION WITH APPLICATIONS TO DNA AND PROTEIN DESIGN
Chenyu Wang" 11, Masatoshi Uchara"* '2, Yichun He 3, Amy Wang", Avantika Lal',
Tommi Jaakkola', Sergey Levine', Aviv Regev?, Hanchen Wang Tommaso Biancalanit?
MITT GGenentech 'Harvard University 4UC Berkeley Stanford University
ABSTRACT
Recent studies have demonstrated the strong empirical performance of diffusion models 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 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 mod- eling 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 dis- crete diffusion models that can generate natural-like sequences, as well as reward models that map sequences to task objectives. We then formulate the reward max- imization 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 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 founda- tion in continuous-timg Markov chains rather than Brownian motion. Finally, we demonstrate the effectiveness our algorithm generating DNA and protein sequences that optimize enhancer activity and protein stability, respectively, impor- tant tasks for gene therapies and baste therapeutics. The code available at /httpss://qitee com/ Chtee/hhnnnnnaaeaeeeeeaaaae MatpruuuWWaanoMooaaaa/a......
INTRODUCTION
Diffusion models have gained widespread recognition as effective generative models in continuous spaces, such as image and video generation (Song et al. (2020a] Ho et al. (022). Inspired by seminal works (c.g., Austin et al. (0021); Campbell et al. (2022); Sun et al. (2022)), recent studies (Lou) et al. 2023 Shi et al. 2024b Sahoo et al. (2224) 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
'Equal contribution: wangchy @ mii.edu, milaarrsmsstooshieecc. gene.comm I Work mainly done during an internship at Genentech.
1 "ooresponndig authors: wang. hanggehhhchnneeeeee hiancalt@gene.com

---

Published as a conference paper at ICLR 2025
Conditioncd Protein Backbones
Sampling
Pacla-rooa
Gumbel Softmax
Stop Gradlent
(Chng G.)AA) Σ, exp (loggy + G,1/A)
Without gradient
With gradient
M
M
M
M
M
M
M
N
Oracle
M
M
A
Reward
arg max R. R[reward] a- KT. term
(4)(
(0)
(m)())
(P7())
(mT())
(7
(T)())
KL. term
(7(0)
('
9
1
4'
Figure 1: DRAKES. We maximize the reward with a penalty term relative to pre-trained discrete diffusion models using the Gumbel-Softmad trick.
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 specific backbone structure in an inverse protein folding design problem (Dauparas) al. 2022; Campbell al. 2024)), 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., (002)); for mRNA vaccine production we desire 5' UTRs that drive high translational efficiency (Castillo-Haiir & Seelig 2021); 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 (Taskiram et al. 2024); and for natural language optimize minimize harmfulness (Touvron et al. 2023).
To address these challenges, our work introduces a fine-tuning approach for well 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 given backbone structure). To achieve this, we formulate the problem as a reward maximization task, analogous reinforcement learning (RL), where the objective function integrates both the reward terms and the KL divergence with respect to the pre- treined discrete diffusion model, which ensures that the generated sequences remain close 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-Softmal trick (Jang et al., 2016 Maddison et al., (016)).
Our main contribution is an RL-based fine-tuning algorithm, Direct Reward bAcKpropagatio with gumbEl Softmax trick (DRAKES), that enables wwward-maximizing finetuning for discrete diffusion models (Fiuure T). We derive theoretical guarantee that demonstrates its ability to generate natural and high-reward designs. and demonstrate its performance empirically
DNA and protein design tasks. While similar algorithms exist for continuous spaces (Fan et al. 2023 Black et al. 2023 Uehara al. 2024b Venkatraman al. 2024 Yuan et al. 2023] Guo et al. 2024), is the first, to the best our knowledge, to address these aspects in (continuous-timm -ime 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. 2024).
work
RELATED WORKS
Discrete diffusion models and their application in biology. Building on the seminal works of Austin et al. (021); Campbell et al. (2022). recent studies masked diffusion models (Lou et al.
2

---

Published as a conference paper at ICLR 2025
2023 Shi et al. 2024b Sahoo et al. 2024) 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 al., 2024 Campbell et et 2024). Compared to autoregressive models, diffusion models may be particularly well-suited for biological sequences, which typically yield molecules that fold into complex structures. In contrast to these works, our study focuses on fine-tuning diffusion models to optimize downstream reward funnctions One application of our approach is the fine-tuning of protein inverse folding generative models to optimize stability, as discussed in Widatalla et al. (2024). However, unlike this prior work, we employ discrete diffusion models as the generative model.
(3D)
Controlled generation in diffusion models. There are three primary approaches:
Guidance: Techniques such as classifier guidance (Song et al. 2020b Dhariwal & Nichol 2021)
and its variants (e.g.. |Banaal et al. (023); Chung et al. (2022a) Ho et al.
(2022)) 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. (024) proposed method specifically designed for discrete diffusion models. Alternative approaches directly applicable to discrete
(e.g..Banssl
diffusion models include sequential Monte Carlo (SMC)-based methods (Wu et al. 2024 Trippe
et al. 2022 Dou & Song, 2024 Cardoso et al. 20233] Phillips et al.) 2024). 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
To maximize reward functions for pretrained diffusion models, numerous recent studies have explored RL-based fine-tuning in continuous diffusion models (i.c.. diffu-
RL-based fine-tuning:
sion models for continuous objectives) (Fan al. 2023 Black et al. 2023 Clark et al. 2023
Prabhudesai et al., (2023). Our work, in contrast, focuses on discrete diffusion models.
Classifier-free fine-tuning (Ho & Salimans] 2022): This approach constructs conditional gen- erative models, applicable in our setting by conditioning on high reward values. Although not originally designed as fine-tuning method, it can also be adapted for fine-tuning (Zhang et al.) 2023 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.) (024c) 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(z[c)) and the offline dataset size consisting of triplets (e, x, r())) is limited, it is challenging to achieve success. Indeed, for this reason, most current RL-based fine-tuning papers (c.g.. Fan et al. (2023); Black et al. (2023); Clark et al. (2023) do not empirically compare their algorithms with -lassifier-free guidance.
3 PRELIMINARY
3.1 DIFFUSION MODELS ON DISCRETE SPACES
In diffusion models, our goal is to model the data distribution Pdata A(X') using the training data, where X' represents the domain. We focus on the case where X' {1, ,,. The fundamental principle (1) introducing 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 al. (2023) Sahoo et al. (2024): Shi et al. (202a))..
First, we consider the family of distributions je E N (a vector summing to 1) that evolves from t 0 to t = T according to continuous-time Markov chain (CTMC):
diz
Q(t),,, Po ~ panta,
(I)
where Q(t) R N
is the generator. Generally, je is designed so that pt approaches a simple limiting
distribution at T. A common approach is to add Mask into X' and gradually mask sequence so
that the limiting distribution becomes completely masked (Shi al. 2024a: Sahoo et al. (024).
Next, we consider the time-reversal CTMC (Sun et al. 2022) that preserves the marginal distribution. This can be expressed as follows:
de(w) jelz)
djr-t dt
Q(T - O)rr-. Qn,(() =
(2)
va (t)
3

---

Published as a conference paper at ICLR 2025
where (t) is (z. -entry of generator Q(t), representing the transition rate matrix from state c to state y. This implies that if we can learn the marginal density ratio pr. (r.yy)ppee (x)), we can sample from the data distribution at = by following the above CTMC controlled by Q(T t). Existing
works (e.g.. Lou et al. 2023)) demonstrate how to train this ratio from the training data. Especially
when we use masked diffusion models (Sahoo et al. 2024 Shi et al. (024a), we get
E u|u = Mask] (y + Mask, = Mask),
Qs.((t)
Σ Mupmurk yB|nu 2/2; = Mask] (y/ = Mask, 2t = Mask)
(3)
(It # Mask)
for certain constant y, 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 hierarchical extension of BERT (Devlin 2018).
Remark 1 (Sequence of multiple tokens). When dealing with sequences of length M, x (1. ... x,(4)), we simply consider the factorized rate matrix, i.e., Qu. : Qeg (Campbell et al. 2022). thereby avoiding exponential blowup.
Remark 2 (Conditioning). We can easily construct conditional generative model for any E C by allowing the generator to a function of e C.
3.2 GOAL: GENERATING NATURAL SAMPLES WHILE OPTIMIZING REWARD FUNCTIONS
In our work, we consider scenario with a pretrained masked discrete diffusion model (x]c) [C + *XX) trained on an extensive dataset and a downstream reward function r X + R. The pretrained diffusion model captures the naturalness validity of samples. For example, in protein design, (...e.))) could be a protein inverse-folding model that generates amino acid sequences that fold back into the given backbone structure (similar Campbell et al., (2024)), 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 (--)(-))) with desirable properties (high r(-)).
Notation. We introduce a discrete diffusion model parameterized by from / to / = =70
due de
(f)ps,
po Plim-
(4)
The parameter 9 from the pretrained model is denoted by Pre and pum denotes the initial distribution.
The distribution time T is denoted as (-). which approximates the training data distribution y pdata, We denote an element of the generated trajectory from i 0 to L T by To.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 single token for notational convenience. In this paper, all of the proofs are in in(ppendix C. We provide an extension to the case when the initial distributions are stochastic in Appendix D
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(xr) is simply maximize the expected value of the reward under the model's distribution, i.c., Ez (r(zT) )], where the expectation is taken over the distribution pa (xo.r) induced by (4) (i.e., the generator Qº). 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 learned reward or violating the physical assumptions of hand-engineered physics- bbased reward (Levine et al., 2020) Clark et al.) 2003 Uchara et al. 2024c). We address this challenge by constraining the optimized model to remain close to 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 Pd (xaer) and the pretrained diffusion model pr (xo:r) in CTMC.
Starting from Section 3.2 simply the notation, we go from t t represent noise data.
4

---

Published as a conference paper at ICLR 2025
Accordingly, our goal during fine-tuning is to solve the following reinforcement learning (RL) problem:
0 argmax B. Poo [r(TT)]
(5)
age
borm
Pu (t)
Qirs (t + (o log
Beng pa
dt
The first term is designed to generate samples with desired properties, while the second term represents the KL divergence. The parameter controls the strength of this regularization term. Finally, after fine-tuning, by using the following CTMC from t = 0 to t
Q" (4)e: Po Plim-
(6)
generate samples at time T. Interestingly, can show the following.
Theorem 1 (Fine-Tuned Distribution). When {Q" 0 E O} is fully nonparametric (i.e., realizability holds), the generated distribution at time T by (6) proportional to
exp(r())/a)p
This theorem offers valuable insights. The first term, exp(r(x)), represents high rewards. Addition- ally, the second term, (.te()), 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 fold back into the target backbone structure.
Remark 3. A similar theorem has been derived for continuous diffusion models (eeaaaa al. 2024b Theorem 1). However, our formulation (5) differs significantly as our framework is based on a CTMC, whereas those works are centered around the Brownian motion. Furthermore, while the use a similar distribution is common in the literature on (autoregressive) large language models (e.g.. Ziegler et al. (2019)), its application discrete diffusson models is novel, considering that pore cannot be explicitly obtained in our context, unlike autoregressive models.
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 (Algo- rithm and Figure 1), which is designed solve the RL problem (5). The core approach involves iteratively (a) sampling from To.r' ~ pd and (b) updating 0 by approximating the objective func- tion with its empirical counterpart and adding its gradient with respect A into the current A. Importantly, for step (b) to be valid, step (a) must retain the gradients from 0. After explaining the representation of st, will provide details on each step.
Representation. To represent representation within R interchangeably. From this perspective, while the original generator corresponds to a map X' x X' + R, we can also regard it as an extended mapping: RM N x IN R. We will use this extended mapping when we consider our algorithm later.
(1,...,,............ N., we often use the N N-dimensional one-hot encoding
Stage 1: Data collection (Step 29
We aim to sample from the distribution induced by the generator Q". the standard discretization of CTMC, for (y, x) X X, time t, we use
p(Z++A1 = y/2, = 1(x y) (t)tt)..
(8)
Thus, by defining n. = (t))),, ... (1 (t) (())), we sample F+ Cat(m)), where Cat(-) denotes the categorical distribution.
However, this procedure is not differentiable with respect to 6, which limits its applicability for optimization. To address this, first recognize that sampling from the categorical distribution can be reduced to a Gumbel-max operation. Although this operation itself remains non- con-differentiable, we can modify it by replacing the operation with a softmax, shown in Line 7, which is also utilized in discrete VAE (Jang et al. 2016). This modification results in a new variable, Ta ~ [o, 1]M, which maintains differentiability with respect 0. As the temperature approaches zero, F. converges to sample from the exact categorical distribution Cat(T,), effectively becoming I;. Thus, we typically set the temperature to a low value to closely approximate the true distribution.
5

---

Published as a conference paper at ICLR 2025
Algorithm DRAKES (Direct Reward bAcKpropagation with gumbEl Softmax trick)
1: Require: Pretrained diffusion models Q
IR, reward r: X - R, learning rate B. Batch size B, Iteration S, Time-step Al, Temperature Regularization parameter
RR R
2: for E [1.... .5] do
3: for E [1, ,1,..., ,B] do
4: for t E [0, At. do
[ +A 1 (t ((in(y) (t=
5:
Set (=) ( (()) where n(t),
6:
Sample & E , G4 Gumbel(0, 1) Set differentiable counterpart of the sample time
7:
exp(( + G1)/7)
(p((T(t))
Σ, exp(((t)) ee)((()))+))))))))))
8:
end for end for
9:
10:
Set the loss:
Q (t)
( B
(
(t: 2r (t) (%2((t) log
T
Ore (t)
1m
TE
VEX grftar
11:
Update a parameter: 0s41 +- 0, + Veg(())
end for
13: Output: 05+1
After approximately sampling from the distribution induced by p we update the parameter 6g by maximizing the empirical objective. Although z, itself may not have a valid gradient, V, retains the gradient with respect to 0. Therefore, use the empirical approximation based on Tt. We offer several remarks below, with details in Appendix
Stage 2: Optimization (Step 10
Validity of Zy: While Hn does strictly belong to X, this is practically acceptable since the generator Q^ is parameterized as map R.N R + R.
SGD Variants: Although Line 11 uses the standard SGD update, any off-the-shelf SGD algorithm, such as Adam (Kingma 2014), can be applied in practice.
Soft Calculation with ãe: Transition probability (t), and the KL divergence term in g(0) are modified to their soft counterparts by using ã, in place of zt.
Straight-Throut Gumbel Softmax: Non-relaxed computations be used in the forward pass (in Line 10 This is commonly known as straight-through Gumbel softmax estimator.
practice, often more effective to backpropagate from intermediate time steps rather than starting from 0. In practice, adopted this truncation approach, as in Clark et al. (202)).
Truncated Backpropagration:
Optimization Objective g0): For the masked diffusion models (3) that we utilized, g(0) can be further simplified reduce computational complexity, detailed in Appendix E.2
THEORY OF DRAKES
In this section, we provide an overview of the proof for Theorem Based on the insights gained from this proof, we reinterpret state-of-the-ar classifier guidance for discrete diffusion models (Nisonoff al. 2024) from new perspective.
5.1 PROOF SKETCH OF THEOREM
We define the optimal value function V, X IR follows:
Qu
B
r(xr)= - Σ
(oinp(s)
Quers + Qu. (s) log
ds Je = x
(9)
Que

---

Published as a conference paper at ICLR 2025
This is the expected return starting from state at time t following the optimal policy. Once the optimal value function is defined, the optimal generator can be expressed in terms of this value function using the -accoi Bellman (HJB) equation in CTMC, shown below.
Theorem (Optimal generator). For x # y (2. E X), the optimal CTMC generator (Tyy(t) can be expressed by the pretrained CTMC generator Quy (t) and the value function
Q (t) = Quyy (t) exp((Vi()) V((x))a)..
(10)
Next, consider 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-Kas Formula in CTMC). The value function follows the Feynman- -ac formulation with respect to the pretrained CTMC distribution pr generated by Quere as follows
exp(Vi(=)/aa) = Paprr (exp(r(=r)/a))x==== =
(11)
With this preparation, we can prove our main theorem, which reduces to Theorem when = T.
Theorem 4 (Marginal distribution induced by the optimal generator Q" (t)). The marginal distribu- tion at time by the CTMC trajectory in (6) with generator Q" (t), i.e, Pf E A(X), satisfies
pi (-) & exp(Ve))/aapp"">),
(12)
where pt E ((X) is a marginal distribution induced by pretrained CTMC at
This is proved by showing the Kolmogorov forward equation in CTMC: dpt /dt = Q (t)pt-
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. for reward maximization. If we can directly obtain the optimal generator Q" we can achieve our objective. Theorem 2 suggests that the optimal generator Q is a product of the generator from the pretrained model and the value functions, Although we don't 't know the exact value functions, they can be learned through regression using Theorem based on
exp(Vi())a) argmin E p (xrl (eep((((()))(((( g(x.)) g:X-R
(13)
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 2; to 2g in pretrained models following DPS (Chung et al. 2022b) or reconstruction guidance (Ho et al. 2022).
Interestingly, similar algorithm was previously proposed by Nisonoff al. (2024). While Nisonoff et al.) (2024) 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) p(=xz) (e.g.. the log-likelihood from classifier) and set a 1. By adapting Theorem and their setting, obtain:
Q4 (t) = Prre (t) x (ell))//.(()))) p.(a(ax) := E PP (p(2 aT) mr).
(14)
Thus, we can rederive the formula Nisonoff et al. (2024). Here, we also note that this type of result is commonly referred to the Doob transform in the literature stochastic processes (Levin & Peres 2017 Section 17).
While this argument suggests that classifier guidance and RL-based fine-tuning approaches theo- retically achieve the same goal in an ideal setting (without function approximation, sampling, or optimization errors), their practical behavior 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 the need to recalculate the generator during inference. Indeed, this classifier guidance requires O MM 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.
differ significantly,
demonstrate in /ecction 6
7

---

Published as a conference paper at ICLR 2025
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. We also experiment in text generation for minimizing toxicity, detailed in Appendix F.4
6.1 BASELINES
We compare DRAKES against several baseline methods discussed in Section 2 which we summarize below with further details in Appenddxx F.
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., 2024), SMC-based guidance methods (e.g.. Wu et al. (2024)):
SMC, where the proposal is a pretrained model and TDS, where the proposal is CG.
Classifier-free Guidance (CFG) (Ho Salimans, 2022). CFG is trained on labeled datasets with the measured attributes we aim to optimize.
Pretrained. We generated sequences using pretrained models without fine-tuning.
DRA BRAKES 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. 2024).
Dataset and settings. We experiment on a publicly available large-scale enhancer dataset Gossa et al. 2023), which measures the enhancer activity of 7000 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.) 2024)
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 (Avsee et al.) (021) 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 al.) 2024 [eehara et al., [024c] Sarkar al.
2024). Detailed information about the pretrained model and reward oracles is in Appendix F.2
Evaluations. To comprehensively evaluate each model's 's performance 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 models are fine-tuned (or guided) with the oracle trained on different subset of the data, splitting based on chromosome following conventions (but in the same cell lines) (Lal et al. 2024). chromatin accessibility (ATAC-Acc). We use an independent binary classification model trained on chromatin accessibility data in the HepG2 cell line (Consortium ct al. 2012) (active enhancers should have accessible chromatin). While this is not used for fine- tuning, we use for evaluation to further validate the predicted activity of the synthetic sequences, following Lal et al. (2024).
Binary classification
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.) (023) with top 0.1% HepG2 activity level. Models that generate sequences that are more natural-like like and in-distribution have higher correlation.
JASPAR motif analysis (JASPAR Corr). We scan the generated sequences of each model with JASPAR transcription factor binding profiles Castro-Mondragoon
al. (2022), 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.
Approximated log- likelihood of sequences (App-Log-Lik), We calculate the log- likeeloodooo of the generated sequences with respect the pretrained model to measure how likely the sequences
8

---

Published as a conference paper at ICLR 2025
are to be natural-like Models that over-optimize the reward oracle generate out-of-distribttion 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. (2024).
Results. DRAKES generates sequences with high predicted activity in the HepG2 cell line, as robustly measured by Pred-Activity and ATAC-Acc (Table [T. 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 tine-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, Ablations on gradient truncation and Gumbel Softmax schedule are in Appendix F.2
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-Ace, 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.
Method
Pred-Activity (median) 1 ATAC-Ace t (%)
3-mer Corr
JASPAR Corr
App-Log-Lik (median) 1
Pretrained
0.17(0.04)
1.5(0.2)
0.249(0.015) -0.061(0.034)
-261(0.6)
CG
3.30(0.00)
0,0(0.0)
-0.065(0.001)
0.212(0.035)
-266(0.6)
SMC
4.15(0.33)
39.9(8.7)
0.840(0.045)
0.756(0.068)
225(2.55)
TDS
4.64(0.21)
45.3(16.4)
0.848(0.008)
0.846(0.044)
.257(1.5)
CFG
5.04(0.06)
92.1(0.9)
0,746(0.001)
0.864(0.01))
-265(0.6)
DRAKES w/o KI.
6.44(0.04)
82.5(2.8)
0,307(0.001)
0.557(0.015)
(281(0.6)
DRAKES
5.61(0.07)
92.5(0.6)
0.887(0.002)
0.911(0.002)
-264(0.6)
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 optimize the stability of these generated sequences. following Widatalla et al. (2024).
Dataset and settings. First, we pretrained an inverse folding model based on the diffusion model (Campbell et al. 2024) and the ProteinMPNN (Dauparas et al., 2022) architecture, using the PDB training set from Dauparas et al. (202)). Next, we trained the reward oracles using a different large-scale protein stability dataset, Megascale (Tsuboyama et al. 2023), which includes stability measurements (i.e., the Gibbs free energy change) for ~.1.MM sequence variants from 983 natural and designed domains. Following dataset curation and a train-validation-tes splitting procedure from Wiidtalla et al.] (2024) (which leads to ~0055M sequences on 333 domains) and using the ProteinMPNN architecture,
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 F3
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 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 vtraet+ 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 al. 2023) 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. 2024 Trippe et al. 2022
Chu et al. 2024).
Following prior works (Campbell et al., 2024; Nisonoff et al. 2024). we calculate the success rate of inverse folding as the ratio of generated sequences with Pred-ddG>> and scRMSD< 2.
9

---

Published as a conference paper at ICLR 2025
Table 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 random seeds, with standard deviations in parentheses.
Method
Pred-ddG (mediaan)
%(ddG> 0) (
scRMSD (median)
StscRMSD 2X(%)
Success Rate (8)
Pretrained
0.544(0.037)
36.6(1.0)
0.849(0.013)
90.9(0.6)
34,4(0.5)
CG
(0.56100.045)
36.9((1))
0.839(0.012)
90.9(0.6)
34.7(0.9)
SMC
0.659(0.044 044)
68.5(3.1)
0.841(0.006)
93.8(0.4)
63.6(4.0)
TDS
0.674(0.086)
68.2(2.4)
0.83410.000)
94,4(1.2)
62.9(2.8)
CFG
-1. 18660000)3
11.000.4)
3.146(0.062)
29.4(1.0)
1.3(0.4)
DRAKES w/o KL
L.00((00004)
100.0(0.0)
7.307(0.054)
34.1(0.2)
34.1(0.2)
DRAKES
1. 1.09((0.026)
86.4(0.2)
0.918(0.006)
91.8(0.5)
78.6(0.7)
7JJK
2KRU
006
DRARIS
IMRAKIS
7756
TDS
SMC
E.8-
SMAC
Density
WEl Syp
Donsity
002-
0.0-
-155
-00
Energy (A.)
Energy (AG)
(a) Conditioning on backbone structure of 7JJK.
(b) Conditioning on the backbone structure 2KRU.
Figure 2: Examples of generated proteins. Red: Wild-type backbone structure (the one we condition on). Yellow: Structure predicted by ESMFold from the wild-type (truc) sequence, Green: Structure predicted by ESMFold from the sequence generated by DRAKES. The structures for sequences generated by DRAKES show good alignment with the original structure (the scRMSDs are 0.768 for 7JJK and 0.492 for 2KRU). Histograms: Gibbs free energy for each generated sequence, calculated using physics-based simulations. In these two cases, the sequences generated by DRAKES appear to be more stable than the baselines.
Results. For inverse protein folding, DRAKES generates high-stability protein sequences capable of folding into the conditioned structure (Table [ achieves the highest Pred-ddG among all methods, while maintaining similar success rate of inverse folding (measured by %(scRMSD 2), the percentage 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 few hundred backbones, making generalization difficult. This is expected,
mention in Section 2 Further details are provided in Apppendix F.T We provide additional results with pLDDT and sequence diversity, as well as ablation studies the gradient truncation number and Gumbel Softmax temperature schedule in Appendix F.3
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, calculate the stability (i.c., Gibbs free energy) of
the generated sequences using physics-based simulations (PyRosetta (Chaudhury et al. 2010) for wild-type protein backbone structures in Figure 2] following (Widatalla et al. (024)). A Although all models are conditioned on the same set of protein backbones, different sets of sequences generated by generative methods can lead significant differences in side chain interactions, which affect folding energies. The results demonstrate that sequences generated by DRAKES are more stable in this validation compared to other baselines. For additional results, refer to Figure 6 Ampppddix F.3
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 demon- strate its effectiveness in generating DNA and protein sequences optimized for task-specific objectives. For future work, plan conduct more extensive in silico validation and pursue wet-lab validation.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
For the theoretical results presented the paper, we provide explanations of assumptions and complete proofs in Appendix C For the proposed algorithm and experimental results, we provide detailed explanations of the algorithm implementations and experimental setup in [ecciinn 4.2) Section 6 Appendix E, and Appendix and the code and data can be found in https::/// github com/ Beenyyuaang-onniek DRRKES For the datasets used in the experiments, we utilize publicly available datasets and elaborate the data processing procedures Section 6 and Appendix
ACKNOWLEDGEMENT
CW and TJ acknowledge support from the Machine Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium, the DTRA Discovery of Medical Countermeasures Against New and Emerging (DOMANE) threats program, and the NSF Expeditions grant (award 1918839) Under- standing the World Through Code.

---

Published as a conference paper at ICLR 2025
REFERENCES
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:17999 19911-1997 2021.
Pavel Avdeyev, Chenlai Shi, Yuhao Tan, Kseniia Dudnyk, and Jian Zhou. Dirichlet diffusion score model for biological sequence generation. arXiv preprini 2rXiv:2305 2023.
Ziga Avsec, Vikram Agarwal, Daniel Visentin, Joseph Ledsam, Agnieszka Grabska-Band Barriiska,, Kyle R Taylor, Yannis Assacl, John Jumper, Pushmeet Kohli, and David R Kelley. Effective gene expression prediction from sequence by integrating long-range interactions. Nature methods, 18 (10): 196-1203, 2021.
Arpit Bansal, Hong-Min Chu, Avi Schwarzschild, Soumyadip Sengupta, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Universal guidance for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 843-852 2023.
Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305. 13301, 2023.
Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Amaud Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266-2822 -2277, 2022.
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications protein co-design. arXiv preprint arXiv:2402.04999 2024.
Gabriel Cardoso, Yazid Janati El Idrissi, Sylvain Le Corff, and Eric Moulines. Monte carlo guided diffusion for bayesian linear inverse problems, arXiv preprint arXiv:2308.070900 07983, 2023.
Sebastian M Castillo-Hair and Georg Seelig. Machine learning for designing next-generation mma therapeutics. Accounts of Chemical Research, 55(1):24-34 2021.
Jaime A Montro-Monnr,on,, Rafael Riudavets-Puig. Ieva Rauluseviciute, Roza Berhanu Lemma, Laura Turchi, Romain Blanc-Mathieu, Jeremy Lucas, Paul Boddie, Aziz Khan, Nicolás Manos- alva Pérez, et al. Jaspar 2022: the 9th release of the open-access database of transcription factor binding profiles. Nucleic acids research, 50(D1):D155-D17,, 2022.
Sidhartha Chaudhury, Sergey Lyskov, and Jeffrey J Gray. Pyrosetta: script-based interface for implementing molecular modeling algorithms using rosetta. Bioinformatics, 26(5):689-66) 2010.
Alexander Chu, Jinho Kim, Lucy Cheng, Gina EI Nesr, Minkai Xu, Richard W Shuai, and Po-Ssu Huang. An all-atom protein generative model. Proceedings of the National Academy of Sciences, 121(27):e231 5501221, 2024.
Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Mare L Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209. 14687, 2022a.
Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. Advances in Neural Information Processing Systems, 35:25683-25696 2022b.
Kevin Clark, Paul Vicol, Kevin Swersky, and David Fleet. Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309. /7400, 2023.
ENCODE Project Consortium et al. An integrated encyclopedia of dna elements in the human genome. Nature, 489(7414):57,2) 2012.
Nicholas Kluge Corrêa. Aira, 2023. URL
Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob de Haas, Neville Bethel, et al. Robust deep saaming-basee protein sequence design using proteinmpnn. Science, 578(6615):49 2022.
12

---

Published as a conference paper at ICLR 2025
Jacob Devlin. Bert: Pre- Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.048 2018.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-877 88-44, 2021.
Carles Domingo-Enrich Michal Drozdzal, Brian Karrer, and Ricky TQ Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. arXiv preprint arXiv:2409. 000800, 2024.
Zehao Dou and Yang Song. Diffusion posterior sampling for linear inverse problem solving: A filtering perspective. In The Twelfth International Conference on Learning Representations, 2024.
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Picter Abbeel, Mohammad Ghavamzadch, Kangwook Lee, and Kimin Lec. DPOK: Reinforcement learning for fine-tuning text-to-image diffusion models. arXiv preprint arXiv:2305. 16381, 2023.
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick. and Stefanie Tellex. Openwebtext corpus. http: //Skylion007 htt.iiinnn7..ittt.ba git.ub..hbb...... htt.oppeeeeeebbeeeooroup... 2019.
Sager Gosai, Rodrigo Castro, Natalia Fuentes, John C Butts, Susan Kales, Ramil R Noche, Kousuke Mouri, Pardis C Sabeti, Steven K Reilly, and Ryan Tewhey. Machine-guided design of synthetic cell type-specific cis cis-regulatory elements. bioRxiv, 2023.
Yingqing Guo, Hui Yuan, Yukang Yang, Minshuo Chen, and Mengdi Wang. Gradient guidance for diffusion models: An optimization perspective. arXiv preprint 2Xiv:2404.14743,, 2024.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598. 2022.
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. Advances in Neural Information Processing Systems, 35:8633-8646, 2022.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv:1611.01144 2016.
Diederik Kingma. Adam: A method for stochastic optimization. arXiv preprint arXhv:1412.6988 2014.
Avantika Lal, David Garfield, Tommaso Biancalani, and Gokcen Eraslan. reglm: Designing realistic regulatory dna with autoregressive language models. bioRxiv, pp. 2024-02, 2024.
David A Levin and Yuval Peres. Markov chains and mixing times, volume 107. American Mathemat- ical Soc., 2017.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 2020.
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, et al. Evolutionary-seale prediction of atomic-level 2023.
protein structure with a language model. Science, 379(663)):
Aaron Lou, Chenlin Meng. and Stefano Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310.76834 2023.
Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv:1611.00711 2016.
Hunter Nisonoff, Junhao Xiong, Stephan Allenspach, and Jennifer Listgarten. Unlocking guidance for discrete state-space diffusion and flow models. arXiv preprint arXiv:2400.001572. 2024.
Angus Phillips, Hai-Dang Dau, Michael John Hutchinson, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. Particle denoising diffusion sampler. arXiv preprint 063302 2024.
13

---

Published as a conference paper at ICLR 2025
Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, and Katerina Fragkiadaki. Aligning text-to-image to-image diffusion models with reward backpropagation. arXiv preprint arXiv:2310.037.9, 2023.
Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. arXiv preprint arXiv:2406.07524 2024.
Anirban Sarkar, Ziqi Tang, Chris Zhao, and Peter Koo. Designing dna with tunable regulatory activity using discrete diffusion. bioRxiv, pp. 2024-05, 2024.
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis K Titsias. Simplified and generalized masked diffusion for discrete data. arXiv preprint arXiv:2406.0433 2024a.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. Advances in Neural Information Processing Systems, 36, 2024b.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.025000 2020a.
Yang Song. Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 1..4466 2020b.
Hannes Stark, Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, and Tommi Jaakkola. Dirichlet flow matching with applications dna sequence design. arXiv preprint arXiv:2402.05881, 2024.
Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous-time discrete diffusion models. arXiv preprint arXiv:2211. 2022.
Ibrahim Taskiran, Katina Spanier, Hannah Dickmänken, Niklas Kempynck, Alexandra Pančíková, Eren Can Ekşi, Gert Hulselmans, Joy N Ismail, Koen Theunis, Roel Vandepoel, al. Cell-type- directed design of synthetic enhancers. Nature, 626(7997) 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprini arXiv:2307.09228. 2023.
Brian Trippe, Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, and Tommi Jaakkola. Diffusion probabilistic modeling of protein backbones in 3d for the motif- scaffolding problem. arXiv preprint arXiv:2206.00 00122006011111 19, 2022.
Kotaro Tsuboyama, Justas Dauparas, Jonathan Chen, Elodie Laine, Yasser Mohseni Behbahani, Jonathan Weinstein, Niall M Mangan, Sergey Ovchinnikov, and Gabriel J Rocklin. Mega-scale experimental analysis of protein folding stability in biology and design. Nature, 620(7973): 434-444, 2023.
Masatoshi Uehara, Yulai Zhao, Tommaso Biancalani, and Sergey Levine. Understanding rein- forcement learning- basedd fine-tuning of diffusion models: A tutorial and review. arXiv preprint MxXi::2407.13334 2024a.
Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali. Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Fine-tuning of continuous-time diffusion models rntropy-regularrized control. arXiv preprint arXiv:2402. 15:400..55599 15194, 2024b.
Masatoshi Uehara, Yulai Zhao, Ehsan Hajiramezanali, Gabriele Scalia, Gökcen Eraslan, Avantika Lal, Sergey Levine, and Tommaso Biancalani. Bridging model-based optimization and generative modeling via conservative fine-tuning of diffusion models. arXiv preprint arXiv:2405.1967. 19673. 2024c.
Siddarth Venkatraman, Moksh Jain, Luca Scimeca, Minsu Kim, Marcin Sendera, Mohsin Hasan, Luke Rowe, Sarthak Mittal, Pablo Lemos, Emmanuel Bengio, et al. Amortizing intractable inference in diffusion models for vision, language, and control. arXiv preprint arXiv:2405. 20971, 2024.
14

---

Published as a conference paper at ICLR 2025
Talal Widatalla, Rafael Rafailov, and Brian Hie. Aligning protein generative models with experimental fitness via direct preference optimization. bioRxiv, pp. 2024-02224000 05, 2024.
Luhuan Wu, Brian Trippe, Christian Naesseth, David Blei, and John P Cunningham. Practical and asymptotically exact conditional sampling in diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
George G Yin and Qing Zhang. Continuous-tine Markov chains and applications: singular perturbation approach, volume 37. Springer, 2012.
Hui Yuan, Kaixuan Huang. Chengzhuo Ni, Minshuo Chen, and Mengdi Wang. Reward-directed conditional diffusion: Provable distribution estimation and reward improvement. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Lvmin Zhang, Anyi Rao, and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 3836-3847, 2023.
Stephen Zhao, Rob Brekelmans, Alireza Makhzani, and Roger Grosse. Probabilistic inference in language models via twisted sequential monte carlo. arXiv preprint vrXiv:2404.17544 2024.
Mingyuan Zhou, Tianqi Chen, Zhendong Wang. and Huangjie Zheng. Beta diffusion. Advances in Neural Information Processing Systems, 36, 2024.
Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv: 08500, 2019.
15

---

Published as a conference paper at ICLR 2025
A MORE RELATED WORKS
Dirichlet diffusion models for discrete spaces. Another approach to diffusion models for discrete spaces has been proposed (Stark et al., 2024 Avdeyev al. 2023; Zhou et al.) 2024). In these models, each intermediate state is represented
vector within a simplex. This is in contrast to masked diffusion models, where each state a discrete variable.
Relative trajectory balance for posterior inference. Venkatraman et al. (2024) proposed to relative trajectory balance to train a diffusion model that sample from a posterior x pmm (z) oc p(x)r(z)) When rxx) is defined the exponential reward, it can be utilized for reward optimization. However, this approach requires estimation the normalizing constant term of the unnormalized density. In contrast, solve control problem with direct backpropagation. Besides, Venkatraman et al. (2024) models continuous diffusion as a Markovian generative process and is not specifically tailored discrete diffusion models.
B POTENTIAL LIMITATIONS
We have formulated the RL problem. (5) in the context of CTMC. The proposed algorithm in our paper 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 enformer borozi. In cases where creating diffecrentiable models is challenging, we recommend using PPO-based algorithms (Black et al. 2023) Fan et al. (003) or reward- weighted Section 6.1).
MIEE (Uehara al. 2024a
C PROOF OF THEOREMS
C.I PREPARATION
We prepare two well-known theorems in CTMC for the pedagogic purpose. For example, refer to Yin & Zhang (002) for the more detailed proof. In these theorems, suppose have the CTMC:
dp: = Q(t)p- dt
(15)
Lemma 1 (Kolmogorov backward equation). We consider g( t) = [R(((r))a:+ = -] where the expectation taken w.e.. (15). Then, this function g : X [0,T] - R is characterized by the
following ODE:
dg(a, t)
- g(x,T) r((T).
dt
Here, we prove that the p.d.f. g satisfies above backward equation. To show the converse, technically require regularity conditions to claim the ODE solution is unique, which often be proved by the contraction mapping theorem under certain regularity conditions. Since this is a well-known result, skip the converse part. If readers interested in details, refer Yin & Zhang 2012).
Proof.
When T, the statement 9αx,T) r(xr) is obvious. For the rest of the proof, aim show result when t + T. We have
g(F,, t) = / g(retde, + dt dt) lo dd
The above implies
g(z,t) = {1+ (1 dt) + (t)t(()) dt).
16

---

Published as a conference paper at ICLR 2025
Now combined with the property of the generator follows
Σ Qu.y (t + dt),
=
With some algebra,
g(z, t) g0, + dt) - Σ (t) +
+ + Ei (dddtt + dt).
Then, have
g(x, t) (a, t dt)
Σ (een((oolu.tttd) +t) gla, dt))
dl
v
Finally, by setting di obtain
dg(x,t)
-
dt
Lemma (Kolmogorov forward equation). The density pt AXX) is characterized the following ODE:
dpe(x) dt
(v.((()))
(eu((t))(()), yZz
Pu Pini-
Proof. Here, we prove that the p.d.f. Pt 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
Petat (x) = f (etta((//ps)ppp
This implies
(x) Σ (1)=tt + (1 Qu, (t)d(t)))
Σ (Q (1)dpp (1- Σ (mn(()))))
Hence,
Pitdi (x) (x))
Σ (m.a(())(())) Σ (np(((p(z))
dl
By taking dl + O, we obtain
dpr(x)
Σ (tep(()) - (t))
dt
Then, the proof is completed.
C.2 PROOF OF THEOREM 2
derive the Haamiloon-Jcobi-Bellln (HJB) equation CTMC. For this purpose, we consider the
recursive equation:
{
(xx,t) t) max
Eet.. Qu
Σ
Quo (t) Q" (t) log
(t. (t) dt
Quee (t)
+ Σ (d) + (0
17

---

Published as a conference paper at ICLR 2025
Using EvEX (t,(()) O, this is equal to
on dt
V(a,t) max
-t) 2z. (t) log
Opre Q (8)
Mattrr
Σ
(( -x
By taking dt to O, the above is equal to
22 (t)
dV(z.))
(t
Qu. (t) log
+ De
(V((())) -V v(z,t)
dt
max dEe
pro Q 22
(16)
This is the HJB equation in CTMC.
Finally, with simple algebra (i.e., taking functional derivative under the constraint 0 Σyex Voy (t)), we can show
Ve + ys Q2,y (t) = Queg (t)))((((())))) (xx,
C.3 PROOF OF THEOREM 3
This theorem is proved by invoking the Kolmogorov backward equation.
First, by plugging
Va # ys Q2ny(t
(())((()))V (xx,t)))
into 16). get
dV (xxt)
Σ
By multiplying (vp(V(x, t)) to both sides, reduces to
dexp(V(<,, ())
Σ
(17)
dt
Furthermore, clearly, Vzx, 7) = r(rT). Then, the statement is proved by invoking the Kolmogorov backward equation.
C.4 PROOF THEOREM 4
We define
H.(z) = exp(V(z,t)))
We aim to prove that the above satisfies the Kolmogorov forward equation:
dH((x)
Ee
Eeh (1),,,),
Pini Ho(-).
dl
v
1hs.
rh.s.
Remark 4. Here, note that Piui 8(- Mask) Ho(-) holds because have considered scenario where the initial is the Dirac delta distribution. When the original initial distribution is stochastic, will see in Appendix D that still have Pini Ho(.).
18

---

Published as a conference paper at ICLR 2025
First, calculate the Lhss. Here, recall
d dexp(V ((x, 1))
Σ ((xxp((()) exp(((x,,))
dt
using 17 and
dpc(x)
Σ (t)er(()rr()))
(t)pr(y) Σ Deg (t)pe(2) (t)p(t))((z))
di
holds, using the Kolmogorov forward equation. Then, we obtain
dH)(())
(xxx((((x1)) pc(x) + exp(V(xx, dpr(a))
dt
C
dt
dt
x
Σ Br
exppl
+ c x exp(V(ax
Eoe (Up((1)
use
Σ (x)p(((()e
1
Eo'g exp( (v ()))
C
C
On the other hand, the r.h.s. is
1
Σ (1) EoT(()))))
C
Σ
f (t) exp( (() π Σ Saa (t) exp((V V(a, (z)()))))+((
x
C
Σ
Qar (t)
Σ
Puy (t))xxp(((
C
Here, from the first line to the second line, we use
Yx + y: 2% (t) = Q27 exp( V(x,t))..
Finally, can see that l.h.s. r.h.s. Furthermore, recalling we have an assumption that Pini Dirac delta distribution, we clearly have pini Ho(.). Hence, the statement is proved by the Kolmogorov forward equation.
D EXTENSION WITH STOCHASTIC INITIAL DISTRIBUTIONS
When initial distributions are stochastic, we need to modify algorithms. Although there are several strategies take stochastic initial distributions into account (Uchara al. 2024b; Domingo-Enrich
et al. 2024), we consider the strategy of learning the initial distributions (Uchara et al. 2024b).
Hence, we consider the following control problem:
0" argmax E P P8 (r(2T)) )
(18)
Reward
Qgrn Q2 (1) Q2. (1) log (
E Eyog P8 P0
di
KL phim)
Quers (t)
v
KL term
18 we parameterize not only generators but also initial distributions. In this case, Theorem still holds.
Compared to the control problem (5), in our new formulation
19

---

Published as a conference paper at ICLR 2025
D.1 PROOF
Recalling the definition of value functions:
Bo r(zr) - a [ Σ
T
Sley Pui (t) + Qu (t) log
(e (t)
(kk,t) t)
dt
t= yy
Q (t)
the optimal initial distribution needs to satisfy
argmax BegmaxEao-aaa -o pli)).
This is equal to the distribution proportional to
The rest of the proof holds in the proof of Section C.4 without any change, referring Remark
DETAILS OF ALGORITHM
STRAIGHT- THROUGH GUMBEL SOFTMAX
We apply the straight-through Gumbel softmax estimator to the last time step, i.e.
ST(2f) (0) TT SG(ET
(5)
where (1) is the corresponding Gumbel-mat variable, i.c. 27 argmax (1) as and SG denotes stop gradient. Then, (T (0) and backward propagation.
(6)
is input into the reward function r(.) instead of (1)) for forward
We observe boost in fine-tuning performance with the straight-through Gumbel softmax, as convert- ing the input r(-) into one-hot vector makes it better aligned with the reward oracle's training distribution.
E.2 SIMPLIFIED FORMULA OF g())
The key objective function in DRAKES, g(0), can be further simplified for the masked diffusion models that we utilized the experiments.
La
(
r
Σ (2)(()) + Qry (t) +2 (t)
(t) log
T
Open (t)
yEX
We denote the second term estimating the KL divergence with the i-th sample as (f):
1 T
Qay (t)
A: (0))
Σ
Q2,y (feg)))++)+()))) (1) + Qu (t) log Que (t)
T
1=1 VEN
DEX
When Mask, the value of Q (t) is irrelevant to the parametrization 0, i.e.
10.0 y * Mask
22.g (1)
Qur (1)
7 Mask
where is constant related the forward process schedule (Sahoo et al. 2024). In particular, when applying a linear schedule (as in our experiments), = 1/t. Thus, the corresponding KL divergence component equals O.
20

---

Published as a conference paper at ICLR 2025
When * Mask,
Qa = [0, # Mask
(Boolr 2/21-1 = Mask], Mask
Denote Es (Tr
2 [2 = Mask] as [angl.. The KL divergence term 1:() (0) can be simplified as
T
Q (t)
1 (0) = T
1a
Σ
+2) + (t) + Qzy (t) log
Qz,y (t)
[=1 mEX
BEX 85
T
(1) + Sove Musk (1) + Qu Ra,MMak (1) log Qg Mask (t)
2-
Qu. (t)
DEX
[+0]
-100l
+ +100 [=00000= 1 + [=0] log
Mask
The simplified formula reduces the computational complexity of calculating k (0) O(NT).
E.3 SCHEDULE OF GUMBEL SOFTMAX TEMPERATURE
We linear schedule for the Gumbel softmax temperature т, decreasing over time 1/ In early time steps, the temperature is higher, introducing more uncertainty, while later steps have 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 reduced.
F EXPERIMENTAL DETAILS AND ADDITIONAL RESULTS
detailed explanation of the baseline. We then diseuss additional results and present more detailed settings for both regulatory DNA sequence design and protein sequence design. Finally, we present results of DRAKES in the natural language domain for loxicity-controlled text generation. We use NVIDIA A100 80GB GPU for all experiments.
In this section, first provide
F.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., 2024). CG adjusts the transition rate of CTMC by calculating the predictor guidance:
p(rly,t)
Quay)r(t)
p(rlæ, t) (t.y(t)
where r is the target property, and the predictor guidance is further approximated using Taylor expansion, i.e.
p(rly.t) t) p(r
log
(y - x)" V log p()lz,t)
The predictor p(r|z, f) is estimated using the posterior mean approach (Chung et al, 2022a), where the pretrained model is first utilized to estimate the clean data from the noisy input ay and then the reward oracle is applied to the predicted clean sequence. We remark that the above Taylor approximation doesn't it have formal theoretical guarantees, considering that is discrete, This could be a reason why it does not work well in the case protein- inverse folding in Section 6.3
21

---

Published as a conference paper at ICLR 2025
SMC (Wu et al. pretrained model as the proposal distribution. While it was originally designed for condition- ing rather than reward maximization, it can be adapted for reward maximization by treating rewards as classifiers. In our experiment, we use this adapted version.
2024). SMC is sequential Monte Carlo-based approach that uses the
TDS (Wu et al. 2024). Similar to SMC, TDS also applies sequential Monte Carlo, but utilizes CG rather than the pretrained model as the proposal.
Classiffer-free Guidance (CFG) (Ho & Salimans| 2022). Unlike guidance-based methods, CFG trains a conditional generative model from scratch and does rely the pretrained model. To generate sequences x with desired properties r(x), CFG incorporates r(æ) as an additional input to the diffusion model and generates samples conditioning on high r(z) 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(z).
is important to note that CFG requires labeled data pairs {ar, r(r)) for training, which can limit its performance in cases with limited labeled data, especially when the pretrained model already a conditional diffusion model p(z|e). For example, in the protein inverse folding task, where æ is the protein sequence, e is the protein structure, and r(x:) is the protein stability, CFG struggles, as shown in T This is due to the small size of the Megascale dataset (containing only a few hundred different protein structures), which reduces its capability and generalizabilttyf While data augmentation can be applied to construct additional training data, it is resource-intensivve, requires significant case-by-care bycase 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 such issue.
F.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.
We train reward oracles to predict activity levels of enhancers in the HepG2 cell line using the dataset from Gosai et al. (2023). Following standard practice (Lal et al.) (024). 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., (02T) architecture initialized with its pretrained weights. One oracle is used for fine-tuning, while the other is reserved for evaluation (i.c. Pred-Activity in Table Denote the subset used for training the fine-tuning oracle 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 high Pearson correlation (> 0.85) on their respective held-out sets (Eval for the fine-tuning oracle and FT for the evaluation oracle).
Reward Oracle.
Table 3: Performance of the reward oracles for predicting HepG2 activity of enhancer sequences.
Model
Eval Dataset
MSE J
Pearson Corr 1
Fine-Tuning Oracle
FT 0.149
0.938
Eval
0.360
0.869
Evaluation Oracle
FT
0.332
0.852
Eval
0.161
0.941
Pretrained Model. We pretrain the masked discrete diffusion model (Sahoo et al. 2024) on the full dataset of Gosai et al. (2023), using the same CNN architecture as in Stark et al. (2024) and a
linear noise schedule. Other hyperparameters are kept identical to those in Sahoo al. (0224). 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 prreentts 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.
In contrast, other methods (guidance-hused methods and fine-tuning methods) leverage the pretrained model trained the much larger PDB dataset (~ 23, 0000 structures) and achieve better performance.
22

---

Published as a conference paper at ICLR 2025
Distribution of HepG2 Activity for Generated and Original Sequenoes
Comparison of Distributions
Gencnaal((T)
077
Geacnued (Eval)
Original (FT)
Original (Pval)
Origitall (Trac)
HepG2 Activity
Density
0.5
0.4
0.3
03-
611
0.9
Generated (PT)
Geserated (Eiva)
Original (PT)
(Eval)
Original (Eval)
Original (Trae)
Hep Activity
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.
Comparison of 3-mer Counts
Comparison of 4-mer Counts
k-mer Count (Generated)
k-mer Count (Generated)
Pearson, 0.957
Carson Corr:
k-mer Count (Original)
k-mer Count (Original)
Figure 4: 3-mer and 4-mer Pearson correlation between the original and generated 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 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.
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 data collection, sequences generated from the model over 128 steps. We set 0.001 to govern the strength of the KL step 50. The base Gumbel Softmax temperature
regularization and truncate the backpropagation is set to 1.0. The model is fine-tuned with 128 samples as 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 to zero. For evaluation, we generate 640 sequences per method (with a 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 Table Figure 5) presents the full distribution of Pred-Activity for each method, which shows consistent patterns as Table
Ablation Study on Gradient Truncation Number. To show the impact of the gradient truncation module, we conduct an ablation study on the gradient truncation number. As shown Table 4 when the truncation number is small, the model cannot effectively update the early stage of the sampling process, leading to suboptimal performance. Meanwhile, when the truncation number is large, the memory cost of the model fine-tuning increases, and the optimization is harder due to the gradient accumulation through the long trajectory. Therefore, an intermediate level of gradient truncation leads to the best performance.
23

---

Published as a conference paper at ICLR 2025
Distribution of Predicted HcpG2 Activity for Generated Sequenees
HenG2 Activity
2-
Pretrained
CG
SMC
TDS Methods
CFG
DRAKES w/ KL DRAKES
Figure 5: Distribution of Pred-Activity for the generated sequences of each method.
Table 4: Ablation study on gradient truncation number on regulatory DNA sequence design.
Truncation Number Pred-Activity (median)
ATAC-Ace1 (%)
3-mer Corr
JASPAR Corr
App-Log-1 Lik (median) f
6.17
95.9
0.569
0.798
-278
50
5.61
92.5
0.887
0.911
-264
75
4.61
52.8
0.238
0.644
-268
Ablation Study Gumbel Softmax Temperature Schedule. We conduct ablation study with different temperature schedules of Gumbel Softmax, i.c. a linear schedule and a constant schedule. As shown in Table 5 both schedules achieve similar performance, indicating the robustness of our method. In the performance reported in Table we utilize the linear schedule.
Table 5: Ablation study on Gumbel Softmax temperature schedule on regulatory DNA sequence design.
Temporature Schedule
Pred-Activity (median)
ATAC-Ace f (%)
3-mer Corr
JASPAR Corr
App-Log Log-Lik---- (median) f
Constant
5,89
91.6
0.852 0.887
0.914 0.911
-258
Linear
5.61
92.5
-264
F.3 PROTEIN INVERSE FOLDING
We first discuss the setup of datasets used for training reward models. Next, present the perfor- mance 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 (T'ubooyama et al. 2023) for the protein inverse folding experiment, which contains stability measurements for 1.8M sequence variants (for example, single mutants and double mutants) from 983 protein domains. We follow the dataset curation and crain-validation-tes splitting procedure from Widatalla al. (024)) 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 AG 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 (2024). We further exclude proteins where significant proportion of the corresponding variants' AG measurements fall outside the experimental range. The final dataset consists of 438,540 sequence variants from 311 proteins the training set, 18, 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. 2024) and the ProteinMPNN (Dauparas al. (022)) architecture encode both sequence and structure as model input. The model is trained on the PDB training set used Dauparas al. 2022), containing 23,349 protein structures and their ground truth sequences, which is distinct
24

---

Published as a conference paper at ICLR 2025
from the dataset in Tuuboyama et al. (023). We first evaluate the effectiveness of the inverse folding model on the PDB test set in Dauparas et al. (2022), which has 1,539 different proteins. As in Nisonoff et al. (2024), we set the temperature during sampling 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 the pretrained model performs similarly to ProteinMPNN, achieving comparable sequence recovery rate.
Table 6: Model performance of protein inverse folding on PDB test set.
Method
Sequence Recovery Rate (
ProteinMPNN
47.9
Discrete Flow Model
48.6
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 7 We calculate the self-consistency RMSD (seRMSD) assess how well a generated sequence folds into the desired structure. Specifically, the generated sequences are folded into 3D structures using ESMFold (Lin jt al., 2023). and scRMSD) calculated as their RMSD relative to the original backbone structure we are conditioning on. An scRMSD lower than 2 typically considered a successful inverse folding (Nisonoff et al. 2024: Campbell et al. (024). As shown in Table 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 the inverse folding task.
Table 7: Model performance of protein inverse folding on Megascale proteins.
Eval Dataset
Sequence Recovery Rate (%)
scRMSD (Å) 4
8(scRMSD< 2)(%) 1
Megascale-Train
47.0
0.825
95.0
Megascale-Test
44.0
0.849
90.9
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 A AAG (calculated as the difference in AG between the variant and the wild-type from the dataset). Similar to Nisonoff et al. (2024), the final layer of the ProteinMPNN architecture is poon-- and mapped single scalar with 2-layer MLP, and the model weights before the mean-post 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 in Table 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 hyper- parameters, we randomly split the full dataset into two subsets, an in-distribution set for training, denoted and an out-of-distribtiion 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 & degascale-Test-I) and the new structures. The Pearson correlation each subset is presented in Table 8 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 we train it on the full dataset using the best hyperparameters selected as discussed. It achieves Pearson correlation of 0.951 Megascale training set and 0.959 on Megascale test set (both being in-distribution for the evaluation oracle).
set contains unseen sequence variants, but no
Finetuning Setup. We utilize the pretrained inverse folding model and the fine-tuning oracle described above for fine-tuning, During DRAKES's stage I data collection, we generate sequences from the model over 50 steps. We set 0.0003 and truncate the backpropagation step 25. The
25

---

Published as a conference paper at ICLR 2025
Table Performance the reward oracles for predicting stability conditioned on protein sequence and structure, across variety of Megascale subsets.
Model
Eval Dataset
Pearson Corr
Fine- Fine-Tuning Oracle
Megascale-Trail
0.828
Megascale- Test
0.685
Megascale- Train-I
0.948
Evaluation Oracle
Megascale- Train-OO
0,942
Megascale-
0,955
Megascale-Teste 0.920
base Gumbel Softmax temperature set to 0.5. 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 a to zero. The model is evaluated on Megascale test set, where we generate 128 sequences conditioned on each protein structure for every method (with a 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.
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 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 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.
Evaluation Oracle Accounts for OOverOOtimizante
Table 9: Model performance protein inverse folding, with Pred-ddG calculated using either the evaluation oracle (Eval) or the fine-tuning oracle (FT).
Method
Pred-ddG- Eval (median)
Pred-ddG-FT (median) 1
Pretrained
-0.544(0.037)
0.161(0.012)
CG
-0.56(
0.158(0.017)
SMC
0.659(0.044))
0.543(0.013)
TDS
0.674(0.086)
0.557(0.005)
CFG
-1.159(0.035)
-1.243(0.013)
DRAKES w/o KL
1.108(0.004)
0.833(0.000)
DRAKES
1.095(0.026)
0.702(0.002)
Ablation Study on Gradient Truncation Number. To show the impact of the gradient truncation module, we conduct an ablation study on the gradient truncation number. The results are shown in Table While the model performance is robust with different truncation numbers, an intermediate level of gradient truncation leads to the best performance.
Table 10: Ablation study on gradient truncation number on inverse protein folding.
Truncation Number
Pred ddG (median) t
Siddi> 0) (S)
ScRMSD (median).
SisscMSDD< 2
Success Rate (%)
15
0.977
83.2
0,852
92.4
76.0
25
1.095
86.4
0.918
91.8
78.6
35
1.033
84.8
0.868
92.7
77.7
Alation Study on Gumbel Softmax Temperature Schedule. We conduct an ablation study with different temperature schedules of Gumbel Softmax, i.e. linear schedule and constant schedule.
26

---

Published as a conference paper at ICLR 2025
As shown in Table I1 both schedules achieve similar performance, indicating the robustness of our method. In the performance reported in Table we utilize the linear schedule.
Table 11: Ablation study on Gumbel Softmax temperature schedule on inverse protein folding.
Temperature Schedule
Prod-déG (median)
&dddG> (%)
scRMSD (median)
&(ccMMSD 2x%)
Success Rate (%)
Constant
1,178
86.5
0.897
93.2
80.3
Lincar
L,095
86.4
0.918
91.8
78.6
Results with pLDDT. In addition to ScRMSD, we further measure the self-consistency of the generated sequences using pLDDT. The results are shown in Table [2] Following the common practice as in Widatalla et al. (2024), we utilize 80 as the cutoff threshold and define a success generation having positive predicted stability (i.e., Pred-ddG>-0) and confident ESMFFold-prediceed structure (i.e.. pLDDT >> 80). The results align well with those in Table using scRMSD. DRAKES significantly outperforms all baseline methods in terms the overall success rate.
Table 12: Model performance on inverse protein folding. DRAKES generates protein sequences
that have both high stability and high confidence in ESMFoll predicted structures, outperforming baselines in the overall success rate. We report the mean across random seeds, with standard deviations in parentheses.
Method
Pred-ddG (median)
SiddG>> 0) (%)
PLDDT (median)
*(pLDDT>> 80)(S))
Success Rate (S)
Pretrained
0.544(0.033)
36.6(1.0)
87.9(0.0)
87.5(0.1)
33.8(0.8)
-0.66100.045)
36.9(1.1)
87.9(0.0)
87.5(0.2)
34.3(0.5)
SMC
0.639(0.044)
68.5(3.1)
88.3(0.1)
84.1(1.0)
53.7(4.3)
TDS
0.674(0.086)
68.2(2.4)
88.3(0.1)
85.6(0.7)
54.4(2.8)
CFG
-1.186(0.035)
11.0(0.4)
65.7(0.0)
7.8(0.1)
0.0(0.0)
DRAKES w/o KL
1.108(0.004)
100.0(0.0)
73.4(0.2)
40.4(0.2)
40.4(0.2)
DRAKES
1.095(0.026)
86.4(0.2)
87.0(0.0)
85.6(0.7)
72.3(0.8)
Diversity of Generated Sequences. We evaluate the diversity of the protein sequences generated by different methods using the average sequence entropy of each protein backbone in the test set. The results shown in Table 13 DRAKES achieves comparable level of diversity the pretrained model, significantly outperforming SMC and TDS. This further justifies the efficacy of DRAKES to achieve the optimization objective of generating stable sequences, while maintaining high diversity. Among all baselines, DRAKES is the only method that have both a high success rate and high diversity.
Table 13: Average sequence entropy of the generated sequences on protein inverse folding.
Method
Sequence Entropy
Pretrained
34.7(0.2)
CG
34.6(0.1)
SMC
24.9(1.2)
TDS
24.9(0.5)
CFG
8.4(0.1)
DRAKES w/o
25.7(0.1)
DRAKES
33.3(0.2)
Additional Results. We provide more examples of the generated proteins in Figure 6 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.
27

---

Published as a conference paper at ICLR 2025
7JJK
0.08
DRAKES
Wild type Energy: -113.4
For the visualized protein:
Protnaraal
scRMSD: 0.768
Density
Wild
Pred-ddG; 1.950
Energy: -131.0
0.02
Across proteins genernted
by DRAKES:
Energy- median: -121.2
-150
-124
-100
-74
-50
-55
Energy (AC)
2KRU
Wild type Energy: -132.0
0.10
TDS
0.08
SMC
Pocraincd
the visualized protein:
scRMSD: 0.492
Donsity
Wild-opp
Pred-dilG: 2.224
0.06
Energy: -142.0
0.04
Across proteins genernted by DRAKES: Energy-median: -133.3
0.02
0.00
Energy (AG)
IFOM
Wild type Energy- -159.5
DRAKES
0.10
For the visualized protein:
Donsity
0.08
Wild-ypp
Pormunnd
scRMSD: 0,771
Pred-ddG: 1.318
Energy: -190.6
0.04
Acress proteins generated
by DRAKES:
Ebergy-median: -191.5
-003
-120 Energy (AC)
160
21.09
Wild type Energy: -133.1
DRAKES
0
T7S
SMC
the visualized protein:
Donsity
Wild-ype
Pred-diG: 1.625
0.06
Energy: -146.4
0.04
Across proteins generated by DRAKES:
Energy-median: -143.7
0.00
-129
Energy (AC)
2M2JJ
Wild type Energy: -176.8
0.7
DRAKES
SMU
For the visualized protein:
Density
0.50
Widdoye
Pestniaed
Pred-ddG: 2.096
ScRMSD: 0.676
Energy: -188.2
0.00
Across proteins generated
DRAKES:
001
Enargy-median:
-187.6
200
190
100
Energy (AG)
Figure 6: Additional examples of generated proteins.
28

---

Published as a conference paper at ICLR 2025
F.4 MOXICITY-COTTOLLIE TEXT GENERATION
To demonstrate the general applicability of DRAKES beyond its use in the biological domain, we conduct experiments in toxicity-controlledd text generation. The corresponding details and results presented below.
Reward Oracle. We follow the setting Zhao al. (024)) using toxicity prediction model trained by Corrra (2023) as the reward oracle, which is a fine-tuned RoBERTa model with 125M parameters that can be used to score the toxicity sentence.
Pretrained Model. We utilize the pretrained masked diffusion language model from Sahoo al. (2024) with 130M non-embedding parameters, which is trained on the Open WebText cor- pus Gokaslan et al. 2019) for one million steps and involves the processing of 33 billion tokens.
Fine-tuning Setup. Following (Zhao al. 2024). we generate 20 output tokens. During DRAKES's stage 1 data collection, sequences are generated from the model over 20 diffusion steps. We set 0.0001 and truncate the backpropagation at step The base Gumbel Softmax temperature is set to 1.0. The model is fine-tuned with batch size 128 (32 samples per iteration, with gradient accumulated over iterations) for 100 steps. For DRAKES w/o KL, we follow the same setup, but set to zero. For evaluation, generation 320 sequences per method (with batch size of 32 over 10 batches).
Results. As shown in Table 14 DRAKES generates sequences with the highest reward (i.c., low toxicity) compared to all bascline models, measured by the mean and median Pred-Score. This demonstrates the effectiveness of DRAKES in the natural language domain.
Table 14: Model performance toxicity controlled text generation.
Method
Pred-Score (mean)
Pred-Score (median) f
Pretrained
6.996
8.003
CG
8.488
9.009
SMC
9.286
9.367
TDS
9.817
9.882
DRAKES
10.044
10.049
29


Reviews:
[{"review_id": "gB2hwvcNOj", "reviewer": "Reviewer_LsP3", "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.", "questions": "See weaknesses above.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. Using Gumbel-softmax in non-differentiable reinforcement learning in training discrete diffusion model is a novel approach.", "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.\n\n\n2. 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. \n\n\n3. 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?\n\n[1] Venkatraman et al. \"Amortizing Intractable Inference in Diffusion Models for Vision, Language, and Control,\" NeurIPS 2024\n\n[2] Fan et al. \"Reinforcement Learning for Fine-Tuning Text-to-Image Diffusion Models,\" NeurIPS 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cQBkNpH8gj", "reviewer": "Reviewer_1d7y", "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.", "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.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 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.", "weaknesses": "Despite some strengths in this paper, I share several concerns regarding the method, its scalability, and evaluation protocol. I will outline this below:\n\n**Theoretical questions**\n\nSo 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.\n\n**Technical concerns**\n\nDRAKES 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.\n\nRegarding, 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 \n as all the theory relies on not doing the Gumbel-softmax approximation so doesn't apply to the practical settings considered in the paper.\n\n**Experimental questions**\n\nI 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.\n\nMoreover, 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.\n\nAt 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.\n\nGiven the suspected similarity of DRAKES and RTB I would encourage the authors to include RTB as a baseline in there experimental settings.\n\n**Presentation weakness**\n\nAll 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.\n\nLemma 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.\n\nLog Likelihood in Table 1 is actually misleading. You are likely computing an upper bound to Perplexity as done in MDLM.\n\n**Typos and minor details**\n\nLine 848 \"competed\" -> completed\nConcrete distribution (Maddison et. al 2016) should be cited for Gumbel-softmax as well.\n\n\n**Concluding remarks**\n\nI am open to increasing my score if all the points in my review are sufficiently addressed, especially the experimental ones.\n\n**References**\n\nLiu, Liyuan, et al. \"Bridging discrete and backpropagation: Straight-through and beyond.\" Advances in Neural Information Processing Systems 36 (2024).\n\nVenkatraman, Siddarth, et al. \"Amortizing intractable inference in diffusion models for vision, language, and control.\" arXiv preprint arXiv:2405.20971 (2024).\n\nZhao, Stephen, et al. \"Probabilistic inference in language models via twisted sequential monte carlo.\" arXiv preprint arXiv:2404.17546 (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GAEoDkFxZy", "reviewer": "Reviewer_X1rg", "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.", "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. \n- 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.\n- Could the authors make sure that they have an equation number for all of their equations?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 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.\n- 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.\n- 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.", "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$.\n- 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)$?\n- 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?\n- 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?\n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DGxwTfVUoI", "reviewer": "Reviewer_aA7C", "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.", "questions": "**Target Distribution)**\n\nIt 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?\n\n**Comparison with other works)**\n\nDNA 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.\n\n[1] Uehara, Masatoshi, et al. \"Fine-tuning of continuous-time diffusion models as entropy-regularized control.\" *arXiv preprint arXiv:2402.15194* (2024).\n\n[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).\n\n[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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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].\n\n[1] Fan, Ying, et al. \"Reinforcement learning for fine-tuning text-to-image diffusion models.\" *Advances in Neural Information Processing Systems* 36.\n\n[2] Clark, Kevin, et al. \"Directly Fine-Tuning Diffusion Models on Differentiable Rewards.\" *The Twelfth International Conference on Learning Representations*.", "weaknesses": "**Gumbel softmax temperature)**\n\nIt 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.\n\n**Position of the paper)**\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)