PRISM-Astro/mt-static-best-m

MT-Static (the lowest validation M)

A multi-task CNN that predicts the host position and the redshift from the same transient-centred PS1 grizy stamps of the P19 sample, with a shared DELIGHT backbone and two heads, trained with static loss weights (rho = w_z / w_pos = 25: w_pos = 0.0769, w_z = 1.923) for 80 epochs.

Model

• Backbone: the DELIGHT convolutional backbone (Förster et al. 2022), shared by the five stamp levels.
• Input: PS1 stamps of 30 x 30 pixels at five levels (0.25 arcsec per pixel at level 0, doubling at each level), band(s) g, r, i, z, y.
• Output: the host offset (x, y) and the redshift distribution.

The host offset is predicted in level-0 pixels of 0.25 arcsec from the stamp centre. Each prediction is the average of the eight rotations and flips of the stamp, transformed back to the original orientation. M is the mean of min(r, 1), with r the residual distance in units of the host's SEP ellipse scaled by three (semi-axes A = 3a, B = 3b); RMSE_d is the root mean square of the residual distance. Following Pasquet et al. (2019), the redshift is treated as a classification: the head gives the probabilities of 180 equal bins in 0 <= z <= 0.4, and the photometric redshift is the mean of the probability distribution averaged over the eight rotations and flips of the stamp. sigma_NMAD = 1.4826 median(|dz - median(dz)|), with dz = (z_phot - z_spec)/(1 + z_spec); bias is the mean dz and eta the fraction with |dz| > 0.05. g_pos and g_z are the relative gains over the single-task baselines P19-Host (position, lowest validation M of the same fold) and P19-z (redshift, lowest validation sigma_NMAD of the same fold); Delta_m is their mean.

Training data and cross-validation

Trained on PRISM-Astro/h2f_ps1_pasquet_autolabeling_radius_norm_p6 with five-fold cross-validation: there is one model per fold, trained on the other four folds and validated on its own fold. The host locators are tested on the real host-transient pairs of PRISM-Astro/h2f_delight_test.

For each fold we keep the epoch with the lowest validation M (best host localisation). The same runs with the other two selections are in PRISM-Astro/mt-static-best-delta-m (the largest multi-task gain Delta_m) and PRISM-Astro/mt-static-best-nmad (the lowest validation sigma_NMAD).

Validation results

Validation metrics of each fold model on its own validation fold.

Fold	Epoch	Step	M	RMSE_d (arcsec)	sigma_NMAD	eta (%)	g_pos (%)	g_z (%)	Delta_m (%)
1	75	955425	0.02607	0.628	0.02056	3.56	-2.91	+3.73	+0.41
2	76	972572	0.02611	0.565	0.02067	3.75	-2.81	+3.06	+0.13
3	78	1004094	0.02585	0.568	0.02071	3.73	-2.08	+3.39	+0.66
4	80	1025600	0.02575	0.601	0.02053	3.75	-1.61	+3.64	+1.02
5	75	956175	0.02615	0.598	0.02064	3.78	-2.83	+1.83	-0.50
Mean ± std			0.02599 ± 0.00016	0.592 ± 0.024	0.02062 ± 0.00007	3.71 ± 0.08	-2.44 ± 0.52	3.13 ± 0.69	0.34 ± 0.51

Files

config.yaml            Hydra configuration of the run (model, dataset, training arguments)
selection.json         checkpoint, epoch, step and validation metrics of each fold
fold_k/model.safetensors
fold_k/trainer_state.json

Usage

Access is gated: request access on this page and log in (huggingface-cli login or HF_TOKEN) before downloading.

from hydra.utils import instantiate
from huggingface_hub import hf_hub_download
from omegaconf import OmegaConf
from safetensors.torch import load_file

repo = "PRISM-Astro/mt-static-best-m"
cfg = OmegaConf.load(hf_hub_download(repo, "config.yaml"))
model = instantiate(cfg.model, _convert_="partial")   # needs the PRISM package
model.load_state_dict(load_file(hf_hub_download(repo, "fold_1/model.safetensors")))
model.eval()

config.yaml is the Hydra configuration of the run, with dataset.fold_idx: null; the model of fold_k was trained on the other four folds of the dataset's fold column and validated on fold k. All fold models were trained with an effective batch of 40 stamps (per-device batch times the number of GPUs, which varied between folds), so the per-device batch size in the file does not apply to every fold.

The input stamps must be preprocessed as in training (band selection and order, inverse hyperbolic sine of the g band when it is used, five levels); see prism.dataset and scripts/inferencia_test.py in the code.