Can Molecular Foundation Models Know What They Don't Know? A Simple Remedy with Preference Optimization

Langzhou He, Junyou Zhu, Fangxin Wang, Junhua Liu, Haoyan Xu, Yue Zhao, Philip S. Yu, Qitian Wu

Keywords: Foundation model, Preference Optimization, Out-of-distribution Detection

Abstract: Molecular foundation models are rapidly advancing scientific discovery, but their unreliability on out-of-distribution (OOD) samples severely limits their application in high-stakes domains such as drug discovery and protein design. A critical failure mode is chemical hallucination, where models make high-confidence yet entirely incorrect predictions for unknown molecules. To address this challenge, we introduce Molecular Preference Aligned Instance Ranking (Mole-PAIR), a simple, plug-and-play module that can be flexibly integrated with existing foundation models to improve their reliability on OOD data through cost-effective post-training. Specifically, our method formulates the OOD detection problem as a preference optimization over the estimated OOD affinity between in-distribution (ID) and OOD samples, achieving this goal through a pairwise learning objective. We show that this objective essentially optimizes the AUROC, which measures how consistently ID and OOD samples are ranked by the model. Extensive experiments across five real-world molecular datasets demonstrate that our approach significantly improves the OOD detection capabilities of existing molecular foundation models, achieving up to $\mathbf{45.8%}$, $\mathbf{43.9%}$, and $\mathbf{24.3%}$ improvements in AUROC under distribution shifts of size, scaffold, and assay, respectively. Our code is available at: https://anonymous.4open.science/status/Mole-PAIR-61B5.

Primary Area: alignment, fairness, safety, privacy, and societal considerations

Submission Number: 4195