A Systematic Evaluation of Co-folding Model Representations for Small-Molecule Learning¶

会议: NeurIPS 2026
arXiv: 2602.13249
OpenReview: https://openreview.net/forum?id=9cpiiqZlW0
代码: https://github.com/hsjang0/boltz-as-FM
领域: 其他
关键词: 分子表示学习, 蛋白质-配体共折叠, 基础模型, ADMET性质预测, 分子生成与优化

一句话总结¶

本文系统评估并将先进蛋白质-配体共折叠模型 Boltz2 的原子级配体表示迁移到独立小分子下游任务中，证明结合态相互作用监督赋予了小分子强大的结构表征力，并在 ADMET 预测、分子生成蒸馏和强化学习结构引导配体优化中取得了优异表现。

研究背景与动机¶

在小分子与药物发现领域，基础模型已成为预测 ADMET 性质和加速分子设计的核心工具。然而，现有绝大多数小分子基础模型（如 MolE、MiniMol、QIP、UniMol 等）均基于孤立分子数据进行自监督或监督预训练，其训练监督信号主要来源于掩码原子预测、离线 3D 构象优化、高通量生物活性筛选标签或量子化学性质回归。虽然这些方法在特定属性预测上展现了优良泛化性，但它们普遍将小分子作为完全孤立的物理实体建模，缺乏与生物大分子相互作用的上下文感知。

与自然语言和计算机视觉领域常常受益于多模态对齐或关系型弱监督不同，小分子领域天然存在一种富含交互物理约束的数据形态——蛋白质-配体共折叠结构。近年来以 AlphaFold3 和开源的 Boltz2 为代表的新一代共折叠模型，将预测粒度从残基级细化到原子级，在学习预测复合体三维坐标的过程中，其 Pairformer 网络骨干隐式编码了极其丰富的结合态氢键、空间位阻、静电吸引和疏水相互作用模式。然而，由于学界对共折叠模型究竟是真正习得物理相互作用还是仅拟合进化多序列比对（MSA）规律尚存争议，此前从未有人系统探究过共折叠模型学到的原子级配体表示能否脱离蛋白质直接迁移到独立小分子任务中。

本文正是在这一空白下展开研究：如果直接截取共折叠模型中已经预训练好的配体原子与原子对表征，它能否作为高质量的开箱即用小分子基础模型？核心 idea：将蛋白质-配体共折叠模型 Boltz2 重塑为小分子基础模型，提取其 Pairformer 骨干中的原子级单体与原子对表征，在无蛋白质输入的独立小分子预测、扩散模型表征对齐蒸馏以及强化学习结构引导配体优化三类任务中实现多维度迁移与表征互补。

方法详解¶

整体框架¶

本文的方法流程涵盖表征提取适配、静态下游属性探测、生成扩散模型表征对齐蒸馏以及在线结构引导强化学习优化四大环节。首先，针对独立小分子任务，作者修改了 Boltz2 的输入接口（利用伪蛋白占位符并切除相应索引），仅输入小分子的 SMILES，经 UFF 初始化三维构象后送入 64 层 Pairformer 主干网络中。随后，抽取多层原子对表示，经由混合统计池化层构建分子级特征向量，输入 MLP 探测头预测 ADMET 属性。在生成任务中，利用冻结的 Boltz2 表征作为监督目标，指导扩散生成模型去噪骨干的隐层特征对齐；在线结构引导生成中，则将 Boltz2 作为白盒奖励教师，同步输出结合亲和力标量奖励与配体原子级表征，以密集对齐信号联合指导 GFlowNet 策略学习。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入小分子 SMILES"] --> B["构象初始化与前向传播<br/>UFF 构象 + Boltz2 64层 Pairformer"]
    B --> C["分层原子对提取与混合统计池化<br/>抽取第 16/32/48/64 层并统计聚合"]
    C --> D["属性预测线性探测<br/>拼接输入特征过 MLP 预测 ADMET"]
    B --> E["生成模型表征对齐蒸馏<br/>余弦相似度监督扩散去噪隐层"]
    B --> F["白盒结构引导强化学习优化<br/>标量亲和力奖励 + 配体隐层密集表征蒸馏"]

关键设计¶

1. 分层原子对提取与混合统计池化：突破共折叠模型无全局池化的表征适配 Boltz2 最初设计用于原子级三维结构预测，其 Pairformer 内部维护 384 维的单实体表示和 128 维的原子对表示，但模型本身并不包含用于提取固定长度分子级向量的池化层。若直接采用简单的全局均值池化，容易将局部关键化学环境（如关键药效团或代谢反应位点）平均化稀释。为此，本文设计了跨深度的分层抽取与混合统计池化方案：从 64 层 Pairformer 的第 $16, 32, 48, 64$ 层（分别对应网络深度 1/4、2/4、3/4 与末层）提取原子对表示矩阵，捕捉从低级局部拓扑到高级全局交互的多层次几何特征。对提取出的原子对矩阵，分别针对对角元（单个原子的自相互作用）、成键原子对（分子化学键约束拓扑）以及全矩阵所有条目三种区域计算统计池化向量（均值与方差），最终拼接成高维固定长度向量送入探测网络。探测网络采用带有残差跳跃拼接特征的 MLP，在保留多尺度原始结构信息的同时兼顾下游分类与回归任务的拟合。

2. 生成模型表征对齐蒸馏：以共折叠几何先验加速分子扩散训练 在小分子图扩散生成模型（如基于 Pairformer 参数化的 GruM）训练过程中，从纯噪声中恢复分子拓扑结构通常面临收敛缓慢、化学有效性与骨架多样性难以兼顾的问题。本文将表征对齐蒸馏范式引入分子图生成领域，将预计算的干净分子 Boltz2 拼接表征 $f(m)$ 作为师资信号。在去噪网络中间层（第 2 层）引入轻量级双层 MLP 蒸馏适配头 $h_\theta(\cdot)$，将加噪分子 $m_t$ 的隐层原子与原子对表征映射到 Boltz2 空间，并通过余弦相似度施加联合训练损失： $$\mathcal{L}(m_t, m) = \mathcal{L}_{\text{GruM}}(s_\theta(m_t), m) - \lambda \cdot \cos\big(h_\theta(m_t), f(m)\big)$$ 实验中设定 $\lambda = 4$。这种密集表征层面的几何与化学环境对齐，为去噪轨迹提供了连续、平滑的梯度引导，在不破坏扩散生成多样性的前提下使模型训练速度提升 2 倍，且同步提升了分子的有效性（Validity）与 Fréchet ChemNet 距离（FCD）。

3. 白盒结构引导强化学习优化：密集表征监督弥补稀疏标量奖励 在基于蛋白质靶点的结构引导配体生成（SynFlowNet-Boltz）中，传统范式仅将 Boltz2 视为黑盒打分器，仅回传标量亲和力奖励 $r$。标量奖励信号在离散大化学空间中高度稀疏，策略网络极难进行精确的信用分配（credit assignment）。本文首次提出白盒结构引导强化学习，在 Boltz2 评估生成分子 $m$ 亲和力的同时，直接导出副产物——受靶标蛋白口袋调制的 Boltz2 配体单原子表征 $f(m)$。随后，在生成策略网络（4 层 Graph Transformer）的第 2 层隐状态 $h_\theta(m)$ 与 $f(m)$ 之间构建辅助对齐损失： $$\mathcal{L}_{\text{align}}(m, f(m)) = -\lambda \cdot \cos\big(h_\theta(m), f(m)\big)$$ 在异步强化学习框架下，策略网络最小化 $\mathcal{L}_{\text{SynFlowNet}} + \mathcal{L}_{\text{align}}$。密集的原子级表征监督将配体在口袋中诱导结合的化学特征直接注入生成策略，大幅提升了对高亲和力分子及其骨架模式（Modes）的采样效率。

4. 基于局部邻域的表征互补性表征分析与融合策略 为了解开共折叠表征与传统单体分子模型之间的内在联系，本文采用中心化核最近邻对齐度量（Centered Kernel Nearest-Neighbor Alignment, CKNNA）定量刻画不同表征空间的局部几何一致性： $$\text{CKNNA} = \frac{\text{ALIGN}\big(\{f(m_i)\}, \{g(m_i)\}\big)}{\sqrt{\text{ALIGN}\big(\{f(m_i)\}, \{f(m_i)\}\big) \cdot \text{ALIGN}\big(\{g(m_i)\}, \{g(m_i)\}\big)}}$$ 其中对齐核函数仅在两个表征空间中互为 $k$ 近邻的分子对上激活。分析发现，Boltz2 与 MolE、UniMol2、MiniMol 等表现出相对较低的 CKNNA 对齐度，证明共折叠学习到了截然不同的特征空间。更有趣的是，当进行表征集成时，与 Boltz2 表征对齐度更弱的 MolE（$\text{Boltz2}+\text{MolE}$）在 CYP 酶抑制等多个下游基准上获得的增益显著超过了单模性能更强但对齐度较高的 MiniMol，揭示了在分子模型集成中“弱对齐带来强互补”的实用准则。

损失函数 / 训练策略¶

对于属性预测，采用带有早停（Patience = 25 轮）和余弦退火学习率调度的 AdamW 优化器，学习率在 $\{1\text{e-}4, 3\text{e-}4, 5\text{e-}4\}$ 中搜索，权重衰减为 $1\text{e-}5$。分子扩散生成训练时批大小设为 512，训练 100 轮，余弦对齐系数 $\lambda=4$。在结构引导在线强化学习中，采用 4 张 NVIDIA B200 GPU 训练策略网络，配合 16 个异步 Boltz2 计算进程实时打分与提取表征，初始重放池预填充 5000 条轨迹，对齐系数 $\lambda=10$。

实验关键数据¶

主实验¶

在 Therapeutics Data Commons (TDC) 的 22 个 ADMET 任务基准中，评估 Boltz2 与专门针对小分子预训练的基础模型对比。下表展示了代表性吸收、分布、代谢、排泄与毒性任务的核心对比：

数据集	任务类别 (指标)	UniMol2 (1B)	MolE	MiniMol	Boltz2 (本文单模)	Boltz2+MiniMol (集成)
Caco2	肠道细胞通透性 (MAE ↓)	0.41	0.31	0.35	0.30	0.30
HIA	人体肠道吸收 (AUROC ↑)	0.87	0.96	0.99	0.99	0.99
Lipophilicity	亲脂性 LogD (MAE ↓)	0.61	0.47	0.46	0.45	0.41
Solubility	水溶性 (MAE ↓)	0.81	0.79	0.74	0.66	0.64
BBB	血脑屏障穿透 (AUROC ↑)	0.86	0.90	0.92	0.93	0.94
PPBR	血浆蛋白结合率 (MAE ↓)	8.95	8.07	7.70	7.65	7.59
VDss	表观分布容积 (Spearman ↑)	0.70	0.65	0.54	0.74	0.75
CYP2C9 V.	代谢酶抑制 (AUPRC ↑)	0.77	0.80	0.82	0.82	0.86
Half Life	药物半衰期 (Spearman ↑)	0.48	0.55	0.50	0.62	0.65
Clearance H.	肝细胞清除率 (Spearman ↑)	0.45	0.38	0.45	0.62	0.60
LD50	急性毒性 (MAE ↓)	0.50	0.82	0.59	0.40	0.40
AMES	诱变性 (AUROC ↑)	0.88	0.88	0.85	0.91	0.91

在所有 22 个非集成评估任务中，Boltz2 在 9 项上斩获第一、4 项位列第二，总体表现显著优于同样采用 1B Pairformer 架构但在 884M 独立分子构象上预训练的 UniMol2。

在 ZINC250k 分子图无条件生成任务中，将不同基础模型表征蒸馏对齐至 GruM*（Pairformer 骨干）的对比实验如下表所示：

生成模型 / 蒸馏源	有效性 (Valid ↑)	分子分布距离 (FCD ↓)	图核距离 (NSPDK ↓)	唯一性 (Unique ↑)	新颖性 (Novel ↑)	骨架相似度 (Scaffold ↑)
GruM* (基线)	99.36%	1.49	0.0007	100.00%	99.99%	0.4923
+ MolE	99.51%	1.41	0.0006	100.00%	100.00%	0.4932
+ MiniMol	99.48%	1.44	0.0007	100.00%	99.99%	0.4812
+ UniMol2	99.60%	1.40	0.0006	100.00%	100.00%	0.4864
+ Boltz2 (本文)	99.65%	1.31	0.0005	100.00%	100.00%	0.5064

消融实验¶

1. Pairformer 主干各层原子对表征对下游 ADMET 性能的消融（单一层 vs. 拼接）：

抽取配置	水溶性 (Solubility ↓)	血脑屏障 (BBB ↑)	代谢抑制 (CYP2C9 V. ↑)	半衰期 (Half Life ↑)	急性毒性 (LD50 ↓)
第 16 层 (1/4 深度)	0.67	0.92	0.81	0.60	0.41
第 32 层 (2/4 深度)	0.67	0.92	0.82	0.63	0.41
第 48 层 (3/4 深度)	0.66	0.92	0.82	0.61	0.40
第 64 层 (末层)	0.64	0.91	0.81	0.62	0.39
分层拼接 (Concat)	0.66	0.93	0.82	0.62	0.40

2. 引入真实蛋白质序列上下文对代谢酶抑制预测的增益消融：

代谢酶抑制任务	无蛋白质输入 (孤立分子模式)	引入真实酶蛋白序列上下文	性能增益
CYP2C9 V. (AUPRC ↑)	0.82	0.83	+0.01
CYP2D6 V. (AUPRC ↑)	0.69	0.72	+0.03
CYP3A4 V. (AUPRC ↑)	0.85	0.87	+0.02

关键发现¶

共折叠物理监督显著优于无约束构象预训练：Boltz2 仅在约 75 万个复合体结构上预训练，数据量远小于 UniMol2（8.84 亿独立小分子构象），但在同为 1B Pairformer 架构下全面超越 UniMol2，证明受蛋白质结合口袋约束的原子相互作用数据质量极高。
生成模型 2 倍加速与质量提升：在分子扩散模型训练中，Boltz2 蒸馏对齐使得模型仅用一半训练步数即达到并超越了原先 100 轮的收敛质量，且在 Valid、FCD、NSPDK 三项指标上同步领先。
代谢底物（Substrate）预测的共有局限：Boltz2 与 UniMol2 在代谢酶底物任务（CYP2C9/2D6/3A4 S.）上均逊于依赖海量生物活性表型预训练的 MiniMol。这表明现有纯三维结构表征在捕捉精细的局部催化反应机理上仍存短板，且部分基线存在预训练数据标签泄漏效应。

亮点与洞察¶

将蛋白质共折叠模型成功逆向迁移为小分子基础模型：突破了以往蛋白质模型与小分子模型井水不犯河水的壁垒，首次系统证明共折叠大模型在无需下游微调的情况下即可作为高质量的分子特征抽取器。
白盒强化学习表征对齐：跳出以往生成模型只能把外部打分函数当作黑盒标量奖励的局限，利用打分器自身的中间激活实施密集表征蒸馏，为离散分子强化学习探索提供了显著更高的样本效率。
“弱对齐、强集成”的表征互补法则：利用 CKNNA 指标证明了共折叠模型与传统独立小分子模型之间的正交性，指导了更高效的模型集成策略（Boltz2 + MolE 胜过 Boltz2 + MiniMol）。

局限与展望¶

局部原子反应机理建模不足：共折叠模型侧重于整体结合构象匹配，对局部酶促反应过渡态或底物特异性原子位点感知较弱，导致代谢底物类任务泛化欠佳。
冻结骨干与开箱评估：本文所有评估均基于冻结的 Boltz2 骨干网络（Linear Probing / Distillation），尚未探索全参数微调或参数高效微调（PEFT）在小分子任务中的潜力。
蛋白质上下文提示词工程的前景：虽然初步验证了加入靶点序列能够提升代谢性质预测，但如何设计系统化的“蛋白质 Prompt”来条件化引导分子表征仍有待深入探索。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次系统发掘并将大尺度蛋白质-配体共折叠模型逆向适配为小分子基础模型，理念前瞻。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 22 项 ADMET 基准、扩散图生成蒸馏、7 个靶点的在线结构强化学习以及系统的表征几何分析，证据链完备。
写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严密，实验消融深入，对模型局限性与泄漏问题剖析诚实客观。
价值: ⭐⭐⭐⭐⭐ 为小分子表征学习开辟了利用复合体结构预训练模型的新范式，对 AI 药物研发具有直接指导价值。