Towards out-of-distribution generalization in large-scale astronomical surveys: robust networks learn similar representations
This addresses the problem of improving model reliability for astronomers dealing with distribution shifts in large-scale surveys, but it is incremental as it applies an existing interpretability method to a specific domain.
The paper tackles the challenge of out-of-distribution generalization in astronomical surveys by analyzing representation similarity using Centered Kernel Alignment, finding that robust models produce substantially different representations across layers on OOD data, while non-robust models show less change.
The generalization of machine learning (ML) models to out-of-distribution (OOD) examples remains a key challenge in extracting information from upcoming astronomical surveys. Interpretability approaches are a natural way to gain insights into the OOD generalization problem. We use Centered Kernel Alignment (CKA), a similarity measure metric of neural network representations, to examine the relationship between representation similarity and performance of pre-trained Convolutional Neural Networks (CNNs) on the CAMELS Multifield Dataset. We find that when models are robust to a distribution shift, they produce substantially different representations across their layers on OOD data. However, when they fail to generalize, these representations change less from layer to layer on OOD data. We discuss the potential application of similarity representation in guiding model design, training strategy, and mitigating the OOD problem by incorporating CKA as an inductive bias during training.