[论文解读] Meta-learning Adaptive Deep Kernel Gaussian Processes for Molecular Property Prediction

Wenlin Chen, Austin Tripp+1 more|arXiv (Cornell University)|May 5, 2022

Metabolomics and Mass Spectrometry Studies被引用 8

一句话总结

该论文提出 ADKF-IFT，一种元学习框架，通过隐函数定理（IFT）进行双层优化，联合优化元学习的特征提取器和任务特定的核参数，以实现高效的梯度计算，从而训练深度核高斯过程。该方法在少样本分子性质预测任务中达到最先进性能，在分布内和分布外任务上均优于先前方法，且无需敏感的超参数调优即可实现泛化与适应之间的平衡。

ABSTRACT

We propose Adaptive Deep Kernel Fitting with Implicit Function Theorem (ADKF-IFT), a novel framework for learning deep kernel Gaussian processes (GPs) by interpolating between meta-learning and conventional deep kernel learning. Our approach employs a bilevel optimization objective where we meta-learn generally useful feature representations across tasks, in the sense that task-specific GP models estimated on top of such features achieve the lowest possible predictive loss on average. We solve the resulting nested optimization problem using the implicit function theorem (IFT). We show that our ADKF-IFT framework contains previously proposed Deep Kernel Learning (DKL) and Deep Kernel Transfer (DKT) as special cases. Although ADKF-IFT is a completely general method, we argue that it is especially well-suited for drug discovery problems and demonstrate that it significantly outperforms previous state-of-the-art methods on a variety of real-world few-shot molecular property prediction tasks and out-of-domain molecular property prediction and optimization tasks.

研究动机与目标

• 为解决在数据有限条件下分子性质预测中可靠不确定性估计的挑战，特别是在药物发现中的贝叶斯优化场景。
• 克服在小样本数据集上训练深度核高斯过程（DKGPs）时常见的过拟合问题，该问题是药物发现中的常见挑战。
• 将现有方法——深度核学习（DKL）与深度核迁移（DKT）——统一并推广为一个单一、合理的框架。
• 开发一种方法，利用元学习获取可泛化的特征表示，同时避免先前方法中存在的强假设或超参数敏感性。

提出的方法

• ADKF-IFT 将深度核高斯过程的训练表述为双层优化问题：元学习特征提取器参数，以最小化跨任务的平均预测损失。
• 利用隐函数定理（IFT）计算内层优化的梯度，实现通过反向传播的端到端训练，并支持高效的二阶更新。
• 将元学习参数（如特征提取器权重）与任务自适应参数（如核超参数）分离，实现解耦优化。
• 一种实用的实现方式是元学习所有特征提取器参数，并仅对每个任务自适应调整核超参数，避免海森矩阵近似，提升可解释性。
• 该框架将 DKL（单任务学习）和 DKT（纯元学习）作为特例统一，提供统一的训练范式。
• 该方法支持精确的不确定性量化，特别适用于少样本和分布外的分子性质预测。

实验结果

研究问题

• RQ1 能否开发一个统一框架，结合元学习与深度核学习的优势，用于分子性质预测？
• RQ2 如何在不依赖强先验或敏感超参数的前提下，将元学习应用于深度核高斯过程，以防止小样本数据集上的过拟合？
• RQ3 联合元学习特征提取器与任务特定的核参数自适应，是否能带来比现有方法更优的泛化性能与不确定性校准？
• RQ4 所提出的方法是否能在分布内与分布外分子性质预测基准上均超越当前最先进方法？

主要发现

• 在 FS-Mol 基准的 111 个少样本回归任务中，ADKF-IFT 显著优于先前最先进方法，在所有支持集大小下均实现更低的均方误差。
• 在 FS-Mol 基准的 157 个少样本分类任务中，ADKF-IFT 表现优于所有基线模型，且随着支持集大小增加，性能增益保持一致。
• 该方法在分布外分子性质预测中展现出强大的泛化能力，表明其在训练分布之外也具有鲁棒性。
• ADKF-IFT 避免了对敏感超参数调优的需求（例如，对比损失中无需 γ），而先前方法如 ADKL-GP 则需精细平衡以避免过拟合或崩溃。
• 该框架具有可解释性且可扩展，其提出的实例化方式可实现精确的海森矩阵计算，从而支持稳定高效的训练。
• 实证结果证实，ADKF-IFT 统一了 DKL 与 DKT，消融研究进一步表明，元学习的特征提取器是性能提升的关键因素。