Lyman Break Galaxy selection and redshift measurement with supervised contrastive learning

Technical Summary: Lyman Break Galaxy selection and redshift measurement with supervised contrastive learning

Problem Statement
The Dark Energy Spectroscopic Instrument (DESI) and future Stage IV/V surveys aim to probe the high-redshift universe ($z \sim 2$ to $4.5$) using Lyman Break Galaxies (LBGs). While LBGs offer significant potential for precision cosmology (e.g., measuring primordial non-Gaussianity and mapping the circum-galactic medium), their faintness and low signal-to-noise ratios (SNR) pose challenges for spectroscopic redshift measurement and sample decontamination.

Current target selection methods, relying on color dropout techniques, suffer from contamination by low-redshift emission line galaxies (ELGs, where [O II] mimics Ly$\alpha$) and high-redshift quasars (QSOs). Existing machine learning pipelines for DESI, specifically lbgNET (a modified QuasarNET), rely on identifying specific spectroscopic lines to infer redshift and classify objects. However, lbgNET faces limitations due to the scarcity of confidently visually inspected (VI) LBG spectra available for training, making it difficult to generalize to faint, noisy samples and distinguish between subtle LBG subtypes (LBGa, LBGme, LBGse) versus contaminants.

Methodology
The authors propose zlbg, a new pipeline based on Supervised Contrastive Learning (CL) to address these limitations. The methodology diverges from the line-finding approach of lbgNET by learning a continuous embedding space where spectral similarity is defined by both object class and redshift.

• Architecture: The backbone encoder is identical to QuasarNET and lbgNET, consisting of four 1D convolutional layers followed by batch normalization and ReLU activations. The input spectra are linearly binned to 775 pixels ($d\lambda = 8$ Å).
• Contrastive Learning: Unlike lbgNET, which uses a "line finder" head, zlbg projects the encoder representations into a 128-dimensional embedding space via a projection head. The network is trained using a Continuously Weighted Contrastive Loss (CWCL).
  • Relationship Weights: Instead of hard binary labels, the loss function employs continuous weights ($r_{ij} \in [0, 1]$) to define similarity.
    • Class Weighting: A similarity matrix $C$ is used where contaminants (ELG, QSO) are pure negatives (weight 0) relative to LBGs. LBG subtypes are treated as partially similar (e.g., LBGme and LBGse share a weight of 0.5), acknowledging their physical continuity.
    • Redshift Weighting: A Gaussian kernel weight is applied to pairs of LBGs based on their redshift difference, normalized by the mean redshift. This encourages the embedding space to preserve redshift gradients.
• Downstream Tasks: After training, the projection head is discarded. The encoder representations are used for:
  • Classification: A Multi-Layer Perceptron (MLP) predicts probabilities for five classes (LBGa, LBGme, LBGse, ELG, QSO).
  • Redshift Inference: A K-Nearest Neighbors (KNN) regressor operates on the encoder representations to predict redshift. A quality metric ($q_z$) is derived from the scatter of the KNN neighbors to quantify redshift confidence.
• Data and Augmentation: The training utilizes small VI datasets (VI-0, VI-1, and VI-contaminants). To compensate for the limited data size, the authors employ data augmentations:
  • SNR Augmentation: Simulating varying exposure times.
  • Redshift Augmentation: Shifting spectra within a $\pm 0.8$ redshift window and filling missing wavelength coverage using stacked LBG templates.

Key Contributions

1. Novel Architecture: Introduction of zlbg, the first application of supervised contrastive learning with continuous relationship weighting for LBG redshift and classification in the DESI context.
2. Improved Decontamination: The method demonstrates superior outlier rejection (ELGs and QSOs) compared to the state-of-the-art lbgNET.
3. Multi-Task Representation: The framework successfully learns a unified representation that simultaneously handles classification and redshift regression, leveraging the small, balanced nature of the VI training set more effectively than previous single-task or line-finding approaches.
4. Quality Metrics: Development of a combined confidence score ($\tau_{zlbg} = s_{zlbg} \cdot q_z$) that integrates classification confidence with redshift estimation uncertainty.

Results
Evaluated on a test set derived from DESI pilot programs:

• Classification: zlbg achieves an Area Under the Curve (AUC) of 0.997 for contaminant selection, compared to 0.988 for lbgNET. It effectively separates contaminants from LBGs, though distinguishing between LBG subtypes (e.g., LBGme vs. LBGa) remains challenging due to the continuous nature of the physical definitions.
• Redshift Performance: The redshift performance (AUC $\approx$ 0.790) is comparable to lbgNET (AUC $\approx$ 0.799). The authors note that zlbg is less prone to catastrophic outliers caused by line misidentification, as it relies on global spectral features rather than specific line detection.
• Sensitivity to Training: The inclusion of redshift augmentations significantly improves redshift performance for zlbg, whereas contaminant selection is robust across different training configurations.
• Complementarity: Correlation analysis between zlbg and lbgNET confidence scores shows low correlation for contaminants, indicating the two pipelines leverage different features and could be combined for further refinement.

Significance and Claims
The paper claims that zlbg provides a robust alternative to the current lbgNET pipeline for the upcoming DESI Run 2, which will target high-redshift LBGs. The primary significance lies in the ability of supervised contrastive learning to extract meaningful representations from small, noisy datasets where traditional supervised methods might overfit or fail to generalize.

The authors modestly state that while redshift performance is comparable to the existing method, the classification performance is significantly improved, particularly in decontaminating the sample. They emphasize that the continuous weighting scheme allows for a more natural handling of the physical continuum between LBG subtypes and contaminants. The work suggests that combining the distinct strengths of zlbg (better outlier rejection) and lbgNET (line-based precision) could yield further improvements in sample purity and completeness for high-redshift cosmology. The pipeline is made available to the community to facilitate these future studies.