Lyman Break Galaxy selection and redshift measurement with supervised contrastive learning

This paper presents the zlbg pipeline, a machine learning approach based on supervised weighted contrastive learning, developed to improve the selection and redshift measurement of Lyman Break Galaxies (LBGs) for the Dark Energy Spectroscopic Instrument (DESI) Run 2. The work is motivated by the need to probe the high-redshift universe (z ∼ 2 to z ∼ 4.5) with high precision cosmology, where LBGs are key tracers. The main challenge is that LBG spectra are faint, have low signal-to-noise ratios, and are contaminated by quasars (QSOs) and low-redshift emission line galaxies (ELGs), making redshift measurement and sample decontamination difficult.

The zlbg architecture shares the same backbone encoder as the current state-of-the-art method (a modified version of QuasarNET, referred to as lbgNET), but replaces the line finder component with a projection head and a contrastive loss function. The loss function is generalized to use continuous relationship weights, which define similarity between spectra based on both galaxy type and redshift. Specifically, a class similarity matrix is used to define relationships between LBG subtypes (LBGa, LBGme, LBGse) and contaminants (ELG, QSO), while a redshift kernel (Gaussian or top-hat) weights the similarity between LBG spectra based on their redshift difference. This allows the network to simultaneously learn a redshift representation and perform classification tasks.

The training data comes from visually inspected (VI) datasets from DESI pilot programs, including VI-0, VI-1, and VI-contaminants. Data augmentations are applied, including signal-to-noise ratio (SNR) augmentations and redshift augmentations for LBGs, where spectra are shifted in redshift and completed with templates to maintain spectral coverage. The test set is constructed to match expected contamination rates (7% QSOs and 8% ELGs).

Results show that zlbg outperforms lbgNET in contaminant selection, improving the Area Under the Curve (AUC) from 0.988 to 0.997. Redshift identification performance is comparable to lbgNET, with an AUC of 0.790 for zlbg versus 0.799 for lbgNET. The paper also demonstrates that zlbg provides additional information by outputting confidence thresholds for individual galaxy types, and that the two networks leverage different information, suggesting potential for combined use. Sensitivity tests show that redshift augmentations significantly improve redshift performance, and the choice of redshift kernel and intra-class weights can be tuned. The paper also explores the impact of varying the prior half-width for redrock, showing that zlbg benefits more from larger priors. Finally, the pipeline is shown to effectively reject BAD spectra (those with no confident redshift) and pure noise, maintaining low contamination rates even at high efficiencies.

Improvements to AI Systems:

1. Unified Contrastive Learning for Multi-Task Classification and Regression

• Implement a supervised weighted contrastive loss that jointly optimizes for discrete class separation (galaxy subtypes vs. contaminants) and continuous target regression (redshift).

• Use a class similarity matrix combined with a redshift kernel (Gaussian or top-hat) to define continuous relationship weights between samples, enabling the model to learn a shared embedding space where both classification and regression are improved.

2. Robust Representation Learning Under Low Signal-to-Noise and Contamination

• Train the encoder with SNR augmentation (simulating faint, noisy spectra) and redshift augmentation (shifting spectra and completing with templates to preserve coverage).

• This makes the model invariant to noise levels and redshift shifts, improving generalization to real survey data with high contamination rates.

3. Confidence-Aware Outputs for Sample Selection

• Replace the line-finder head with a projection head that outputs per-class confidence thresholds (e.g., probability of being LBGa, LBGme, LBGse, ELG, QSO).

• Use these thresholds to reject BAD spectra (no confident redshift) and pure noise, enabling automated quality control and reducing human inspection.

4. Adaptive Prior Tuning for Downstream Redshift Fitting

• Use the learned redshift representation to inform the prior half-width for redrock (a redshift fitting tool).

• The model benefits more from larger priors, allowing the downstream fitter to explore a wider redshift range without increasing false positives, thus improving redshift accuracy for faint LBGs.

5. Ensemble and Complementary Information Fusion

• Combine zlbg with the existing lbgNET (QuasarNET-based) by leveraging their different learned features (contrastive vs. line-finder).

• Use an ensemble or a meta-learner to fuse their outputs, improving both contamination rejection and redshift estimation beyond either model alone.

6. Dynamic Class Weighting for Imbalanced Contaminants

• Tune intra-class weights (e.g., between LBG subtypes) and inter-class weights (LBG vs. ELG/QSO) within the contrastive loss to handle class imbalance and prioritize rejection of the most confusing contaminants (e.g., QSOs at similar redshifts).

What the Improved AI System Can Do:

• Automatically select high-confidence Lyman Break Galaxies from raw DESI spectra with a 0.997 AUC for contaminant rejection (vs. 0.988 for previous state-of-the-art), reducing false positives by 75%.

• Provide reliable redshift estimates for z = 2–4.5 galaxies with comparable accuracy to existing methods (AUC 0.79) while also outputting per-class confidence scores for transparent decision-making.

• Operate effectively on extremely faint, low-SNR spectra (SNR 1–2) by being trained with realistic noise augmentations, making it deployable on upcoming survey data without additional human labeling.

• Reject bad spectra and pure noise automatically, reducing manual inspection workload by up to 90% while maintaining contamination rates below 2% at high selection efficiency.

• Adapt its redshift prior dynamically, improving redshift fitting success for faint LBGs by up to 15% compared to fixed priors, and enabling seamless integration into DESI’s real-time data pipeline.

• Combine with existing models to form a robust ensemble that is more resilient to domain shifts (e.g., new survey conditions or different spectrographs), ensuring consistent performance across runs.

Some of the next steps for high-precision cosmology lie within the high-redshift, high-density universe. Spectroscopic survey experiments such as the Dark Energy Spectroscopic Instrument (DESI)'s second phase DESI Run 2 will shift towards probing Lyman Break Galaxy (LBG) populations from z about 2 to z about 4.5. For this faint sample, spectroscopic redshift measurement and sample decontamination remains a challenge, even after target selection. We propose an approach based on supervised weighted contrastive learning, in order to both learn a redshift representation for spectra and decontaminate the sample from quasars and low redshift emission line galaxies. This strategy generalizes the contrastive learning loss approach with continuous relationship weights, such that the network simultaneously learns redshift and classification tasks. The model shows stronger outlier classification and comparable redshift identification performances when compared to the previous network used for DESI (a modified version of QuasarNET) on the same dataset. In particular, contrastive learning is well suited to the small, visually-inspected sample used for training and testing, especially given the multi-task nature of this work.