From b302892e7d4e970b7cee61a2f1f3f9fc6d112212 Mon Sep 17 00:00:00 2001 From: sadamov Date: Fri, 5 Jun 2026 09:23:15 +0200 Subject: [PATCH] fix: Scale pred_std by diff_std at init in BaseGraphModel (#523) Re-applies @Debadri-das's PR #523 onto current main. The target file `base_graph_model.py` was moved to `models/step_predictors/graph/base.py` by #208; the surrounding code is otherwise unchanged. Expanded the NOTE-comment to explain WHY the scaling is needed (initial pred_std otherwise sits at softplus(0) = ln(2), independent of variable scale, which blows up NLL/CRPS in early training for any variable whose physical step-diff std is much larger). Co-Authored-By: Debadri Das Co-Authored-By: Claude Opus 4.7 (1M context) --- CHANGELOG.md | 2 ++ neural_lam/models/step_predictors/graph/base.py | 11 +++++++++-- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ee3166b..1260de5e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -64,6 +64,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Log every figure passed to `CustomMLFlowLogger.log_image` instead of silently dropping all but the first, using per-figure indexed keys (`{key}_{i}`) when more than one is supplied [\#499](https://github.com/mllam/neural-lam/pull/499) @Raj-Taware +- Scale predicted `pred_std` by `self.diff_std` in `output_std=True` graph step-predictors so it starts on the empirical one-step difference scale, avoiding early-training NLL/CRPS blow-up for variables whose physical step-diff std is much larger than `softplus(0) = ln(2)` ([#347](https://github.com/mllam/neural-lam/issues/347)) [\#523](https://github.com/mllam/neural-lam/pull/523) @Debadri-das + ### Maintenance - Add comprehensive type hints to `neural_lam/metrics.py` [\#447](https://github.com/mllam/neural-lam/pull/447) @sidhantpande diff --git a/neural_lam/models/step_predictors/graph/base.py b/neural_lam/models/step_predictors/graph/base.py index 4522d307..82263082 100644 --- a/neural_lam/models/step_predictors/graph/base.py +++ b/neural_lam/models/step_predictors/graph/base.py @@ -272,10 +272,17 @@ def forward(self, prev_state, prev_prev_state, forcing): pred_delta_mean, pred_std_raw = net_output.chunk( 2, dim=-1 ) # both (B, num_grid_nodes, d_f) - # NOTE: The predicted std. is not scaled in any way here + # Scale predicted std. with one-step difference std. so that + # the initial pred_std is on the empirical scale of the data, + # mirroring the diff_std scaling of pred_delta_mean below. + # Without this, NLL/CRPS explode in early training for + # variables whose physical step-diff std is much greater than + # softplus(0) = ln(2). # linter for some reason does not think softplus is callable # pylint: disable-next=not-callable - pred_std = torch.nn.functional.softplus(pred_std_raw) + pred_std = ( + torch.nn.functional.softplus(pred_std_raw) * self.diff_std + ) else: pred_delta_mean = net_output pred_std = None