ICML 2024 | 离散状态空间上的生成流：实现多模态流及其在蛋白质共同设计中的应用

DRUGAI

今天为大家介绍的是来自Tommi Jaakkola团队的一篇论文。结合离散数据和连续数据是生成模型的重要能力。作者提出了离散流模型（DFMs），这是一种新的基于流的离散数据模型，弥补了在多模态连续和离散数据问题中应用基于流的生成模型的缺失环节。作者的关键见解是，可以使用连续时间马尔可夫链实现连续空间流匹配的离散等价形式。DFMs从一个简单的推导出发，包括离散扩散模型作为特定实例，同时在性能上优于现有的基于扩散的方法。作者利用DFMs方法构建了一个多模态的基于流的建模框架。作者将这一能力应用于蛋白质共同设计任务，在其中作者学习一个联合生成蛋白质结构和序列的模型。作者的方法在共同设计性能上达到了最先进的水平，同时允许同一多模态模型用于灵活生成序列或结构。

科学领域经常涉及连续的原子相互作用和离散的化学描述。扩展生成模型处理离散和连续数据（作者称之为多模态）的能力，是使其在科学应用中广泛采用的基础问题之一。蛋白质共同设计是需要多模态生成模型的一个应用，其目标是联合生成连续的蛋白质结构和相应的离散氨基酸序列。蛋白质已经被广泛研究：蛋白质的功能由其结构赋予，而序列是结构形成的蓝图。这种相互作用激发了作者联合生成结构和序列，而不是单独生成。为此，作者工作的重点是开发一个能够进行共同设计的多模态生成框架。

扩散模型有潜力作为多模态框架，因为它们可以在连续和离散空间上定义。然而，它们的采样时间不灵活使其不适合多模态问题。另一方面，基于流的模型通过更简单的框架改善了扩散模型，允许通过采样灵活性实现更高性能。不幸的是，目前无法在离散空间上定义基于流的模型，这阻碍了实现多模态流模型的进展。

为了解决这个问题，作者引入了一种新型的基于流的离散数据模型，称为离散流模型（DFMs），从而解锁了基于流的多模态生成建模的完整框架。作者的关键见解是，离散的基于流的模型可以通过连续时间马尔可夫链（CTMCs）实现。DFMs是一种新的离散生成建模范式：比扩散模型限制更少，允许在不重新训练的情况下实现采样灵活性，并能与连续状态空间流简单结合，形成多模态流模型。

离散流模型DFM

图 1

图1A概述了离散流模型（DFMs）。离散流模型（DFM）是一种围绕从噪声到数据插值的概率流构建的离散数据生成模型。为了采样新的数据点，作者模拟一个匹配从噪声到数据概率流的序列轨迹。该流的构建使得能够将DFM与连续数据流模型结合起来，从而定义一个多模态生成模型。具体来说，首先定义一个概率流，它从噪声线性插值到数据。然后，通过模拟一个随时间跟随的序列轨迹来生成新数据，这需要训练一个带有交叉熵的去噪神经网络。序列轨迹可能会有很多或很少的转换，作者称之为CTMC随机性。先前的离散扩散模型相当于在训练时选择特定的随机性，而作者可以在推理时调整它：增强样本质量并控制样本分布特性。DFM采样步骤伪代码如算法1所示，训练目标函数如式1所示。

算法 1

式 1

多模态的蛋白生成模型

现在，作者使用DFM来创建一个多模态蛋白质生成模型。为了生成多模态数据，作者将定义一个多模态生成流。我们定义在不同模态上进行因子化，从而可以为每个模态单独定义。作者的训练损失只是每个模态的标准流损失之和。在推理时，作者还可以在每个模拟步骤中单独更新每个模态，使用ODE处理连续数据，使用CTMC处理离散数据。现在作者将这一功能应用于蛋白质结构-序列生成。

蛋白质可以被建模为一个线性残基链，每个残基都有一个指定的氨基酸和3D原子坐标。蛋白质共同设计的目标是联合生成氨基酸（序列）和坐标（结构）。先前的工作使用了一个模态（序列或结构）的生成模型，并通过一个单独的模型来预测另一模态。而作者的方法使用一个生成模型来联合采样两种模态：DFM用于序列，FrameFlow用于结构。作者称之为共同生成序列和结构，因此该方法称为Multiflow。

结构表示为SE(3)的元素，以捕捉沿主链局部框架的刚性。为残基碳-α原子的平移，是残基局部框架相对于全局参考框架的旋转矩阵，是20种氨基酸之一或掩码状态M。在训练过程中，作者使用每种模态的条件流对数据进行破坏。是SO(3)上的均匀分布。

式 2

在训练过程中，作者的网络将以噪声蛋白质作为输入，预测去噪后的平移x、旋转r和氨基酸分布p。作者最小化损失函数（式3）。

式 3

然后作者将预测转换为向量场和速率矩阵（式4）。

式 4

为了使用Multiflow进行采样，作者沿着平移和旋转的ODE轨迹进行积分，同时跟随氨基酸序列的CTMC。采样期间的每个欧拉步都有更新（式5）。

式 5

在采样氨基酸时，作者发现使用纯度（purity）来选择每一步要取消遮蔽的索引是有益的。使用解耦时间表进行训练的优点在于，作者可以自由地使用任意组合的 进行采样。作者利用这一点进行条件修补，通过将 或设为1来固定其中一个模态。作者在图1C和表2中总结了这些功能。

表 2

文本建模性能

图 2

图2绘制了不同η值和采样温度下的结果。为了比较，作者还包括了使用吸收态损坏（absorbing state corruption）的离散扩散D3PM方法的结果。作者发现，DFM由于额外的采样时间灵活性，表现优于D3PM。作者能够选择在采样时优化帕累托前沿的η值（这里η=15），而D3PM没有这种灵活性。作者在图2中展示了η=0的前沿。当η=0时，性能相似，因为在这种设置下，DFMs是D3PM的连续时间泛化。作者还在图2中包括了自回归模型的结果以作参考；然而，注意到这并不是一个完全的对比，因为自回归模型训练所需的计算量比基于扩散的模型少得多。

共设计结果

按照RFdiffusion的基准，作者为每种长度（70、100、200和300）采样100个蛋白质。使用Multiflow时，作者采用500个时间步长，温度设为0.1（PMPNN也使用0.1），随机性水平设为η = 20。作者将模型的结构质量与最先进的结构生成方法RFdiffusion进行比较。在共同设计方面，作者与Protpardelle和ProteinGenerator进行比较。所有方法均使用其公开发布的代码并以相同的方式进行评估。

表 3

作者的结果见表3。作者发现，Multiflow的共同设计能力超越了之前的共同设计方法，这些方法均未使用联合多模态生成过程。Multiflow生成的序列在与生成的结构一致性方面达到了与PMPNN（作者对于ProteinMPNN的缩写）相当的水平，通过比较Co-design 1和PMPNN 1的可设计性可以看出。在纯结构生成方面，作者发现Multiflow在结构质量（通过PMPNN 8可设计性测量）上超过了所有基线。Multiflow还达到了与之前方法相当的多样性和新颖性。作者对蒸馏的使用进行了消融研究，发现蒸馏不仅改善了整体可设计性，还提高了多样性。最后，作者使用相同的架构进行训练，仅在蒸馏数据集上对结构建模，使用Yim等人提出的损失函数。作者发现，联合结构-序列模型在结构质量上与仅结构版本相同，但在生成过程中额外包含序列则增加了结构的多样性。

编译 | 黄海涛

审稿 | 曾全晨

参考资料

Campbell, A., Yim, J., Barzilay, R., Rainforth, T., & Jaakkola, T. (2024). Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design. arXiv preprint arXiv:2402.04997.