[论文解读] Enforcing Predictive Invariance across Structured Biomedical Domains

Wengong Jin, Regina Barzilay+1 more|arXiv (Cornell University)|Jun 6, 2020

Machine Learning in Materials Science参考文献 33被引用 8

一句话总结

该论文提出后悔最小化（RGM）及其结构化扩展（SRGM），通过在分子骨架和蛋白家族等复杂结构化环境中强制预测不变性，以提升生物医学AI中的模型泛化能力。通过将泛化问题建模为对拥有事后访问未见环境能力的“预言者”的后悔最小化，并利用基于梯度的扰动来模拟未见的领域变化，RGM/SRGM在分子性质预测（QM9数据集上MAE为41.7 vs. 52.3）和蛋白质稳定性预测（Spearman相关系数为0.793 vs. 0.73）任务上达到最先进性能。

ABSTRACT

Many biochemical applications such as molecular property prediction require models to generalize beyond their training domains (environments). Moreover, natural environments in these tasks are structured, defined by complex descriptors such as molecular scaffolds or protein families. Therefore, most environments are either never seen during training, or contain only a single training example. To address these challenges, we propose a new regret minimization (RGM) algorithm and its extension for structured environments. RGM builds from invariant risk minimization (IRM) by recasting simultaneous optimality condition in terms of predictive regret, finding a representation that enables the predictor to compete against an oracle with hindsight access to held-out environments. The structured extension adaptively highlights variation due to complex environments via specialized domain perturbations. We evaluate our method on multiple applications: molecular property prediction, protein homology and stability prediction and show that RGM significantly outperforms previous state-of-the-art baselines.

研究动机与目标

• 解决生物医学模型在训练数据分布之外的泛化挑战，特别是在环境具有结构化且稀疏表示的场景下（如分子骨架、蛋白超家族）。
• 克服现有基于不变性的方法（如IRM）在单样本或高度结构化环境中的局限性，这些方法因对表征学习的约束较弱而表现不佳。
• 开发一种后悔最小化框架，量化相对于拥有事后访问未见环境能力的预言者的泛化性能，从而实现更强的不变性诱导。
• 通过使用基于梯度的表征扰动来模拟未显式定义的领域变化，将后悔框架扩展至结构化领域。
• 在真实世界生物医学基准上验证所提方法的优越性，包括分子性质预测、蛋白质稳定性预测与同源分类任务。

提出的方法

• 将不变风险最小化（IRM）重构为后悔最小化问题：模型最小化在有无未见环境访问条件下训练的预测器之间的差异（即后悔）。
• 通过从训练数据中采样子集来定义未见环境 $E_e$，并计算在有无 $E_e$ 情况下训练的预测器之间的损失差异作为后悔值。
• 对于结构化环境（如分子骨架），通过使用独立分类器 $g$（用于预测领域，如骨架或拓扑）对表征 $\phi$ 进行基于梯度的扰动，来模拟未见的领域变化。
• 扰动 $\nabla_\phi \mathcal{L}_{\text{cls}}(g, \phi)$ 突出显示骨架特异性变化，对扰动样本计算的后悔值衡量对未见骨架类型的泛化能力。
• 通过端到端优化训练主模型，以最小化在多个未见环境采样和扰动方案下的后悔值。
• 使用聚类将稀有或稀疏的环境（如少于10个样本的蛋白超家族）划分为两个粗粒度环境 $E_0, E_1$，以稳定标准基线的训练；而SRGM则直接通过基于梯度的扰动在结构化领域上运行。

实验结果

研究问题

• RQ1 在环境稀疏或结构化设置下，后悔最小化是否能为模型提供强于IRM的不变性归纳偏置？
• RQ2 在未显式建模每个环境的前提下，基于梯度的表征扰动如何提升在结构化生物医学领域中的泛化能力？
• RQ3 利用结构化领域信息（如分子骨架、蛋白拓扑）在多大程度上能提升分子与蛋白质建模任务中的泛化性能？
• RQ4 在分布外与骨架划分评估协议下，所提方法是否优于现有领域泛化基线（如IRM、MLDG、CrossGrad）？
• RQ5 当训练与测试分布显著偏离时，该方法在不同水平的领域偏移下表现如何？

主要发现

• 在QM9数据集的骨架划分设置下，RGM将12种分子性质的平均MAE从最佳基线的52.3降低至41.7，显著提升了泛化性能。
• SRGM在蛋白质稳定性预测任务（Rocklin et al., 2017）中达到0.793的Spearman相关系数，优于此前最先进方法（0.73）与ERM基线（0.736）。
• 在同源预测任务中，SRGM达到23.8%的准确率，超过最佳基线（ERM）的22.3%，并展现出对大进化差距的鲁棒性。
• 消融研究显示，SRGM在不断增加的领域偏移水平下（如在≤8、7或6个原子的分子上训练）始终优于RGM与CrossGrad，证实其在严重分布偏移下的有效性。
• 在HIV数据集上，SRGM将准确率提升至0.735，显著优于次佳基线（CrossGrad为0.708），凸显结构化领域建模的优势。
• 该方法在多种生物医学任务中均表现出一致的性能提升，包括分子性质预测、蛋白质稳定性预测与同源性检测，验证了其在结构化、低数据环境中的泛化能力。