Towards Mitigating Systematics in Large-Scale Surveys via Few-Shot Optimal Transport-Based Feature Alignment
This addresses a major challenge for using pre-trained models in astronomy and cosmology by mitigating systematics in large-scale surveys, though it is incremental as it builds on existing feature alignment techniques.
The paper tackles the problem of systematic errors contaminating observables in large-scale surveys, which cause distribution shifts that hinder pre-trained models, by proposing a few-shot optimal transport-based feature alignment method; results show optimal transport effectively aligns out-of-distribution features even with limited data, as validated on MNIST and applied to neutral hydrogen maps.
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.