[论文解读] Improving AlphaFlow for Efficient Protein Ensembles Generation

Shaoning Li, Mingyu Li+5 more|arXiv (Cornell University)|Jul 8, 2024

Cell Image Analysis Techniques被引用 5

一句话总结

Alpha Flow -Lit 通过在特征条件化下生成蛋白质集合，同时保持与 Alpha Flow 相近的准确性，实现约 47x 更快的采样。

ABSTRACT

Investigating conformational landscapes of proteins is a crucial way to understand their biological functions and properties. AlphaFlow stands out as a sequence-conditioned generative model that introduces flexibility into structure prediction models by fine-tuning AlphaFold under the flow-matching framework. Despite the advantages of efficient sampling afforded by flow-matching, AlphaFlow still requires multiple runs of AlphaFold to finally generate one single conformation. Due to the heavy consumption of AlphaFold, its applicability is limited in sampling larger set of protein ensembles or the longer chains within a constrained timeframe. In this work, we propose a feature-conditioned generative model called AlphaFlow-Lit to realize efficient protein ensembles generation. In contrast to the full fine-tuning on the entire structure, we focus solely on the light-weight structure module to reconstruct the conformation. AlphaFlow-Lit performs on-par with AlphaFlow and surpasses its distilled version without pretraining, all while achieving a significant sampling acceleration of around 47 times. The advancement in efficiency showcases the potential of AlphaFlow-Lit in enabling faster and more scalable generation of protein ensembles.

研究动机与目标

• 了解 Alpha Flow 在蛋白质集合采样效率方面的局限性。
• 开发一个更轻量、特征条件化的变体以降低计算负担。
• 在 MD 轨迹数据上对比评估 Alpha Flow -Lit、Alpha Flow 及其蒸馏版本。
• 分析生成的集合在蛋白质动力学、局部残基排列，以及长程相关性方面。

提出的方法

• 保留 AlphaFold 嵌入器和 Evoformer 处于冻结状态，以预计算单特征和成对特征。
• 将轻量级 StructureModule 以计算得到的单特征和成对特征（扭转角和接触图）为条件，而不是 MSAs。
• 使用受 ControlNet 启发的预计算特征条件输入嵌入，以尽量不干扰预训练权重。
• 在 ALTAS MD 轨迹上训练 Alpha Flow -Lit，未进行 PDB 预训练，使用完整的去噪步骤和轻量级去噪器。
• 使用基于 MD 的度量（RMSD、RMSF、JSD、DCCM 等）比较运行时间和集合质量，与 Alpha Flow -Full 和 Alpha Flow -Distilled 相比。
• 维持基于 AlphaFlow 的 CNF/流匹配框架，x_t 作为特征，使用学习得到的向量场。

Figure 1: Model architecture of sequence-conditioned Alpha Flow ( left ) and feature-conditioned Alpha Flow -Lit ( right ). $T$ : Denoising steps; $x_{t}$ : Noisy structure; $\tilde{x}_{0}$ : Predicted structure.

实验结果

研究问题

• RQ1 特征条件化、更加轻量的结构模块是否能够取代以 MSAs 为锚点的完整条件化，从而在不牺牲集合质量的前提下加速 Alpha Flow？
• RQ2 相较于 Alpha Flow 及其蒸馏版本，Alpha Flow -Lit 如何随蛋白长度扩展？
• RQ3 Alpha Flow -Lit 的集合是否能保留与 ground-truth MD 集合相当的基本动力学、局部残基行为和长程相关性？
• RQ4 Alpha Flow -Lit 与 Alpha Flow -Full/Distilled 在多样性和准确性上的权衡是什么？

主要发现

• 与 Alpha Flow -Full 相比，Alpha Flow -Lit 实现了平均约 47x 的更快采样。
• Alpha Flow -Lit 在 ATLAS MD 轨迹上的集合质量与 Alpha Flow 相当，在若干指标上优于 Alpha Flow -Distilled。
• Alpha Flow -Lit 在一般动力学（基于 RMSD）的地面真实 MD 集合上显示出最强的皮尔逊相关性，同时保持与 Alpha Flow -Full 相当的多样性。
• 在基于 PCA 的本质动力学中，Alpha Flow -Lit 的表现与 Alpha Flow -Full 相当，优于 Alpha Flow -Distilled。
• 局部残基分析（RMSF、接触概率、二面角分布）表明 Alpha Flow -Lit 能更接近真实模式，相对于 Distilled 有显著改进。
• 长程耦合分析（DCCM）显示 Alpha Flow -Lit 比蒸馏版本更好地捕获残基相关性。

Figure 2: Visualization of MD evaluation from MD, Alpha Flow -Lit and Alpha Flow . ( A ) Runtime comparison corresponding to the sequence length and their fitted curves. ( B ) Principal components analysis (PCA) for 6q9c_A ensembles. The representative structures are pointed out. ( C, D ) Ensembles