CausalBind: Causal Modeling and Learning for Protein-Molecule Virtual Screening¶

会议: NeurIPS 2026
arXiv: 2610.07340
OpenReview: https://openreview.net/forum?id=bURAYAEF84
代码: https://github.com/lokali/CausalBind
领域: 计算生物
关键词: 虚拟筛选, 蛋白质-小分子结合, 因果表示学习, 稀疏交互, 可辨识性

一句话总结¶

针对基于双塔检索的虚拟筛选容易过度拟合全局冗余相关性的缺陷，CausalBind 将仅观测到结合对的数据建模为 Heckman 风格的选择撞节 V-结构，在理论上证明了密集交互不可辨识而稀疏反链/低秩交互可辨识，并通过稀疏概念对瓶颈在 DUD-E 和去偏基准 LIT-PCBA 上显著提升了早期富集性能。

研究背景与动机¶

在现代药物研发流程中，从包含数百万至数十亿化合物的超大规模分子库中筛选活性先导化合物是一项核心挑战。传统的基于物理的分子对接方法（如 Glide-SP、AutoDock Vina）通过力场和构象采样模拟复合物空间构象，虽具备良好的物理可解释性，但对单个靶点-配体对的计算耗时达数秒至数分钟，无法扩展至亿级分子库的初筛；而结构感知的深度学习评分函数（如 Gnina、RTMScore、EquiScore、TankBind）在推理时依然重度依赖预先生成的复合物三维构象，推理延迟在数十至数百毫秒，更适用于小规模候选分子的重打分。为了突破通量瓶颈，近期兴起的检索型双塔表征学习方法（如 DrugCLIP、LigUnity、HypSeek、S2Drug）将分子与蛋白质分别独立编码到共享嵌入空间中，将亲和力预测转化为亚毫秒级的向量内积或余弦相似度检索，使得超大规模化合物库的初筛成为可能。

然而，现有检索型方法在构建表征时普遍采用“全局全表面稠密对齐”（dense holistic alignment）策略，即将分子的全局嵌入与蛋白质表面的整体嵌入直接做稠密匹配。这种全局稠密对齐在结构层面上与蛋白质-分子结合的真实生物物理机制存在严重脱节。在生物体系中，配体与受体的结合高度受控于局部的“热点”（hot spots）相互作用，整个结合自由能绝大部分由接触界面上极少数关键的局部位点相互作用（如氢键、疏水接触、盐桥和 $\pi$-$\pi$ 堆积）决定，而非取决于整体的全局折叠轮廓。全局全表面匹配使得模型极易纠缠不变量结合决定因素与全局非因果冗余特征，进而退化为利用靶点家族偏置或小分子骨架假相关进行捷径学习（shortcut learning），严重损害了模型跨蛋白质家族和未知化学骨架的外推泛化能力。

本文的核心洞见在于：既然真实的结合由稀疏的局部相容性主导，那么在表征学习中显式提取概念级局部表征并施加稀疏/低秩交互约束，是否能够消除冗余相关性并带来理论保证？针对训练数据仅包含阳性结合对（$B=1$）的特点，本文的核心 idea 是将仅含活性对的结合数据建模为 Heckman 风格选择撞节（collider）的 V-结构因果图，严格证明了未约束的稠密交互在理论上不可辨识，而稀疏反链交互具备分量级可辨识性、低秩交互具备子空间可辨识性；据此提出了 CausalBind 框架，通过 Perceiver 概念提取器与可学习的稀疏/低秩交互瓶颈，迫使对比表征学习通过稀疏的因果交互回路完成跨模态匹配。

方法详解¶

整体框架¶

CausalBind 继承了检索型双塔架构的高吞吐推理特性，能够支持化合物库特征离线预编码。整体流程包含三个主要阶段：多模态局部标记提取、潜概念分解（Concept Extraction）、以及跨模态稀疏概念交互与池化打分。

在输入端，模型接收小分子三维构象 $x_{\mathrm{mol}}$、蛋白质结合口袋结构 $x_{\mathrm{poc}}$ 和蛋白质序列 $x_{\mathrm{seq}}$ 三种模态。Uni-Mol 分别提取分子和口袋的局部三维空间特征矩阵 $\mathbf{H}_{\mathrm{mol}}$ 和 $\mathbf{H}_{\mathrm{poc}}$，ESM-2 提取全局残基序列特征 $\mathbf{H}_{\mathrm{seq}}$。随后，每个模态通过 Perceiver 架构的交叉注意力模块，从连续的 token 序列中提取 $K$ 个潜在原子概念向量 $\mathbf{C}_v \in \mathbb{R}^{K \times d_c}$。接着，在（分子-口袋）与（分子-序列）两个检索分支上，CausalBind 通过可学习的稀疏交互掩码矩阵对概念对相容性矩阵进行逐元素门控滤波，滤除不相关的全局噪声，仅保留关键的局部交互证据；最后经池化汇聚生成标量检索得分，并利用对比损失与几何正则化进行端到端优化。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["多模态输入<br/>分子结构 + 口袋结构 + 蛋白序列"] --> B["阶段 1：因果可辨识性建模<br/>V-结构选择撞节与稀疏/低秩交互先验"]
    B --> C["阶段 2：Perceiver 概念提取器<br/>跨模态分解为 K 个原子概念向量"]
    C --> D["阶段 3：稀疏交互瓶颈与双分支门控<br/>Straight-Through 硬阈值过滤非因果概念对"]
    D --> E["阶段 4：低秩分解与嵌入空间约束变体<br/>CausalBindSP / LR / EMB 多粒度实现"]
    E --> F["联合训练优化<br/>双向对比损失 + 掩码 L1 稀疏约束 + 层次锥正则"]
    F --> G["输出：结合亲和力排序得分<br/>高通量虚拟筛选早期富集检索"]

关键设计¶

1. 基于 V-结构选择撞节的因果可辨识性建模：从非可辨识到分量与子空间可辨识 针对虚拟筛选训练数据中“仅观测到活性结合对”（binder-only selection）导致的数据偏置，本文将结合事件建模为一个伯努利选择撞节（collider）变量 $B \in \{0, 1\}$，构建了典型的 V-结构因果生成模型： $$x^M \leftarrow \mathbf{c}^M \to B \leftarrow \mathbf{c}^P \to x^P$$ 在先验上，小分子概念 $\mathbf{c}^M$ 与蛋白质概念 $\mathbf{c}^P$ 相互独立（$\mathbf{c}^M \perp\!\!\!\perp \mathbf{c}^P$），但条件化在阳性结合事件 $B=1$ 上时，后验分布 $p^{\mathcal{D}}(\mathbf{c}^M, \mathbf{c}^P) \propto p(\mathbf{c}^M) p(\mathbf{c}^P) \sigma\left(\sum_{i,j} \mathcal{M}_{ij} \phi(c^M_i, c^P_j)\right)$ 诱导了跨模态的统计相关性。这种撞节选择诱导的相关性是恢复潜在概念和交互矩阵的唯一信号源。作者深入分析了该机制下的可辨识性，建立了三个渐进的理论定理： - 定理 1（稠密非可辨识性）：若不对交互掩码 $\mathcal{M}$ 施加结构约束，存在无限多组概念变换与交互矩阵均能生成相同的选择后联合分布 $p^{\mathcal{D}}(x^M, x^P)$，稠密全表面匹配在原理上不可辨识。 - 定理 2（稀疏反链分量可辨识性）：当交互掩码 $\mathcal{M}$ 满足稀疏反链（antichain）支撑条件（即不存在任意两行或两列其非零元素集合存在严格包含关系），且分子侧有限探针诱导的交叉混合二阶偏导数向量线性无关时，概念与稀疏交互图可在置换、分量单调变换和尺度放缩的平凡等价下实现分量级可辨识（component-wise identifiability）。 - 定理 3（低秩子空间可辨识性）：当精确元素级稀疏过于严格时，若交互核具备最小秩 $r$ 分解 $S(\mathbf{c}^M, \mathbf{c}^P) = \mathbf{u}(\mathbf{c}^M)^\top \mathbf{v}(\mathbf{c}^P)$ 且模态评估矩阵满秩，则分子与蛋白质两侧的秩-$r$ 交互子空间可在 $\mathrm{GL}(r)$ 可逆变换下实现子空间可辨识（subspace identifiability）。这为后续神经网络的归纳偏置设计提供了坚实的理论基石。

2. Perceiver 概念提取器与双分支稀疏交互瓶颈 为了在神经网络中落实上述因果概念，模型必须能够从高维输入中解耦出潜在概念，并对概念间交互实施稀疏门控。对于每个模态 $v \in \{\mathrm{mol}, \mathrm{poc}, \mathrm{seq}\}$，CausalBind 使用 $L$ 层包含可学习查询向量 $\mathbf{Q}_v \in \mathbb{R}^{K \times d_c}$ 的 Perceiver 交叉注意力块，将骨干网络提取的 token 特征序列逐步压缩映射为 $K$ 个具有语义代表性的概念表征 $\mathbf{C}_v \in \mathbb{R}^{K \times d_c}$。在核心的主变体 CausalBindSP 中，针对（分子-口袋）与（分子-序列）两条配对分支，分别分配未受限的可学习交互权重矩阵 $\widehat{\mathcal{M}}_{a,b} \in \mathbb{R}^{K \times K}$。首先通过平滑激活函数 $\mathbf{Z}_{a,b} = \tanh(\widehat{\mathcal{M}}_{a,b}) + 1$ 映射到 $[0, 2]^{K \times K}$，随后采用 Straight-Through 离散硬阈值算子 $H_\tau^{\mathrm{ST}}$ 进行门控（默认阈值 $\tau=0.5$）： $$[\mathcal{M}_{a,b}^{\mathrm{SP}}]_{ij} = [\mathbf{Z}_{a,b}]_{ij} \cdot \mathbf{1}\{[\mathbf{Z}_{a,b}]_{ij} > \tau\}$$ 前向传播中低于阈值的概念对被严格置零断开，反向传播中梯度沿连续张量 $\mathbf{Z}_{a,b}$ 正常回传。概念对相容性矩阵由向量内积定义 $\mathbf{A}_{a,b}[i,j] = \langle \mathbf{C}_a[i,:], \mathbf{C}_b[j,:] \rangle$。分支标量检索分数计算为保留交互项的平均汇聚：$u_{a,b}^{\mathrm{SP}} = K^{-1} \sum_{i,j} [\mathcal{M}_{a,b}^{\mathrm{SP}} \odot \mathbf{A}_{a,b}]_{ij}$。该稀疏掩码扮演了跨模态因果信息瓶颈，强迫模型仅依赖少数关键相容性进行结合判断。

3. 低秩分解与多粒度交互约束变体 为了验证理论中定理 3 的低秩子空间先验并探索计算效率，作者进一步设计了低秩概念掩码变体 CausalBindLR 与嵌入层掩码变体 CausalBindEMB： - CausalBindLR（低秩概念分解）：直接依据定理 3，将原始未激活掩码参数因子化为低秩矩阵乘积 $\widehat{\mathcal{M}}_{a,b} = \mathbf{U}_{a,b} \mathbf{V}_{a,b}^\top$，其中 $\mathbf{U}_{a,b}, \mathbf{V}_{a,b} \in \mathbb{R}^{K \times r}$（默认 $r=1$）。这使得单分支掩码参数量由 $K^2$ 骤降至 $2Kr$（当 $K=128, r=1$ 时从 16,384 减少到 256）。低秩因子化在概念轴上施加了显式的低维相互作用模态约束，在无需额外稀疏损失的情况下自发实现了子空间可辨识性。 - CausalBindEMB（嵌入空间掩码）：将稀疏瓶颈后置。先对概念向量沿概念轴进行均值池化得到模态嵌入 $\bar{\mathbf{h}}_v = \frac{1}{K} \sum_{i=1}^K \mathbf{C}_v[i,:]$，随后在嵌入特征维度施加掩码变换 $\widetilde{\mathbf{h}}_v = \bar{\mathbf{h}}_v \mathcal{M}_v^{\mathrm{EMB}}$（$\mathcal{M}_v^{\mathrm{EMB}} \in \mathbb{R}^{d_c \times d_c}$），最后计算投影相似度。该变体作为消融对照，探讨了约束作用在概念轴与池化嵌入轴之间的本质差异。 - CausalBindATOM（原子级注意力）：作为高分辨率变体，概念提取器直接在 Uni-Mol 的全原子级 token 序列上执行注意力，使得提取的概念可回溯至具体原子位点，兼顾检索精度与原子级可解释性。

损失函数 / 训练策略¶

模型采用多任务联合目标函数在由 ChEMBL、BindingDB 和 PDBBind 构建的非冗余数据集上进行端到端优化（训练集已剔除与测试集 UniProt ID 重叠的测定项）： $$\mathcal{L} = \mathcal{L}_{\mathrm{cont}} + \lambda_{\mathrm{spa}} \mathcal{L}_{\mathrm{spa}} + \lambda_{\mathrm{reg}} \mathcal{L}_{\mathrm{reg}}$$ 1. 跨模态对比检索损失 $\mathcal{L}_{\mathrm{cont}}$：在小批量（batch size $N$）内，针对分子-口袋与分子-序列分支分别计算对称式 InfoNCE 损失： $$\ell_{\mathrm{cont}}(\mathbf{U}^{a,b}) = -\frac{1}{2N} \sum_{i=1}^N \left[ \log \frac{\exp(U^{a,b}_{ii}/\tau_{\mathrm{c}})}{\sum_{j=1}^N \exp(U^{a,b}_{ij}/\tau_{\mathrm{c}})} + \log \frac{\exp(U^{a,b}_{ii}/\tau_{\mathrm{c}})}{\sum_{j=1}^N \exp(U^{a,b}_{ji}/\tau_{\mathrm{c}})} \right]$$ $$\mathcal{L}_{\mathrm{cont}} = \ell_{\mathrm{cont}}(\mathbf{U}^{\mathrm{poc,mol}}) + \ell_{\mathrm{cont}}(\mathbf{U}^{\mathrm{seq,mol}})$$ 2. 交互稀疏惩罚 $\mathcal{L}_{\mathrm{spa}}$：对平滑激活后的掩码 $\mathbf{Z}$ 施加连续 $L_1$ 范数惩罚以保持回传梯度： $$\mathcal{L}_{\mathrm{spa}} = \|\mathbf{Z}_{\mathrm{mol},\mathrm{poc}}\|_1 + \|\mathbf{Z}_{\mathrm{mol},\mathrm{seq}}\|_1$$ 默认超参数设置为 $\lambda_{\mathrm{spa}} = 10^{-4}$。 3. 层次几何正则化项 $\mathcal{L}_{\mathrm{reg}}$：继承自 HypSeek 的双曲空间几何结构正则，包括防止角度坍缩的角度边界正则 $R_{\mathrm{ang}}$（margin 设为 0.15）以及利用测定内部活性分子排序加权的测定内异质性负对数似然损失 $R_{\mathrm{het}}$，迫使模型在同靶点活性化合物内部区分高亲和力与弱亲和力。

实验关键数据¶

主实验¶

评估在经典虚拟筛选基准 DUD-E（102 个靶点，侧重靶点多样性）以及经过严格物理化学属性去偏的 LIT-PCBA（15 个靶点，严格消除人工捷径，评测真实结合识别）上进行。模型评估使用 50 epoch 最终检查点，评价指标涵盖全局判别力 AUROC，以及更加契合工业界实际高通量筛选需求的高灵敏早期富集指标 BEDROC（$\alpha=80.5$，指数加权聚焦于前 2% 排名）与 EF@1%（前 1% 富集倍数）。

方法	分类类型	DUD-E AUROC	DUD-E BEDROC80.5	DUD-E EF@1%	LIT-PCBA AUROC	LIT-PCBA BEDROC80.5	LIT-PCBA EF@1%
Glide-SP	物理力场对接	0.767	0.407	16.18	0.532	0.040	3.41
Gnina	结构评分模型	0.782	0.299	17.73	0.609	0.054	4.63
RTMScore	结构评分模型	0.753	0.434	27.10	0.525	0.039	2.94
TankBind	结构评分模型	0.751	0.330	13.00	0.597	0.039	2.90
DrugCLIP	检索型双塔	0.809	0.505	31.89	0.572	0.062	5.51
LigUnity	检索型双塔	0.897	0.674	44.20	0.599	0.075	6.50
HypSeek	双曲检索双塔	0.909	0.605	37.83	0.603	0.061	4.62
S2Drug	两阶段检索 (序列+结构)	0.925	0.793	43.06	0.582	0.087	7.38
CausalBindSP	本文稀疏概念对 ($K=128$)	0.935	0.744	47.69	0.639	0.097	8.56
CausalBindLR	本文低秩概念 ($K=64, r=1$)	0.935	0.704	44.87	0.632	0.103	8.84
CausalBindEMB	本文嵌入层掩码 ($K=64$)	0.939	0.754	48.56	0.636	0.102	9.60
CausalBindATOM	本文全原子注意力 ($K=128$)	0.943	0.768	49.66	0.621	0.079	6.25

消融实验¶

表 2 汇总了 CausalBindSP、CausalBindLR 与 CausalBindEMB 在 5 次独立重复运行下的均值与标准差消融分析，系统考察了稀疏损失 $\mathcal{L}_{\mathrm{spa}}$ 以及低秩约束施加位置的作用：

配置	DUD-E AUROC	DUD-E BEDROC80.5	DUD-E EF@1%	LIT-PCBA AUROC	LIT-PCBA BEDROC80.5	LIT-PCBA EF@1%	说明
CausalBindSP (完整模型)	0.939 ± 0.004	0.746 ± 0.004	47.99 ± 0.39	0.635 ± 0.006	0.088 ± 0.007	7.55 ± 0.73	稀疏掩码完整模型
CausalBindSP $-\mathcal{L}_{\mathrm{spa}}$	0.935 ± 0.004	0.732 ± 0.014	46.93 ± 1.07	0.623 ± 0.005	0.084 ± 0.004	7.17 ± 0.36	移除 $L_1$ 稀疏惩罚，早期富集明显衰减
CausalBindLR (秩 $r=1$)	0.937 ± 0.004	0.732 ± 0.011	46.87 ± 0.91	0.630 ± 0.006	0.085 ± 0.006	7.22 ± 0.31	纯显式秩 1 因子分解，无需稀疏正则
CausalBindLR $+\mathcal{L}_{\mathrm{spa}}$	0.935 ± 0.003	0.727 ± 0.007	46.39 ± 0.58	0.623 ± 0.007	0.083 ± 0.005	7.51 ± 0.94	额外施加稀疏正则未带来增益
CausalBindEMB	0.938 ± 0.001	0.752 ± 0.005	48.47 ± 0.44	0.634 ± 0.008	0.101 ± 0.007	9.52 ± 0.96	嵌入层连续掩码基线
CausalBindEMB $+\mathbf{M}_{\mathrm{emb}}^{r=1}$	0.723 ± 0.022	0.319 ± 0.013	20.96 ± 0.63	0.539 ± 0.003	0.018 ± 0.001	1.38 ± 0.10	嵌入空间强加秩 1 产生毁灭性坍缩

关键发现¶

去偏基准上的早期富集提升极其显著：在去除了假阴性与简单理化偏置的严格基准 LIT-PCBA 上，CausalBindSP 相比于检索基线 LigUnity 和 HypSeek，EF@1% 从 6.50 和 4.62 跃升至 8.56，BEDROC80.5 从 0.075 和 0.061 提升至 0.097。这有力证实了稀疏因果交互有效过滤掉了全局理化伪相关，捕获了真正的特异性结合驱动因素。
低秩约束的位置至关重要：在概念轴上施加秩 1 分解（CausalBindLR）性能稳定且极大精简了参数量（保持了 7.22 的 LIT-PCBA EF@1%）；但在池化后的嵌入特征轴强加秩 1 投影（$\mathbf{M}_{\mathrm{emb}}^{r=1}$）会导致两基准的全面性能崩溃（DUD-E EF@1% 暴跌至 20.96，LIT-PCBA 跌至 1.38）。这表明低秩归纳偏置必须作用于跨模态的相互作用模式（概念轴），而非压缩单模态表征特征空间。
稀疏强度与反链理论的验证：超参数扫描（Table 3）显示，$\lambda_{\mathrm{spa}} = 10^{-4}$ 时掩码未出现硬零点（零值率为 0%），但有效秩显著降低，处于定理 3 的低秩子空间体制；而 $\lambda_{\mathrm{spa}} = 10^{-3}$ 时零值率达 20.5%，反链违背次数降为 0（完全满足定理 2 的稀疏反链支撑条件），DUD-E 达到最高 EF@1% 49.73；当 $\lambda_{\mathrm{spa}} \ge 10^{-2}$ 时掩码发生过度稀疏崩溃（零值率超 78%）。

亮点与洞察¶

将撞节选择偏置严格形式化为表征学习先验：论文突破性地指出虚拟筛选训练集本质上是条件化在结合事件 $B=1$ 上的选择样本，并在该 V-结构因果图下首次证明了稠密交互的不可辨识性与稀疏/低秩交互的可辨识性，架起了因果推断理论与计算生物物理实践之间的严谨桥梁。
不牺牲推理速度的高泛化归纳偏置：与依赖 3D 构象生成的缓慢对接和复合物网络不同，CausalBind 保留了双塔检索架构，模态概念可离线独立编码，毫秒级稀疏门控瓶颈直接过滤非因果假相关，完美兼顾了高通量与抗过拟合能力。
理论预测与实证指标的高度自洽：通过对学到的交互掩码进行奇异值分析与反链违背统计，成功在实际神经网络中复现了定理 2 所预测的无违背反链稀疏区与定理 3 的低秩子空间区，为深度模型解释提供了绝佳范例。

局限与展望¶

作者承认的局限：模型评估主要基于跨数据集泛化协议（DUD-E 和 LIT-PCBA），尚未在时间切分（time-split）或更严格的蛋白质序列家族完全隔离切分上开展大规模验证；此外，案例分析提供的是模型层面的归因，而非严格湿实验物理因果干预验证。
实验与方法设计的局限：概念数量 $K$ 与 Perceiver 层数需要仔细权衡，过大的 $K$（如 256）会导致过参数化，而过小的特征维度会导致欠拟合；目前稀疏掩码是在训练集上全局共享学习的，尚未实现根据蛋白质口袋微环境动态自适应调节的实例级稀疏路由。
未来改进方向：可将固定稀疏掩码拓展为输入自适应的条件因果路由网络（instance-dependent routing）；同时可将稀疏因果交互先验推广至多肽-蛋白质或蛋白质-蛋白质复合物的相互作用预测。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [首次在双塔虚拟筛选中建立 Heckman 选择撞节 V-结构模型，并证明了稀疏反链与低秩交互的可辨识性理论]
实验充分度: ⭐⭐⭐⭐⭐ [覆盖 DUD-E、LIT-PCBA、DEKOIS 2.0 及 Merck-FEP-8，提供详尽的消融、奇异值分析与反链违背统计]
写作质量: ⭐⭐⭐⭐⭐ [理论推导严密层层递进，从因果图假说到神经架构实现逻辑丝丝入扣]
价值: ⭐⭐⭐⭐⭐ [解决了虚拟筛选领域捷径拟合的根本痛点，为超大规模 AI 药物筛选算法设计提供了理论指引]