[论文解读] Learning Invariant Molecular Representation in Latent Discrete Space

Xiang Zhuang, Qiang Zhang+6 more|arXiv (Cornell University)|Oct 22, 2023

Computational Drug Discovery Methods被引用 5

一句话总结

该论文提出 iMoLD，一种通过先使用图神经网络（GNN）编码分子，再利用残差向量量化（RVQ）模块在离散潜在空间中分离不变特征，从而学习分子表示不变性的新框架。该方法采用与任务无关的自监督目标来识别因果性、对分布偏移具有鲁棒性的特征，在18个分子数据集上实现了最先进（SOTA）的分布外（OOD）泛化性能，尤其在协变量偏移和概念偏移场景下表现优异。

ABSTRACT

Molecular representation learning lays the foundation for drug discovery. However, existing methods suffer from poor out-of-distribution (OOD) generalization, particularly when data for training and testing originate from different environments. To address this issue, we propose a new framework for learning molecular representations that exhibit invariance and robustness against distribution shifts. Specifically, we propose a strategy called ``first-encoding-then-separation'' to identify invariant molecule features in the latent space, which deviates from conventional practices. Prior to the separation step, we introduce a residual vector quantization module that mitigates the over-fitting to training data distributions while preserving the expressivity of encoders. Furthermore, we design a task-agnostic self-supervised learning objective to encourage precise invariance identification, which enables our method widely applicable to a variety of tasks, such as regression and multi-label classification. Extensive experiments on 18 real-world molecular datasets demonstrate that our model achieves stronger generalization against state-of-the-art baselines in the presence of various distribution shifts. Our code is available at https://github.com/HICAI-ZJU/iMoLD.

研究动机与目标

• 为解决分子表示学习中在不同数据源引起的分布偏移下表现不佳的分布外（OOD）泛化问题。
• 克服现有“先分离后编码”方法在处理纠缠分子结构时的局限性，并避免对环境的假设。
• 开发一种与任务无关、鲁棒的框架，适用于回归和多标签分类等多种分子任务。
• 通过识别在分布偏移下保持稳定的不变因果分子特征，提升泛化能力。
• 设计一种自监督目标，实现对不变性的精确识别，而无需任务特定的监督信号。

提出的方法

• 提出“先编码后分离”策略：首先使用 GNN 将完整分子图编码为潜在表示，再进行特征分离。
• 引入残差向量量化（RVQ）模块对潜在空间进行离散化，以减少过拟合，同时保持编码器的表达能力。
• 使用评分 GNN 从离散潜在表示中识别并提取不变子结构。
• 设计一种与任务无关的自监督学习目标，以鼓励识别在不同环境中保持不变的特征。
• 使用一致性损失将编码器输出与学习到的码书对齐，确保训练过程稳定。
• 该框架可直接应用于多种下游任务（如回归和多标签分类），无需修改网络架构。

Figure 1: (a) First-Separation-Then-Encoding: the input is separated using a subgraph generator, then each subgraph is encoded respectively. (b): First-Encoding-Then-Separation: the input is encoded, then the representation is separated by a scorer.

实验结果

研究问题

• RQ1 ‘先编码后分离’策略是否在捕捉不变分子特征方面优于传统的‘先分离后编码’方法？
• RQ2 残差向量量化（RVQ）在不损失表示能力的前提下，是否能有效提升 OOD 泛化性能？
• RQ3 所提出的与任务无关的自监督目标在多种分子任务中对不变性识别的增强程度如何？
• RQ4 该模型在各种分布偏移（包括协变量偏移和概念偏移）下的表现如何？
• RQ5 各组件（如 RVQ、不变性目标）对最终 OOD 泛化性能的贡献程度如何？

主要发现

• 所提出的 iMoLD 框架在 18 个真实世界分子数据集上，于多种分布偏移场景（包括协变量偏移和概念偏移）下均实现了最先进性能。
• 在 GOOD-HIV-Scaffold 数据集上，移除 VQ 模块导致性能显著下降，证明其在 OOD 泛化中起关键作用。
• 在更复杂的 GOOD-PCBA 数据集上，移除残差连接导致性能明显下降，表明其对维持模型容量至关重要。
• 消融实验确认，不变性学习目标 $\mathcal{L}_{\textrm{inv}}$ 和一致性损失 $\mathcal{L}_{\textrm{cmt}}$ 均为实现鲁棒性能所不可或缺。
• t-SNE 可视化显示，iMoLD 生成的特征分布更均匀，且环境间距离更小（Wasserstein 距离），证实其具有更好的不变性。
• iMoLD 在验证集上得分更高，且训练得分与验证得分更接近，表明其过拟合更少，泛化能力优于 ERM 基线方法。

Figure 2: An overview of distribution shifts in molecules. Distribution shifts occur when molecules originate from different scaffold, size or assay environments.