训练 ADMET 多任务模型：一次预测多个性质

ADMET 建模的典型困境是每个终点的数据都不多，但终点之间有相关性——溶解度与 logP 相关、通透性与分子量相关。多任务学习让这些终点共享表示，理论上数据少的终点能从数据多的终点借力。

数据格式：稀疏标签

import pandas as pd

# 每列一个终点，缺失值留空 —— Chemprop 会自动只在有标签处回传梯度
data = pd.DataFrame({
    "smiles":     ["CCO", "CCN", "CCC", "CCCl"],
    "logS":       [0.5, None, -1.2, None],
    "logD":       [-0.3, 0.8, None, 1.5],
    "hERG_pIC50": [None, None, 4.2, 5.1],
    "HLM_CLint":  [12.0, 45.0, None, None],
})
data.to_csv("admet_multitask.csv", index=False)

# 检查标签稀疏度
for col in data.columns[1:]:
    n = data[col].notna().sum()
    print(f"{col:12s} {n:5d} 条 ({n/len(data):.1%})")

训练多任务模型

chemprop train \
  --data-path admet_multitask.csv \
  --task-type regression \
  --split-type scaffold_balanced \
  --num-folds 5 --ensemble-size 5 \
  --epochs 60 \
  --descriptors-generator rdkit_2d_normalized \
  --save-dir ckpt_admet/ \
  --metrics rmse mae

Chemprop 原生支持稀疏标签：某个分子在某任务上没有标签时，该任务的损失不计入梯度。这是它做 ADMET 多任务的主要便利。

混合回归与分类任务

ADMET 中常有回归（溶解度）与分类（是否为 P-gp 底物）混合。Chemprop 单次训练只支持一种任务类型，需分开处理：

# 方案一：分开训两个模型（最简单可靠）
chemprop train --data-path admet_regression.csv --task-type regression ...
chemprop train --data-path admet_classification.csv --task-type classification ...

# 方案二：把分类转成回归（若有连续的原始值）
# 如「是否 hERG 阳性」→ 直接用 pIC50 回归，更多信息

# 方案三：自己用 PyTorch 实现共享编码器 + 多个任务头
import torch, torch.nn as nn

class MultiTaskHead(nn.Module):
    def __init__(self, enc_dim, reg_tasks, clf_tasks):
        super().__init__()
        self.shared = nn.Sequential(
            nn.Linear(enc_dim, 512), nn.ReLU(), nn.Dropout(0.2))
        self.reg_head = nn.Linear(512, reg_tasks)
        self.clf_head = nn.Linear(512, clf_tasks)
    def forward(self, h):
        z = self.shared(h)
        return self.reg_head(z), self.clf_head(z)

def masked_loss(pred, target, mask, loss_fn):
    """只在有标签处计算损失"""
    if mask.sum() == 0:
        return torch.tensor(0.0, device=pred.device)
    return loss_fn(pred[mask], target[mask])

多任务真的有帮助吗：必须验证

多任务不是必然有益的——任务间若无相关性，共享表示反而会互相干扰（负迁移）。必须做对照：

results = {}

# 对照一：每个终点单独训练
for task in tasks:
    single_df = data[["smiles", task]].dropna()
    single_df.to_csv(f"single_{task}.csv", index=False)
    # chemprop train --data-path single_{task}.csv ...
    results[f"single_{task}"] = rmse_single

# 对照二：多任务训练后，逐任务评估
# results[f"multi_{task}"] = rmse_multi

# 比较
for task in tasks:
    s, m = results[f"single_{task}"], results[f"multi_{task}"]
    delta = (s - m) / s * 100
    flag = "✓ 受益" if delta > 3 else ("✗ 受损" if delta < -3 else "— 无差异")
    print(f"{task:12s} 单任务 {s:.3f} → 多任务 {m:.3f}  ({delta:+.1f}%) {flag}")

典型结果：数据量小的终点（几百条）常从多任务受益明显；数据量大的终点（上万条）通常无差异甚至略有下降。据此可以决定哪些终点放进多任务、哪些单独训。

任务权重与量纲

# 不同终点的量纲差别很大（logS 范围约 -8~2，CLint 可能 0~300），
# 不做归一化会让大量纲任务主导梯度

from sklearn.preprocessing import StandardScaler

scalers = {}
for task in reg_tasks:
    mask = data[task].notna()
    sc = StandardScaler()
    data.loc[mask, task] = sc.fit_transform(data.loc[mask, [task]])
    scalers[task] = sc          # 保存，预测后要反变换

# Chemprop 会自动做目标归一化，自建模型时必须自己处理

# 预测后反变换回原尺度
pred_original = scalers[task].inverse_transform(pred_scaled.reshape(-1, 1))

用哪些终点组合

组合	相关性	多任务效果
logP / logD / 溶解度	强	通常有益
Caco-2 / MDCK / PAMPA 通透性	强	通常有益
各 CYP 亚型抑制	中	常有益
微粒体 / 肝细胞清除率	中强	常有益
hERG 与溶解度	弱	可能无益
毒性终点与理化性质	弱	需验证

延伸资源

• 工具：174《Chemprop》；数据集：244《TDC ADMET Group》、173《TDC》；
• 上一步：333《训练 Chemprop 活性预测模型》；评测：335《用 TDC Benchmark 评估模型》；ADMET 概念见「成药性」模块。