EMA_WEIGHTS.md

EMA checkpoints (cBottle-3d-ema/, cBottle-3d-tc-ema/)

Exponential-moving-average (EMA) checkpoints for the cbottle-3d-moe denoiser and its paired cbottle-3d-moe-tc tropical-cyclone classifier, released to reproduce the results in "Towards accurate extreme event likelihoods from diffusion model climate emulators" (egusphere-2026-2610). These are the weights used for every figure and number in the revised manuscript; the originally submitted version used the raw (non-EMA) optimizer weights in cBottle-3d/ and cBottle-3d-tc/ instead.

Inference code is public at NVlabs/cBottle (scripts/inference_coarse.py for AMIP-style sampling, scripts/inference_odds_ratio.py for the guided-sampling / importance-sampling pipeline used in the likelihood paper).

Contents

• cBottle-3d-ema/ema-checkpoint-000512000.checkpoint — MoE expert A (σ ≥ 100 active range), 512k images seen
• cBottle-3d-ema/ema-checkpoint-002048000.checkpoint — MoE expert B (10 ≤ σ < 100), 2048k images seen
• cBottle-3d-ema/ema-checkpoint-009856000.checkpoint — MoE expert C (σ < 10), 9856k images seen
• cBottle-3d-tc-ema/ema-checkpoint-002176000.checkpoint — paired TC classifier, 2176k images seen

Same training run and same image counts (nimg) as cBottle-3d/ and cBottle-3d-tc/ — these are an EMA-averaged view of the identical training run, not a separately trained model. Do not mix these EMA weights with the classifier from the raw-weight run, or vice versa: the likelihood-paper reproduction used the EMA classifier throughout, paired only with the EMA MoE denoiser.

Why EMA weights

The manuscript's original submission used net_state.pth from the raw training-state checkpoints (the optimizer's live weights, i.e. what's in cBottle-3d/ and cBottle-3d-tc/). A later evaluation found the EMA weights close part of the gap between cBottle's tropical-cyclone climatology and ERA5's, at no cost elsewhere. We gated the switch by re-running the original cBottle paper's own AMIP evaluation suite (spectra, seasonal cycle, precipitation PDF, variance maps, global-mean biases, Arctic Oscillation, ENSO teleconnection, TC frequency) comparing raw vs. EMA weights against ERA5. All figured evidence below is from that gate; diagnostics we evaluated but aren't shown as figures here are called out explicitly rather than asserted without a plot.

Validation loss / log-probability

80 held-out samples per checkpoint — EMA is equal to or better than the raw weights at all three checkpoint steps, with the largest gain at the least-trained checkpoint (512k):

Breaking this down by noise level confirms EMA is not worse anywhere, and is consistently at least as good within each expert's own MoE-active σ range (shaded):

Tropical cyclone frequency

AMIP inference, 1980–2000, TempestExtremes detections, same pipeline as the original cBottle paper — EMA closes roughly half the gap to ERA5 without any guidance or classifier involved:

	detections/year	ρ̄ (detections vs. ERA5 ratio)
ERA5 (truth)	1007	1.3
cBottle raw MoE	707	0.95
cBottle EMA MoE	873	1.2

Global-mean biases

This is the most notable improvement outside of TC frequency. Global-mean, deseasonalized anomaly time series (1980–2000, shading = ensemble spread) against ERA5 (black):

Precipitation — the raw-weight model has a persistent ~+0.3 mm/d global bias relative to ERA5 that does not appear in the EMA weights, which track ERA5's inter-annual variability directly instead of just running parallel to it offset upward:

raw weights	EMA weights

Z500 — similarly, the raw-weight model runs persistently above ERA5 for most of the 1980s and misses the 1984–85 dip entirely; the EMA weights track ERA5 much more closely throughout, including that dip:

raw weights	EMA weights

2 m temperature — both track ERA5's inter-annual variability similarly well; no meaningful difference between raw and EMA here:

raw weights	EMA weights

Modes of variability (no regression)

Arctic Oscillation — daily AO index against ERA5, train-period CRPS is identical (5.0) for both raw and EMA weights; switching to EMA does not change AO skill:

raw weights	EMA weights

ENSO precipitation teleconnection — regression of precipitation anomaly onto the ENSO index, against ERA5 (a); raw (b, left) and EMA (b, right) both reproduce the tropical Pacific dipole pattern with no visible difference between them:

raw weights	EMA weights

Precipitation PDF and variance maps

Effectively unchanged, and if anything marginally closer to ERA5 truth for precipitation and near-surface temperature RMS:

Power spectra

Spherical-harmonic power spectra against ERA5 truth (black) for Z500, pr, T850, tcwv, tas, and uas, raw weights (dashed) and EMA weights (teal) overlaid directly. The two are visually indistinguishable across the full wavenumber range for all six variables — EMA does not shift spectral content at any scale:

Seasonal cycle

DJF-minus-JJA seasonal contrast maps for Z500, precipitation, 2 m temperature, and sea-ice concentration — truth (top row), raw weights (middle row), EMA weights (bottom row). All three rows are qualitatively very similar, including the Southern Ocean sea-ice edge and the Southern Hemisphere storm track pattern in panel (d)/(h)/(l):

Conclusion: EMA weights are climatologically as good as or better than the raw weights on every metric checked, with clear improvements in TC frequency and in the global-mean precipitation and Z500 biases, and no regression in validation loss, power spectra, seasonal cycle, precipitation PDF, variance maps, Arctic Oscillation skill, or ENSO teleconnection. This is why the revision uses EMA weights throughout.

Usage

Unconditional / AMIP-forced sampling (reproduces the climatology plots above):

python3 scripts/inference_coarse.py \
  --dataset amip --sample.mode sample --sample.sigma_max 200 \
  --sample.sampler all --sample.batch_gpu 8 --sample.bf16 \
  --end_time 2017-12-31 \
  cBottle-3d-ema/ema-checkpoint-000512000.checkpoint,cBottle-3d-ema/ema-checkpoint-002048000.checkpoint,cBottle-3d-ema/ema-checkpoint-009856000.checkpoint \
  <output_dir>

Guided sampling / odds-ratio calculation for the TC likelihood paper (uses the classifier checkpoint too):

python3 scripts/inference_odds_ratio.py \
  --moe cBottle-3d-ema/ema-checkpoint-000512000.checkpoint,cBottle-3d-ema/ema-checkpoint-002048000.checkpoint,cBottle-3d-ema/ema-checkpoint-009856000.checkpoint \
  --classifier cBottle-3d-tc-ema/ema-checkpoint-002176000.checkpoint \
  ...  # see scripts/inference_odds_ratio.py --help for the full sampler/guidance CLI

Provenance

• Converted from EMA network snapshots (network-snapshot-*.pkl) to inference-ready checkpoints via scripts/cmip_eval_scripts/convert_ema_snapshot_to_checkpoint.py in the cBottle training repo.
• Denoiser: three-expert mixture (SongUNet, HEALPix HPX64), each expert responsible for a disjoint noise-level range (σ<10, 10≤σ<100, σ≥100).
• Classifier: cBottle-3d-tc, a SongUNet with an HPX3-resolution classification head, trained jointly with a denoising objective, binary cross-entropy term (weight 0.1) on IBTrACS-derived tropical-cyclone labels.

License

Same terms as the rest of this repository: the NVIDIA Software and Model Evaluation License Agreement.

Citation

@article{manshausen2026towards,
  title={Towards accurate extreme event likelihoods from diffusion model climate emulators},
  author={Manshausen, Peter and Brenowitz, Noah and Berner, Julius and Kashinath, Karthik and Pritchard, Mike},
  journal={arXiv preprint arXiv:2605.03802},
  year={2026}
}

This is the arXiv preprint; update to the camera-ready EGUsphere/journal citation once egusphere-2026-2610 is formally published.