Atom-level Protein Representation Learning Improves Protein Structure Prediction¶

会议: NeurIPS 2026
arXiv: 2605.22133
OpenReview: https://openreview.net/forum?id=iyl2sAzrHx
代码: https://holymollyhao.github.io/TriProRep/
领域: 计算生物
关键词: 蛋白质表征学习、全原子几何、结构预测、VQ-VAE、同源二聚体对接

一句话总结¶

TriProRep 提出了融合氨基酸序列、主链局部几何与残基内全原子几何的三视角蛋白质表征预训练框架，借助轻量生成器构造跨视角扰动并由判别器进行纠错恢复，并在全新提出的 RepSP 结构预测基准上显著提升了同源二聚体柔性对接、残基结合特征预测与单体结构预测蒸馏精度。

研究背景与动机¶

蛋白质语言模型（如 ESM2、SaProt、ProstT5）在酶类别（EC）预测和基因本体（GO）注释等经典功能分类基准上取得了令人瞩目的进展。然而，这些传统任务仅能间接衡量表征对宏观生物功能的捕捉能力，却无法直接检验表征中是否蕴含了下游三维几何建模所必需的残基级精细空间信息。在实际评测中，纯序列模型在 EC/GO 分类上往往能与显式融入三维结构的表征模型打成平手，这导致研究者难以确切评估现有的结构监督是否真正赋予了表征模型用于预测复合物界面、引导折叠或辅助多聚体组装的结构几何能力。

与此同时，现有结构感知蛋白质分词方案（如 Foldseek 3Di 结构字母表、ESM3 结构轨道或 AminoAseed 主链码本）主要将注意力集中在局部主链骨架几何的离散化上，普遍忽略了残基内部重原子的排布与侧链旋转异构态（rotamers）。在生物大分子的相互作用界面和结合口袋中，侧链原子的微观空间构型与堆积紧密程度直接决定了分子识别与结合特异性。缺少侧链全原子几何信息的表征在面对复合物界面建模、接触预测与结构微调时存在天然的信息表达瓶颈。

受到计算机视觉中自监督表征既能充当生成模型条件信号、又能作为特征对齐蒸馏目标的启发，本文探索将预训练蛋白质表征直接迁移并赋能于结构预测任务。核心 idea：联合建模氨基酸序列、主链几何与全原子几何三个对齐的残基视角，通过轻量生成器构造跨视角自洽性破坏并执行稠密纠错恢复预训练，构建直接面向结构预测与复合物对接的高分辨率蛋白质表征模型 TriProRep 及评测基准 RepSP。

方法详解¶

整体框架¶

对于任意输入的蛋白质单体三维结构，TriProRep 首先在残基级别对齐抽取三个视角的离散分词序列：反映序列身份的氨基酸 token、反映局部骨架空间形态的主链几何 token，以及刻画残基内部微观空间排布的全原子几何 token。在预训练过程中，模型引入一个轻量级的 Transformer 生成器，分别在三个视角中独立对掩码位置进行在分布采样替换，制造出“单视角看似合理但跨视角物理冲突”的破坏序列；随后由表征编码器接收该破坏序列，并在全部残基位置和三个视角上计算纠错恢复交叉熵损失，迫使模型学习不同层次结构视角之间的物理约束与互补几何特征。预训练完成后舍弃生成器与预测头，冻结表征模型提取残基隐状态，直接输入 RepSP 基准中的同源二聚体柔性对接模型、残基级属性探测头以及单体结构预测蒸馏框架。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["单体蛋白质输入<br/>氨基酸序列与 3D 原子坐标"] --> B["三视角残基分词<br/>氨基酸 / 主链 / 全原子"]
    B --> C["生成器跨视角扰动<br/>各视角独立掩码与采样替换"]
    C --> D["纠错恢复预训练<br/>表征模型全位置重构原始 token"]
    D --> E["冻结蛋白质表征<br/>残基级多视角融合隐状态"]
    E --> F["RepSP 结构预测任务<br/>二聚体对接 / 结合性质探测 / 单体蒸馏"]

关键设计¶

1. 全原子几何 VQ-VAE 分词：捕获残基侧链旋转异构与精细局部原子排布 以往的主链分词器（如 Foldseek 3Di、ESM3 结构轨道）无法区分不同侧链异构体。本文针对残基内 Atom37 重原子坐标构建局部规范坐标系：以 $\mathrm{N}, \mathrm{C}_\alpha, \mathrm{C}$ 主链原子构建具备 $\mathrm{SE}(3)$ 不变性的局部参考系，提取所有重原子相对 $\mathrm{C}_\alpha$ 的三维空间位移、重原子与 $\mathrm{C}_\alpha$ 的标量距离、离散化主链键角以及侧链二面角 $\chi_1, \ldots, \chi_4$（在 $[-\pi, \pi]$ 区间内均匀离散为 21 个 bin）。随后通过 6 层 Pairformer 架构联合更新单原子表示与 $37 \times 37$ 原子间配对表示，经残基内重原子均值池化后获得连续向量，并映射至大小为 512 的可学习离散码本。分词器训练目标结合了标准 VQ-VAE 损失、成对原子截断距离回归损失 $\mathcal{L}_{\mathrm{pair}}$ 与侧链二面角交叉熵损失 $\mathcal{L}_\chi$。实证分析表明，全原子 token 能在 $\chi_1$ 单纯形空间中精准聚类区分 Trans、Gauche+ 与 Gauche- 旋转异构态，提供了主链分词完全缺失的高分辨率侧链微环境信息。

2. 基于生成器破坏视角的纠错预训练：利用跨视角张力施加稠密结构物理约束 直接对离散 token 应用掩码语言模型（MLM）仅利用了显式 [MASK] 标记，且仅在被掩码的稀疏位置反向传播梯度。TriProRep 借鉴 ELECTRA 机制但将其重构为多类别纠错恢复任务：使用轻量级 Transformer 生成器 $\mathcal{G}$ 对三个视角独立以 0.6 概率掩码的位置预测概率分布，并采样替换候选 token $\hat{x}_i^v$。由于三个视角的替换是彼此独立进行的，生成器输出往往“单视角内符合各自先验但跨视角自相矛盾”（例如替换后的氨基酸种类与当前主链曲率或侧链空间排布产生严重立体碰撞）。表征模型 $\mathcal{E}$ 接收融合投影后的破坏序列 $x^{\mathrm{corrupt}}$，在所有 $L$ 个残基位置和全部 3 个视角上计算原始 token 的恢复交叉熵： $$\mathcal{L}_{\mathrm{rec}} = \mathbb{E}\left[\sum_{v} \sum_{i=1}^L \mathcal{L}_{\mathrm{CE}}\left(p_{\mathcal{E}}^v(x^{\mathrm{corrupt}})_i, x_i^v\right)\right]$$ 这种机制迫使表征编码器深度挖掘序列信息与三维空间原子排布的互补约束，并在整条蛋白质链上提供了稠密的监督信号。

3. RepSP 结构预测表征评测基准：全面考察单体表征对三维复合物推理与生成的支撑能力 针对传统 EC/GO 任务无法评测几何信息利用率的缺陷，RepSP 构建了涵盖三大下游维度的结构化评测体系。第一项为同源二聚体柔性对接，修改基于流匹配的 SimpleFold 架构，仅将单体冻结表征输入去噪骨架，通过 100M 参数的条件流匹配模型预测同源二聚体全原子结构；第二项为残基级二聚体结合属性探测，通过双层 MLP 线性探测残基结合位点、溶剂可及表面积变化 $\Delta\mathrm{SASA}$、Levy 区域分级（core, rim, support 等）及 5 种非共价键相互作用类型；第三项为单体结构预测蒸馏，将单体表征作为密集对齐目标（REPA），最大化流匹配骨架中间层表示与预训练表征的余弦相似度： $$\mathcal{L}_{\mathrm{REPA}} = -\lambda_{\mathrm{REPA}} \frac{1}{|\mathcal{V}|} \sum_{i \in \mathcal{V}} \frac{\langle \hat{h}_i, h_{\mathrm{tgt}, i} \rangle}{\|\hat{h}_i\|_2 \|h_{\mathrm{tgt}, i}\|_2}$$ 从而直接检验预训练表征是否能作为指导结构生成的有效先验。

损失函数 / 训练策略¶

TriProRep 预训练共分两阶段：分词器预训练采用 AdamW 优化器，学习率 $5 \times 10^{-4}$，在 AlphaFoldDB 单体结构上完成码本训练；表征模型预训练在 83.7M 蛋白质结构（18.2M AFDB UniRef30 代表结构 + 65.5M ESMAtlas 高置信结构）上进行，序列裁剪长度为 512，有效批量大小为 1024，在 NVIDIA B200 集群上训练 500K 步。优化采用 AdamW，峰值学习率 $4 \times 10^{-4}$，权重衰减 0.01，$\beta = (0.9, 0.95)$，结合 bf16 混合精度与梯度截断 1.0。

实验关键数据¶

主实验¶

主实验包含同源二聚体柔性对接评测（RepSP Task 1）以及单体结构预测蒸馏评测（RepSP Task 3）。

模型	参数量	DockQ ↑	iRMSD (Å) ↓	LRMSD (Å) ↓	Fnat ↑	TM-score ↑	LDDT ↑	RMSD (Å) ↓
ESM2 (Small)	35M	0.278	11.549	25.942	0.417	0.531	0.603	16.176
SaProt (Small)	35M	0.296	10.380	25.308	0.414	0.567	0.712	14.660
TriProRep (Ours)	35M	0.371	8.915	21.964	0.486	0.633	0.771	12.778
ESM2 (Medium)	150M	0.310	10.743	25.328	0.450	0.563	0.667	15.170
TriProRep (Ours)	150M	0.419	8.170	20.292	0.512	0.666	0.812	11.874
ESM2 (Large)	650M	0.374	9.517	22.803	0.533	0.613	0.721	13.913
SaProt (Large)	650M	0.443	7.449	18.878	0.548	0.690	0.807	11.196
S-PLM (Large)	704M	0.366	9.350	22.678	0.507	0.619	0.708	13.591
MIF-ST (Large)	643M	0.299	10.598	25.420	0.384	0.563	0.818	14.521
TriProRep (Ours)	650M	0.477	6.883	18.243	0.588	0.705	0.829	10.554
ESM2 (Huge)	3B	0.387	8.679	21.455	0.547	0.635	0.732	13.025
ESM3 (Huge)	1.4B	0.448	7.060	18.554	0.579	0.695	0.835	10.719
ProstT5 (Huge)	1.2B	0.409	7.790	20.403	0.541	0.661	0.782	11.843
TriProRep (Ours)	2.8B	0.499	6.370	17.087	0.612	0.724	0.834	9.875

单体结构预测蒸馏对齐（REPA）实验（Large 规模基线对比）：

REPA 对齐目标	参数量	TM-score ↑	GDT-TS ↑	GDT-HA ↑	LDDT ↑	LDDTbb ↑	RMSD (Å) ↓
无对齐基线 (SimpleFold)	–	0.5422	0.3432	0.1990	0.6208	0.6746	12.98
SaProt	650M	0.5659	0.3677	0.2192	0.6410	0.6947	12.26
S-PLM	704M	0.5569	0.3559	0.2081	0.6301	0.6838	12.40
MIF-ST	643M	0.5491	0.3520	0.2075	0.6310	0.6827	12.73
TriProRep (Ours)	650M	0.5822	0.3858	0.2344	0.6513	0.7051	11.91

消融实验¶

预训练数据与预训练目标消融（150M 规模残基级二聚体结合属性探测）：

模型配置	结合位点 AUPRC	结合位点 AUROC	$\Delta\mathrm{SASA}$ Pearson	$\Delta\mathrm{SASA}$ Spearman	Levy tier Macro F1	Levy tier Acc.	作用力类型 AUPRC	作用力类型 AUROC
纯序列基线 (ESM2)	0.7104	0.9066	0.6376	0.3701	0.5701	0.7972	0.3534	0.9542
仅 AFDB 数据 ($\text{Ours}_{\mathrm{AF}}$)	0.7759	0.9314	0.7285	0.4344	0.6458	0.8470	0.3818	0.9576
标准 MLM 目标 ($\text{Ours}_{\mathrm{MLM}}$)	0.7677	0.9297	0.7309	0.4380	0.6424	0.8471	0.3748	0.9588
完整模型 (TriProRep)	0.8156	0.9432	0.7797	0.4662	0.6819	0.8668	0.4184	0.9644

关键发现¶

结构感知优于纯参数缩放：在同源二聚体柔性对接中，TriProRep-650M 的 DockQ（0.477）显著超越了参数量大数倍的纯序列模型 ESM2-3B（0.387）及结构模型 ESM3-1.4B（0.448），证明在缺乏精细几何表征的情况下单靠扩大参数量无法弥补几何归纳偏置的匮乏。
纠错恢复带来关键增益：消融实验显示，将生成器采样的纠错目标替换为传统 MLM 目标后，结合位点 AUPRC 从 0.8156 下降至 0.7677，$\Delta\mathrm{SASA}$ 相关系数从 0.7797 降至 0.7309，证实跨视角自洽性扰动驱动的稠密纠错能迫使模型学习更为本质的几何物理约束。
零样本泛化至真实晶体结构与异源结合：在 RCSB 真实晶体二聚体（2023年6月后存入）上，TriProRep-2.8B 的高质量对接率达到 13.5%，远超 ESM2-3B（6.0%）和 ESM3-1.4B（8.7%）；在 PPIRef50K 异源复合物界面探测中同样保持领先。

亮点与洞察¶

全原子几何分词补全侧链盲区：传统结构模型多采用 3Di 或主链骨架分词，忽视了决定结合特异性的侧链旋转异构态；TriProRep 引入 $\mathrm{SE}(3)$ 不变的全原子 VQ-VAE，让模型在单体阶段便内嵌了高分辨率侧链微环境。
跨视角破坏纠错机制极其巧妙：通过轻量生成器对不同视角独立替换，自然制造出“各视角局部合理但跨视角物理冲突”的困难负例，使表征判别器在全局自洽性恢复中建立起坚实的几何物理关联。
建立结构生成导向的表征评测闭环：突破了仅依赖 EC/GO 粗粒度分类的旧评价范式，提出 RepSP 基准，验证了预训练特征可作为复合物对接的条件输入与生成模型的稠密蒸馏先验。

局限与展望¶

预训练与评测基准主要基于 AFDB 预测结构：虽然在 RCSB 真实结构上验证了趋势的一致性，但预测结构的固有误差仍可能对极端精细构象造成系统性偏差，未来需在更大规模高精度实验解析结构上展开验证。
评测主要集中于同源二聚体：同源二聚体去除了伴侣配对多样性的干扰，是绝佳的受控测试场，但在多聚体与通用异源蛋白复合物上的柔性组装与界面建模仍待进一步扩展。
全原子分词器计算开销相对较大：Pairformer 与多视角对齐引入了额外的离散化预处理流程，对超长链蛋白质的大规模预处理效率提出了更高要求。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次提出融合全原子侧链几何与生成器跨视角纠错的三视角蛋白质表征学习框架。
实验充分度: ⭐⭐⭐⭐⭐ 构建百万级 RepSP 基准，涵盖柔性对接、残基探测与单体蒸馏，并在真实晶体结构上完成全面检验。
写作质量: ⭐⭐⭐⭐⭐ 动机阐述清晰，问题抽象到位，实验设置详实严谨。
价值: ⭐⭐⭐⭐⭐ 为结构生物学与生成式大模型架构的融合提供了重要的表征范式与评测标准。