Upload README.md

README.md

@@ -0,0 +1,127 @@
+# NativeSpecZ-FM-76M
+
+A 76M-parameter unimodal foundation model for DESI spectra. Trained from scratch on 97,332 DESI EDR spectra. No AION weights, no pretrained encoder — everything is ours.
+
+## 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
+
+The augmentation simulates non-DESI instrument characteristics during training, which is the cleanest way to encourage cross-instrument generalization without using non-DESI data.
+
+## Headline results
+
+### DESI held-out (n=2500, deduped from training)
+
+| Metric | Value |
+|---|---:|
+| MAE(z) | **0.0516** |
+| Median AE | 0.00212 |
+| RMSE | 0.189 |
+| Pearson r | **0.936** |
+| NMAD | 0.0019 |
+| Cat \|dz\|/(1+z)>0.01 | 13.5% |
+| Cat \|dz\|/(1+z)>0.05 | 9.0% |
+| Cat \|dz\|/(1+z)>0.15 | 6.8% |
+| Acc \|dz\|<0.10 | 90.6% |
+| Acc \|dz\|<0.20 | 92.3% |
+| Acc \|dz\|<0.30 | 93.6% |
+| Masked recon MSE (mask=0.25) | 0.037 |
+| Masked recon MSE (mask=0.30) | 0.437 |
+
+### Cross-instrument generalization
+
+| Dataset | n | MAE(z) | NMAD | AION-base MAE |
+|---|---:|---:|---:|---:|
+| **DESI held-out** | 2500 | **0.052** | 0.0019 | 0.074 (AION loses) |
+| **SDSS (real non-DESI)** | 2000 | 0.382 | 0.385 | 0.127 |
+| **VIPERS (real non-DESI)** | 2000 | **0.172** | 0.087 | 0.274 (AION loses) |
+
+We beat AION-base on **DESI** by 30% and on **VIPERS** by 37%. SDSS is harder for our raw-flux architecture; for SDSS-style inputs, the AION-tokenizer-based variants in our ensemble are better.
+
+## Folder structure
+
+```
+NativeSpecZ-FM-76M_Submission/
+├── NativeSpecZ-FM-76M.ipynb       ← demo notebook (load model, run eval, plot results)
+├── README.md                       ← this file
+├── REPORT.md                       ← full methodology + results write-up
+├── weights/
+│   ├── best.pt                     ← 292 MB model checkpoint
+│   ├── training_args.json          ← all hyperparameters
+│   └── best_metrics.json           ← training-time eval metrics
+├── code/
+│   └── hybrid_redshift.py          ← model architecture + collator + training loop
+├── eval_results/
+│   └── desi_2500_metrics.json      ← test-set metrics in friend-style format
+└── plots/
+    ├── friend_style_scatter.png            ← Pearson r=0.9358 scatter (matches friend's format)
+    ├── friend_style_reconstruction.png     ← 4-panel masked recon overlay
+    ├── comparison_vs_aion.png              ← MAE bar chart vs AION-base
+    ├── comparison_3way_desi.png            ← 6-metric chart: us vs AION vs friend
+    ├── foundation_evidence.png             ← cross-instrument robustness ratio
+    ├── dashboard.png                       ← 6-panel model summary
+    ├── multi_mask_reconstruction.png       ← rec quality vs mask ratio
+    └── stress_curve.png                    ← instrument-shift robustness
+```
+
+## How to reload the model
+
+```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=True)
+model.eval()
+```
+
+See `NativeSpecZ-FM-76M.ipynb` for the full inference + evaluation pipeline.
+
+## Hugging Face
+
+Model also available at `tempAstro/NativeSpecZ-FM-76M` on Hugging Face.
+
+## Submission checklist
+
+- [x] Approach A (joint z-head training, encoder shaped by z)
+- [x] Approach B (always-masked z token)
+- [x] Redshift prediction working (MAE 0.052 on held-out DESI)
+- [x] Masked spectrum reconstruction working (rec MSE 0.037 at 25% mask)
+- [x] Cross-instrument testing on real non-DESI data (SDSS + VIPERS)
+- [x] Unimodal — DESI spectra + z only, no imaging
+- [x] No AION pretrained weights used
+- [x] Single-model submission with notebook + weights + report