AstroMAE: Redshift Prediction Using a Masked Autoencoder with a Novel Fine-Tuning Architecture

TL;DR

AstroMAE addresses the challenge of redshift prediction with limited labeled data by pretraining a Vision Transformer encoder via a masked autoencoder on SDSS images to learn global patterns without labels. It couples this self-supervised pretraining with a novel fine-tuning architecture that injects local information through Inception blocks and a Magnitude block, while freezing most of the pretrained encoder. Across two experiments, AstroMAE outperforms baselines, including transformer- and CNN-based methods, and demonstrates robustness when scaling from 80% to 100% of the data. The approach offers efficient pretraining, improved locality-aware feature extraction, and strong generalization for photometric redshift estimation in large astronomical surveys.

Abstract

Redshift prediction is a fundamental task in astronomy, essential for understanding the expansion of the universe and determining the distances of astronomical objects. Accurate redshift prediction plays a crucial role in advancing our knowledge of the cosmos. Machine learning (ML) methods, renowned for their precision and speed, offer promising solutions for this complex task. However, traditional ML algorithms heavily depend on labeled data and task-specific feature extraction. To overcome these limitations, we introduce AstroMAE, an innovative approach that pretrains a vision transformer encoder using a masked autoencoder method on Sloan Digital Sky Survey (SDSS) images. This technique enables the encoder to capture the global patterns within the data without relying on labels. To the best of our knowledge, AstroMAE represents the first application of a masked autoencoder to astronomical data. By ignoring labels during the pretraining phase, the encoder gathers a general understanding of the data. The pretrained encoder is subsequently fine-tuned within a specialized architecture tailored for redshift prediction. We evaluate our model against various vision transformer architectures and CNN-based models, demonstrating the superior performance of AstroMAEs pretrained model and fine-tuning architecture.

Figures (34)

• Figure 2: It illustrates the architecture of pretraining the AstroMAE using a masked autoencoder algorithm. Loss involves comparing the generated patches corresponding to masked areas with their original counterparts.
• Figure 3: Inception module
• Figure 4: Learning rate during pretraining.
• Figure 5: It shows the learning rate during fine-tuning, with the yellow section highlighting the changes over two cycles.
• Figure 8: Learning rate for training during the second experiment.