Learning to Select Source Domains: Proxy-Rewarded Policy Optimization for Molecular OOD Generalization

Robust generalization under extreme out-of-distribution (OOD) shifts remains a fundamental challenge for molecular property prediction. Existing scaffold-based evaluation protocols often fail to eliminate subtle structural and semantic overlap between training and test molecules, leading to overly optimistic estimates of model extrapolation capability. Meanwhile, conventional domain adaptation approaches typically perform indiscriminate alignment across heterogeneous source domains, which may introduce irrelevant knowledge and cause negative transfer under severe distribution shifts. In this work, we investigate molecular OOD generalization from the perspective of adaptive source domain selection. We first introduce SCOPE-Bench, a challenging molecular OOD benchmark constructed through cluster-level partitioning in an explicit physicochemical descriptor space, enabling more rigorous evaluation of extrapolation ability beyond conventional scaffold splitting. Based on this benchmark, we propose Policy Optimization for Multi-source Adaptation (POMA), a framework that formulates knowledge transfer as a source selection problem. POMA learns a policy to identify transferable source subsets from a large candidate domain pool, followed by dual-scale adaptation that captures both global molecular topology and local pharmacophore patterns. Extensive experiments across diverse molecular prediction architectures demonstrate that current state-of-the-art models suffer substantial performance degradation under realistic OOD shifts, with prediction errors increasing by up to 8.0× and 5.9× on average compared with conventional evaluations. POMA consistently improves generalization performance across multiple backbone models, achieving up to 11.2% reduction in mean absolute error and an average relative improvement of 6.2%. These results reveal the importance of adaptive source utilization for reliable molecular learning and provide a general strategy for mitigating negative transfer under extreme distribution shifts.