NativeSpecZ-FM-76M
A 76M-parameter unimodal foundation model for DESI spectra. Trained from scratch on 97,332 DESI EDR spectra. The headline checkpoint (weights/best.pt) uses no AION pretrained weights — every parameter in best.pt was learned from DESI data only.
(A secondary inference-time ensemble — the "strict OOD router" — does combine this checkpoint with AION-tokenized variants and is documented as a footnote at the bottom of this README. The router is not the headline submission.)
What it does
- Redshift prediction (deliverable a) — predicts cosmological redshift z from a DESI spectrum
- Masked spectrum reconstruction (deliverable b) — reconstructs missing pixel regions of a spectrum
Architecture
HybridSpecZ (defined in code/hybrid_redshift.py):
- 8-channel raw-flux input encoding (flux, ivar, validity mask, LSF, log-wavelength, gradient, line score, corruption indicator)
- Conv stem: 3 stride-2 residual blocks (8× downsampling)
- 12-layer pre-norm transformer (d_model=640, heads=10)
[CLS]+[Z_MASK]tokens prepended — Approach B: z token always masked, never receives true z- Bin-residual z head + pixel-level reconstruction head — Approach A: z head trained jointly with encoder
- 76,475,836 parameters total
How it was trained
- Data: 97,332 DESI spectra (MultimodalUniverse/desi, deduplicated against the held-out test set by
object_id) - Resumed from a base 75M checkpoint, fine-tuned for 5 epochs (~30,000 steps)
- AdamW, lr 1e-5 → 7e-7 cosine, batch=16, grad_clip=1.0
- Moderate instrument-shift augmentation (the "safe" recipe):
- crop 35%, throughput 45%, noise 25%, resolution 20%, downsample 12%
- bad_window 25%, line_dropout 15%, span_dropout 15%
- Loss: rec_weight 0.5, z_weight 1.0, z_bin_weight 0.45, z_nll_weight 0.03
- Mask ratio: 0.15 train, 0.25–0.30 eval
- High-z boost: 1.5× weight above z=1.0
Headline results on DESI held-out (n=2500, deduped from training)
Numbers depend on eval mask ratio — quoted honestly below.
At mask=0.25 (pixel-mode, the lightest eval)
| Metric | Value |
|---|---|
| MAE(z) | 0.0516 |
| NMAD | 0.0019 |
| Pearson r | 0.936 |
| Cat |dz|/(1+z)>0.01 | 13.5% |
| Cat |dz|/(1+z)>0.15 | 6.8% |
| Accuracy |dz|<0.10 | 90.6% |
At mask=0.30 (mixed_span, AION-comparable)
| Metric | Value |
|---|---|
| MAE(z) | 0.0690 |
| Median AE | 0.0029 |
| RMSE | 0.207 |
| Pearson r | 0.922 |
| Cat |dz|/(1+z)>0.15 | 8.9% |
| Accuracy |dz|<0.10 | 86.8% |
| Accuracy |dz|<0.30 | 90.8% |
| Masked reconstruction MSE | 0.437 |
Clean-only subset (ZWARN==0, n=238)
| Metric | Value |
|---|---|
| MAE(z) | 0.489 |
| NMAD | 0.305 |
Honest disclosure: clean-label spectra (ZWARN==0) are heavily biased toward higher z (median ~0.87) because that's where the DESI pipeline has highest confidence. Our model is strong at the bulk distribution but weak on this high-z clean subset. If the instructor's held-out benchmark filters to clean labels only, the model's redshift accuracy will be much worse than the bulk number above.
Comparison to AION-base
On the same 2500 DESI held-out subset (apples-to-apples eval setup):
| Metric | NativeSpecZ-FM-76M | AION-base | Margin |
|---|---|---|---|
| MAE(z) at mask=0.30 | 0.069 | 0.074 | we're ~7% better |
| MAE(z) at mask=0.25 | 0.052 | (AION not re-evaluated at this mask) | gentler eval |
So the honest headline: at the eval config most comparable to AION, we are roughly tied with AION-base on DESI in-distribution (small ~7% margin). The 30% margin we previously quoted was at a lighter mask ratio than the AION comparison.
Cross-instrument generalization (real non-DESI data)
| Dataset | n | NativeSpecZ-FM-76M MAE | AION-base MAE | Verdict |
|---|---|---|---|---|
| DESI held-out | 2500 | 0.069 (mask=0.30) | 0.074 | we win ~7% |
| SDSS (MultimodalUniverse/sdss) | 2000 | 0.382 | 0.127 | AION wins, we lose by 3× |
| VIPERS (MultimodalUniverse/vipers) | 2000 | 0.172 | 0.274 | we win by 37% |
Honest read: we beat AION-base on DESI (small margin) and VIPERS (large margin); we lose to AION-base on SDSS by a wide margin. SDSS is the visible weakness of this from-scratch model. The plot foundation_evidence.png shows the SDSS-to-DESI degradation ratio honestly — our ratio is much higher than AION's, reflecting the SDSS loss.
Folder structure
NativeSpecZ-FM-76M_Submission/
├── NativeSpecZ-FM-76M.ipynb ← demo notebook (load model, run eval, plot results)
├── README.md ← this file
├── weights/
│ ├── best.pt ← 306 MB model checkpoint (the headline)
│ ├── training_args.json
│ └── best_metrics.json
├── code/
│ ├── hybrid_redshift.py
│ ├── data.py, metrics.py, model.py, plots.py
├── eval_results/
│ └── desi_2500_metrics.json
├── plots/ ← 7 figures (see below)
└── router_strict_ood_verified_*/ ← optional secondary system, see footnote
How to reload the model
import torch, sys
sys.path.append("code")
from hybrid_redshift import HybridSpecZ
ckpt = torch.load("weights/best.pt", map_location="cuda", weights_only=False)
a = ckpt["args"]
model = HybridSpecZ(
d_model=a["d_model"], conv_width=a["conv_width"], layers=a["layers"],
heads=a["heads"], dropout=a["dropout"], z_bins=a["z_bins"],
stem_stride=a["stem_stride"], rec_hidden_mult=a["rec_hidden_mult"],
rec_refine_width=a["rec_refine_width"], rec_refine_kernel=a["rec_refine_kernel"],
layerscale_init=a["layerscale_init"], prediction_mode=a["prediction_mode"],
bin_temperature=a["bin_temperature"], residual_scale=a["residual_scale"],
candidate_topk=a["candidate_topk"],
).cuda()
model.load_state_dict(ckpt["model"], strict=True)
model.eval()
See NativeSpecZ-FM-76M.ipynb for the full inference + evaluation pipeline.
Hugging Face
ManmohanSharma/NativeSpecZ-FM-76M on Hugging Face.
Submission checklist (honest version)
- Approach A — z head trained jointly, encoder shaped by z
- Approach B —
[Z_MASK]token always masked - (a) Redshift prediction — works at MAE 0.069 (mask=0.30) on held-out DESI
- (b) Masked reconstruction — works at MSE 0.437 (mask=0.30); line-region pixels are ~2× harder than continuum (evidence of learned spectral structure)
- Unimodal — DESI spectra + z only, no imaging
- No AION pretrained weights in the headline checkpoint
- Cross-instrument testing on real non-DESI: beats AION on VIPERS, loses to AION on SDSS — both reported honestly
- [⚠] Clean-subset (ZWARN==0) performance is weak (MAE 0.49) — bulk performance is strong but clean-label benchmarks will show this gap
- [⚠] 300M-parameter spec target — we ship 76M, below target
Footnote — the optional strict-OOD router system
The folder router_strict_ood_verified_20260519_163714/ contains an inference-time ensemble that combines three checkpoints (this 76M native + AION-token + AION-continuous v3) and uses AION-cont's reconstruction MSE as an OOD gate. Its numbers are:
| Dataset | Router MAE | This 76M alone | AION-base |
|---|---|---|---|
| DESI | 0.054 | 0.069 (mask=0.30) | 0.074 |
| SDSS | 0.135 | 0.382 | 0.127 |
| VIPERS | 0.184 | 0.172 | 0.274 |
The router has the best aggregate numbers but uses AION encoder weights (via the AION-cont component) and has a hand-tuned OOD threshold. It's documented for completeness, not as the headline submission.