▶ Live demo — try the model on real held-out DESI spectra (or your own) in the browser.

Encoder-only transformer (8 layers, d_model=512, 8 heads, ~26M parameters) trained with masked-token prediction on DESI EDR/SV3 spectra (Multimodal Universe). The redshift token is always masked and a prediction head is trained jointly with reconstruction, so redshift enters the representation space from step one — the redesign of AION-1's redshift handling that this course project asked for.

Current checkpoint: v2.1 — a fine-tune of the v1 encoder (not a from-scratch model) with a new classification head over 100 log(1+z) bins: cross-entropy normalized by log(n_bins), sqrt-inverse class weights estimated from the real 80k training labels, 1:1 loss weighting with reconstruction, and a leak-free train/held-out split. It replaced v1 after passing every release gate of the project plan (comparison.json, decision: promote_v2_1). The official redshift prediction is z_pred_map (posterior argmax); z_pred (posterior expectation) is kept for backward compatibility.

Both models evaluated on the same 2,000 valid-label held-out spectra (the ones following the 80,000 used for training — never seen by either model). v2.1 uses z_pred_map:

The full bias/outlier analysis (with plots) is in the evaluation notebook; the executed v1 "before" picture is preserved in evaluation_v1_baseline.ipynb.

Accepts spectra from any instrument: inputs are interpolated onto an internal log-λ grid (3600–9800 Å, 7081 pixels) and the positional embedding is a sinusoidal encoding of physical log(λ), not a token index, so wavelength coverage different from DESI's is handled transparently.

Trained on 80k spectra on a laptop (Apple MPS) — a course project, not for production science (the DESI pipeline is ~3 orders of magnitude more accurate on z). v2.1 removed v1's z ≈ 2 prediction ceiling and cut catastrophic outliers from 22.6 % to 14.95 % globally (82.7 % → 23.5 % in z ∈ [1.5, 2.5)), but ~15 % of held-out spectra are still catastrophic outliers — dominated by low-S/N spectra and line-misidentification degeneracies — and predictions above z ≈ 3.5 remain unconstrained (few training examples there). Use z_confidence to filter unreliable predictions. See the evaluation notebook for the quantified analysis.

The faiss/ pair powers the find_similar_spectra semantic-search tool in spectra-copilot.