第 8 名私有解决方案 - 只需要一个基于 SMILES 的 Transformer 模型。以及适当的基准测试
背景
业务背景：https://www.kaggle.com/competitions/leash-BELKA/overview
数据背景：https://www.kaggle.com/competitions/leash-BELKA/data
总结
解决方案总结：基于配体、仅 SMILES、多靶点 Transformer 模型，分别在共享和非共享构建块划分上独立训练。对于非共享块，使用了特殊的命中识别数据划分技术和 k 折训练与验证。模型针对测试集的共享和非共享部分独立平均。
数据划分
这场竞赛的关键在于相似分子倾向于具有相似性质——这是初始命中化合物先导优化的主要假设。然而，正确的评估策略很少被应用，这就是为什么最近推出了 Polaris，这是一个为药物发现建立无偏基准的集体努力。此外，很难说什么是特定任务的正确方法。然而，我的同事 Simon Steshin 最近表明，对于新命中化合物识别，应考虑按 0.4 Tanimoto 相似度进行划分策略 (Lo-Hi: Practical ML Drug Discovery Benchmark)。也许我们可以应用类似的方法在其他任务中构建更具泛化能力的模型，特别是在这场竞赛中，66% 的私有分数 dedicated 于分布外分子。
非共享构建块 (Non-shared BBs)
组织者提供的分子由三个构建块组成，因此我们可以考虑按构建块的相似度进行划分。让我们检查测试构建块与训练构建块的分布：
训练集和测试集之间有许多共享的构建块（最大 Tanimoto 为 1.0），我们显然需要避免在我们的非共享训练/验证划分中出现相同的情况。我们看到 BB1 不与 BB2 相交，而 BB2 与 BB3 显著相交。除了明显的重复外，还有许多相似的构建块，根据上述假设，这将给我们的评估引入强烈的偏差。
基于构建块的命中识别划分 (HIBB)
所以，让我们消除这种偏差。我们将使用来自 Lo-Hi 基准的 Hi 划分算法，该算法解决平衡顶点最小 k-割问题，以构建训练集和验证集，使得两个集合之间最近的分子至少相隔截止 Tanimoto 距离 (代码)。为了决定使用什么阈值，让我们查看每个块位置的训练集 Tanimoto 相似度分布，仅考虑相同位置（我们稍后将处理 BB2-BB3 交叉位置）：
设置什么阈值是一个开放性问题，最有可能的是使用多个阈值是最好的方法。我选择了接近每个位置相似度平均值的阈值（别问我为什么）：BB1<->BB2 为 0.7，(BB2+BB3<->BB2+BB3) 为 0.4。是的，我合并了 BB2 和 BB3，因为它们已知相互交叉。好吧，这个划分是最难的。在其上训练模型会很快显示 nothing is working（换句话说，结构相似性驱动了与活性的大量相关性）。
注意构建块的独占性——如果一个块出现在一个集合中，它不能以任何可能的组合出现在另一个集合中。
加权构建块划分 (BB)
虽然 Hi 划分很有用，但对于竞赛问题来说可能过于悲观，并且由于将构建块独占分配给训练/验证集而删除了大量数据。因此，我决定创建一个更简单的划分，接近其他参与者在复制 host 划分时所做的事情，但在几个方面有所不同：
- 块再次对训练/验证集独占
- 与 @thedrcat 的出色 notebook 以及 @roberthatch 和 @hengck23 的 工作 中的 BB 划分相比，每个位置的块数量增加了 3 倍
- 有滚动 5 折，对于每一折，在每个位置选择不同的块
- 训练数据由 `任意阳性` 分子（至少结合一个靶点的分子）和从相应折的构建块中随机采样的相同数量的 `全阴性` 分子（不结合所有靶点）混合组成
- 验证数据由 `任意阳性` 集合中剩余的部分和一定数量的 `全阴性` 化合物混合组成，以大致匹配原始训练数据集的不平衡，即 `任意阳性` 率约为 1.5%
总的来说，折看起来像这样（为简单起见不考虑随机化）：
| 折 (Fold) | BB1 | BB2&BB3 | BB3\BB2 | 训练大小 | 训练阳性率 | 验证大小 | 验证阳性率 |
|---|---|---|---|---|---|---|---|
| 1 | 0-50 | 0-101 | 0-5 | 1 875 494 | 50% | 356 620 | 1.67% |
| 2 | 51-101 | 102-203 | 6-11 | 1 777 354 | 50% | 406 629 | 1.62% |
| 3 | 102-152 | 204-305 | 12-17 | 1 430 522 | 50% | 363 503 | 2.96% |
| 4 | 153-203 | 306-407 | 18-23 | 1 924 372 | 50% | 311 943 | 1.47% |
| 5 | 204-254 | 408-509 | 24-29 | 1 903 992 | 50% | 313 838 | 1.47% |
与 HIBB 划分一样，每折的验证集大小相似，但我们现在有更多的训练数据：
共享构建块 (Shared BBs)
训练集和测试集之间共享构建块的策略很明确——我们只需要过拟合已知的构建块，同时覆盖尽可能多的数据。因此，在这里，我使用了所有 9800 万个分子，其中我随机选择了 1% 用于验证，99% 用于训练。
模型与训练
简单来说，我使用的训练策略如下：
- 共享构建块：过拟合的同时跟踪验证性能，取最后一个检查点
- 非共享构建块：取验证集上的最佳检查点（HIBB 和 BB 划分策略相同）
我尝试了大量的模型，并对以下模型进行了广泛的超参数调优（特别是在非共享构建块上）：
- MolFormer, RoBERTa ZINC, ChemBERTa：都显示出相似的质量，MolFormer 和自定义 RoBERTa 在集成中表现更好
- GPS++：具有原子和键特征的 SOTA 全注意力 GNN（实现 adopted from Graphium）。对于原子，使用以下特征：原子序数、族、周期、总价电子数、度、形式电荷、自由基电子、杂化、手性、隐含价、氢原子数、芳香性、环内、电负性。对于键：键类型 onehot、立体化学、环内
- MPNN++：具有相同原子和键特征的 GNN（实现 adopted from Graphium），见 scaling GNNs paper 中的性能。在我的案例中，MPNN++ 的表现也 substantially better than GPS++，可能是由于没有 attending everything to everything，从而具有更强的归纳偏置
- 构建块与 top-50 最近分子的 Tanimoto 相似度分数放在一起的多头注意力模型，其中 queries/keys 是相似度，values 是真实值。据推测，这种模型（我称之为
SimAttn模型）可以从最近的排名分子列表中学习，它确实做到了，但我无法用它达到足够高的 AP - 基于 RDKit 210 描述符（用标准 scaler 标准化）和 ECFP4 的 XGBoost
- 基于 RDKit + ECFP4 + 相似度特征的 MLP
- 基于 RDKit + ECFP4 + 相似度特征 + 表现最好的 MolFormer 的 Meta MLP 模型：在 BB 划分上性能略有提高，但不足以说服我提交
综上所述，表现最好的模型是微调的 MolFormer（4500 万参数）和自定义 RoBERTa（约 800 万参数），如 我公司最近的论文 所述。自定义主要涉及 500 大小的 BPE SMILES tokenizer，30% 的情况下 15% 的掩码 token，以及 50% 的情况下的 SMILES 重新枚举。
结论
我对未见化学子空间泛化能力的看法：
- 这不仅对于这场竞赛，而且在公共领域的任何地方都不是一个已解决的问题
- 我们可以通过广泛的基准测试来解决这个问题
- 基于 SMILES 的 Transformer 模型的表现略优于 SOTA GNN 模型
- 缩放参数数量改善了分布外分子的质量
- 由于没有预处理，缩放基于 SMILES 的 Transformer 比 GNN 容易得多，因此应考虑缩放至十亿-size 模型和数据集以提高泛化能力，作为后续研究
致谢
我感谢 Copilot 和 Continue.dev 一直作为我的编码队友。我也感谢我的同事 Simon 他在分子基准测试方面的公开工作。我感谢其他参与者的解决方案和讨论；我没有时间参与讨论，但我与其中许多人保持一致。非常感谢 Leash Bio 和 Kaggle 组织这场竞赛。对于私有 leaderboard 上被打乱的团队，我很抱歉。我在比赛结束前十天就有了 0.299 的解决方案，所以我预计许多其他团队也会发生同样的情况。我相信他们都努力破解这个未解决但引人入胜的小分子泛化问题。