PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Chenyu Wang, Masatoshi Uehara, Yichun He, Amy Wang, Tommaso Biancalani, Avantika Lal, Tommi Jaakkola, Sergey Levine, Hanchen Wang, Aviv Regev

arXiv:2410.13643v2cs.LGcs.AI

TL;DR

Design tasks need sequences that remain natural-like while optimizing specific objectives, a requirement not met by modeling conditional distributions alone. DRAKES fine-tunes pretrained discrete diffusion models with reward optimization and a KL penalty, using Gumbel-Softmax for trajectory backpropagation. The paper reports theoretical support and DNA and protein design results.

  • Problem

    Discrete diffusion models can capture conditional sequence distributions, but design tasks also require natural-like sequences optimized for downstream objectives.

    Source: §1 · p. 2

  • Method

    DRAKES fine-tunes pretrained discrete diffusion models by maximizing reward with a KL penalty and backpropagating through trajectories using Gumbel-Softmax.

    Source: §1 · p. 2

  • Results

    DRAKES generates DNA enhancer sequences with high predicted activity and protein sequences with high predicted stability while maintaining naturalness-related measures or folding success.

    Source: §6.2 · p. 9, §6.3 · p. 10

  • Takeaways & Limitations

    The paper demonstrates reward-optimized DNA enhancer and protein stability design with discrete diffusion models.

    Source: §7 · p. 10

    The setup assumes a completely masked initial state; stochastic initial distributions are addressed in an appendix extension.

    Source: §3.2 · p. 4

Abstract

from arXiv · show

Recent studies have demonstrated the strong empirical performance of diffusion models on discrete sequences across domains from natural language to biological sequence generation. For example, in the protein inverse folding task, conditional diffusion models have achieved impressive results in generating natural-like sequences that fold back into the original structure. However, practical design tasks often require not only modeling a conditional distribution but also optimizing specific task objectives. For instance, we may prefer protein sequences with high stability. To address this, we consider the scenario where we have pre-trained discrete diffusion models that can generate natural-like sequences, as well as reward models that map sequences to task objectives. We then formulate the reward maximization problem within discrete diffusion models, analogous to reinforcement learning (RL), while minimizing the KL divergence against pretrained diffusion models to preserve naturalness. To solve this RL problem, we propose a novel algorithm, DRAKES, that enables direct backpropagation of rewards through entire trajectories generated by diffusion models, by making the originally non-differentiable trajectories differentiable using the Gumbel-Softmax trick. Our theoretical analysis indicates that our approach can generate sequences that are both natural-like 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 DRAKES in generating DNA and protein sequences that optimize enhancer activity and protein stability, respectively, important tasks for gene therapies and protein-based therapeutics.

1 INTRODUCTION

Discrete diffusion models can capture conditional distributions, but design tasks also require sequences that optimize downstream objectives while remaining natural-like. DRAKES fine-tunes pretrained models with a reward-and-KL objective and differentiable trajectory backpropagation.

1 INTRODUCTION

  • Discrete diffusion models capture conditional sequence distributions, but practical design also seeks natural-like sequences optimized for objectives such as protein stability or cell-specific gene expression.The introduction also names high translational efficiency and minimizing language-model harmfulness as objectives.Source: §1 · p. 2
  • DRAKES combines downstream reward maximization with a KL penalty against the pretrained model to preserve naturalness.It uses Gumbel-Softmax to make trajectories differentiable and backpropagate rewards through them.Source: §1 · p. 2
  • DRAKES addresses discrete diffusion models’ continuous-time Markov chain trajectories, which are non-differentiable, and provides a theoretical guarantee for natural, high-reward designs.The paper reports empirical demonstrations on DNA and protein design tasks.Source: §1 · p. 2

2 RELATED WORKS

Related work covers discrete diffusion for language and biological sequences, approaches to controlled generation, and reward-oriented fine-tuning. This paper focuses on reward optimization for discrete diffusion models.

2 RELATED WORKS

  • Discrete diffusion models have been applied to natural language and biological sequence generation, while this work focuses on optimizing downstream rewards.Prior biological applications include DNA and protein sequences; the paper applies its approach to protein inverse folding with discrete diffusion rather than the prior work’s generative model.Source: §2 · p. 3
  • Controlled-generation alternatives include inference-time guidance and sequential Monte Carlo, which generally take longer to infer than fine-tuned models.The passage also describes classifier-free fine-tuning and reward-based fine-tuning in continuous diffusion models.Source: §2 · p. 3
  • Prior RL-based fine-tuning studies primarily address continuous diffusion models, whereas this work focuses on discrete diffusion.The cited discussion says most current RL-based fine-tuning papers do not empirically compare their algorithms with classifier-free guidance.Source: §2 · p. 3

3 PRELIMINARY

Discrete diffusion models use a forward continuous-time Markov chain and a learned time reversal to generate samples. The paper considers fine-tuning a pretrained model to balance sample naturalness and downstream reward.

3.1 DIFFUSION MODELS ON DISCRETE SPACES

  • The forward discrete diffusion process evolves a sequence distribution through a continuous-time Markov chain, commonly toward a simple limiting distribution such as a fully masked sequence.Its generator is a time-dependent matrix, and masking is implemented by adding a Mask state and gradually masking the sequence.Source: §3.1 · p. 3
  • A time-reversal CTMC preserves the marginal distribution and can sample from the data distribution when controlled by a learned marginal density ratio.For masked diffusion models, the text also describes a rate formulation based on the forward process.Source: §3.1 · pp. 3–4

3.2 GOAL: GENERATING NATURAL SAMPLES WHILE OPTIMIZING REWARD FUNCTIONS

  • The fine-tuning setup pairs a pretrained masked discrete diffusion model, which captures sample naturalness or validity, with a downstream reward function.The stated objective is to generate samples with both high pretrained-model likelihood and high reward.Source: §3.2 · p. 4
  • The setup assumes a completely masked initial state, while stochastic initial distributions are treated in an appendix extension.The paper denotes the pretrained parameters by θpre and the initial distribution by plim.Source: §3.2 · p. 4

4 ALGORITHM

The algorithm optimizes downstream reward while regularizing the fine-tuned model toward a pretrained diffusion model. DRAKES makes trajectory sampling differentiable so reward gradients can update the model.

4.1 KEY FORMULATION

  • Reward maximization alone can produce unrealistic samples, so the objective penalizes KL divergence from the pretrained model, with α setting regularization strength.Reward models may be exploited by samples outside their training distribution or by violations of physics-based reward assumptions.Source: §4.1 · pp. 4–5
  • Under full nonparametric realizability, the fine-tuned terminal distribution is proportional to the pretrained distribution weighted by exp(r(x)).The reward term favors high-reward sequences, while the pretrained distribution supplies prior information about natural sequences.Source: §4.1 · p. 5

4.2 DIRECT REWARD BACKPROPAGATION WITH GUMBEL SOFTMAX TRICK (DRAKES)

  • DRAKES iteratively samples trajectories and updates model parameters using an empirical objective whose soft samples retain gradients.The update approximates the reward-and-regularization objective with its empirical counterpart.Source: §4.2 · pp. 5–6
  • Replacing categorical sampling with a softmax yields differentiable samples that approach exact categorical samples as temperature τ approaches zero.The method typically uses a low temperature to approximate the true distribution closely.Source: §4.2 · p. 5
  • The implementation can use non-relaxed forward computations and truncate backpropagation at intermediate time steps.For the masked diffusion models used, the objective can also be simplified to reduce computational complexity.Source: §4.2 · p. 6

5 THEORY OF DRAKES

The theory section sketches a CTMC-based analysis of DRAKES and reinterprets discrete classifier guidance through the optimal generator. It also identifies practical differences between guidance and RL-based fine-tuning.

5.1 PROOF SKETCH OF THEOREM 1

  • The proof sketch defines the optimal value as expected return and uses the CTMC Hamilton-Jacobi-Bellman equation to express the optimal generator.The value function also follows a Feynman–Kac formulation under the pretrained CTMC, which is useful for learning it.Source: §5.1 · pp. 6–7
  • The proof establishes the marginal distribution induced by the optimal generator by checking the CTMC Kolmogorov forward equation.The argument is presented as the route to the main theorem, which reduces to Theorem 1 at t = T.Source: §5.1 · p. 7

5.2 RELATION TO CLASSIFIER GUIDANCE FOR DISCRETE DIFFUSION MODELS

  • A fine-tuning-free alternative would use the optimal generator, a product of the pretrained generator and value functions learned through regression.The exact expectation can be replaced by an empirical analogue or approximated using a pretrained-model map from x_t to x_0.Source: §5.2 · p. 7
  • Adapting Theorems 2 and 3 recovers Nisonoff et al. (2024)'s conditional-generation approach for reward maximization; this result is also called a Doob transform.Their conditional-generation setting uses r(x) = log p(z|x) and α = 1.Source: §5.2 · p. 7
  • Despite ideal-setting theoretical equivalence, classifier guidance can require O(NM) value-function computations per step and can be challenging to train in practice.Its inference cost may be significantly longer; Taylor approximation can mitigate this, but the heuristic has no theoretical guarantee for discrete diffusion models.Source: §5.2 · p. 7

6 EXPERIMENTS

The experiments evaluate DRAKES against guidance, classifier-free guidance, pretrained models, and a no-KL ablation on enhancer activity and protein stability. Across these tasks, the reported results pair task performance with measures of sequence naturalness or structural consistency.

6.1 BASELINES

  • The baselines include classifier guidance, SMC and TDS guidance, classifier-free guidance, pretrained sampling, and DRAKES without KL regularization.The no-KL ablation tests whether KL regularization mitigates over-optimization.Source: §6.1 · p. 8

6.2 REGULATORY DNA SEQUENCE DESIGN

  • The enhancer study uses a roughly 700,000-sequence, 200-bp dataset and evaluates predicted activity, ATAC accessibility, sequence correlations, motif correlations, and pretrained-model likelihood.Fine-tuning and evaluation activity oracles are trained on different chromosome-based subsets.Source: §6.2 · pp. 8–9
  • For HepG2 enhancer design, DRAKES produces high predicted activity and ATAC accessibility alongside naturalness indicators, while its higher Pred-Activity without KL is attributed to over-optimization.The naturalness indicators are 3-mer and JASPAR motif correlations and likelihood; analogous regions across chromosomes may keep held-out sequences similar.Source: §6.2 · p. 9, Table 1 · p. 9

6.3 PROTEIN SEQUENCE DESIGN: OPTIMIZING STABILITY IN INVERSE FOLDING MODEL

  • DRAKES achieves the highest Pred-ddG among all methods and a similar inverse-folding success rate to the pretrained model (percentage with scRMSD < 2); considering both, it significantly outperforms all baselines in overall success rate.Source: §6.3 · p. 10
  • Without KL regularization, DRAKES tends to over-optimize: scRMSD rises even when Pred-ddG remains high.The reported high scRMSD indicates failure to fold back to the target backbone.Source: §6.3 · p. 10
  • In physics-based validation, DRAKES-generated sequences are more stable than baseline sequences, while the illustrated structures align well with their conditioned backbones.For the examples shown, scRMSDs are 0.768 for 7JJK and 0.492 for 2KRU.Source: §6.3 · p. 10, Figure 2 · p. 10

7 CONCLUSIONS

The conclusion presents DRAKES as a reward-maximization method for discrete diffusion and reports DNA and protein design demonstrations. Related-work discussion distinguishes this approach from alternative discrete diffusion formulations and trajectory-balance methods.

7 CONCLUSIONS

  • DRAKES incorporates reward maximization into discrete diffusion using Gumbel-Softmax for differentiable reward backpropagation, with demonstrations on DNA and protein design.The conclusion identifies more extensive in silico validation and wet-lab validation as future work.Source: §7 · p. 10

A MORE RELATED WORKS

  • Dirichlet diffusion represents intermediate states as simplex vectors, unlike masked diffusion models, whose intermediate states are discrete variables.Source: A MORE RELATED WORKS · p. 16
  • Relative trajectory balance requires estimating an unnormalized density's normalizing constant, whereas this work uses direct backpropagation for a control problem and is tailored to discrete diffusion.The cited trajectory-balance approach models continuous diffusion as a Markovian generative process.Source: A MORE RELATED WORKS · p. 16

B POTENTIAL LIMITATIONS

The proposed algorithm assumes differentiable reward models; the authors recommend alternative optimization methods when building such models is challenging.

B POTENTIAL LIMITATIONS

  • DRAKES requires differentiable reward models; when these are difficult to create, the authors recommend PPO-based algorithms or reward-weighted MLE.Source: B POTENTIAL LIMITATIONS · p. 16

C PROOF OF THEOREMS

The appendix prepares continuous-time Markov chain results and applies backward and forward equations in theorem proofs. It also extends the control formulation to stochastic initial distributions by parameterizing both the generators and initial distributions.

C.1 PREPARATION

  • The appendix states the Kolmogorov backward and forward equations for CTMCs; it notes that proving converse uniqueness requires regularity conditions and omits those converse proofs.Source: §C.1 · pp. 16–17

C.2 PROOF OF THEOREM 2

  • Theorem 2's proof derives the CTMC Hamilton-Jacobi-Bellman equation from a recursive equation by taking dt to zero.Source: §C.2 · pp. 17–18

C.3 PROOF OF THEOREM 3

  • Theorem 3's proof invokes the Kolmogorov backward equation after relating the value function to the reward at the terminal time.Source: §C.3 · p. 18

C.4 PROOF OF THEOREM 4

  • Theorem 4's proof uses the Kolmogorov forward equation and assumes a Dirac-delta initial distribution for its initial-condition argument.A remark says the equality of initial distributions also holds for stochastic initial distributions, with details deferred to Appendix D.Source: §C.4 · pp. 18–19

D EXTENSION WITH STOCHASTIC INITIAL DISTRIBUTIONS

  • For stochastic initial distributions, the extension learns the initial distributions as well as the generators, and Theorem 1 still holds.The authors select this strategy among several possible strategies.Source: D EXTENSION WITH STOCHASTIC INITIAL… · p. 19

D.1 PROOF

The appendix describes proof reuse, implementation choices, and experimental baselines for discrete diffusion fine-tuning. It reports that straight-through Gumbel softmax boosts fine-tuning, a simplified KL calculation has O(NT) complexity, and baseline methods have distinct data and inference constraints.

D.1 PROOF

  • The proof proceeds as in Section C.4 without change, referring to Remark 4.The preceding proof passages recall value functions and state conditions on the optimal initial distribution.Source: §D.1 · p. 20

E.1 STRAIGHT-THROUGH GUMBEL SOFTMAX

  • Straight-through Gumbel softmax boosts fine-tuning by making the reward-oracle input a one-hot vector better aligned with its training distribution.The estimator is applied at the last time step, and its straight-through representation is fed to the reward function.Source: §E.1 · p. 20

E.2 SIMPLIFIED FORMULA OF g(θ)

  • The simplified KL-divergence calculation reduces the complexity of computing k_i(θ) to O(NT).With the experimental linear schedule, γ = 1/t and the corresponding KL-divergence component is zero.Source: §E.2 · pp. 20–21

E.3 SCHEDULE OF GUMBEL SOFTMAX TEMPERATURE

  • A linear Gumbel-softmax temperature schedule τ ∼ 1/t starts with more uncertainty and approaches the true distribution at later steps.The passage states that this schedule improves fine-tuning as inputs become closer to clean data and prediction uncertainty falls.Source: §E.3 · p. 21

F EXPERIMENTAL DETAILS AND ADDITIONAL RESULTS

  • The additional experiments cover DNA, protein, and toxicity-controlled text generation, using one NVIDIA A100 80GB GPU.The section also promises baseline explanations, additional results, and detailed settings for DNA and protein design.Source: F EXPERIMENTAL DETAILS AND ADDITIONAL… · p. 21

F.1 BASELINES

  • Guidance-based methods adjust sampling and take longer at inference, while SMC uses the pretrained model as proposal and TDS uses CG.SMC and TDS are adapted for reward maximization by treating rewards as classifiers.Source: §F.1 · pp. 21–22
  • CG's discrete-state Taylor approximation lacks formal guarantees and may explain weak protein inverse-folding performance; CFG needs labeled pairs and struggles with the small Megascale dataset.CFG is trained from scratch rather than relying on the pretrained model, which uses a larger PDB dataset of approximately 23,000 structures.Source: §F.1 · pp. 21–22

F.2 REGULATORY DNA SEQUENCE DESIGN

The DNA experiments assess oracle quality and whether the pretrained model generates in-distribution enhancer sequences before evaluating fine-tuning. The reported additional results include activity distributions and ablations of gradient truncation and Gumbel-softmax temperature.

Reward oracle

  • The fine-tuning and evaluation reward oracles achieve similar Pearson correlations above 0.85 on their respective held-out chromosome subsets.The oracles are trained independently on separate halves of the 23 human chromosomes, with the evaluation oracle reserved for Pred-Activity evaluation.Source: §F.2 · p. 22, Table 3 · p. 22

Pretrained model

  • The pretrained model is trained on the full Gosai et al. (2023) dataset, and its generated sequences are compared with randomly selected original sequences for activity distributions.The comparison uses 1,280 generated and 1,280 original sequences, with activity predictions from both the fine-tuning and evaluation oracles.Source: §F.2 · p. 22, Figure 3 · p. 23
  • Generated sequences' activity levels align with the original dataset, and their 3-mer and 4-mer Pearson correlations both exceed 0.95.The pretrained model's activity distributions also match the originals closely; the correlation comparison is shown in Figure 4.Source: §F.2 · p. 23, Figure 4 · p. 23, Figure 3 · p. 23

Additional results

  • Figure 5 shows full Pred-Activity distributions for each method, with patterns consistent with the median values reported in Table 1.The evaluation generates 640 sequences per method for each random seed and reports means and standard deviations across three seeds.Source: §F.2 · p. 23, Figure 5 · p. 24

Gradient truncation ablation

  • An intermediate gradient-truncation number performs best: small values limit early-sampling updates, while large values increase memory cost and complicate optimization.The comparison is reported in Table 4.Source: §F.2 · p. 23, Table 4 · p. 24

Gumbel-softmax temperature ablation

  • Linear and constant Gumbel-softmax temperature schedules achieve similar performance, indicating robustness to this schedule choice.The linear schedule is used for the Table 1 results; the comparison is in Table 5.Source: §F.2 · p. 24, Table 5 · p. 24

F.3 PROTEIN INVERSE FOLDING

The protein experiments evaluate inverse-folding performance and reward optimization, including stability, structure confidence, diversity, and robustness to implementation choices. The appendix also reports toxicity-controlled text-generation results, where DRAKES attains the highest reward among the compared baselines.

F.3 PROTEIN INVERSE FOLDING

  • The pretrained inverse-folding model matches ProteinMPNN's sequence recovery on the PDB test set and achieves over 90% success on Megascale proteins.The Megascale success assessment uses scRMSD below 2 Å; the dataset is curated into 438,540 training, 15,182 validation, and 23,466 test variants.Source: §F.3 · pp. 24–25, Table 6 · p. 25, Table 7 · p. 25
  • The fine-tuning reward oracle's train–test performance gap highlights difficulty generalizing to unseen protein structures, whereas the evaluation oracle is trained on all dataset splits.The evaluation oracle reports Pearson correlations of 0.951 on Megascale training and 0.959 on Megascale test, both in-distribution for that oracle.Source: §F.3 · p. 25, Table 8 · p. 26
  • The evaluation oracle shows similar Pred-ddG for DRAKES with and without KL, despite higher fine-tuning-oracle Pred-ddG without KL, suggesting less evaluation over-optimization for proteins.The comparison is in Table 9; the authors contrast protein structures with more homogeneous enhancer sequences.Source: Table 9 · p. 26, §F.3 · p. 26
  • Intermediate gradient truncation performs best, while protein results remain robust across truncation numbers and linear versus constant temperature schedules.The truncation and temperature comparisons are reported in Tables 10 and 11, respectively.Source: Table 10 · p. 26, §F.3 · pp. 26–27, Table 11 · p. 27
  • DRAKES significantly outperforms all baselines in overall success rate using the criterion Pred-ddG > 0 and pLDDT > 80.Table 12 reports means across three random seeds with standard deviations in parentheses.Source: Table 12 · p. 27, §F.3 · p. 27
  • DRAKES matches pretrained-model diversity, exceeds SMC and TDS, and is the only baseline reported to combine high success rate with high diversity.Diversity is measured by average sequence entropy for each test-set protein backbone.Source: Table 13 · p. 27, §F.3 · p. 27

F.4 TOXICITY-CONTROLLED TEXT GENERATION

  • In toxicity-controlled text generation, DRAKES produces the highest reward, corresponding to low toxicity, among the compared baselines.The reported reward measures are mean and median Pred-Score.Source: §F.4 · p. 29, Table 14 · p. 29
Loading 2410.13643v2…