LICO: Large Language Models for In-Context Molecular Optimization

2024-06-27

International Conference on Learning Representations

AI Summary

Molecular property optimization faces challenges in modeling black-box objective functions and adapting pre-trained large language models (LLMs) to domain-specific tasks. Method: We propose LICO—a lightweight, fine-tuning-free framework that enables any base LLM to perform context-aware surrogate modeling and optimization solely from query history (molecular structures + performance feedback), without relying on molecular text descriptions. LICO introduces a learnable embedding layer and regression head, integrated with SELFIES/SMILES encodings, and is trained via in-context learning across diverse benchmark functions. Results: On the PMO benchmark comprising 20+ objectives, LICO significantly outperforms Bayesian optimization and graph neural networks, achieving state-of-the-art performance. It is the first method to enable generalizable, in-context surrogate modeling by LLMs for pure black-box molecular optimization.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Optimizing black-box functions is a fundamental problem in science and engineering. To solve this problem, many approaches learn a surrogate function that estimates the underlying objective from limited historical evaluations. Large Language Models (LLMs), with their strong pattern-matching capabilities via pretraining on vast amounts of data, stand out as a potential candidate for surrogate modeling. However, directly prompting a pretrained language model to produce predictions is not feasible in many scientific domains due to the scarcity of domain-specific data in the pretraining corpora and the challenges of articulating complex problems in natural language. In this work, we introduce LICO, a general-purpose model that extends arbitrary base LLMs for black-box optimization, with a particular application to the molecular domain. To achieve this, we equip the language model with a separate embedding layer and prediction layer, and train the model to perform in-context predictions on a diverse set of functions defined over the domain. Once trained, LICO can generalize to unseen molecule properties simply via in-context prompting. LICO achieves state-of-the-art performance on PMO, a challenging molecular optimization benchmark comprising over 20 objective functions.

Problem

• Optimizing black-box functions in scientific domains
• Developing LLM-based surrogate models for molecular optimization
• Generalizing to unseen molecule properties via in-context learning

Innovation

• Extends LLMs with separate embedding and prediction layers
• Trains model for in-context predictions on diverse functions
• Enables generalization to unseen properties via prompting