Wavelet-Domain ENSO ForecastingReport 6 of 9 · Probabilistic heads

Software companion

From Factorized Gaussian Noise to Spatially Correlated Generative Ensembles

Overview

All three heads attach to the same adopted deterministic model, which is never retrained. That is the design decision the report rests on: because the centre is frozen, every difference in ensemble behaviour is attributable to the head, and the deterministic skill established earlier is not silently traded away for calibration.

The heads differ in exactly how they generate spread. The Gaussian head predicts a per-coefficient variance, and its independence across coefficients is what produces spatially incoherent members. The latent-noise head draws one shared noise vector per member and maps it through a learned projection, so members are spatially correlated by construction. The diffusion head is a FiLM-conditioned denoiser over correction fields with a cosine noise schedule, generating members by reverse diffusion rather than by sampling a parametric family. Comparing them needs the same evaluation surface, so probabilistic_evaluation.py scores ensembles, spread-skill relationships, rank histograms and regional index distributions identically regardless of which head produced the members.

Implementation

Centre
The adopted deterministic model, frozen; no head retrains it
Heads
A per-coefficient Gaussian, a shared-latent-noise residual head, and a conditional DDPM denoiser over correction fields
Diffusion
FiLM-conditioned convolutional denoiser with a cosine cumulative-alpha schedule
Evaluation
A shared ensemble layer reporting field-level and regional-index calibration separately, because the two disagree
Configurations
Six probabilistic evaluation configurations covering the head variants and the CRPS ensemble comparison

Components

models/uncertainty.py
All three heads and their prediction types: the Gaussian coefficient prediction, the latent-noise residual head and its ensemble, and the conditional diffusion head
probabilistic.py
The probabilistic model wrapper around the frozen deterministic centre
probabilistic_evaluation.py
The shared ensemble evaluation surface: CRPS, spread-skill, reliability, and regional index distributions
conditional_variance.py
Conditional-variance diagnostics used to ask whether width is flow-dependent at all
training/losses.py
Gaussian likelihood and ensemble scoring objectives
pretraining/diffusion_pretrain.py
The diffusion-head training loop over correction fields
scripts/pretrain_diffusion.py
The diffusion training entry point
scripts/phase_g0_correlated_noise.py
The correlated-noise study that follows from the independence failure

Head comparison

  1. freeze Fix the deterministic centre and its selection; no head may move it.
  2. train head Fit the head on residuals or correction fields under its own objective.
  3. sample Generate a member ensemble — parametric draw, shared-latent draw, or reverse diffusion.
  4. score fields Evaluate spatial coherence, reliability and spread-skill on the reconstructed fields.
  5. score index Evaluate the regional index distribution separately, where the calibration wall shows up.

Reproducibility and validation

  • Because the deterministic centre is frozen, ensemble mean skill is held constant across heads and only the distributional behaviour varies.
  • Field-level and regional-index calibration are reported as separate results rather than averaged, since the report’s finding is that a head can fix the first and leave the second untouched.
  • Each head has its own test module covering its sampling path, so a correlated head cannot silently degenerate to independent noise.
  • The negative result is stated at the mechanism level — per-coefficient independence cannot produce a calibrated regional index — rather than as a tuning shortfall, and survived a sixty-four-fold increase in training data.

Availability

Not publicly released. The heads are small relative to the rest of the system, but they are meaningless without the frozen deterministic checkpoint they attach to, which is not distributed.

Companion research

A generative ladder built around a frozen deterministic center: per-coefficient Gaussian heads, spatially correlated latent-noise heads, a conditional diffusion head, and its regularized and simulation-pretrained variants.

Read the research report