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
TL;DR
This work tackles the challenge of generating natural-like discrete sequences while optimizing task-specific objectives. It introduces DRAKES, a method that fine-tunes pretrained discrete diffusion models by solving a reward-maximization objective with a KL penalty to preserve naturalness, and makes diffusion trajectories differentiable through the Gumbel-Softmax trick. The authors provide theoretical guarantees linking the learned distribution to a reward-weighted prior and demonstrate empirical superiority on enhancer DNA design and protein stability tasks, with ablations highlighting the beneficial role of the KL term. The approach addresses fundamental differences between CTMC-based discrete diffusion and continuous diffusion, enabling effective design in biology with potential for broader NLP and therapeutics applications. The results suggest substantial practical impact for generating functional, in-distribution sequences in gene therapy and protein engineering.
Abstract
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.