NativeSpecZ-FM-76M

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

1. Redshift prediction (deliverable a) — predicts cosmological redshift z from a DESI spectrum
2. 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)

At mask=0.30 (mixed_span, AION-comparable)

Clean-only subset (ZWARN==0, n=238)

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):

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)

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:

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.