Sample Efficient Generative Molecular Optimization with Joint Self-Improvement

Serra Korkmaz, Adam Izdebski, Jonathan Pirnay, Rasmus Møller-Larsen, Michal Kmicikiewicz, Pankhil Gawade, Dominik G. Grimm, Ewa Szczurek

Keywords: reinforcement learning, molecule design, self-improvement, joint modeling

TL;DR: Combining joint modeling with an efficient inference-time sampling for molecular optimization

Abstract: Generative molecular optimization aims to design molecules with properties surpassing those of existing compounds. However, such candidates are rare and expensive to evaluate, yielding sample efficiency essential. Additionally, surrogate models introduced to predict molecule evaluations, suffer from distribution shift as optimization drives candidates increasingly out-of-distribution. To address these challenges, we introduce Joint Self-Improvement, which benefits from (i) **a joint generative-predictive model** and (ii) **a self-improving sampling scheme**. The former aligns the generator with the surrogate, alleviating distribution shift, while the latter biases the generative part of the joint model using the predictive one to efficiently generate optimized molecules at inference-time. Experiments across offline and online molecular optimization benchmarks demonstrate that Joint Self-Improvement outperforms state-of-the-art methods under limited evaluation budgets.

Submission Track: Feedback-Based Learning for Materials Design - Full Paper

Submission Category: AI-Guided Design

Submission Number: 16