arXiv 2025 | Protenix-Mini: 高效结构预测器——紧凑架构、少步扩散与可切换pLM

该研究提出 Protenix-Mini，通过紧凑架构与2 步 ODE 采样显著降低生物分子结构预测的计算成本，仅损失 1–5% 精度。作者发现 AF3 模型存在冗余模块，可裁剪 Pairformer 与 Diffusion Transformer 块，并用 ESM 蛋白质语言模型替代高开销的 MSA 模块。实验表明，Mini 在 RecentPDB 与 Posebusters 上表现接近完整模型，Tiny 成本降约 85%。ESM 版本在大部分任务中接近 MSA，但蛋白–蛋白界面预测显著下降。比较 EDM 与 Flow Matching，性能差异极小。针对少步 ODE 的塌陷与碰撞问题，调整步长策略可提升质量并保持多样性。未来将探索稀疏注意力、更优架构及结构增强型 pLM 预训练，实现长序列与多组分体系的高效预测。

获取详情及资源:

论文: https://doi.org/10.48550/arXiv.2507.11839
代码和数据: https://github.com/bytedance/Protenix

0 摘要

轻量化推理对于生物分子结构预测及其他下游任务至关重要，它能够在大规模应用中实现高效的实际部署与推理阶段的可扩展性。本研究针对模型效率与预测精度之间的平衡，提出了几项关键改进：首先，将多步的 AF3 采样器 替换为仅需少量步骤的 常微分方程（ODE）采样器，显著降低了推理时扩散模块的计算开销；其次，在开源 Protenix 框架中发现，一部分 pairformer/diffusion transformer 模块对最终结构预测并无显著贡献，因此可进行架构裁剪与轻量化重设计；最后，引入 ESM 模块 取代传统的 MSA 模块，从而减少多序列比对的预处理时间。

基于这些核心洞见，研究提出了 Protenix-Mini——一种紧凑且优化的高效蛋白质结构预测模型。该模型采用了更加高效的架构设计与两步ODE采样策略，在去除冗余Transformer组件并精简采样流程的同时，显著降低了模型复杂度，仅带来极小的精度损失。在基准数据集上的评估表明，与完整版本相比，其性能下降仅为 1–5%，但推理速度与资源利用率得到大幅提升。这使得 Protenix-Mini 特别适合于计算资源受限但仍需高精度结构预测的应用场景。

图1 展示了在 RecentPDB 数据集上平衡模型计算开销与性能的关键配置。右侧面板中，为了优化效率，Mini 与 Tiny 两个变体均减少了 MSA 模块和 pairformer 模块的数量，并将扩散步骤截断为仅两步。值得注意的是，这种设计带来的性能下降极为有限，并且权衡经过精心校准，以保证功能精度的基本不变。论文中所有数值评估均基于测试集中长度少于768个tokens的蛋白质。更详细的实验信息见表8。右侧面板还比较了不同模型在不同token长度与MSA规模下的FLOPs情况：在左图中，MSA数量固定为2048；而在右图中，token长度固定为384。两种情况下，原子数均固定为8832。这种对比清晰地表明，在保持结构预测可用性的同时，通过减少模块数量和扩散步数，可以显著降低计算量，从而在资源受限的环境下实现更高的推理效率。

1 高效采样与轻量化设计的核心思路

少步ODE的高效采样

近年来，基于分数的生成模型研究表明，扩散过程可用常微分方程（ODE）进行近似，从而实现确定性采样，跳过传统随机微分方程（SDE）求解器中的随机性。分析发现，无论是使用 EDM 框架（如 AlphaFold3）还是 flow matching 框架训练的模型，都对采样步数的大幅削减表现出惊人的鲁棒性。仅使用2步ODE采样，依然能够生成结构精确的蛋白质构象，挑战了“高保真采样必须依赖几十次迭代”的传统认知。

冗余组件的压缩

在对开源 Protenix 框架（基于扩散的结构预测器）分析时发现，其部分 pairformer/diffusion transformer 模块对最终预测结果几乎没有贡献。实验证明，在后训练阶段推理时直接移除这些模块，只会带来轻微的性能下降。通过识别并删除这些非关键组件，重新配置模型架构，可有效降低计算量。正如图1所示，在不同token数量和MSA规模下，FLOPs均有明显下降。

用pLM替代MSA模块

考虑到MSA搜索与计算模块的高计算成本，提出用**预训练蛋白质语言模型（pLM）**生成的嵌入替代MSA特征，以降低开销。具体实现中，使用 ESM2-3B 模型提取蛋白质表示，完全跳过耗时的MSA预处理与计算步骤。

基于以上洞见，提出了Protenix-Mini：在减少Transformer模块数量、采样步数（如1–2步）的同时，实现了对复杂生物分子结构的高效预测。实验结果表明，在 interface LDDT、complex LDDT、ligand RMSD 成功率 等指标上，性能下降仅为 1–5%，但推理速度与资源效率大幅提升。

这一成果证明了轻量化模型在生物分子结构预测中可与大型模型媲美。通过同时优化采样效率与架构开销，Protenix-Mini 向高保真结构预测的普及化迈出了一步。未来，作者计划结合架构设计、蒸馏、量化等方法进一步提升推理效率。

2 推理流程概述

本节简要介绍了 AF3 风格结构预测模型（如 [1, 20, 21, 24]）的架构与推理步骤，其整体流程（见图2）由**条件编码（Conditioning）与扩散生成（Diffusion）**两个核心阶段组成。

① 条件编码（Conditioning）

模型首先计算条件信号。设输入序列为 $S e q$ ，多序列比对结果为 $M S A$ 。将 $M S A$ 与序列嵌入输入到 MSA 模块 与 Pairformer 中，循环 $M$ 次后，得到逐token表示与成对表示，分别记为 $s$ 与 $z$ ：

(s, z) = C o n d (S e q, M S A)

这两个表示将作为后续扩散模块的条件输入。

② 扩散生成（Diffusion）

扩散过程通过逐步去噪，将初始随机样本生成目标的三维原子结构。设 $x_{1}$ 为真实结构，初始状态 $x_{0} \sim N (0, σ_{d a t a}^{2} I)$ 为高斯噪声。与标准扩散模型相同，结构生成遵循正向–反向范式：

x_{1} \overset{forward}{\to} x_{0}, x_{0} \overset{reverse}{\to} {\hat{x}}_{1} \approx x_{1}

其中，正向过程通过逐步加噪扰动真实结构，反向过程则利用训练好的去噪器恢复原始信号。

在 AF3 中，这一过程由采样算法实现（见算法1），遵循 EDM 框架 [6]：在每一步迭代中，按照预设的时间表向当前样本 $x_{t}$ 注入噪声 $ε_{t}$ ，Diffuser 模块预测去噪结果，然后计算速度项（velocity），并按步长参数 $η$ 对 $x_{t}$ 进行更新。更多实现细节见 [1] 的补充材料。

3 轻量化结构预测器

3.1 通过采样器配置实现少步扩散

在基于扩散的结构生成中，最大的性能瓶颈之一是长采样轨迹带来的高计算开销。AF3 风格模型默认使用多达 200 步的采样器 [1]。然而，这样的长采样计划在推理时是否真的必要？

令人意外的是，答案是否定的——只要采样算法配置得当，采样步数就可以大幅减少而不显著影响性能。研究发现，即使在无需重新训练的情况下，AF3 风格模型在采样步数大幅缩减时依然表现出极强的有效性。例如，200 步采样可以直接替换为仅 2 步 ODE 采样，仍能获得接近完整版本的预测效果；甚至在单步推理下，也能给出合理的结构预测，仅伴随轻微的性能下降。

这种结果颠覆了“高保真生成必须依赖长迭代链”的传统认知，为扩散模型的高效推理提供了直接可用的途径。

图2 展示了 AF3 风格模型架构的整体概览。模型分为两个阶段：首先是条件编码阶段（Diffusion Cond），将输入特征转化为潜在表示；随后进入扩散生成阶段（Diffusion Noise），从随机初始化的结构出发，逐步生成最终的三维原子结构。

在直接将默认的 AF3 采样器 用于较少推理步数时，实验发现当步数低于 10 时，模型性能会急剧崩溃——生成的结构往往残缺或不合理。

进一步分析表明，这种退化并非源于模型本身的能力不足，而是由于默认采样算法并不适合少步推理。通过对采样过程的系统分析，研究找出了两项简单但关键的改动，使得少步推理成为可能：

切换为 ODE 采样器：将 $γ_{0} = 0$ ，从而移除额外的噪声注入；
设置步长比例 $η = 1.0$ ：确保与底层基于速度（velocity-based）的公式保持一致性。

图3 展示了步长比例与采样步数对模型性能的影响。图中为 ODE 采样器在不同配置下的表现趋势，纵轴为 complex LDDT 分数，横轴为推理阶段的扩散步数。

影响少步推理性能的核心因素之一是步长比例 $η$ ，该参数最初在 AF3 中作为对 EDM 采样器 [6] 的改动引入。EDM 默认使用 $η = 1$ ，而 AF3 将其提高到 $η = 1.5$ 。虽然原始论文并未明确解释为何选择 $η > 1$ ，但先前研究（如 Auto Guidance [7]）提出，在低噪声阶段（ $t \to 1$ ）去噪器往往会低估速度（velocity）幅值，因此适当放大步长比例有助于提高长采样链中的生成质量。然而，本研究发现这一策略在少步采样中反而有害：当采样步数很少时，较大的 $η$ 会导致更新过程不稳定，从而降低预测精度。实验结果表明，在少步采样场景下， $η = 1.0$ 是必要条件。这一交互效应在图3中得到了直观体现。

图4 对比了不同采样器下的模型性能。AF3 采样器：参数设定为、、；ODE 采样器：参数设定为、、。

图4 对比了默认 AF3 采样器（ $η = 1.5, λ = 1.003, γ_{0} = 0.8$ ）与改进后的 ODE 采样器（ $η = 1.0, λ = 1.0, γ_{0} = 0$ ）在不同采样步数下的性能表现。结果表明，默认配置在 200 步或 20 步时表现优异，但一旦采样步数降至 10 以下就会性能崩溃，在 5 步时几乎完全失效。相比之下，ODE 配置在极少步数下依然保持稳定，即使只有 1–2 步也能生成高质量结构。例如，2 步 ODE 采样器在配体–蛋白界面上的 LDDT 得分为 0.645，几乎与 200 步基线的 0.65 持平；而单步版本也能达到 0.64，表现仍处于合理范围。

虽然少步推理在平均性能上表现优异，但仍会偶尔出现结构问题，例如物理合理性下降或原子碰撞等现象，具体示例见第4.2.3节。为了完整性，研究还在第4节评估了使用 flow matching 训练的模型，结果显示其在少步采样中同样具有较强的鲁棒性。

3.2 压缩 Protenix 模型中的冗余模块

表1 表明，在进一步微调后，小规模模型依然可以获得与原模型接近的性能。具体而言，微调过程使用批量大小 256、学习率，并进行 10K 次迭代，即可在保持较低计算成本的同时恢复大部分预测精度。

AF3 的架构包含 48 个 Pairformer 块（建模 token 两两间的距离关系）以及 24 个 Diffusion Transformer 块（生成三维坐标）。虽然这种深层网络在捕捉复杂结构依赖上表现出色，但其计算开销巨大，也引发了架构冗余的疑问——是否每一层都对最终预测有实质贡献。类似现象在先前研究中也有报道。

图5 展示了 Protenix 在后训练阶段进行模块裁剪的性能表现。结果表明，即使在不进行微调的情况下，移除部分模块也仅会带来轻微的性能下降，与完整模型相比差距很小，这进一步验证了模型架构中存在一定的冗余空间。

为此，作者进行了逐块消融实验，按顺序移除 Pairformer 网络的早期模块。结果显示，移除前 4 个 Pairformer 块在基准数据集上仅带来轻微精度下降（见图5）；进一步如表1所示，在对剩余参数进行微调后，即使去掉若干块也不会显著影响性能。这一发现促使作者尝试更紧凑的架构配置。

基于此，提出了两种训练方案：

剪枝 + 微调：移除部分模块后微调剩余模型；
直接缩小架构并从零训练。

通过架构搜索，发现 Protenix-Mini 在 LDDT 分数与效率之间达成了较优平衡（见图1与表8）。此外，作者还进行了更激进的裁剪——移除最靠近输入的 前 8 个 Pairformer 块，在 800 次迭代微调后，complex LDDT 下降约 10%（从 0.80 降至 0.72），并由此得到更小的 Protenix-Tiny 模型。详细训练配置见表8。

3.3 基于 pLM 的 Protenix-Mini

由于 MSA 搜索与计算往往需要大量内存与并行计算资源，作者提出用**预训练蛋白质语言模型（pLM）**替代传统的 MSA 流程，以显著降低计算负担。

在架构上，首先将输入序列转化为嵌入表示，再通过线性变换层调整维度。在实现中，采用 ESM2-3B 模型生成上下文相关的序列嵌入，并将其添加到结构预测流水线中 input embedder 模块的 $s_{inputs}$ 中；训练与推理时均传入空的 MSA 输入字典，彻底屏蔽 MSA 模块的作用。

表2 展示了模型在 RecentPDB（≤768）测试集上的性能表现。结果为随机采样多次取平均值，在样本选择时未使用置信度分数进行筛选，并分别报告了不同界面类型的 interface LDDT 得分。

表3 则给出了 ESM 模型在同一 RecentPDB（≤768）测试集上的性能表现，同样是对多次采样结果取平均，并且在样本选择过程中不依赖置信度分数，以保证结果的客观性与可比性。

为促进 MSA 模型与 ESM 架构之间的知识迁移，作者设计了混合训练策略：在每次训练迭代中，以 50% 概率选择使用 MSA 模块或 ESM 特征，两条路径共享所有核心组件（如原子编码器、原子解码器、Transformer 块、Diffusion Transformer 块），保证特征表示一致性。该方法可视作一种隐式知识蒸馏，即 ESM 模块在无显式监督的情况下学习逼近 MSA 模块的输出分布，从而生成紧凑但信息丰富的嵌入。

在推理阶段，完全移除 MSA 模块，仅使用 ESM 模块 进行预测，从而实现高效的端到端结构生成。

4 实验结果

4.1 模型配置与训练细节

基于前文分析，作者训练了 Protenix-Mini 模型，配置为 16 个 Pairformer 块、8 个 Diffusion Transformer 块、1 个 MSA 模块块、1 个原子解码器块、1 个原子编码器块，其余训练与架构参数与开源 Protenix [20] 保持一致。模型分别在 EDM 和 flow matching 框架下训练，性能相近。

Protenix-Mini 的训练参数为 batch size 64、200K 次迭代、学习率 $10^{- 3}$ 。在获得训练完成的 Mini 模型后，进一步裁剪最靠近输入的 8 个 Pairformer 块并微调剩余权重，得到更小的 Protenix-Tiny 模型（batch size 64、100K 次迭代）。

此外，将 Mini 模型加载权重后，将 MSA 置空，引入 ESM2 表示作为输入的一部分，得到 Protenix-Mini-ESM 模型，并额外训练 100K 次迭代。推理时，骨干网络循环 4 次，扩散模块采用 2 步 ODE 采样（Alg. 2）以节省计算。评估指标为 complex LDDT 与 interface LDDT，数据集为 token 数少于 768 的 RecentPDB 子集。

表4 展示了 Flow 模型在 RecentPDB（≤768）测试集上的性能，其中 Flow 表示使用 flow matching 损失训练的模型。结果表明，Flow Matching 与 EDM 框架在各指标上的差异极小。

表5 给出了各模型在 Posebusters（≤768）测试集上的表现。结果为随机采样取平均，样本选择不使用置信度分数。成功率定义为 RMSD ≤ 2 的比例。结果显示，Mini 模型与完整 Protenix 模型接近，Tiny 模型略有下降。

4.2 实验结果

在 RecentPDB（≤768） 与 Posebusters（≤768） 数据集上的表现如下：

Mini 模型：在 complex 与 interface LDDT 上比标准 Protenix 略低；
Tiny 模型：各类界面 LDDT 降低约 2–3%，但相比 Protenix 推理计算量减少约 85%（见图1）；
Posebusters 数据集：以平均配体 RMSD 与 RMSD ≤ 2 的成功率为指标，Mini 与完整 Protenix 表现接近，Tiny 略差（见表5）。

表6 统计了 Posebusters 数据集上的原子碰撞（Clash）数量。每个模型生成 5 × 5 个样本，并验证是否存在碰撞：符号 ∃ 表示在样本集中存在碰撞，∀5×5 表示所有样本均有碰撞。结果显示，少步 ODE 与多步 AF3 采样器在 Mini 模型上都会出现碰撞问题，但完整 Protenix 在多步采样下表现更好。

表7 展示了 RecentPDB（≤768）测试集上 Protenix 模型的采样多样性。该表量化了 5 × 5 样本中表现最佳与最差结果的差距，分析结果确认 ODE 采样器在保持结构多样性方面并无退化。

4.2.1 ESM 模型性能

表3显示，Protenix-Mini-ESM 在多数测试指标上略低于 MSA 版本，例如 complex LDDT 从 0.80 降至 0.775；但在蛋白–蛋白界面上降幅显著，从约 0.50 降至 0.40，下降超 10%。这凸显了成对 MSA 特征在蛋白–蛋白界面预测中的关键作用。未来计划开发结构感知型蛋白语言模型以提升此类任务精度。

4.2.2 Flow Matching 模型性能

在 flow matching 框架下，将 $x$ 预测网络替换为速度预测网络，并调整 $t$ 加权函数。训练损失包括：

L_{m s e} = | x_{1} - x_{0} - D i f f u s e r (x_{t}, t ∣ C) |^{2}

并通过

{\hat{x}}_{1} = x_{t} + (1 - t) \cdot D i f f u s e r (x_{t}, t ∣ C)

计算附加损失（平滑 LDDT、键长损失等）。训练时 $t \sim B e t a (2.5, 2.5)$ ，在 $t \approx 0.5$ 处密度最高，以重点学习中期扩散阶段的去噪能力，并减少极端噪声或平凡映射下的不稳定性。

表4显示，Flow Matching 与 EDM 在不同交互类型上的性能差异极小（差值 -0.005 至 +0.009，Complex LDDT 差仅 0.013），说明二者对最终性能影响不显著。

此外，在 2 步 ODE 与 多步 AF3 采样器对比中（5 个随机种子 × 每种 5 个扩散样本）：

Mini EDM 模型：两种采样器均会出现原子碰撞（clash），但增加采样数量可缓解；
完整 Protenix 模型：2 步 ODE 仍有碰撞问题，而多步 AF3 在这方面表现更好（见表6）。

4.2.3 案例分析

可视化对比发现，2 步 ODE 相比多步 AF3 更易生成塌陷的配体结构（见图6）。将 ODE 步数增至 10 可明显缓解此问题，推测原因是 $t \leq 0.5$ 区域下模型欠拟合。实际操作中，前 5 步设置 $η = 1$ ，后 5 步设置 $η = 1.5$ 可在无碰撞的情况下得到高质量结果。

在多样性评估中，两种采样器的最佳与最差样本差异几乎一致（见表7），说明少步 ODE 在保持采样多样性的同时，只需合理调节步长就能取得与多步采样相当的性能。

图6 显示了 2 步 ODE 采样器偶尔生成塌陷结构的案例，例如结构 “7bnh”。

5 结论与未来方向

本文提出了 Protenix-Mini，一种将紧凑架构设计与 2 步 ODE 采样器相结合的轻量化框架，用于高效的生物分子结构预测。在关键基准上仅带来 1%–5% 的性能损失，却显著提升了推理效率。通过裁剪冗余 Transformer 组件并引入稳定的损失函数，验证了在扩散模型中平衡精度与速度的可行性，为其在真实应用中的部署提供了重要途径。同时，提出了基于 ESM 的变体以替代 MSA 搜索模块，在效果与效率间取得新的平衡。

未来的优化方向包括：

探索稀疏或自适应注意力机制，在处理大型生物分子复合物时优化 Transformer 主干，缓解其二次复杂度问题；
深入搜索架构细节配置中的最佳平衡点；
在蛋白质语言模型（pLM）预训练中引入结构信息。

这些工作旨在实现可扩展、低延迟的长序列与多组分系统结构预测，推动药物发现与合成生物学的进一步发展。