Learning Molecular Semantic Invariant Representation with Prototype Constraint

Molecular representation learning has achieved remarkable progress in molecular property prediction, yet out-of-distribution (OOD) generalization remains challenging. In practice, training data typically cover only a limited portion of the chemical space, causing models to rely on environment-dependent factors that fail to transfer when scaffold structures or functional compositions shift. To address this issue, we propose MoSIR, a framework for learning molecular semantic invariant representation with prototype constraint, which projects entangled molecular embeddings into a learnable semantic prototype space to extract semantic invariant representation while isolating environment-sensitive variations. Building upon this decomposition, we optimize a bi-level min-max objective that introduces representation perturbations to simulate plausible environment shifts and enforce semantic stability. We further provide theoretical guarantees for MoSIR by deriving an OOD generalization bound under distribution shifts. Extensive experiments on multiple molecular OOD benchmarks demonstrate that MoSIR consistently outperforms strong baselines across diverse shift settings, and qualitative analyses confirm that the learned prototypes capture meaningful chemical semantics.

AI models are increasingly used to predict molecular properties, which can help researchers discover useful new molecules for drugs, materials, and chemistry. However, these models often learn from limited training data and may fail when they encounter molecules with new core structures or different functional groups. This is a serious challenge because real-world molecular discovery usually requires making predictions beyond the chemical patterns seen during training. We propose MoSIR, a method designed to help molecular AI models learn more stable and transferable molecular knowledge. Instead of using the original molecular representation directly, MoSIR maps molecules into a set of learned semantic prototypes, which can be viewed as reusable chemical concepts shared across different environments. This allows the model to separate stable molecular semantics from unstable variations that may only appear in specific training conditions. MoSIR further trains the model by simulating possible distribution changes, encouraging its predictions to remain consistent under such shifts. We also provide theoretical analysis showing why this strategy can improve generalization to unseen molecular distributions. Experiments on multiple molecular benchmarks show that MoSIR performs better than strong existing methods and learns prototypes that reflect meaningful chemical patterns.