JunuChoi/astra-ka-ablations

Astra KA ablation checkpoints

Checkpoints of the Kantorovich Attention (KA) ablations for the journal version of Efficient probabilistic weather forecasting at quarter-degree resolution with two-sided transport attention (Astra).

• Network: W-Net-L/4 (217 M parameters) at 0.25 degrees, trained on ERA5 2018-2019 (anchors at 00 and 12 UTC).
• MeanFlow stage: one H200 per run, 4 rows per step (2 JVP rows, 2 rows at r = t), 7,000 steps, a checkpoint every 1,000 steps; seeds 1-3. Each arm changes only the tiles that every KA round selects at the 1-degree and 2-degree levels:
  • ka: the learned two-sided (pull and push) selection.
  • ka_local: each tile's great-circle-nearest tiles.
  • ka_random: fixed, degree-balanced random tables.
  • pull_only: the push branch converted to a second pull branch.
• CRPS stage: from the ka MeanFlow checkpoint at 4,000, 3,000 steps of one case with 4 members: field CRPS (crps_field), field and feature CRPS (crps_feat), and field and feature CRPS with the feature network refreshed every 1,000 steps (crps_feat_refresh, in segments seg0-seg2); seeds 1-3.
• Layout: wnet_l_quarter/<arm>_s<seed>/ckpt_<step>.*. .sf_x.eqx holds the averaged weights used for evaluation, .meta.json the run identity and counters; SHA256SUMS lists the files of a run.
• assets/: the normalisation statistics, masks, statics, model priors and SST forcing that the evaluation reads (assets/ASSETS_PUB.json).