[论文解读] Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design

Andrew T. Campbell, Jason Yim+3 more|PubMed|Feb 7, 2024

Cancer-related gene regulation被引用 9

一句话总结

引入离散流模型（DFMs），实现对离散与连续数据的多模态生成；提出 Multiflow，用于联合结构-序列蛋白质设计，达到最先进的结果。

ABSTRACT

Combining discrete and continuous data is an important capability for generative models. We present Discrete Flow Models (DFMs), a new flow-based model of discrete data that provides the missing link in enabling flow-based generative models to be applied to multimodal continuous and discrete data problems. Our key insight is that the discrete equivalent of continuous space flow matching can be realized using Continuous Time Markov Chains. DFMs benefit from a simple derivation that includes discrete diffusion models as a specific instance while allowing improved performance over existing diffusion-based approaches. We utilize our DFMs method to build a multimodal flow-based modeling framework. We apply this capability to the task of protein co-design, wherein we learn a model for jointly generating protein structure and sequence. Our approach achieves state-of-the-art co-design performance while allowing the same multimodal model to be used for flexible generation of the sequence or structure.

研究动机与目标

• 通过在单一框架内统一离散数据与连续数据来激发多模态生成建模。
• 基于连续时间马尔可夫链发展离散流模型（DFMs），以建模离散数据流。
• 通过将 DFMs 与连续状态流模型结合，创建多模态框架。
• 将该框架应用于蛋白质协同设计，以联合生成蛋白质结构与序列。
• 展示在多模态生成中的最先进性能和采样灵活性。

提出的方法

• 定义一个概率流 p_t，通过数据条件流 p_{t|1}(·|x_1) 在噪声与数据之间插值。
• 将 p_t 表示为条件流的期望，从而实现封闭形式构造和可控训练（Eq. 6）。
• 实例化一个速率矩阵 R_t(x_t, j|x_1)，其在 p_{1|t} 的期望下生成 p_t（Prop. 3.1）。
• 使用交叉熵训练一个神经去噪器 p_{1|t}^θ(x_1|x_t) 来逼近去噪分布（Eq. 10）。
• 引入起始速率 R_t^*，并通过带有细节平衡的分量 R_t^{DB} 扩展，形成控制 CTMC 随机性的族 R_t^η（Eq. 14）。
• 将 DFMs 与连续流结构模型结合，形成用于联合结构与序列生成的 Multimodal Flow（Multiflow）（Section 4）。
• 通过将序列建模为 DFM、结构建模为类似 FrameFlow 的连续流，应用于蛋白质，启用协同设计（Section 4）。

实验结果

研究问题

• RQ1 基于连续时间马尔可夫链构建的离散流模型，是否能在保留连续流模型灵活性的同时对离散数据进行采样？
• RQ2 将 DFMs 与连续流模型结合，是否能够在多模态任务中实现有效的多模态生成与采样控制？
• RQ3 与先前将结构与序列分开建模的方法相比，所提的 Multiflow 框架在蛋白质协同设计上的表现如何？
• RQ4 CTMC 随机性（通过 η）在蛋白质协同设计中的样本质量、多样性和可设计性方面的影响程度如何？
• RQ5 该框架是否支持跨模态的条件化与修复（如在给定序列的情况下生成结构，或反之）？

主要发现

方法	协同设计 1（Des.）	协同设计 1（Div.）	协同设计 1（Nov.）	PMPNN 8（Des.）	PMPNN 8（Div.）	PMPNN 8（Nov.）	PMPNN 1（Des.）	PMPNN 1（Div.）	PMPNN 1（Nov.）
Protpardelle	0.05	6	0.75	0.92	46	0.67	0.63	33	0.68
ProteinGenerator	0.34	31	0.74	0.88	73	0.71	0.75	56	0.72
RFdiffusion	N/A	N/A	N/A	0.90	161	0.69	0.69	120	0.70
Multiflow	0.88	143	0.68	0.99	156	0.68	0.87	142	0.69
Multiflow w/o distillation	0.41	73	0.68	0.89	126	0.68	0.75	110	0.69
Multiflow w/o sequence	N/A	N/A	N/A	0.99	118	0.69	0.86	95	0.69

• DFMs 通过 CTMC 随机性实现对离散数据的采样时间灵活性，在适当的 η 下，优于文本数据上的离散扩散基线。
• Multiflow 在联合生成蛋白质结构与序列方面实现了最先进的协同设计性能。
• 模型支持对任一模态的条件化与灵活的多模态采样，包括 inpainting 和前向/逆向折叠等任务。
• CTMC 随机性（由 η 控制）影响样本多样性与跳变频率，提供一个平衡探索与保真度的调节钮。
• 初步结果表明 Multiflow 在逆向与前向折叠方面具有潜力，为通用蛋白质生成建模指明了一条路径。