# 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`:

```python
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

- [x] Approach A — joint z-head training
- [x] Approach B — always-masked z token
- [x] ~300M parameter target — **296M, meets the spec target**
- [x] (a) Redshift prediction — MAE 0.067 on held-out DESI (beats AION 0.074)
- [x] (b) Masked reconstruction — rec MSE 0.066 at mask=0.25
- [x] Unimodal DESI + z only, no imaging
- [x] No AION pretrained weights
- [x] 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