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
Primary area: generative models
Keywords: Discrete Diffusion Models, Reward Optimization, Fine-Tuning, AI for science, Reinforcement learning

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.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729777283845

thread_type:
official_review

<thread_object>
creator: reviewer_aa7c

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

strengths:
While most of diffusion fine-tuning methods conduct experiments on continuous spaces such as images, this paper tackles more practical and difficult setting, fine-tuning discrete diffusion models using reward function. The proposed method, gumbel-softmax trick for tackling non-differentiablility of discrete diffusion, is simple but have shown effectivenss in various experiments. Moreover, authors utilize various observations of prior methods for fine-tuning diffusion models, such as KL regularization from DPOK [1], truncated backpropagation from DRAFT [2].

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

[2] Clark, Kevin, et al. "Directly Fine-Tuning Diffusion Models on Differentiable Rewards." *The Twelfth International Conference on Learning Representations*.

weaknesses:
**Gumbel softmax temperature)**

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

**Position of the paper)**

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

questions:
**Target Distribution)**

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

**Comparison with other works)**

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

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

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

[3] Uehara, Masatoshi, et al. "Bridging Model-Based Optimization and Generative Modeling via Conservative Fine-Tuning of Diffusion Models." *Advances in Neural Information Processing Systems* 37.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730500453748

thread_type:
official_review

<thread_object>
creator: reviewer_x1rg

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

strengths:
- The proposed work is well motivated theoretically and tackles a challenging problem of fine-tuning discrete diffusion models to allow for realistic generation while respecting certain desirable downstream properties, defined by a reward model.
- The algorithm, DRAKES, outperforms some of the baselines considered in the work (eg. guidance and SMC-based approaches) and shows significant promise in terms of generating realistic protein and DNA sequences while optimizing for better stability / activity.
- While similar approaches have been considered in continuous space diffusion models, this work tackles the challenging problem of discrete space. In particular, it constructs theory to be able to formally do this finetuning and shows alternate ways of looking at guidance approaches already proposed for discrete space models.

weaknesses:
- I found the draft a bit hard to follow. If I understood it correctly, the preliminary section deals with generation from mask to noise as a trajectory from $t = T\ldots 0$, whereas this story is flipped from Section 3.2 onwards where the generation is now considered as a process from $t = 0\ldots T$.
- In line with the above problem, it was quite confusing to switch between $Q_{x, y}$ to $Q_{y, x}$, without really having some notion of what the $(i, j)$ entry of $Q$ represents. This might be because of my lack of understanding, but does it mean that $p(x_{t+dt} = i | x_t = j) = \sum_{k} Q_{i, k} p(x_t = k)$?
- The authors use truncated backpropagation to train their setup and the bias introduced due to that is unclear. Can the authors provide some kind of ablation into how the amount of truncation affects task performance?
- For the equation under Section 5.2, the left hand side of the equation is a function of $x$ and $t$ whereas the right side of the equation is not. Is there a typo?
- It is not quite clear whether the theory proposed in the work is very novel and insightful or a direct extension of existing theory in Uehara et. al 2024, for example. In particular, could the authors clarify that in their setup, what exactly is different and how is it not covered in prior work?

questions:
- Could the authors write Equation 3 in terms of the KL divergence terms, to better understand what is the exact KL divergence that they want to regularize with. It is not immediately clear where the terms not associated with the $\log$ come from. 
- For the equation under Stage 1, how does the probability that the authors write sum to 1? In particular, the sum over y of that equation would be greater than or equal to 1 always, unless $Q$ is negative, which will lead to other issues.
- Could the authors make sure that they have an equation number for all of their equations?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730636150777

thread_type:
official_review

<thread_object>
creator: reviewer_1d7y

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

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:

**Theoretical questions**

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

**Technical concerns**

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

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

**Experimental questions**

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

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

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

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

**Presentation weakness**

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

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

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

**Typos and minor details**

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


**Concluding remarks**

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

**References**

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

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

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

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730755057296

thread_type:
official_review

<thread_object>
creator: reviewer_lsp3

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

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.


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


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

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

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

questions:
See weaknesses above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732078266126

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Consolidated response to all reviewers and the AC

comment:
We thank all the reviewers for their constructive feedback that helped us improve our paper. Based on these comments, we have conducted additional experiments and revised our manuscript, with the changes highlighted in magenta. Below we provide a brief summary of the key revisions and clarifications. Please refer to the responses to individual reviewers for more details.

**Revisions**
* We included additional metrics (pLDDT and sequence entropy) to showcase the superior performance of our method in the protein inverse folding experiment. In addition, we conducted ablation studies on the gradient truncation and Gumbel Softmax temperature schedule.
* We improved the notations and added clarifications to our problem setting and formulas.
* We included discussions on additional related work RTB[a].

**Clarifications**
* **Position of this paper (reviewer 1d7y and aA7C):** The paper is motivated by and focuses on the practical applications of biological sequence design problems, essential in real-world gene therapies and protein-based therapeutics, as emphasized in the abstract and introduction sections. We want to emphasize that biological sequences are the unique domains where masked diffusion models shine since they are really discrete (unlike image) and non-autoregressive (unlike text) by nature.
* **The initial value bias problem (reviewer 1d7y and aA7C):** We use masked diffusion models for this paper and the initial distribution is a Dirac delta distribution (completely masked state) following existing works. In this case, there is no initial value bias problem. We have updated the paper to clarify this further. Additionally, our method can be expanded to settings where the initial distribution is not a Diract delta and the initial value bias problem arises.
* **Relationship with RTB[a] (reviewer LsP3 and 1d7y):** [a] uses losses from trajectory balance, while we solve a control problem with direct backpropagation. Hence, the actual algorithms are significantly different. Also, theoretically, our result has more insights tailored to discrete diffusion models.

[a] Amortizing Intractable Inference in Diffusion Models for Vision, Language, and Control. NeurIPS 2024.
</content>
</thread_object>

</thread>


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