[论文解读] In-Context Learning for Few-Shot Molecular Property Prediction

Christopher Fifty, Jure Leskovec+1 more|arXiv (Cornell University)|Oct 13, 2023

Machine Learning in Materials Science被引用 5

一句话总结

本文提出上下文感知分子性质预测（CAMP），一种新颖的少样本分子性质预测上下文学习框架，无需微调即可直接从（分子，性质）示范样本的上下文中推断分子性质。CAMP 使用 Transformer 编码器按顺序处理分子和标签嵌入，实现了小支持集下的最先进性能，且推理速度比原型网络快 4.5 倍。

ABSTRACT

In-context learning has become an important approach for few-shot learning in Large Language Models because of its ability to rapidly adapt to new tasks without fine-tuning model parameters. However, it is restricted to applications in natural language and inapplicable to other domains. In this paper, we adapt the concepts underpinning in-context learning to develop a new algorithm for few-shot molecular property prediction. Our approach learns to predict molecular properties from a context of (molecule, property measurement) pairs and rapidly adapts to new properties without fine-tuning. On the FS-Mol and BACE molecular property prediction benchmarks, we find this method surpasses the performance of recent meta-learning algorithms at small support sizes and is competitive with the best methods at large support sizes.

研究动机与目标

• 将上下文学习——此前仅限于自然语言任务——拓展至分子性质预测领域，其中少样本学习对药物发现至关重要。
• 开发一种无需微调或梯度更新即可快速适应新分子性质的方法。
• 通过实现低延迟、高精度的推理，克服现有少样本学习方法在分子科学中的局限性，适用于大规模筛选任务。
• 探究上下文学习原理是否可成功迁移至非语言领域（如分子科学），并利用序列建模实现。
• 评估位置嵌入与示范顺序对分子少样本学习中模型泛化能力的影响。

提出的方法

• CAMP 使用分子编码器（如 MPNN）编码查询和示范分子，使用标签编码器编码性质测量值。
• 将分子和标签嵌入拼接为序列，并输入 Transformer 编码器以预测查询分子的性质。
• 模型通过单次前向传播完成推理，直接从示范样本的上下文中预测标签，无需微调或中心点计算。
• 该架构对示范顺序不变，消融研究显示位置嵌入会降低性能并增加过拟合风险。
• 该方法在小样本少样本数据集（约 5 GB）上从随机初始化训练，与基于大语言模型的上下文学习不同，无需大规模预训练。
• 消融研究对比了朴素 ICL（交错嵌入）与 CAMP，结果表明分离嵌入拼接可提升训练稳定性和性能。

实验结果

研究问题

• RQ1 上下文学习原理能否成功应用于分子性质预测这一非语言领域？
• RQ2 在分子设置中，使用分子与标签嵌入的序列建模方法是否优于基于梯度和基于度量的少样本学习基线？
• RQ3 上下文中的示范顺序如何影响模型性能？位置编码是有利还是有害？
• RQ4 在小数据集上从随机初始化训练的模型是否能在无微调的情况下实现有竞争力的性能？
• RQ5 嵌入拼接与位置编码等架构选择对训练稳定性和泛化能力有何影响？

主要发现

• 在 FS-Mol 和 BACE 基准测试中，CAMP 在小支持集下优于近期元学习算法，展现出更优的少样本泛化能力。
• 在大支持集下，CAMP 性能与最佳基线（包括原型网络）相当，且推理速度提升约 4.5 倍。
• 无位置嵌入的模型验证损失更低（0.5746），低于含位置嵌入的变体（0.5817），表明位置编码会损害性能。
• 使用位置嵌入训练导致更快过拟合，最终训练损失为 0.3181，而无位置嵌入的变体为 0.3374。
• 朴素 ICL（交错拼接分子与标签嵌入）表现出训练收敛不稳定且性能劣于 CAMP 的序列拼接方法。
• CAMP 对示范顺序具有鲁棒性，支持了分子预测中上下文学习应具备排列不变性的假设。