Dirichlet Flow Matching with Applications to DNA Sequence Design

📄 arXiv: 2402.05841v2 📥 PDF

作者: Hannes Stark, Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, Tommi Jaakkola

分类: q-bio.BM, cs.LG

发布日期: 2024-02-08 (更新: 2024-05-30)

备注: Published at ICML 2024. (Proceedings of the 41st International Conference on Machine Learning, Vienna, Austria. PMLR 235, 2024)

🔗 代码/项目: GITHUB

💡 一句话要点

提出Dirichlet流匹配以解决DNA序列设计问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱四：生成式动作 (Generative Motion)

关键词: DNA序列设计 流匹配 Dirichlet分布 生成模型 机器学习 生物信息学

📋 核心要点

  1. 现有的线性流匹配方法在训练目标上存在不连续性，导致生成序列的控制性不足。
  2. 提出的Dirichlet流匹配利用Dirichlet分布混合作为概率路径，克服了现有方法的不足。
  3. 在复杂DNA序列生成任务中，所提方法在性能上显著优于基线，且实现了高效的序列生成。

📝 摘要（中文）

离散扩散或流模型能够比自回归模型更快速且可控地生成序列。我们展示了简单的线性流匹配在此目标上不足之处，因其在训练目标上存在不连续性等问题。为此，我们开发了基于Dirichlet分布混合的Dirichlet流匹配框架，建立了混合得分与流向量场之间的联系，支持分类器和无分类器引导。此外，我们提供了蒸馏Dirichlet流匹配，实现了一步序列生成，性能损失最小，相较于自回归模型实现了$O(L)$的加速。在复杂的DNA序列生成任务中，我们在分布度量和实现设计目标方面表现优于所有基线。最后，我们展示了无分类器引导方法在无条件生成中的改进效果，并有效生成满足设计目标的DNA序列。

🔬 方法详解

问题定义：本论文旨在解决现有线性流匹配方法在DNA序列生成中的不连续性和控制性不足的问题。现有方法在训练目标上存在缺陷，导致生成序列的质量和可控性较差。

核心思路：论文提出Dirichlet流匹配，通过使用Dirichlet分布的混合来构建概率路径，从而实现更平滑的流动匹配，增强序列生成的可控性和效率。

技术框架：整体架构包括Dirichlet分布混合的构建、流向量场的计算以及分类器和无分类器引导的实现。主要模块包括数据预处理、模型训练和生成阶段。

关键创新：最重要的技术创新在于建立了混合得分与流向量场之间的联系，使得生成过程更加平滑且可控。这一设计与传统的线性流匹配方法有本质区别。

关键设计：在模型设计中，采用了特定的损失函数以优化流动匹配过程，并通过蒸馏技术实现了一步序列生成，确保在性能损失最小的情况下提高生成速度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，所提Dirichlet流匹配方法在复杂DNA序列生成任务中显著优于所有基线模型，尤其在分布度量和设计目标实现方面表现突出，达到了$O(L)$的速度提升，展示了其在实际应用中的有效性。

🎯 应用场景

该研究在DNA序列设计领域具有广泛的应用潜力，能够为生物工程、合成生物学等领域提供高效的序列生成工具。未来，该方法可能推动基因合成和基因编辑技术的发展，助力个性化医疗和生物制药的进步。

📄 摘要（原文）

Discrete diffusion or flow models could enable faster and more controllable sequence generation than autoregressive models. We show that naïve linear flow matching on the simplex is insufficient toward this goal since it suffers from discontinuities in the training target and further pathologies. To overcome this, we develop Dirichlet flow matching on the simplex based on mixtures of Dirichlet distributions as probability paths. In this framework, we derive a connection between the mixtures' scores and the flow's vector field that allows for classifier and classifier-free guidance. Further, we provide distilled Dirichlet flow matching, which enables one-step sequence generation with minimal performance hits, resulting in $O(L)$ speedups compared to autoregressive models. On complex DNA sequence generation tasks, we demonstrate superior performance compared to all baselines in distributional metrics and in achieving desired design targets for generated sequences. Finally, we show that our classifier-free guidance approach improves unconditional generation and is effective for generating DNA that satisfies design targets. Code is available at https://github.com/HannesStark/dirichlet-flow-matching.