🤖 AI Summary
In large-scale astronomical surveys, systematic errors induce distributional shifts in observations, severely degrading the generalization of pre-trained models under label-scarce conditions. To address this, we propose an optimal transport (OT)-based few-shot feature alignment method that directly aligns in-distribution (ID) and out-of-distribution (OOD) sample distributions in the pre-trained feature space—without requiring paired samples or prior knowledge of systematic errors. Our approach jointly optimizes mean squared error and OT distance to achieve robust transfer learning. We validate the method on both MNIST and real-world neutral hydrogen (HI) large-scale sky maps. Results demonstrate substantial improvements in downstream task performance, particularly in realistic astronomical settings where systematic errors are difficult to model and labeled data are extremely scarce. The method exhibits strong adaptability and generalization capability under such challenging conditions.
📝 Abstract
Systematics contaminate observables, leading to distribution shifts relative to theoretically simulated signals-posing a major challenge for using pre-trained models to label such observables. Since systematics are often poorly understood and difficult to model, removing them directly and entirely may not be feasible. To address this challenge, we propose a novel method that aligns learned features between in-distribution (ID) and out-of-distribution (OOD) samples by optimizing a feature-alignment loss on the representations extracted from a pre-trained ID model. We first experimentally validate the method on the MNIST dataset using possible alignment losses, including mean squared error and optimal transport, and subsequently apply it to large-scale maps of neutral hydrogen. Our results show that optimal transport is particularly effective at aligning OOD features when parity between ID and OOD samples is unknown, even with limited data-mimicking real-world conditions in extracting information from large-scale surveys. Our code is available at https://github.com/sultan-hassan/feature-alignment-for-OOD-generalization.