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

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.

• The paper presents DRAKES, a reward optimization approach that fine-tunes discrete diffusion models for generating biologically viable DNA and protein sequences.
• It balances task-specific objectives with natural sequence properties by minimizing KL divergence from pretrained models.
• Empirical results show enhanced enhancer activity and improved protein stability, underscoring its potential for gene therapy and protein design.

Fine-Tuning Discrete Diffusion Models for Biological Sequence Design

The paper "Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design" presents a methodology to enhance the capabilities of discrete diffusion models in generating biological sequences, specifically DNA and protein sequences, that are both natural-like and optimized for specific task objectives. The paper leverages the reinforcement learning (RL) paradigm, introducing an algorithm named DRAKES to effectively fine-tune diffusion models for reward maximization.

Overview and Contributions

Diffusion models have proven their efficacy across various domains, including natural language processing and biological sequence generation. The research extends these models' application within discrete spaces, aiming to optimize specific objectives beyond mere generative quality. For instance, in protein design, the goal is not only to generate viable sequences but also to optimize attributes like stability—crucial for therapeutic interventions.

The key contribution of the paper lies in formulating the reward maximization problem analogous to RL while maintaining the sequence's natural essence by minimizing the KL divergence against pretrained diffusion models. This equilibrium allows for balancing task-specific optimization and the preservation of intrinsic sequence properties. The proposed algorithm, DRAKES, enables the optimization process by making use of the Gumbel-Softmax trick to treat ordinarily non-differentiable trajectories as differentiable, facilitating effective backpropagation of rewards.

Theoretical Analysis

The authors offer theoretical guarantees demonstrating that their approach can generate sequences that maintain high probability within a pretrained model's distribution and achieve significant reward scores. The derivation of a theoretical framework that parallels advancements in classifier guidance further underpins the robustness of the presented method.

The work differentiates itself from previous approaches by tackling unique algorithmic and theoretical challenges inherent in discrete diffusion models. Unlike continuous diffusion models that utilize Brownian motion, discrete models are grounded in continuous-time Markov chains, necessitating distinct methodological adaptations.

Empirical Results

The efficacy of DRAKES is demonstrated through applications to DNA and protein sequence designs. DNA sequences are generated with enhanced enhancer activity, which is essential for gene therapy, while protein sequences that optimize stability are crucial for protein-based therapeutic solutions. Empirical evaluations highlight that the sequences generated not only fit naturally into expected distributions but also achieve high task-specific rewards.

Implications and Future Directions

Practically, the implications of this research are profound for fields like gene therapy and protein engineering, offering a refined tool for designing sequences that meet stringent biological criteria. Theoretically, it reinforces the potential of integrating RL frameworks with generative diffusion models, expanding the modeling capabilities beyond traditional applications.

Future research could explore further algorithmic enhancements, integrations with more complex biological objectives, and validation through in silico or experimental wet-lab studies. The understanding and management of trade-offs between naturalness and task-specific optimization will continue to be a pivotal aspect in the evolution of such models.

In summary, this paper advances the field of computational biology by providing a more nuanced control mechanism within discrete diffusion models, ensuring they create biologically viable and optimally functional sequences for complex applications.