Large Language Models are In-Context Molecule Learners

TL;DR

ICMA introduces In-Context Molecule Adaptation to enable LLMs to learn molecule-text alignment for the molecule-caption translation task from retrieved context examples, without domain-specific pre-training. The framework combines Hybrid Context Retrieval, Post-retrieval Re-ranking (Random Walk and Sequence Reversal), and In-Context Molecule Tuning to ground molecule-caption mappings and adapt model behavior. Across ChEBI-20 and PubChem324k, ICMA delivers state-of-the-art or competitive Mol2Cap and Cap2Mol results with backbone models such as Galactica-125M and Mistral-7B, including BLEU-4 up to $0.581$ and EM up to $0.460$. The findings show LLMs are inherently in-context molecule learners, and ICMA provides a model-agnostic, data-efficient path to biomolecule tasks including molecule property prediction.

Abstract

Large Language Models (LLMs) have demonstrated exceptional performance in biochemical tasks, especially the molecule caption translation task, which aims to bridge the gap between molecules and natural language texts. However, previous methods in adapting LLMs to the molecule-caption translation task required extra domain-specific pre-training stages, suffered weak alignment between molecular and textual spaces, or imposed stringent demands on the scale of LLMs. To resolve the challenges, we propose In-Context Molecule Adaptation (ICMA), as a new paradigm allowing LLMs to learn the molecule-text alignment from context examples via In-Context Molecule Tuning. Specifically, ICMA incorporates the following three stages: Hybrid Context Retrieval, Post-retrieval Re-ranking, and In-context Molecule Tuning. Initially, Hybrid Context Retrieval utilizes BM25 Caption Retrieval and Molecule Graph Retrieval to retrieve similar informative context examples. Additionally, Post-retrieval Re-ranking is composed of Sequence Reversal and Random Walk selection to further improve the quality of retrieval results. Finally, In-Context Molecule Tuning unlocks the in-context learning and reasoning capability of LLMs with the retrieved examples and adapts the parameters of LLMs for better alignment between molecules and texts. Experimental results demonstrate that ICMA can empower LLMs to achieve state-of-the-art or comparable performance without extra training corpora and intricate structures, showing that LLMs are inherently in-context molecule learners.

Figures (7)

• Figure 1: An illustration of three similar molecules alongside their molecule captions. The molecules are represented as both SMILES strings and graphs, while the molecule captions elucidate their structures and functions. Here, the three molecules are similar, considering their 2D graph embeddings, and the overlaps in their captions are highlighted in blue and pink.
• Figure 2: Framework of In-Context Molecule Adaptation (ICMA). Generally, ICMA consists of three stages, Hybrid Context Retrieval, Post-retrieval Re-ranking, and In-Context Molecule Tuning.
• Figure 3: The model performance with the change of context settings, including the number of refined examples (i.e., $n$) and cutoff length. Mol2Cap Results (Left) and Cap2Mol Results (Right).
• Figure 4: The scaling law of ICMA. Three models with different levels of parameters are selected, including the Galactica series (Blue) and the Llama-3 series (Green). Mol2Cap Results (Left) and Cap2Mol Results (Right).
• Figure 5: The performance of Galatcica-125M with the maximum skip probability $p_{max}$. Mol2Cap Results (Left) and Cap2Mol Results (Right).