Towards out-of-distribution generalization: robust networks learn similar representations

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 here use the 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 robust models, i.e., those that score high accuracy on both in-distribution (ID) and OOD data, learn similar representations, whereas non-robust models do not. We observe a strong correlation between similarity and accuracy in recovering cosmological parameters from three fields across the IllustrisTNG and SIMBA simulations. 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.