自监督学习基本原理

做分子ADMET预测这两年，我最直观的感受就是：标注数据太难搞了。一个分子体外实验做下来，成本高、周期长，能拿到的样本通常只有几百到几千条。但无标注的分子结构，PubChem、ZINC上一抓一大把，几千万个躺着没人用。自监督学习正好解决这个问题——先让模型自己从大量无标签数据里学分子表征，再用小样本微调。今天我们就来拆解三大主流自监督范式的核心思想。

对比学习（Contrastive Learning）

对比学习的思路其实很朴素。举个例子：你手里拿着两个SMILES序列——CCO（乙醇）和CCCO（丙醇），它们很相似。再拿一个c1ccccc1（苯），跟乙醇明显不同。我们希望模型学会把乙醇和丙醇的表征拉近，把乙醇和苯的表征推远。这就是对比学习的基本假设。

我个人习惯把它拆成三步：

1. 构造正负样本对 —— 对同一个分子做两次不同的数据增强，得到两个“视图”，它们互为正样本；不同分子之间互为负样本。
2. 编码器提取特征 —— 把增强后的分子图（或SMILES）送入GNN或Transformer，得到表征向量。
3. 对比损失拉近推开 —— 最简单的是InfoNCE损失，最大化正样本对之间的一致性，最小化负样本之间的一致性。

我之前在项目中用过MolCLR框架。它对分子图做原子掩码、边删除和子图裁剪这三种增强。一开始我不知道顺序敏感，随便乱裁，结果正样本对里包含完全不同的官能团，模型学到的表征一塌糊涂。后来改成保留核心骨架、只扰动外围侧基，效果才稳定下来。

对比学习的核心代码框架其实很简洁。以PyTorch为例，一个典型的最小实现大致长这样：

class ContrastiveModel(nn.Module):
    def __init__(self, encoder, projector):
        super().__init__()
        self.encoder = encoder          # GNN或Transformer
        self.projector = projector      # MLP映射到头空间

    def forward(self, x1, x2):
        h1 = self.encoder(x1)
        h2 = self.encoder(x2)
        z1 = self.projector(h1)
        z2 = self.projector(h2)
        # 计算InfoNCE损失，这里省略具体实现
        loss = info_nce_loss(z1, z2)
        return loss

你可能会问：为什么需要一个额外的projector？我刚开始也想不通。后来看SimCLR的论文才知道，projector的作用是过滤掉与增强方式相关的噪声信息，迫使编码器保留真正对语义判别有用的特征。嗯，这个细节挺值得注意。

生成式预训练（Generative Pretraining）

生成式预训练的核心逻辑是：让模型学会预测序列中的下一个元素。在NLP领域GPT已经验证了这个方法的威力。在分子领域，我们通常把SMILES当作一个字符串语言模型来处理。

例如分子“乙醇”，其规范SMILES是CCO。模型从左到右依次看到字符C → CC → CCO，每一步的任务都是预测下一个字符。本质上就是一个因果语言模型（Causal LM）。

我试过在几百GB的ZINC SMILES数据上预训练一个8层Transformer。训练完成之后，用它的中间层表征去做ADMET分类，效果比随机初始化好了一截。但要注意：SMILES的语法是严格定义的，比如括号必须成对、环必须正确闭合。模型如果生成非法SMILES，预训练毫无意义。所以训练时我额外加了一层合法性校验——虽然慢点，但值。

生成式预训练的损失函数就是标准的交叉熵，和语言模型一模一样。下面是一个训练循环的伪代码示例：

for batch in dataloader:
    # batch是填充后的SMILES序列 [batch_size, seq_len]
    logits = model(batch)                # 输出 [batch_size, seq_len, vocab_size]
    shift_logits = logits[:, :-1, :]
    shift_labels = batch[:, 1:]
    loss = cross_entropy(shift_logits, shift_labels)
    loss.backward()
    optimizer.step()

掩码建模（Masked Modeling）

掩码建模的思想来自BERT。随机遮住输入的一部分，让模型去预测被遮住的内容。在分子领域，我们不仅可以遮SMILES中的字符，更流行的是在分子图上遮节点或边。

举个例子：一个苯环分子图有6个碳原子和6条键。我们随机遮住其中2个碳原子的特征（比如原子类型设为[MASK]），模型需要根据周围原子的环境推断被遮原子的类型。这强迫模型学到原子之间的局部环境和长程依赖关系。

我自己做过一个实验，分别用原子层级掩码和边层级掩码去预训练图神经网络。结果发现：原子掩码对预测分子极性（logP）有帮助，边掩码对预测溶解度更有效。单独用原子掩码时，模型在溶解度任务上比边掩码低3个点。后来我干脆两者都做——同时遮节点和边，让模型同时推断——ADMET五项指标平均提升了约1.5%。

掩码建模的训练流程也很直观：

1. 对分子图进行数据增强，随机选择15%的节点。
2. 以80%概率用[MASK]替代，10%概率随机替换为其他原子类型，10%概率保持不变。
3. 用图编码器（如GIN）提取每个节点的表征。
4. 在[MASK]位置接一个分类头，预测该原子类型。

下面是用PyTorch Geometric实现分子图掩码建模的核心片段：

import torch
from torch_geometric.utils import mask_feature

def mask_node_features(data, mask_ratio=0.15):
    # data.x: [num_nodes, num_features]
    mask = torch.rand(data.x.size(0)) < mask_ratio
    masked_x = data.x.clone()
    # 80%替换为[MASK]向量（这里是全零向量）
    masked_x[mask] = 0.0
    # 10%随机替换
    random_idx = torch.where(torch.rand(data.x.size(0)) < 0.1 & mask)[0]
    if len(random_idx) > 0:
        random_vec = torch.randn_like(masked_x[random_idx])
        masked_x[random_idx] = random_vec
    # 剩余10%保持原样（已经没有被mask的x）
    return masked_x, mask

三种方法的对比

从上表可以看出来，对于ADMET预测这种分子性质回归/分类任务，对比学习和掩码建模是当前的主流打法。生成式预训练更适用于分子生成、逆合成预测等场景。我个人的策略是：如果预算有限，优先跑对比学习；如果手头图数据多且结构复杂，掩码建模往往能带来更稳定的提升。