基于奖励优化的离散扩散模型精细调优：面向 DNA 与蛋白设计

最近的研究表明，扩散模型在从自然语言到生物序列生成等多个领域的离散序列上都表现出强大的经验性能。例如，在蛋白逆折叠任务中，条件扩散模型已在生成能够折叠回原始结构的自然样子序列方面取得了令人瞩目的成绩。然而，实际的设计任务往往不仅需要建模条件分布，还需要优化特定的任务目标。例如，我们可能更倾向于选择具有高稳定性的蛋白序列。为此，我们考虑这样的场景：已有能够生成自然样子序列的预训练离散扩散模型，以及将序列映射到任务目标的奖励模型。我们在离散扩散模型中构建奖励最大化问题，类似于强化学习（RL），同时最小化与预训练扩散模型之间的 KL 散度以保持自然性。为求解该 RL 问题，我们提出了一种新颖算法 DRAKES，使奖励能够通过扩散模型生成的整个轨迹进行直接反向传播，通过使用 Gumbel-Softmax 技巧将原本不可微的轨迹变为可微的。我们的理论分析表明，该方法能够生成既自然又能获得高奖励的序列。虽然类似任务在连续域的扩散模型中已有探索，但我们的工作解决了离散扩散模型特有的算法和理论挑战，这些挑战源于其基于连续时间马尔可夫链而非布朗运动的基础。最后，我们在生成优化增子活性的 DNA 序列和优化蛋白稳定性的蛋白序列方面，证明了 DRAKES 方法的有效性，后者对于基因疗法和基于蛋白的治疗药物具有重要意义。