通过奖励优化微调离散扩散模型及其在DNA和蛋白质设计中的应用
Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design
AI总结：
本文提出DRAKES算法，通过Gumbel-Softmax使离散扩散轨迹可微并结合KL约束进行奖励优化，实现高奖励且类天然的DNA与蛋白质序列设计。
摘要：
近期研究表明，扩散模型在从自然语言到生物序列生成等领域的离散序列上具有出色的实证表现。例如，在蛋白质逆折叠任务中，条件扩散模型在生成能够折叠回原始结构的类天然序列方面取得了令人瞩目的结果。然而，实际设计任务通常不仅需要对条件分布建模，还需要优化特定任务目标。例如，我们可能更偏好具有高稳定性的蛋白质序列。为解决这一问题，我们考虑这样一种场景：已有能够生成类天然序列的预训练离散扩散模型，以及将序列映射到任务目标的奖励模型。随后，我们在离散扩散模型中构建奖励最大化问题，类似于强化学习（RL），同时最小化相对于预训练扩散模型的KL散度，以保持自然性。为求解该RL问题，我们提出了一种新算法DRAKES，它利用Gumbel-Softmax技巧使原本不可微的轨迹变得可微，从而使奖励能够通过扩散模型生成的完整轨迹直接反向传播。我们的理论分析表明，该方法能够生成既类天然又具有高奖励的序列。尽管近期已有研究在连续域扩散模型中探索类似任务，但我们的工作解决了离散扩散模型特有的算法和理论挑战，这些挑战源于其基础是连续时间马尔可夫链而非布朗运动。最后，我们证明了DRAKES在生成分别优化增强子活性和蛋白质稳定性的DNA与蛋白质序列方面的有效性，这两项任务对基因治疗和基于蛋白质的治疗均具有重要意义。
英文摘要：
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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
