4. 分子图数据集划分与泄漏防范

在分子图模型训练中，数据集划分不是简单的“切一刀”，而是决定模型评估是否可信的关键工程环节。分子数据具有天然的结构相似性：同一骨架的衍生物、同一实验批次测得的化合物、同一文献来源的类似物，往往在特征空间中高度接近。如果划分时把这些相似样本分散到训练集和测试集，模型在测试集上的表现会被高估，这就是数据泄漏。本节围绕四种主流划分策略——随机划分、骨架划分、时间划分、相似性划分——讲清它们各自适用的场景、实现方式，以及如何检测和处理泄漏。

四种划分策略的适用逻辑

随机划分是最直观的方式：把全部样本打乱后按比例切分，例如 8:1:1 得到训练、验证、测试集。它假设样本独立同分布，适合数据来源单一、分子多样性高、没有明显批次或时间结构的场景。实现上只需固定随机种子，用 sklearn.model_selection.train_test_split 或 random_split 即可。但随机划分对分子数据往往过于乐观：同一骨架的多个变体可能同时落入训练和测试，导致测试集“见过类似题”。

骨架划分（scaffold split）按 Murcko 骨架分组，保证同一骨架的所有分子只出现在一个子集中。这是药物发现领域评估模型泛化能力的常用标准。实现时先用 RDKit 计算每个分子的 Murcko 骨架，再以骨架为分组键做 GroupShuffleSplit。它的核心价值是回答“模型能否对全新骨架做出预测”，适合先导化合物优化、虚拟筛选等任务。代价是当骨架种类少或分布极不均衡时，测试集可能过小或代表性不足。

时间划分按数据产生或收录的时间顺序切分，训练集用早期数据，测试集用后期数据。它模拟真实部署场景：用历史数据训练，预测未来新出现的分子。实现时按时间戳或数据版本排序后取前若干比例作训练。时间划分能有效避免“未来信息泄漏到过去”，适合持续更新的化合物库、专利数据、实验记录等。需要注意的是，时间字段必须真实可靠，且要检查时间分布是否均匀。

相似性划分基于分子指纹（如 Morgan 指纹、MACCS 键）计算样本间相似度，把高相似样本聚到同一侧，或按相似度阈值构造“冷启动”测试集。常见做法是：对每个测试分子，要求其与训练集的最大 Tanimoto 相似度低于阈值（如 0.7），从而构造出真正“陌生”的测试集。它适合评估模型对新颖化学空间的探索能力，但计算量较大，且阈值选择需要结合任务调整。

数据泄漏的常见来源与检测方法

数据泄漏并不总是显式的重复样本。在分子图任务中，它常以更隐蔽的形式出现。第一类是骨架泄漏：训练集和测试集包含同一 Murcko 骨架的分子，模型只需记住骨架相关的子结构就能“猜中”测试标签。检测方法是提取两侧分子的骨架集合，计算交集比例；若交集占比超过 5%～10%，就应警惕。第二类是相似性泄漏：即使骨架不同，分子也可能因取代基模式高度相似而接近。可以计算测试集每个分子与训练集的最大 Tanimoto 相似度，若大量测试样本的相似度高于 0.8，说明测试集并不“陌生”。第三类是时间泄漏：训练集包含了测试集之后才出现的信息，例如用 2023 年的数据训练、用 2022 年的数据测试。检测时按时间排序，确认训练集时间戳全部早于测试集。第四类是批次或来源泄漏：同一实验批次、同一文献的分子被分到不同子集，导致模型学到批次特征而非化学规律。可以按来源字段做分组统计，检查各子集的来源分布是否重叠。

此外还有标签泄漏：某些特征在真实预测时不可获得，却出现在训练特征中。例如用“是否通过临床二期”作为特征去预测“是否通过临床三期”，而前者在预测时点尚未知。这类问题需要结合业务时间线逐特征审查，不能只靠统计检测。

处理策略与工程落地

检测到泄漏后，处理方式取决于泄漏类型和任务目标。对于骨架泄漏，最直接的做法是改用骨架划分重新切分；若必须保留随机划分，则应在训练集中剔除与测试集骨架相同的分子，或把测试集限制为“训练集未见骨架”。对于相似性泄漏，可以设定相似度阈值，把与训练集过于接近的测试分子移出测试集，或将其归入训练集。对于时间泄漏，严格按时间切分，并确保特征工程只使用训练集时间点之前可获取的信息。对于批次泄漏，按批次或来源做分组划分，保证同一批次只出现在一个子集中。

工程落地时，建议把划分与检测写成可复现的脚本，固定随机种子并记录划分版本。下面是一个骨架划分加泄漏检查的示例流程，参数为教学示例，实际使用时需根据数据集调整：

from rdkit import Chem
from rdkit.Chem.Scaffolds import MurckoScaffold
from sklearn.model_selection import GroupShuffleSplit
import numpy as np

# 1. 计算 Murcko 骨架作为分组键
def get_scaffold(smiles):
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return "INVALID"
    return MurckoScaffold.MurckoScaffoldSmiles(mol=mol)

scaffolds = [get_scaffold(s) for s in smiles_list]

# 2. 按骨架分组划分，8:1:1
gss1 = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
train_idx, temp_idx = next(gss1.split(smiles_list, groups=scaffolds))
gss2 = GroupShuffleSplit(n_splits=1, test_size=0.5, random_state=42)
val_idx, test_idx = next(gss2.split(
    [smiles_list[i] for i in temp_idx],
    groups=[scaffolds[i] for i in temp_idx]))

# 3. 泄漏检查：训练集与测试集骨架交集
train_scaf = set(scaffolds[i] for i in train_idx)
test_scaf  = set(scaffolds[i] for i in test_idx)
overlap = train_scaf & test_scaf
print(f"骨架重叠数: {len(overlap)}, 占比: {len(overlap)/len(test_scaf):.2%}")
# 期望输出：骨架重叠数: 0, 占比: 0.00%

运行后若骨架重叠占比为 0，说明骨架划分生效；若不为 0，需要检查分组键是否计算正确，或是否存在无效 SMILES 被归入同一组。对于相似性划分，可以进一步计算测试集与训练集的最大相似度分布，若 95 分位数低于设定阈值（如 0.7），则认为测试集足够“冷”。

最后，划分不是一次性的动作。当数据更新、任务变化或发现新的泄漏模式时，应重新执行划分与检测，并保留每次划分的版本记录。把划分脚本、随机种子、骨架重叠率、相似度分布等指标纳入实验日志，才能让模型评估结果可追溯、可比较、可复现。