In-Context Learning of Physical Properties: Few-Shot Adaptation to Out-of-Distribution Molecular Graphs

TL;DR

This paper explores leveraging in-context learning in large language models combined with geometry-aware graph neural networks to predict out-of-distribution molecular properties, demonstrating improved performance over traditional models.

Contribution

It introduces a novel compound model that integrates GPT-2 with graph neural networks for effective in-context learning on molecular data, especially out-of-distribution examples.

Findings

Model surpasses traditional graph neural networks on out-of-distribution molecules

Combining GPT-2 with graph neural networks enhances geometric feature utilization

Significant performance improvement on QM9 dataset sequences with shared substructures

Abstract

Large language models manifest the ability of few-shot adaptation to a sequence of provided examples. This behavior, known as in-context learning, allows for performing nontrivial machine learning tasks during inference only. In this work, we address the question: can we leverage in-context learning to predict out-of-distribution materials properties? However, this would not be possible for structure property prediction tasks unless an effective method is found to pass atomic-level geometric features to the transformer model. To address this problem, we employ a compound model in which GPT-2 acts on the output of geometry-aware graph neural networks to adapt in-context information. To demonstrate our model's capabilities, we partition the QM9 dataset into sequences of molecules that share a common substructure and use them for in-context learning. This approach significantly improves the performance of the model on out-of-distribution examples, surpassing the one of general graph neural network models.