Add per-sample spectral whitening option to the energy-score loss

[mcgibbon]

Small-scale (high-l) spectral power converges very slowly in SFNO training: the energy score lives in a red spectrum, so high-l modes carry tiny amplitude, contribute little to the score, and their dhconv filter weights see proportionally small, low-SNR gradients that the energy score's uniform per-(l,m) mode_weights do not correct.

This PR adds an opt-in per-sample spectral whitening mode to EnergyScoreLoss. It reweights the per-(l,m) energy score by a power of the inverse per-degree RMS amplitude of each target sample, flattening the target's angular power spectrum so high-l modes are no longer starved and their gradient SNR rises. The reweight is computed from the detached target coefficients (no new gradient path) and is magnitude-preserving (rescaled per (sample, channel), so the overall energy-score magnitude — and the meaning of energy_score_weight — is unchanged; only the per-scale balance shifts). A white-spectrum target yields a uniform factor (no-op), and the default (whitening off) is bit-for-bit backward compatible.

Whitening strength is the exponent γ ∈ (0, 1] (per-degree factor (1/amp_l)**γ), interpolating between no reweighting (γ→0) and full flattening (γ=1). Enabling whitening (kind='per_sample') defaults to γ=0.5, the validated stable knee; full flattening (γ=1) over-upweights noise-dominated low-amplitude degrees and destabilizes residual rollouts, so it is opt-in. Options are grouped in a SpectralWhiteningConfig (kind, eps_frac, exponent); kind='none' (default) disables whitening, and eps_frac/exponent left unset resolve to 0.02/0.5 when enabled (setting either with kind='none' raises). In a loss config it is a nested block:

loss:
  type: EnsembleLoss
  kwargs:
    crps_weight: 0.9
    energy_score_weight: 0.1
    energy_score_whitening:
      kind: per_sample
      exponent: 0.5

Changes:

• fme.core.loss.SpectralWhitening / SpectralWhiteningConfig — new operator and config for per-sample whitening. The config (validated in __post_init__) build()s the operator, or None when kind='none'; the operator owns the reweight (per-(sample, channel, l), detached, floored at eps_frac of the per-sample mean degree amplitude, raised to exponent, magnitude-preserving).

• fme.core.loss.EnergyScoreLoss / EnsembleLoss — take the built whitening: SpectralWhitening | None operator (not the config) and apply its factor after mode_weights.

• fme.core.loss.LossConfig.build — constructs SpectralWhiteningConfig from the opaque loss kwargs dict, builds the operator, and passes it down.

• fme.core.test_loss — no-op default, white-target invariance, small-scale boost, magnitude preservation, exponent range, symmetric per_sample-required validation of both knobs, resolved per_sample defaults, partial-whitening interpolation, and config wiring via dict and dataclass.

• Tests added

• If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated

[mcgibbon]

Full analysis and figures: whitening-γ-selection report

whitening_exponent (γ) knob added + validation

The latest commit adds the whitening_exponent γ knob ((1/amp_l)**γ, γ ∈ (0, 1], default 1 = prior full-whitening behavior). Motivation: full whitening (γ=1) maximizes the small-scale gradient boost but over-upweights noise-dominated low-amplitude high-l degrees, which destabilizes residual-prediction rollouts (stratospheric-temperature-led blowup). Partial whitening keeps most of the spectral win while staying stable.

γ selected: 0.5. A two-seed sweep of γ ∈ {0, 0.2, 0.5, 0.8} on the 4°/daily v2 ERA5-only residual recipe (matched epoch 120, seed-pair noise sizing) found γ=0.5 the unique arm clearing all four adoption conditions: (a) large smallest-scale spectral win well outside seed noise (h500 −51% / >7σ, TMP2m −86%; the win saturates by γ=0.5, γ=0.8 adds only −5 pts), (b) 10-yr rollout health within seed noise of baseline, (c) 1-step val CRPS within seed noise, (d) no regression on the 46-yr baseline surface. γ=0.8 fails (d) with a reproducible two-seed 46-yr stratospheric regression; γ=0.2 under-delivers (a) and destabilizes one seed.

Sweep runs (γ, seed):

γ	seed 0	seed 1
0 (baseline)	wandb oshj5u79	beaker 01KWW0J0KWPGDXT2X5PNXY2SBG / wandb eyfz92l8
0.2	beaker 01KWW0EPJZKKE25P6ZSWHEPS7S / wandb bwbg9r2k	beaker 01KWW0FBKYYNA7R9ET1Z6T1JF4 / wandb jjprdcdt
0.5 (selected)	beaker 01KWHS1DHTR858DC51720J1Q5C / wandb erug2fqs	beaker 01KWW0HBG0509A1C6Q86YQA803 / wandb 4dltgowi
0.8	beaker 01KWW0G0SN0383BM8FSD3SR940 / wandb 8vgzlp54	beaker 01KWW0GNXKVWRBKGY3DMAY12VA / wandb xen5ofv9

Caveat carried by the verdict: the γ arms stop at epoch 120 while the baseline (extended to 259) drifts past ~ep 150, so γ=0.5 is verified clean through 120 but not yet proven not to merely delay the same intrinsic residual-recipe drift; a decisive hardening run (extend the two γ=0.5 seeds past ep 150) is tracked separately.

---

Note: this commit also fixes a pre-existing pre-commit (mypy) CI failure — the test_energy_score_invalid_whitening_raises test passes a deliberately-invalid literal, now marked # type: ignore[arg-type].

[mcgibbon]

Isn't this always non-None? But you're checking for None below? Choose one.

[mcgibbon]

Claude: Good catch on the readability — the (whitening or SpectralWhiteningConfig()) config is always non-None, but .build() returns SpectralWhitening | None (None when kind='none'), so self._whitening genuinely can be None and the check in forward is deliberate: it skips the whole reweight when whitening is disabled (rather than building a no-op operator that still runs the SHT/factor math every step). I kept the None path and made the intent explicit — annotated self._whitening: SpectralWhitening | None with a comment saying build() returns None when disabled (801a719).

[mcgibbon]

Claude: Applied self-review follow-ups (pushed as 7df72cf9), updating the earlier γ-knob note above:

• Default γ flipped 1 → 0.5. Enabling kind='per_sample' now defaults to the sweep-selected stable knee (0.5), not full whitening — shipping the opt-in defaulted to the value the sweep found destabilizing was a footgun. eps_frac/exponent are now None-sentinel and resolve to the validated defaults only when enabled; setting either with kind='none' raises symmetrically (previously only exponent did).
• Build moved to the config layer. LossConfig.build now builds the SpectralWhitening operator and passes it down; EnergyScoreLoss/EnsembleLoss take the built operator instead of the config, so the nn.Modules no longer build their own sub-config.
• factor cleanup. Drops the target's singleton ensemble dim up front (y_hat[:, 0]) instead of a trailing conditional squeeze, and computes |y_hat| once.

New tests cover the eps_frac-with-none guard and the resolved per_sample defaults. Full loss suite (72) + mypy/ruff pass locally.

[mcgibbon]

Report for reviewer: https://github.com/ai2cm/reports/pull/2/changes#diff-a2d53a810a58b87bed5ab372f8611c88b4041df7d758ef0afb29b6112ca3c0e5

[mcgibbon]

Claude: Empirical motivation and validation, moved out of the description so it reads as a clean commit message (detail kept here):

Motivation — gradient probe. On a trained checkpoint, a per-l gradient probe showed an ~8× monotone decline of high-l gradient magnitude across degrees, confirming the high-l dhconv filter weights are gradient-starved under the red-spectrum energy score.

Validation. On 4°/daily residual runs, base whitening delivered a large small-scale spectral-fidelity gain at neutral CRPS. A two-seed γ ∈ {0, 0.2, 0.5, 0.8} sweep selected γ=0.5 as the stable knee — the default this PR ships for kind='per_sample'. Full run table, adoption criteria, and the whitening-γ-selection report are in the sweep comment above.