Importance-Weighted Operator Learning Under Probability Measure Shifts

Operator learning has attracted significant interest for its capability to approximate mappings between function spaces. A common assumption in this field is that training and test samples are drawn from the same probability measure. However, this assumption is often violated, leading to probability measure shifts. Under such shifts, standard operator learning methods often lead to degraded performance because the expected training error does not match the expected test error. In this paper, we propose importance-weighted operator learning (IWOL), a framework for learning operators under probability measure shifts in function spaces. We prove that, when the test measure is absolutely continuous with respect to the training measure, the expected test error is equivalent to an importance-weighted expected training error, with weights given by the Radon--Nikodym derivative. This result provides a theoretical basis for correcting measure shifts in operator learning, but importance weighting relies on this absolute-continuity condition and may yield unbounded weights. To overcome these limitations, we introduce relative importance weights defined with respect to a mixture of the training and test measures. The resulting weights remain well-defined even when absolute continuity fails, and they are uniformly bounded. These weights induce a relative importance-weighted training objective for neural operators under probability measure shifts. To estimate these weights, we present binary classifiers that take functions as inputs, defined through functionals modeled by kernel integral operators. This design enables discretization-invariant weight estimation in function space. Experiments on multiple PDE benchmarks and neural operator architectures demonstrate that the proposed method consistently improves performance under probability measure shifts.