Training

Every published checkpoint is reproducible from a config file.

specsr-roman train sr1   --config configs/sr1.yaml
specsr-roman train zhead --config configs/zhead.yaml     # needs SR1
specsr-roman train sr2   --config configs/sr2.yaml       # needs SR1 + ZHead

Or from Python:

from specsr_roman.config import SR1Config, load_config
from specsr_roman.training import train_sr1

summary = train_sr1(load_config(SR1Config, "configs/sr1.yaml"))

The shipped configs are the published chain — configs/sr1.yaml produces sr1_ou2024_v6_independent — and use split_mode: three_way. Any field can be overridden (--epochs 50, --no-augment); an unknown key raises rather than silently doing nothing.

The three stages

SR1 maps the LR spectrum to a mean and a heteroscedastic log-variance. It sees [flux, err] on a shared scale, so the channel ratio is the per-pixel S/N.

ZHead freezes SR1 and predicts P(z) over 310 bins from [LR flux, LR err, SR1 mean, SR1 log-σ] plus photometry.

SR2 freezes both and learns a delta on SR1, with one attention token per rest-frame line, run once per redshift hypothesis.

Two settings that must not drift

phot_tier: medium         # bands that ship with the grism
phot_eval_mag_err: 0.05   # never score against noiseless truth photometry

medium is the three Roman bands of the published chain. medium_rubin adds Rubin ugrizy with the sky-limited noise of a coadd rubin_depth_offset magnitudes shallower than the ten-year one (configs/zhead_rubin_y10.yaml, configs/zhead_rubin_y1.yaml). Each redshift head has its own SR2 (configs/sr2.yaml, sr2_rubin_y1.yaml, sr2_rubin_y10.yaml, sr2_photonly_rubin_y10.yaml), identical except for the head it reads; scripts/run_chain.sh rubin sr2_bands trains the ones beyond the published three-band chain. grids.MAX_PHOT_BANDS refuses an explicit band list longer than those nine. Feeding a model more bands than a survey delivers with the grism hands it an effectively complete, noiseless SED, from which the redshift can be read without the spectrum contributing anything — specsr-roman evaluate ablation exists to keep that visible. Setting phot_eval_mag_err: 0 reports an accuracy nobody will reproduce on sky.

Checkpoint selection is not the validation loss

This is the part most worth understanding before changing anything.

SR1 monitors line-flux recovery plus the hallucination penalty, smoothed by an EMA — not the total loss. The validation loss is dominated by the NLL, which rises as the model commits to line amplitudes: a confident half-amplitude line costs more than a hedged flat one. An early run had every line-recovery metric improving to epoch 150 while the loss picked a mid-run “best”.

SR2 minimises -recov_amp + lam_hallu * hallu_amp — integrated predicted-over-true flux on recoverable strong lines against the same on undetectable ones. Plain NLL is continuum-dominated and reliably selects the SR2 that draws nothing; a line-region MSE ratio is diluted across 98 windows and picks the timidest epoch available.

Note

The Roman SR2 best checkpoint is epoch 8 of 150. Selection uses only the 2,327-spectrum validation partition; the 2,384-spectrum test partition is reserved until all nine models are frozen.

Logging

Runs sync to Weights & Biases when it is installed, logging example spectra and residual histograms alongside scalars. That is not decoration: a scalar loss curve cannot tell you a super-resolution model has converged to the prior mean. The loss looks fine; only a picture of a predicted spectrum next to its target shows the flat line.

Set wandb_mode: disabled to turn it off; training never fails because a metrics service is unreachable.

Publishing checkpoints

push_to_hub: true
hub_repo: your-org/your-repo

Uploads the best checkpoint with a card.json recording metrics, dataset, upstream checkpoints, resolved config and the W&B run URL. A .pth without that record is an unreproducible artefact within a week.