[Paper Review] AstroCLIP: A Cross-Modal Foundation Model for Galaxies

Liam Parker, François Lanusse+13 more|arXiv (Cornell University)|Oct 4, 2023

CCD and CMOS Imaging Sensors30 references4 citations

TL;DR

AstroCLIP introduces a cross-modal contrastive learning framework that aligns galaxy images and optical spectra into a shared embedding space, enabling zero-shot prediction of physical properties like redshift and stellar mass. The method achieves state-of-the-art performance on redshift prediction (R² = 0.99) using only spectral embeddings, demonstrating that self-supervised pretraining can yield highly informative, general-purpose astronomical foundation models without fine-tuning.

ABSTRACT

We present AstroCLIP, a single, versatile model that can embed both galaxy images and spectra into a shared, physically meaningful latent space. These embeddings can then be used - without any model fine-tuning - for a variety of downstream tasks including (1) accurate in-modality and cross-modality semantic similarity search, (2) photometric redshift estimation, (3) galaxy property estimation from both images and spectra, and (4) morphology classification. Our approach to implementing AstroCLIP consists of two parts. First, we embed galaxy images and spectra separately by pretraining separate transformer-based image and spectrum encoders in self-supervised settings. We then align the encoders using a contrastive loss. We apply our method to spectra from the Dark Energy Spectroscopic Instrument and images from its corresponding Legacy Imaging Survey. Overall, we find remarkable performance on all downstream tasks, even relative to supervised baselines. For example, for a task like photometric redshift prediction, we find similar performance to a specifically-trained ResNet18, and for additional tasks like physical property estimation (stellar mass, age, metallicity, and sSFR), we beat this supervised baseline by 19\% in terms of $R^2$. We also compare our results to a state-of-the-art self-supervised single-modal model for galaxy images, and find that our approach outperforms this benchmark by roughly a factor of two on photometric redshift estimation and physical property prediction in terms of $R^2$, while remaining roughly in-line in terms of morphology classification. Ultimately, our approach represents the first cross-modal self-supervised model for galaxies, and the first self-supervised transformer-based architectures for galaxy images and spectra.

Motivation & Objective

• To develop a foundation model that bridges multi-band images and optical spectra of galaxies using self-supervised contrastive learning.
• To overcome the limitations of single-modality self-supervised learning by aligning complementary observational modalities into a shared embedding space.
• To enable zero- and few-shot prediction of physical properties like redshift and stellar mass using only the learned embeddings.
• To introduce a novel transformer-based architecture and pretraining strategy for galaxy spectra, advancing beyond traditional 1D CNNs.

Proposed method

• The model uses a contrastive learning objective to align image and spectral embeddings of the same galaxy in a shared latent space, minimizing contrastive loss between positive pairs.
• It employs a vision transformer for processing multi-band galaxy images and a dedicated transformer-based encoder for optical spectra, both trained end-to-end via contrastive pretraining.
• The training process leverages large-scale, weakly labeled data from the DESI survey, using augmented views of images and spectra as positive pairs.
• The model is trained using a contrastive loss function that pulls embeddings of the same galaxy closer while pushing apart embeddings from different galaxies.
• After pretraining, the frozen embeddings are used directly for downstream tasks such as cross-modal retrieval and regression, without further fine-tuning.
• A k-nearest neighbors (k-NN) regressor is used for zero-shot prediction of redshift and stellar mass, with performance evaluated via R² scores.

Experimental results

Research questions

• RQ1 Can cross-modal contrastive learning between galaxy images and optical spectra produce well-aligned, informative embeddings in a shared space?
• RQ2 Can these embeddings enable accurate zero-shot prediction of physical properties like redshift and stellar mass without fine-tuning?
• RQ3 Does the spectral encoder learn to encode redshift information even when images do not contain strong redshift signals?
• RQ4 Can a transformer-based architecture outperform traditional 1D CNNs for galaxy spectrum representation learning?
• RQ5 Does contrastive pretraining improve the intrinsic structure of embeddings within each modality, even when cross-modal alignment is the primary goal?

Key findings

• The spectral embeddings achieve an R² of 0.99 for redshift prediction using a k-NN regressor, outperforming all other methods tested.
• The image embeddings achieve an R² of 0.71 for redshift prediction via k-NN, significantly outperforming the self-supervised baseline from Stein et al. (2021b) with R² = 0.39.
• The model demonstrates that cross-modal contrastive learning leads to emergent structure in both modalities, with spectral embeddings showing near-perfect redshift retrieval despite images lacking strong redshift information.
• The k-NN zero-shot approach outperforms the few-shot MLP fine-tuning baseline, suggesting that the k-NN method benefits from tighter, though slightly biased, predictions.
• The spectral encoder learns to organize its embedding space around redshift, even without explicit supervision, indicating that redshift is a fundamental organizing principle in the learned representation.
• The proposed transformer-based spectral encoder enables high-quality pretraining and paves the way for broader adoption of transformers in galaxy spectrum modeling.