[论文解读] TaxDiff: Taxonomic-Guided Diffusion Model for Protein Sequence Generation

Lin Zongying, Hao Li+6 more|arXiv (Cornell University)|Feb 27, 2024

Microbial Metabolic Engineering and Bioproduction被引用 4

一句话总结

TaxDiff 是一种用于可控蛋白质序列生成的新型扩散模型，通过在每个 Transformer 模块中整合分类学信息，引导生成结构稳定的蛋白质。通过采用分块机制结合全局与局部注意力，其在无条件生成和分类学引导生成任务中均达到最先进性能，结构度量指标超越以往模型，且推理时间仅为基线扩散模型的 25%。

ABSTRACT

Designing protein sequences with specific biological functions and structural stability is crucial in biology and chemistry. Generative models already demonstrated their capabilities for reliable protein design. However, previous models are limited to the unconditional generation of protein sequences and lack the controllable generation ability that is vital to biological tasks. In this work, we propose TaxDiff, a taxonomic-guided diffusion model for controllable protein sequence generation that combines biological species information with the generative capabilities of diffusion models to generate structurally stable proteins within the sequence space. Specifically, taxonomic control information is inserted into each layer of the transformer block to achieve fine-grained control. The combination of global and local attention ensures the sequence consistency and structural foldability of taxonomic-specific proteins. Extensive experiments demonstrate that TaxDiff can consistently achieve better performance on multiple protein sequence generation benchmarks in both taxonomic-guided controllable generation and unconditional generation. Remarkably, the sequences generated by TaxDiff even surpass those produced by direct-structure-generation models in terms of confidence based on predicted structures and require only a quarter of the time of models based on the diffusion model. The code for generating proteins and training new versions of TaxDiff is available at:https://github.com/Linzy19/TaxDiff.

研究动机与目标

• 为解决无条件蛋白质序列生成在生物应用中的局限性，即随机生成需经过大量过滤。
• 通过将生物分类学作为条件信号，实现对蛋白质设计的细粒度、可控化。
• 通过混合注意力机制提升生成蛋白质序列的结构一致性和建模准确性。
• 在保持或提升性能的前提下，减少推理时间，相较现有基于扩散的蛋白质生成模型更具优势。

提出的方法

• 将分类学控制特征集成到去噪 Transformer 模块的每一层中，实现基于物种或家族分类的条件生成。
• 提出一种分块注意力机制，结合对完整序列的全局注意力与对固定大小氨基酸分块（如 16 个残基）的局部注意力。
• 在家族和物种层级重新分类蛋白质序列，以整合过于细粒度的 UniProt 分类，提升训练稳定性与控制保真度。
• 采用与 EvoDiff 类似的基于扩散的生成框架，实现在共享潜在空间中联合生成序列与结构。
• 为无条件与条件生成采用统一架构，实现对不同生物设计任务的无缝适配。
• 采用 OmegaFold 进行结构预测与评估，使用 pLDDT、TM-score、RMSD 和 Fident 作为关键指标，评估生成序列的质量。

实验结果

研究问题

• RQ1 分类学信息能否有效用于引导生成生物上合理的蛋白质序列？
• RQ2 全局与局部注意力机制的结合如何影响生成蛋白质的结构与序列质量？
• RQ3 与标准自注意力相比，分块注意力机制在生成效率与结构一致性方面改善程度如何？
• RQ4 TaxDiff 在无条件与分类学引导的蛋白质序列生成基准测试中，相较于最先进模型表现如何？
• RQ5 TaxDiff 是否能生成结构置信度更高、推理时间更短的蛋白质序列，相较现有基于扩散的模型？

主要发现

• 在无条件蛋白质序列生成中，TaxDiff 相较基线模型在 Fident 得分上提升 7.13%，TM-score 提升 11.93%，表现优于甚至超过基于结构的模型。
• 在分类学引导生成中，TaxDiff 的 pLDDT 得分显著高于其他序列模型，接近天然蛋白质水平，表明具有高结构置信度。
• 模型仅用 24 分钟即可生成 1,000 条蛋白质序列，耗时仅为其他基于扩散模型的 1/4 至 2/3，展现出卓越的推理效率。
• 全局与局部注意力结合（方法 A）在所有指标上均表现最佳，仅使用局部注意力（方法 E）也优于标准注意力配置。
• 将序列划分为 16 个氨基酸分块时，结构度量指标（如 TM-score 与 RMSD）表现最优；而更大分块（64 个残基）则更优序列级指标（如 Fident）。
• 实证结果证实，相较于全序列注意力，局部注意力在蛋白质序列上更有效，提示其内在具有类似语言中“短句”的局部结构模式。