NativeSpecZ-296M — Wavelength-Aware Scale-Up Experiment

A 296M-parameter unimodal foundation model for DESI spectra, trained from scratch with no AION pretrained weights. This is the project's scale-up ablation — it hits the spec's ~300M-parameter target and is competitive with AION-base on DESI in-distribution, but it does NOT beat AION on out-of-distribution data. The smaller NativeSpecZ-FM-76M is the recommended headline submission; this 296M model is documented here as the larger wavelength-aware scaling experiment.

What it is

• 296,269,116 parameters
• Conv-stem transformer: 20 layers, d_model=1024, conv_width=512, 16 heads, stem_stride=8
• 8-channel raw-flux input including log-wavelength channels
• [CLS] + [Z_MASK] tokens — Approach B: z token always masked
• Bin-residual z head + pixel-level reconstruction head — Approach A: joint training
• Trained on 97,332 DESI spectra with wavelength-grid jitter and OOD-style augmentation (the "wavelength-aware" recipe designed to improve cross-instrument transfer)
• Mixed-span masking at 30%
• Resumed from an earlier 300M joint checkpoint, gentle fine-tune (3 epochs, lr 2e-6, 1500 steps)

Results (held-out, TTA)

Dataset	n	MAE(z)	NMAD	Cat>0.01
DESI held-out	2500	0.0674	0.0048	0.213
SDSS (real non-DESI)	2000	0.314	0.278	0.692
VIPERS (real non-DESI)	2000	0.316	0.154	0.906

Clean-subset (ZWARN==0) MAE: DESI 0.457, SDSS 0.306 — same high-z clean-label weakness as the 76M.

Masked reconstruction (DESI, mask=0.25): rec MSE 0.066, line-region ~2× harder than continuum.

Three-way comparison

Dataset	NativeSpecZ-296M	NativeSpecZ-76M (headline)	AION-base
DESI	0.067	0.069	0.074
SDSS	0.314	0.382	0.127
VIPERS	0.316	0.172	0.274

Honest read of the scale-up:

• On DESI in-distribution, the 296M is the best of the three (ties the 76M, beats AION). Scaling helped in-distribution.
• The wavelength-jitter training homogenized cross-instrument behavior — SDSS improved over the 76M (0.314 vs 0.382), but VIPERS regressed badly (0.316 vs the 76M's 0.172).
• Critically, the 296M beats AION-base on no OOD dataset (loses on both SDSS and VIPERS). The 76M, by contrast, beats AION on VIPERS by 37% — the strongest foundation-model claim in the project.
• Conclusion: bigger + wavelength-aware did not improve the foundation-model criterion. This model is a legitimate scale ablation, not the headline.

Why it's still worth submitting as an ablation

1. It hits the spec's ~300M-parameter target (the 76M is below it).
2. It is genuinely competitive with AION-base on DESI in-distribution (0.067 vs 0.074), from scratch, no AION weights.
3. It demonstrates the scaling + wavelength-jitter direction was tried and honestly evaluated — the negative OOD result is a real finding, not a gap.

Folder structure

NativeSpecZ-296M_Submission/
├── README.md
├── NativeSpecZ-296M.ipynb
├── weights/
│   ├── best.pt                     (1.18 GB checkpoint)
│   ├── training_args.json
│   ├── best_metrics.json
│   └── final_metrics.json
├── code/
│   ├── hybrid_redshift.py          (model architecture; load with strict=False — see note)
│   ├── data.py, metrics.py, model.py, plots.py
│   └── run_eval_296m.py            (evaluation script, strict=False load)
├── eval/
│   ├── desi_heldout/  sdss/  vipers/   (summary.json + NPZ predictions + multi_mask.json)
└── plots/
    ├── scatter_redshift.png            (predicted vs true z, DESI)
    ├── scatter_3datasets.png           (z scatter on DESI/SDSS/VIPERS)
    ├── spectrum_reconstruction.png     (4-panel masked-region reconstruction overlay)
    ├── comparison_296m_76m_aion.png    (MAE bars vs 76M and AION)
    ├── multi_mask_reconstruction.png   (rec MSE vs mask ratio, line vs continuum)
    └── stress_curve.png                (instrument-shift robustness)

Note on loading

The current hybrid_redshift.py includes z_rerank_head and z_calib_head modules that were added AFTER this checkpoint was trained. Load with strict=False:

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=False)  # rerank/calib heads unused in bin_residual mode
model.eval()

The z_rerank_head and z_calib_head are not used in the bin_residual prediction path — they were exploratory and left in the code. The bin-residual z prediction and the reconstruction head are fully loaded.

Hugging Face

ManmohanSharma/NativeSpecZ-296M on Hugging Face.

Submission checklist

• Approach A — joint z-head training
• Approach B — always-masked z token
• ~300M parameter target — 296M, meets the spec target
• (a) Redshift prediction — MAE 0.067 on held-out DESI (beats AION 0.074)
• (b) Masked reconstruction — rec MSE 0.066 at mask=0.25
• Unimodal DESI + z only, no imaging
• No AION pretrained weights
• Cross-instrument tested on SDSS + VIPERS
• [⚠] Does NOT beat AION-base on any OOD dataset (the 76M does, on VIPERS) — this is why the 76M is the headline
• [⚠] Clean-subset MAE 0.46 (DESI) — high-z clean-label weakness shared with the 76M