Understanding protein function with a multimodal retrieval-augmented foundation model

本文介绍了 PoET-2，一种结合检索增强、多模态及家族中心建模的蛋白质基础模型，它通过层次化 Transformer 和双解码器架构，在零突变效应预测和基于小数据集的序列 - 功能关系学习中均取得了最先进性能。

要点

这篇论文介绍了一个名为 PoET-2 的人工智能模型，它的任务是帮助科学家理解和设计蛋白质。为了让你更容易理解，我们可以把蛋白质想象成乐高积木搭建的复杂机器，而 PoET-2 则是一位超级天才的“乐高大师”。

以下是用通俗语言和生动比喻对这篇论文的解读：

1. 背景：蛋白质是什么？为什么我们需要 AI？

• 蛋白质 = 生命的乐高积木：蛋白质是由氨基酸（就像乐高小颗粒）串成的链条。它们折叠成特定的 3D 形状，就像搭好的乐高模型，负责身体里的各种工作（比如消化食物、抵抗病毒）。
• 突变 = 换掉几个积木：有时候，蛋白质里的几个氨基酸会发生变化（突变）。这就像把乐高模型里的一个红色积木换成了蓝色。
  • 如果换得不好，机器可能坏了（导致疾病）。
  • 如果换得好，机器可能变得更强大（用于开发新药或工业酶）。
• 难题：以前，科学家要预测换掉一个积木会发生什么，需要像大海捞针一样做大量实验。现在的 AI 模型（叫“蛋白质语言模型”）虽然能猜个大概，但面对复杂的改动（比如一次换好几个积木，或者插入/删除积木），它们往往就“晕”了，或者需要海量的数据才能学会。

2. PoET-2 的三大绝招

PoET-2 之所以厉害，是因为它用了三个独特的“超能力”：

绝招一：像“查字典”一样学习（检索增强）

• 传统 AI：像是一个死记硬背的学生，只靠脑子里背过的知识做题。如果题目太偏，它就答不上来。
• PoET-2：像是一个带着百科全书的专家。当它遇到一个新蛋白质时，它不会只靠死记硬背，而是会立刻去“图书馆”（数据库）里查找和这个蛋白质长得像的“亲戚”（同源蛋白）。
• 比喻：这就好比你要修一辆从未见过的新型自行车。普通修车师傅只能靠经验猜；而 PoET-2 会立刻拿出手机，搜索“这种自行车的亲戚们是怎么设计的”，然后结合这些亲戚的维修手册来修你的车。这让它即使没见过这个具体型号，也能猜得很准。

绝招二：既看“文字”又看“图纸”（多模态）

• 传统 AI：通常只看氨基酸的“文字序列”（比如 A-B-C-D...），就像只看乐高的说明书文字，却不去看拼好的 3D 模型长什么样。
• PoET-2：它既看文字，又看 3D 结构。它不仅能读懂氨基酸的排列顺序，还能“看见”这些氨基酸在空间里是怎么折叠的。
• 比喻：就像你不仅知道乐高的说明书（序列），还能直接看到拼好的模型（结构）。如果它发现某个积木的位置在空间上会打架，它就知道这个设计肯定不行。

绝招三：双核大脑（双重训练目标）

• 传统 AI：通常只擅长一种任务，要么擅长“填空”（猜中间缺什么），要么擅长“续写”（猜后面是什么）。
• PoET-2：它有两个“大脑”同时工作：
  1. 创造性大脑（生成式）：擅长从头开始设计新的蛋白质，或者计算某个改动后的蛋白质“存活率”有多高。
  2. 理解性大脑（双向理解）：擅长深入分析，提取蛋白质的特征，用来做精细的分类或预测。
• 比喻：就像一位艺术家，左手能即兴创作新画作（生成），右手能像艺术评论家一样深度分析画作的精妙之处（理解）。

3. PoET-2 做成了什么？（主要成就）

• 搞定“乱改”的难题：以前的 AI 最怕一次改好几个地方，或者插入/删除几个积木（这叫 Indels）。PoET-2 在这方面打破了世界纪录，准确率比以前的冠军高了 20% 以上。
  • 比喻：以前 AI 只能猜“把第 5 块积木换成红色”会怎样；PoET-2 能猜“把第 5 块换成红色，同时把第 10 块拿掉，再在第 15 块后面加两块”会发生什么。
• 小样本学习（少即是多）：在需要科学家做实验收集数据时，PoET-2 只需要很少的数据（比如几十个样本）就能学会预测，而以前的模型需要成千上万个样本。
  • 比喻：以前的老师要教学生 1000 道题才能学会解题；PoET-2 这位天才学生，看 10 道题就能举一反三，学会解所有类似的题。这对那些很难获取数据的罕见病研究特别有用。
• 预测疾病突变：它能更准确地判断人类基因突变是否会导致疾病，帮助医生更快找到致病原因。

4. 总结：为什么这很重要？

PoET-2 就像是一个超级高效的蛋白质设计助手。

• 以前：设计一个新蛋白质或预测突变效果，像在大雾里摸索，需要大量试错，耗时耗力。
• 现在：有了 PoET-2，科学家可以像有了“透视眼”和“超级计算器”，能更快地设计出更稳定的药物、更高效的酶，或者更快地找到治愈疾病的钥匙。

一句话总结：PoET-2 通过“查亲戚资料”（检索增强）、“看图说话”（多模态）和“双核驱动”（双重目标），让 AI 在理解和设计蛋白质这件事上，从“小学生”进化成了“博士”，而且只需要很少的练习就能达到顶尖水平。

技术摘要

以下是关于论文《Understanding protein function with a multimodal retrieval-augmented foundation model》（PoET-2）的详细技术总结：

1. 研究背景与问题 (Problem)

蛋白质语言模型（PLMs）通过从数亿条天然蛋白质序列中学习概率分布，在蛋白质理解和设计方面展现出巨大潜力。然而，现有的 PLMs 在解决以下关键挑战时仍存在局限性：

• 突变预测的局限性：大多数基于掩码语言模型（MLM）的方法仅限于预测单点替换突变，无法有效处理插入/缺失（Indels）以及高阶突变中的上位效应（Epistatic effects）。
• 零样本与监督学习的差距：虽然零样本预测（Zero-shot）能评估序列适应度与功能的关联，但在蛋白质工程中，从业者更需要从有限的突变实验数据中学习（Few-shot/Supervised），以优化特定功能。现有模型在小数据集上的泛化能力和数据效率仍有待提高。
• 模型扩展的边际效应：单纯增加模型参数量（Scaling）似乎主要提升了结构预测能力，但对适应度建模和功能预测的提升有限，甚至可能因过拟合导致泛化能力下降，且训练和推理成本高昂。
• 多模态与检索增强的缺失：现有研究要么结合多模态信息（序列 + 结构），要么结合检索增强（Retrieval Augmentation），但鲜有将两者结合以同时利用进化约束和结构信息的模型。

2. 方法论 (Methodology)

作者提出了 PoET-2，这是一个多模态、检索增强的蛋白质基础模型。其核心架构和训练策略包括：

A. 核心架构

• 编码器 - 双解码器结构：PoET-2 采用一个编码器（Encoder）和两个解码器（Decoders）的 Transformer 架构。
  • 编码器：处理用户提供的提示（Prompt），包含一组进化相关的蛋白质（Context）和可选的查询（Query，指定部分序列或结构约束）。编码器采用层级注意力机制（Hierarchical Attention），确保对蛋白质输入顺序的等变性（Equivariance），即无论输入蛋白质的顺序如何，输出表示保持一致。
  • 自回归解码器 (CLM Decoder)：基于因果语言建模目标训练，用于生成新序列和计算序列的对数似然（Log-likelihood）。这使得模型能够处理可变长度的序列（如 Indels）并建模高阶突变的全联合概率分布。
  • 双向解码器 (MLM Decoder)：基于掩码语言建模目标训练，用于生成强大的上下文感知嵌入（Embeddings），捕捉全局依赖关系，适用于监督学习任务。
• 多模态输入：模型同时处理蛋白质序列（Sequence）和结构信息（Structure）。
  • 结构信息包括：残基间的 $C_\alpha$ 距离（离散化为 128 个桶）、局部骨架距离以及预测结构置信度（pLDDT）。
  • 引入了基于结构的注意力偏置（Structure-based Attention Bias），在注意力机制中根据 3D 空间距离而非线性序列位置来调整注意力分数。
• 检索增强 (Retrieval-Augmentation)：通过“上下文（Context）”机制，模型可以动态地“检索”并学习特定蛋白质家族的进化约束，而无需在训练数据中见过该特定序列。这实现了真正的上下文学习（In-context Learning）。

B. 训练目标

模型通过联合优化三个损失函数进行训练：

1. $L_{MLM\_encoder}$：编码器上的标准掩码语言建模损失。
2. $L_{CLM\_decoder}$：自回归解码器的因果语言建模损失（用于生成和评分）。
3. $L_{MLM\_decoder}$：双向解码器的掩码语言建模损失（用于表示学习）。

C. 提示工程 (Prompt Engineering)

• Context：包含同源蛋白质家族的代表性样本，引导模型学习特定家族的进化约束。
• Query：允许用户指定部分序列或结构约束（如全长、信号肽、活性位点或骨架结构），模型仅生成满足这些约束的蛋白质。

3. 关键贡献 (Key Contributions)

1. 首个结合多模态与检索增强的蛋白质基础模型：PoET-2 成功将序列、结构信息和家族特定的进化约束（通过检索增强）统一在一个框架内。
2. 突破 Indels 和高阶突变预测瓶颈：利用自回归解码器，PoET-2 能够直接处理插入和缺失突变，并有效建模高阶突变中的上位效应，这是传统 MLM 模型难以做到的。
3. 极高的数据效率：在监督学习（Few-shot）场景下，PoET-2 的嵌入表示仅需极少量的训练数据（如 10-100 个样本）即可达到甚至超越需要数千样本的现有最先进方法（如 Kermut）。
4. 轻量化与高效性：模型参数量仅为 1.82 亿 (182M)，远小于许多百亿参数模型，但在性能上实现了 SOTA，且推理速度快，对 GPU 资源要求低。

4. 实验结果 (Results)

实验基于 ProteinGym 基准测试，涵盖深度突变扫描（DMS）数据集和临床变异数据集。

• 零样本变异效应预测 (Zero-shot)：

  • Indels 预测：在 DMS Indels 基准上，PoET-2 显著优于所有现有模型，Spearman 相关系数 ($\rho$) 比上一最佳模型 PoET-1 提高了约 0.05，比非 PoET 类最佳模型提高了约 0.10（提升超过 20%）。
  • 高阶突变：在 3 个及以上突变的高阶变异预测中，PoET-2 表现卓越，显著优于 VenusREM 等集成模型。
  • 临床变异：在区分致病和良性突变方面，PoET-2 在替换和 Indels 任务上均建立了新的 SOTA（AUROC 分别提升 0.008 和 0.018）。
  • 单点突变：在单点替换任务上，PoET-2 与当前 SOTA 的集成模型（VenusREM）表现相当，且在识别有益突变的指标（NDCG）上略胜一筹。
  • 集成效果：PoET-2 与 VenusREM 的简单集成模型在所有指标上均超越了单一模型，表明两者捕捉了互补的适应度信号。

• 监督学习 (Supervised Learning)：

  • 数据效率：在 DMS 监督基准上，基于 PoET-2 的高斯过程（GP）模型在所有交叉验证方案（随机、模数、连续）中均大幅超越 Kermut 和 ESM 系列模型。
  • 小样本表现：在训练数据极少（如仅 10 个样本）的情况下，PoET-2 的表现甚至优于使用最大训练集（~2600 样本）的 ESM C 模型，展示了极强的泛化能力。

• 结构条件的作用：

  • 在零样本预测中，引入结构信息（作为 Context 或 Query）能显著提升性能，特别是在稳定性相关的任务中。
  • 在监督学习中，显式的结构条件带来的提升有限，表明 PoET-2 的嵌入表示已经隐式编码了关键的结构信息。

5. 意义与影响 (Significance)

• 填补技术空白：PoET-2 解决了现有 PLMs 无法有效处理 Indels 和高阶突变的问题，为蛋白质工程中的复杂设计任务提供了新工具。
• 降低门槛：其轻量化（182M 参数）和高数据效率特性，使得在计算资源有限或实验数据稀缺的场景下进行高精度的蛋白质功能预测和设计成为可能。
• 范式转变：证明了“检索增强 + 多模态 + 双解码器”架构在蛋白质科学中的巨大潜力，为未来基础模型的设计提供了新的方向，即不再单纯依赖堆叠参数，而是通过更聪明的架构利用进化上下文和结构信息。
• 应用前景：该模型可加速新型酶、治疗性蛋白质和更稳定蛋白质的设计，在药物开发和合成生物学领域具有广泛的应用价值。

总结：PoET-2 是一个高效、多功能的蛋白质基础模型，它通过创新的架构设计，在零样本和高阶突变预测、以及小样本监督学习方面均取得了突破性进展，为蛋白质功能理解和设计树立了新的标杆。