当你的模型不知道它不知道什么：一项药物发现的机器学习项目
使用图注意力网络、MC Dropout不确定性量化和贝叶斯超参数优化构建分子结合亲和力预测管道
在药物发现过程中，一个错误但自信的预测会将研究人员引入价值百万美元的死胡同。一个知道自己不知道什么的模型，比一个不知道自己不知道什么的模型更有价值。
这是一个关于我如何构建端到端机器学习管道的故事，该管道用于预测分子对EGFR（一种驱动癌症的激酶）的结合亲和力，并对每个预测进行不确定性量化。
问题
EGFR（表皮生长因子受体）是一种蛋白质，当它过度表达或突变时，会在非小细胞肺癌、结直肠癌和几种其他肿瘤类型中驱动不受控制的细胞生长。像厄洛替尼(Erlotinib)和吉非替尼(Gefitinib)这样的药物已经针对它。本项目提出的问题是：对于一个我们从未测试过的分子，它与EGFR的结合强度如何，以及我们对这个答案的置信度如何？
数据集来自ChEMBL，这是世界上最大的开放生物活性数据库。经过清理后，我们拥有13,286个具有测量结合亲和力值的唯一分子。
从SMILES到数字
ChEMBL中的每个分子都以SMILES字符串形式存储。例如，厄洛替尼(Erlotinib)如下所示：
COCCOC1=C(OCCOC)C=C2C(=C1)NC(=N2)NC3=CC=CC(=C3)C#C
神经网络无法直接处理文本。我们需要将分子转换为数值表示。我使用了两种方法。
Morgan指纹
Morgan指纹将分子转换为2048位二进制向量，其中每一位编码特定化学子结构是否存在于任何原子的半径2范围内。可以将其视为分子条形码。计算速度快，长度固定，在化学信息学中广泛用作基线。
python
from rdkit.Chem import rdFingerprintGenerator
_MORGAN_GEN = rdFingerprintGenerator.GetMorganGenerator(
radius=2, fpSize=2048
)
def smiles_to_morgan(smiles):
mol = Chem.MolFromSmiles(smiles)
if mol is None:
return None
fp = _MORGAN_GEN.GetFingerprint(mol)
return np.array(fp, dtype=np.uint8)
分子图
分子图保留了实际的分子拓扑结构。原子成为节点，键成为边。每个原子携带特征，如原子序数、度数、芳香性、形式电荷、氢原子计数和环成员资格。这使模型能够直接访问分子结构，而不是哈希指纹。
python
def smiles_to_graph(smiles):
mol = Chem.MolFromSmiles(smiles)
node_feats = []
for atom in mol.GetAtoms():
node_feats.append([
atom.GetAtomicNum(),
atom.GetDegree(),
int(atom.GetIsAromatic()),
atom.GetFormalCharge(),
atom.GetNumExplicitHs(),
atom.GetNumImplicitHs(),
int(atom.IsInRing()),
])
x = torch.tensor(node_feats, dtype=torch.float)
return Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
为什么骨架分割很重要
大多数机器学习项目使用随机训练/测试分割。对于分子来说，这是一个错误。
具有相似结构的分子在化学空间中会聚在一起。随机分割允许几乎相同的分子出现在训练/测试边界的两侧，使模型进行插值而不是泛化。你的指标看起来很好，但模型实际上是在记忆数据。
Bemis-Murcko骨架分割根据核心环骨架对分子进行分组。所有共享相同骨架的分子都进入同一分割。这迫使模型泛化到真正新颖的结构。
python
from rdkit.Chem.Scaffolds import MurckoScaffold
def get_scaffold(smiles):
mol = Chem.MolFromSmiles(smiles)
if mol is None:
return None
return MurckoScaffold.MurckoScaffoldSmiles(
mol=mol, includeChirality=False
)
两种模型：MLP与图注意力网络
MLP基线
一个基于Morgan指纹的3层全连接网络。简单、快速，并且是一个出人意料的强力基线。
python
class MLP(nn.Module):
def init(self, in_dim, hidden_dims, dropout):
super().init()
layers = []
prev = in_dim
for h in hidden_dims:
layers += [nn.Linear(prev, h), nn.ReLU(), nn.Dropout(dropout)]
prev = h
layers.append(nn.Linear(prev, 1))
self.net = nn.Sequential(*layers)
def forward(self, x):
return self.net(x).squeeze(-1)
图注意力网络
GAT直接在分子图上操作。三个GATConv层使用学习到的注意力权重从邻近原子聚合信息，后接全局平均池化和回归头。注意力机制是关键区别。每个原子在聚合信息时学会对邻居赋予不同的权重。对于预测结合亲和力，某些邻居比其他邻居更重要，模型学会了哪些是重要的。
python
class GATRegressor(nn.Module):
def init(self, in_dim, hidden_dim, heads, dropout):
super().init()
self.conv1 = GATConv(in_dim, hidden_dim, heads=heads, concat=True)
self.conv2 = GATConv(hidden_dim * heads, hidden_dim, heads=heads, concat=True)
self.conv3 = GATConv(hidden_dim * heads, hidden_dim, heads=heads, concat=False)
self.mlp = nn.Sequential(
nn.Linear(hidden_dim, 64), nn.ReLU(),
nn.Dropout(dropout), nn.Linear(64, 1)
)
def forward(self, x, edge_index, batch):
x = F.elu(self.conv1(x, edge_index))
x = F.elu(self.conv2(x, edge_index))
x = F.elu(self.conv3(x, edge_index))
x = global_mean_pool(x, batch)
return self.mlp(x).squeeze(-1)
令我惊讶的发现
MLP击败了GAT
这实际上是分子机器学习中一个有据可查的现象。Morgan指纹MLP异常难以超越。这里有两个因素解释了差距：GAT在v1版本中原子特征极少，而且在CPU上训练意味着到第200轮时模型尚未收敛。验证集均方误差仍在下降。
将原子特征从3个扩展到7个，无需更改任何模型代码就改善了所有三个指标。特征化是一个高杠杆杠杆。
骨架测试集上的结果：
- MLP基线：RMSE 0.93，R平方 0.65，Spearman 0.78
- GAT v1（3个原子特征）：RMSE 1.10，R平方 0.52，Spearman 0.67
- GAT v2（7个原子特征）：RMSE 1.05，R平方 0.56，Spearman 0.72
不确定性量化：知道自己不知道什么的模型
标准神经网络在推理时是确定性的。它们产生一个单一数字，没有置信度指示。在药物发现中，这是危险的：对于模型从未见过类似结构的分子，自信的预测看起来与对经过充分研究的骨架的自信预测相同。
MC Dropout以低成本解决了这个问题。Gal和Ghahramani（2016年）的关键见解是，带有dropout的神经网络相当于高斯过程的变分近似。在推理时，我们不关闭dropout，而是保持其激活并运行多次随机前向传播。由于随机dropout掩码，每次传播产生略有不同的预测。均值等于预测值，标准差等于不确定性。
python
@torch.no_grad()
def mc_predict(model, loader, n_passes=50):
model.train() # 保持dropout激活
all_passes = []
for _ in range(n_passes):
preds = []
for x_batch, _ in loader:
preds.append(model(x_batch).cpu().numpy())
all_passes.append(np.concatenate(preds))
samples = np.stack(all_passes) # 形状: (n_passes, n_molecules)
return samples.mean(axis=0), samples.std(axis=0)
不确定性估计携带真实信号。高不确定性分子（按标准差排名前25%）在测试集上的平均绝对误差为0.72。低不确定性分子（底部25%）的MAE为0.68。该模型在很大程度上知道何时可靠性较低。
测试集预测，按预测标准差着色。黄色表示高不确定性，紫色表示低不确定性。
按实际pChEMBL排序的分子，显示平均预测值和正负1个标准差的阴影。
使用Optuna进行贝叶斯超参数优化
我使用Optuna与Tree-structured Parzen Estimator算法，而不是手动调整GAT。与网格搜索或随机搜索不同，TPE构建一个概率模型，显示哪些超参数区域产生良好结果，并将采样集中在这些区域。
python
def objective(trial):
lr = trial.suggest_float("lr", 1e-4, 1e-2, log=True)
dropout = trial.suggest_float("dropout", 0.1, 0.5)
hidden_dim = trial.suggest_categorical("hidden_dim", [64, 128, 256])
heads = trial.suggest_categorical("heads", [2, 4, 8])
batch_size = trial.suggest_categorical("batch_size", [32, 64, 128])
训练并返回验证集MSE
50次试验。每次试验训练50轮。最佳验证集MSE：1.11。
最有趣的发现：学习率以0.61的重要性得分占据主导地位。隐藏维度在测试的大小下几乎无关紧要，小于0.01。模型在测试的大小下不受容量限制。正确设置学习率比几乎所有其他因素都更重要。
一个注意事项：每次试验在CPU上大约需要20分钟。完整研究花了19个小时。我运行了两次，因为第一次不小心关闭了终端。
最佳验证集MSE在50次试验中稳步改善，最终试验时达到约1.11。
学习率以0.61的重要性得分占据主导地位。隐藏维度几乎无关紧要。
我会如何做不同
MC Dropout成本低廉但校准不佳。高不确定性和低不确定性MAE之间的差距是真实的但很小。深度集成（训练多个独立模型并聚合预测）在更多计算成本下产生更好校准的不确定性。保形预测将提供严格的、与分布无关的覆盖保证。
GAT在GPU上经过更多轮次和更丰富的特征化（涵盖手性、杂化和部分电荷）后可能会优于MLP。像GPS或Graphormer这样的图Transformer架构，结合局部消息传递和全局注意力，是自然的下一步。
代码与进一步阅读
完整代码、结果和文档可在GitHub上获取。该项目包括从ChEMBL获取数据、骨架分割、Morgan指纹和分子图特征化、MLP和GAT训练、MC Dropout不确定性量化以及使用Optuna进行贝叶斯超参数优化。
GitHub:
参考文献
Mendez等。ChEMBL：迈向生物测定数据的直接沉积。Nucleic Acids Research，2019。
Velickovic等。图注意力网络。ICLR，2018。
Gal和Ghahramani。Dropout作为贝叶斯近似。ICML，2016。
Akiba等。Optuna：下一代超参数优化框架。KDD，2019。
Bemis和Murcko。已知药物的性质。Journal of Medicinal Chemistry，1996。
【全文结束】