Upload README.md

README.md

@@ -0,0 +1,114 @@
+# 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
+    ├── scatter_3datasets.png
+    ├── comparison_296m_76m_aion.png
+    ├── multi_mask_reconstruction.png
+    └── stress_curve.png
+```
+
+## 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