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