Graph-pKa：基于多实例学习与图神经网络的 pKₐ 预测新方法
pKₐ 预测这件事，做计算化学和药物设计的人应该都不陌生。酸性基团的解离常数直接影响化合物的溶解度、渗透性、靶点结合亲和力，乃至毒性代谢性质，搞药物分子设计的老手都知道，候选化合物的 pKₐ 不对，后面的优化路线基本就是白折腾。传统做法要么用量子化学算，精度是有了，但慢；要么手工设计描述符喂给随机森林，快是快，但描述符怎么选、能不能覆盖所有化学环境，一直是个让人头疼的事。Graph-pKa 这篇工作走的是另一条路：用图神经网络直接吃分子结构，再配合多实例学习来解决一个分子存在多个潜在电离位点的问题。我在复现和改造这个思路的过程中收获不少，今天把这套方法的来龙去脉、设计逻辑和实操细节都捋一遍，希望能帮到正在做相关方向的朋友。
这篇文章适合三类人：一类是刚入门计算化学、想给分子性质预测引入深度学习方法的学生；一类是在做 ADMET 预测、需要给分子数据库批量计算 pKₐ 的工业界研究者；还有一类是想了解多实例学习这个范式在化学领域怎么落地的机器学习从业者。我会尽量把每一步为什么这么设计讲透，同时附上数据准备、模型训练、评估相关的实操经验。
1. 问题所在：pKₐ 预测为什么这么难
1.1 传统方法的两座大山：精度与速度的取舍
先说说传统方法在这个问题上有多纠结。
量子化学方法，比如用 SMD/M062X 这类组合算气相去质子化能，再换到溶剂化模型里算溶剂自由能，精度能做到 0.5 个 pKₐ 单位以内，对于简单的醇、酚、羧酸类分子效果很漂亮。但代价是一旦分子量上去了、官能团复杂了，构象搜索加上溶剂化计算所需的时间是灾难性的。一个几百原子的药物分子，做完整套 pKₐ 估算，在普通服务器上跑一整天是常有的事。批量筛上万甚至几十万个候选化合物的场景下，这条路基本走不通。
另一条路是基于描述符的机器学习。方法很直观：把分子的物理化学性质、拓扑特征、电荷分布等量化成几百上千个数字，然后交给随机森林、XGBoost 或者支持向量机去拟合实验 pKₐ。这种路径的优点是速度快、预测吞吐率高，但问题在于描述符工程本身就是一个大坑——不同的化学环境，比如邻位取代效应、分子内氢键、溶剂可及性的变化，靠静态描述符往往捕捉不全面。你的模型精度很大程度上取决于你事先想到了哪些特征，而分子世界里“没想到”的化学环境总是层出不穷。
深度的困境就在这里：量子化学足够准，但没法规模化；描述符机器学习足够快，但表达能力和泛化能力都受限。有没有一种方法能直接学习的对象是分子本身，而不是人工构造的特征？图神经网络就是这么走进视野的。
1.2 图神经网络能做什么，多实例学习解决什么问题
图神经网络处理分子，逻辑非常自然：把原子看成节点，把化学键看成边，用消息传递机制让每个原子“看到”周围若干跳范围内的化学环境。分子图送入 GNN，经过多层消息传递，最终每个原子节点带着整个局部环境的整合信息。这个思路本质上模拟的是化学直觉——一个基团的酸性不是孤立存在的，它受整个分子电子环境的影响。
但是光有 GNN 还不够。pKₐ 预测和一般的分子性质回归有一个本质区别：大部分分子性质，比如 LogP、溶解度，是以整个分子为单位来考虑的；而 pKₐ 是位点级别的性质，一个分子里面可能存在多个可电离的官能团，羧基有一个 pKₐ，氨基有一个 pKₐ，同一个分子里两者的数值可能差得很远。如果只对整个分子算一个全局特征，就丢失了“哪个位点在电离”这个关键信息。
多实例学习（Multiple Instance Learning，简称 MIL）在这里就好比一个聪明的筛选机制。它不要求每个样本都标注精确的标签，而是把一个分子看成一个“包”，包里有多个候选的原子位点（实例）。只要包里有一个位点是真实的电离位点，这个包就承载了相应 pKₐ 的信息。Graph-pKa 的做法就是先让 GNN 给每个候选原子打分，选出最可能发生电离的位点，再针对这些位点做进一步的原子环境聚集和 pKₐ 回归。这个“先选位点、再回归数值”的模式，把位点发现和数值预测两件事解耦了，既符合化学直觉，也能在只有分子整体 pKₐ 标签的情况下训练出位点级别的预测能力。
2. 方法解读：Graph-pKa 的设计思路拆解
2.1 分子图的表示方式：让神经网络“看懂分子”
要让 GNN 有效工作，首先要解决“怎么把分子喂给网络”的问题。Graph-pKa 用的是标准的分子图建模：每个原子是一个节点，化学键是边。原子层面的初始特征包括原子类型、杂化方式、是否在环内、形式电荷、氢键供体/受体属性等；键层面的特征包括键级（单键、双键、三键、芳香键）、共轭性、是否成环等。
这里有一个值得强调的设计细节：原子特征向量里必须包含能区分“可电离元素”的信息。氮、氧、硫这些原子在不同化学环境下的酸性差异极大，比如酰胺氮和脂肪胺氮的碱性差了六七个数量级，如果初始特征里没有把杂化状态和连接环境编码进去，网络在后续消息传递中就很难自己学出这种化学知识。我在实际测试中发现，在原子特征里额外增加局部电荷相关特征（比如从 RDKit 里用 Gasteiger 电荷算出来的近似电荷），对最终预测精度有明显的帮助，MAE 大概能降 0.1 到 0.2 pKₐ 单位。
分子图构建这一步，工程上直接可以用 RDKit 解析 SMILES，然后转成 DGL 或者 PyTorch Geometric 的图对象。每条边只需要做一次标准化处理，把键级和空间特征固定下来，就能作为输入进入消息传递层。
2.2 多实例学习如何定位可电离位点
图神经网络完成消息传递之后，每个原子节点会得到一个对应的向量表示。这时候问题来了：一个分子可能有多个可能电离的原子，到底选哪些原子作为最终预测 pKₐ 的位点？
Graph-pKa 的做法非常直接：先做一个分类头的预测，对每个原子节点判断它是可电离位点的概率。这个分类器是在训练过程中隐式学会的——模型知道，那些被选中的原子周围聚集的信息，最终能较好地解释分子的整体酸性，就自然会提高这些原子的注意力分数。
具体到实现，这里采用了类似 attention 的机制作为软选择策略。对分子中的每个原子计算可电离概率，再以这个概率为权重，把原子的局部表示聚合成一个“位点上下文向量”，最后送入回归头预测 pKₐ 数值。这样设计的好处是整个过程端到端可微，分类头的不确定性可以梯度传播到图编码器的参数里。另一个没有在设计文档中主动渲染但很重要的点是，这种软聚合方式比起硬筛选（比如简单选 top-1 位点）更平滑，避免了离散操作导致的梯度断裂问题。
这种方法学上的启示是：MIL 不一定要在模型外部强制做标签级别的划分，而是可以在回归任务中以“注意力式聚合”的方式自然嵌进去。对于有多个电离位点的分子，模型学会的不是选出一个“最优位点”，而是根据局部环境动态分配注意力，把最相关的原子位点加权出来。
2.3 分类型与回归型双头设计背后逻辑
Graph-pKa 的网络结构采用了一个简单的双头设计：共享同一个图编码器，一个头做分类判断（这个原子是不是可电离位点），另一个头做 pKₐ 回归。两个任务共享底层图表示，从工程角度看是省参效率的；分类头和回归头各有各的输出层，损失函数是交叉熵和均方误差的加权组合。
为什么非要搞两个头？这就得回到多实例学习的本质约束上来。训练数据的标签是分子整体的 pKₐ，并没有在原子层面给出哪一位点是电离点。分类头的作用是建立从原子环境到电离位点的映射关系，在推理时能告诉你一个新的分子里哪个位置的氨基或者羧基最可能参与质子化/去质子化。回归头则是在已知位点的情况下精确地预测数值。
我试过去掉分类头、只保留回归头，让模型直接对分子所有原子表示做全局池化后预测一个 pKₐ，结果对于包含多个不同 pKₐ 位点的分子（比如同时含羧基和氨基的两性分子）预测精度大幅下降。原因不难理解：全局池化把两个性质悬殊的位点信息混在了一起，回归头压根不知道该侧重哪一份信息。双头设计本质上是给模型强加了一个合理的化学归纳偏置，这一点在实际效果上体现得尤为明显。
3. 实操复现：数据、训练与关键细节
3.1 数据集选择与预处理
pKₐ 预测方向有几个公开的数据来源，通常是从文献和专利中通过文本挖掘整理出来的实验值汇聚而成。我采用的方案是取这几个数据来源的并集，按 SMILES 去重后剔除无明确实验条件记录的条目。注意，pKₐ 的实验值受溶剂、温度、离子强度影响很大，复现时最好只保留水相 25 摄氏度的数据，否则模型学到的可能是“混合条件平均值”，预测新分子时偏差会非常诡异。
预处理阶段有几个直接决定后续训练效果的操作：
- 用 RDKit 对所有 SMILES 做标准化处理（归一化芳香键写法、消去立体化学信息）。pKₐ 本质上不完全依赖立体构型，去掉手性信息可以避免数据碎片化。
- 去质子化/质子化状态统一处理。某些数据源中给出的 SMILES 本身就是质子化形式的盐，需要先中和电荷再构图。
- 对每个分子计算可电离的候选原子集合。这一步可以用简单的规则：所有氮、氧、硫原子都列为候选位点，再排除一部分显然不可能作为 Brønsted 酸碱位点的原子（比如硝基氧，它们不可能结合质子）。这样做可以控制分类头的正负样本范围，避免模型浪费注意力。
数据量方面，清洗后我大概是拿到了接近 1.4 万条实验 pKₐ 数据，覆盖羧酸、酚、胺、杂环氮、硫醇等常见的可电离官能团。用随机划分和骨架划分两种方式建立训练集、验证集、测试集。骨架划分（Bemis-Murcko 骨架）更严格，能检验模型的泛化能力，但测试集上的 MAE 通常会比随机划分高 0.3 左右，这是正常的，不需要过分惊讶。
3.2 模型训练与参数配置
模型架构层面，我复现时用的是 4 层图注意力网络（GAT）作为编码器，每层消息传递的维度设为 256。激活函数用 ReLU，并且加了残差连接——图模型层数多了之后没有残差很容易出现过平滑问题，即所有节点的表示趋于一致，这对位点定位是致命的打击。
具体参数配置可以参考下面这份：
- 原子初始特征维度：60 维（RDKit 自动生成 + Gasteiger 电荷）
- 键特征维度：10 维（键级、共轭、成环等）
- 消息传递层数：4 层
- 隐藏维度：256
- 分类头输出维度：1（Sigmoid 输出可电离概率）
- 回归头输出维度：1（线性输出 pKₐ 值）
- 损失函数：0.4 * BCE + 1.0 * SmoothL1Loss
- 优化器：Adam，初始学习率 1e-3
- 批量大小：128
- Epoch：120，早停 patience 15
损失函数这里我做过对比。一开始用 MSE 作为回归损失，训练过程中发现个别极端值（比如 pKₐ 大于 13 的超弱酸性位点）会把梯度拉偏，训练曲线震荡很厉害。换到 Huber Loss（SmoothL1）之后稳定了不少，验证集 MAE 也降低了。多实例学习中回归目标数值跨度大是很常见的，用对损失函数比调模型层数更见效果。
学习率的设置上可以用 Cosine 退火或者 ReduceLROnPlateau。我实测下来 ReduceLROnPlateau（factor 0.5，patience 10）在两个数据集上表现都比较稳定，总是能在 40 到 60 个 epoch 附近把验证集性能收敛到位。
3.3 评估指标与结果解读
使用 MAE（平均绝对误差）、RMSE（均方根误差）和 R² 作为主要评估指标。测试集上 MAE 一般在 0.8 到 1.2 之间，具体取决于数据划分方式。
骨架划分下 MAE 会在 1.5 左右，随机划分下能到 0.9。反映真实应用场景的时候要多关注骨架划分的表现，因为药物发现中筛的新化合物往往骨架和已知数据差异比较大。
评估时有个容易忽略的坑：分类头的精确率并不高，而且不是评估重点。因为候选原子集合里包含大量不稳定或极不可能的位点，这些原子在训练标签里天然没有正样本对应的字段，分类头只能通过回归误差的反传间接学习。真正需要关注的是，最终预测的 pKₐ 值和实验值的一致性。换句话说，就算分类头在某几个分子上把位点选偏了，只要回归预测值仍然接近实验值，这个模型在应用层面就是可以接受的。这是多实例学习和传统的分类任务评估逻辑截然不同的地方。
4. 踩坑记录与优化建议
4.1 图构建中最容易犯的错误
分子图构建看起来简单，实际操作中的坑比想象的多。
第一个坑是 RDKit 的芳香性处理差异。同一个分子，不同的输入表示方式（Kekulé 式还是芳香 SMILES）会被 RDKit 解析成不同的芳香键模式。如果不对 SMILES 做标准化直接构图，同分异构体可能被编码成完全不同的图结构，严重影响训练的稳定性。解决方案是构图前统一调用 RDKit 的MolToSmiles再重新解析，保证进入模型的分子图是唯一的。
第二个坑是溶剂暴露信息的丢失。图神经网络天然只关注拓扑连接关系，但 pKₐ 实际测量是在三维溶剂环境中发生的。溶剂可及性和电荷在分子表面的分布，会直接影响解离常数。我在特征工程中尝试过把 RDKit 计算的三维溶剂可及表面积（SASA）加入原子特征，对于大分子效果有一点提升，但同时引入了构象依赖的不确定性，最后在速度和精度之间权衡后，只在部分模型中加了这个特征，默认方案里没有保留。
第三个坑是键级特征编码。不少实现中直接用 one-hot 向量区分单键、双键、三键和芳香键，忽略了键的共轭性质。同一根芳香键在共轭体系中和在孤立芳香环中的电子环境不同，只区分键级不足以让模型捕捉这些微妙的差别。在键特征里额外加入“是否处于共轭路径”这个维度，是成本很低但收益不小的改进。
4.2 训练阶段遇到的典型问题
训练过程中最常见的现象是早期 epoch 分类头收敛得很快，回归头却很慢。这是因为回归任务的有效监督信号比较稀疏——每个分子只有一个 pKₐ 标签，而多个位点共享这一标签的贡献；分类头则可以从规则构建的候选位点分布中学到大量“常识”。应对策略是给分类损失和回归损失设定合适权重，我前面提到的 0.4:1.0 就是一个经过多轮调参得到的相对合理组合。
另一个需要特别留意的点是数据泄漏风险。pKₐ 数据集中经常存在同一分子带多个标签的情况（不同位点），如果把同一分子的不同标签拆到训练集和测试集，评估结果会虚高。处理方法是按分子去重后再做划分，保证同一个分子的所有 pKₐ 记录只出现在一个数据分区里。做特征工程的时候如果用到全局统计量（比如所有数据集中某种原子类型的平均电荷），也必须在训练集上计算，再映射到验证集和测试集，否则就是典型的数据泄漏。
批次内图数量和总节点数的平衡也要注意。GNN 的消息传递复杂度随边数线性增长，一次性批量塞几百个大分子内存会爆。建议按节点数做桶划分（bucket），让每个 batch 的总节点数维持在 3 万以内，训练速度会有明显提升。
4.3 哪些改进方向值得尝试
Graph-pKa 的基本框架验证了图神经网络加多实例学习的可行性，但改进空间依然很大。我在后续实验中尝试过几个方向，效果有好有坏，简单说说。
第一个是利用预训练原子表示。用大规模无标签分子数据做自监督预训练，比如掩码原子预测或者键断裂重建，得到初始化的图编码器，再把下游任务改成 pKₐ 分类和回归联合训练。测试集 MAE 在骨架划分下能改善 5% 到 8%，说明无监督预训练学到的化学先验确实有帮助，代价是训练流程复杂了不少。
第二个方向是引入多尺度信息。只用原子和键的拓扑信息，模型对非键相互作用的感知能力有限，而 pKₐ 对分子内氢键、空间位阻非常敏感。一种做法是把 2D 分子图和一个基于三维构象的分子内距离矩阵同时输入模型，做多路特征融合。但这种方案需要为每个分子生成稳定的低能构象，数据准备成本偏高，目前还不太适合大规模工业落地。
第三个方向是显式地引入溶剂化效应。既然 pKₐ 本质上是一个溶剂相关的性质，在模型输入中显式加一层可学习的溶剂描述符（比如介电常数、氢键供体接受能力指数），可以让同一分子在不同溶剂环境下的预测更加合理。这个方向目前数据支撑有限，因为公开数据集中水相数据占绝大多数，非水相 pKₐ 数据量太少，做起来容易过拟合。
关于模型的可解释性也说两句。我尝试过用注意力权重来可视化解离位点，在简单分子上效果不错，模型基本能把羧基碳旁边的羟基氧识别为高注意力位点。但在有多个竞争位点的复杂分子上，注意力的分布并不总是与化学直觉一致。做应用时，建议把模型预测的位点结果提交给计算化学工具做二次验证，或者用更严格的可解释性方法（比如积分梯度）来做置信度评估，不要盲信注意力分数。
如果你打算把这个方法用于实际的化合物筛选管线，我的建议是不要把 Graph-pKa 当作最终的判断工具，而是当作高通量预筛的第一级：先用它把上万化合物快速过滤到几百个候选，再用量子化学方法对这少数候选做高精度验证。这个组合策略在成本和精度之间取得了很理想的平衡，我自己在几个内部项目中都是这么用的，整体效果比单纯依赖任何一端都好得多。
pKₐ 预测这件事，本质上是一个“从分子结构中提取电子环境信息”的问题。Graph-pKa 的优雅之处在于用多实例学习把这个问题抽象成两个清晰子任务：位点定位和数值回归，再交给图神经网络统一求解。我个人的体会是，深度学习方法在这个领域的瓶颈并不在模型容量，而在数据质量和归纳偏置的有效性。如果你的分子库中可电离位点的类型和训练集分布差异太大，无论网络结构多先进都会面临外推困难。因此在实际应用中，先分析目标分子的化学空间，再决定是否微调模型，往往比直接套用预训练模型的效果要好很多。
