arXiv:2410.13643cs.LGcs.AI2024-10ICLR被引 78

用奖励优化提升生成序列的生物功能,保持自然性。

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

  • 通过Gumbel-Softmax让离散扩散模型轨迹可微,实现奖励直接反向传播。
  • 在DNA和蛋白质设计中,显著提升增强子活性与蛋白稳定性。
  • 适合需生成高功能性生物序列的研究者,如基因治疗与药物设计。

近期研究表明,扩散模型在自然语言到生物序列生成等离散序列任务中表现优异。例如,在蛋白质逆折叠任务中,条件扩散模型能生成可折叠回原始结构的自然序列。然而,实际设计任务不仅需建模条件分布,还需优化特定目标(如蛋白稳定性)。本文针对已有预训练离散扩散模型及奖励模型的场景,将奖励最大化问题形式化为类似强化学习的优化任务,同时最小化与预训练模型的KL散度以保持序列自然性。为此提出新算法DRAKES,利用Gumbel-Softmax使原本不可微的生成轨迹可微,实现奖励在整个生成路径上的直接反向传播。理论分析表明该方法能生成既自然又高奖励的序列。相较于连续域扩散模型的类似工作,本文解决的是离散扩散模型特有的挑战——其基于连续时间马尔可夫链而非布朗运动。实验验证了DRAKES在优化DNA增强子活性和蛋白稳定性方面的有效性,对基因治疗和蛋白类疗法具有重要意义。

原文摘要 · Abstract (English)

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.

序列生成扩散模型蛋白质设计奖励优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。