feat : a general probabilistic forecasting interface#700
feat : a general probabilistic forecasting interface#700Sir-Sloth-The-Lazy wants to merge 7 commits into
Conversation
Add abstract Forecaster.compute_training_loss returning a finished (loss, loss_components) pair, so each forecaster owns its complete training objective. ForecasterModule.training_step now only injects the configured scoring rule, interior mask and per_var_std, and logs the result. The deterministic ARForecaster loss is unchanged in value. Add the abstract ProbabilisticForecaster (sample_ensemble capability), ProbabilisticARForecaster (sequential sampled rollouts, trains on the configured score of the ensemble mean) and a minimal ProbabilisticForecasterModule whose validation samples an ensemble and logs the RMSE of the ensemble mean. Interface design from mllam#685.
|
@joeloskarsson @observingClouds , this is the implementation of the dicussion on issue #685 hope this is what you wanted ! 👀. |
observingClouds
left a comment
observingClouds
left a comment
There was a problem hiding this comment.
Hi @Sir-Sloth-The-Lazy thanks for drafting this. I had a look at the discussion around the proposal in #685 and this PR. It looks well aligned. So far, I only have a few minor comments.
| trajectory. This class adds ensemble forecasting on top: unrolling | ||
| several trajectories and stacking them along an ensemble dimension. | ||
| The default training objective scores the ensemble mean with the | ||
| injected scoring rule; forecasters with model-specific objectives |
There was a problem hiding this comment.
"injected scoring rule". Can this be written more explicitly? So that it is clear to users where the scoring rule is set.
| states at each predicted step, used both as the prediction | ||
| targets and to overwrite boundary nodes during the rollouts. | ||
| Dims: same as one ensemble member. | ||
| score_fn : Callable |
There was a problem hiding this comment.
fn always reminds me as an abbreviation for filename. Maybe use score_func or score_metric?
There was a problem hiding this comment.
Or even just score.
There was a problem hiding this comment.
renamed to score_metric as per the suggestion.
|
@observingClouds , Hope the latest commit solves the concerns 😁 |
joeloskarsson
left a comment
joeloskarsson
left a comment
There was a problem hiding this comment.
This looks quite nice I think, and is still pretty easy to follow. I had some outstanding thoughts regarding the connection to the deterministic version + some small things.
| init_states, | ||
| forcing_features, | ||
| target_states, | ||
| score_metric=self.loss, |
There was a problem hiding this comment.
If the loss calculation is now entirely handled by the Forecaster, I think it could make more sense that self.loss is instantiated and sits on the Forecaster, rather than being passed to it? What do you think, is there any problem with that?
There was a problem hiding this comment.
Good question, I don't think there's a hard problem with it, but it's a bit bigger than a one-liner so wanted to lay out the scope before doing it.
Right now self.loss isn't only feeding compute_training_loss ForecasterModule also calls it directly in validation_step/test_step (and for the spatial loss map) to report val/test loss, independent of whatever the forecaster's training objective ends up being. So "loss calculation entirely handled by the Forecaster" isn't quite true yet for the deterministic path the Module still owns a second, separate use of the same metric for reporting.
It also isn't uniform across the two modules we have: ProbabilisticForecasterModule.validation_step doesn't touch self.loss at all, it scores the ensemble mean with raw metrics.mse, and test_step isn't implemented there yet.
If we move it, the plan would be:
loss: str becomes a ctor arg on Forecaster (or ARForecaster), instantiated there as self.loss = metrics.get_metric(loss)
compute_training_loss drops the score_metric param entirely and reads self.loss internally
ForecasterModule.validation_step/test_step read self.forecaster.loss instead of self.loss
loss gets threaded through the forecaster constructors in train_model.py (currently only passed to ForecasterModule) and the ~15 test call sites that build ARForecaster/ProbabilisticARForecaster directly
None of that is a blocker, just confirming that's the scope before I make the change and whether ForecasterModule should keep a loss kwarg at all (for backward-compat / overriding the reporting metric independently of the training metric) or drop it entirely in favor of always reading it off the forecaster.
| target_states, | ||
| score_metric=self.loss, | ||
| interior_mask_bool=self.interior_mask_bool, | ||
| per_var_std=self.per_var_std, |
There was a problem hiding this comment.
Similarly for per_var_std, the main use case of this seems to be computing the loss, so would it also belong in the Forecaster then, that handles the loss?
There was a problem hiding this comment.
-
On score_metric/self.loss: Moved. ARForecaster (and ProbabilisticARForecaster through it) now takes aloss: str = "wmse"constructor argument and buildsself.loss = metrics.get_metric(loss)itself.compute_training_lossno longer takes ascore_metricparam, it readsself.lossinternally.ForecasterModule.training_stepnow just injects the interior mask; validation_step/test_step readself.forecaster.lossfor their reporting metric instead of owning a separate copy. -
On
per_var_std: Moved too, same reasoning. ARForecaster also takes an optionalconfig: NeuralLAMConfig = Noneconstructor argument and computesself.per_var_stdfrom it (same formula as beforediff_std / sqrt(feature_weights)), only when notself.predicts_std. Made config optional rather than required: since per_var_std is only needed for scoring, a forecaster built purely for inference (or the several existing tests that only exercise forward()) doesn't need to supply aNeuralLAMConfigat all it staysNoneand is simply never used, same as today'spredicts_std=Truecase.
| time_step_loss = torch.mean( | ||
| self.loss( | ||
| prediction, | ||
| target_states, | ||
| pred_std, | ||
| mask=self.interior_mask_bool, | ||
| ), | ||
| dim=0, | ||
| ) |
There was a problem hiding this comment.
This loss needs to also be computed by the Forecaster. It might be a good idea to have a read over the refactor in #675, as that should make some of this simpler. Not saying to rebase on top of that (although I think we can merge that very soon), just anticipate that change.
There was a problem hiding this comment.
On this specific comment: those lines already changed in 987fecd (pushed earlier today), which moved self.loss onto the forecaster. It's self.forecaster.loss(...) now, not self.loss(...).
On #675: Read through it. It consolidates training_step/validation_step/test_step's duplicated prediction+loss logic into one _compute_prediction_and_loss helper (plus a _log_step_loss helper). That's genuinely the right anticipation, once #675 merges and we rebase, our change gets simpler: instead of patching the loss/std fallback in 3 places, there'll be exactly one call site (inside _compute_prediction_and_loss) to point at self.forecaster.loss/self.forecaster.per_var_std instead of self.loss/self.per_var_std. Not rebasing now per your note, just flagging that the interaction is small and one-directional (their structural refactor, our data-ownership change).
One thing I found while reading it and fixed (511a6d5): #675 adds assert pred_std is not None after the per_var_std fallback. Combined with our change, config is now optional on ARForecaster/ProbabilisticARForecaster, so it's possible to build a forecaster without it, someone could build a forecaster for a predictor that doesn't output its own std, forget to pass config, wrap it in a ForecasterModule, and previously that would only surface as a crash at the first validation/test step. Added a check in ForecasterModule.__init__ that raises immediately with a clear message if forecaster.per_var_std is None and not forecaster.predicts_std, so it fails at construction instead.
| Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``. | ||
| The sampled forecasts, stacked along the ensemble dimension | ||
| ``S``. | ||
| ensemble_std : torch.Tensor or None |
There was a problem hiding this comment.
We should call this something different, as it is not the ensemble std, rather an std associated with each member. So maybe per_member_std.
When this is predicted, our ensemble is practically a mixture of Gaussians. Maybe good to explain somewhere.
There was a problem hiding this comment.
Pushed as 56b3d6b. Summary:
- Renamed
ensemble_std→per_member_stdthroughoutprobabilistic.py(the abstractsample_ensemblecontract onProbabilisticForecaster, ProbabilisticARForecaster's implementation, and its use incompute_training_loss), plus the matching test variable. - Added the mixture-of-Gaussians explanation to ProbabilisticForecaster's class docstring: when
per_member_stdis present, the predictive distribution isp(x) = mean_s N(x; ensemble[:, s], per_member_std[:, s]**2), not a single Gaussian and its variance isn't just the average of the per-member variances, it also picks up the spread between the member means. - Flagged the direct consequence in
ProbabilisticARForecaster.compute_training_loss's docstring: since it currently scores the ensemble mean usingper_member_std.mean(dim=1), that's a known simplification of the true mixture variance (drops the between-member spread term), not the exact predictive std.
| self, | ||
| predictor: StepPredictor, | ||
| datastore: BaseDatastore, | ||
| ensemble_size: int, |
There was a problem hiding this comment.
Making this ensemble size default part of the state seems unnecessary to me, but let me know if I am missing something. Can't we just always require an explicit number of ens members?
|
|
||
| def test_step(self, batch, batch_idx): | ||
| """ | ||
| Not supported: ensemble test evaluation is not implemented. |
There was a problem hiding this comment.
Go ahead and implement also this, should look quite similar to the validation step.
| ) | ||
| time_step_rmse = torch.sqrt(time_step_mse) | ||
| mean_rmse = torch.mean(time_step_rmse) | ||
| self._warn_skipped_val_steps(len(time_step_rmse), "val") |
There was a problem hiding this comment.
I think these types of warnings refer to the validation loss logging, not the ensemble mean MSE. View the loss as something separate, and the ensemble RMSE as a separate metric.
| self.eval_ensemble_size = eval_ensemble_size | ||
| self.val_metrics = {"ens_mse": []} | ||
|
|
||
| def validation_step(self, batch, batch_idx): |
There was a problem hiding this comment.
For the validation_step, now the approach seems to be to rewrite this fully separated from the validation_step of the deterministic module. It seems like we do a lot of the same thing, so I wonder if it's not possible to re-use the deterministic validation_step with a super call here? It could then compute things for 1 ensemble member, and then here in the extended validation_step (for the probabilistic extension) we do ensemble-specific things?
There was a problem hiding this comment.
This is up for discussion, maybe it would get very messy to try to use the validation_step of the ForecastModule with a probabilistic forecaster?
| - Add a general probabilistic forecasting interface: an abstract | ||
| `ProbabilisticForecaster` capable of sampling ensemble forecasts | ||
| (`sample_ensemble`, members stacked along a new dimension after batch), | ||
| its auto-regressive implementation `ProbabilisticARForecaster` (samples | ||
| independent trajectories through a stochastic step predictor and by | ||
| default trains on the configured scoring rule applied to the ensemble | ||
| mean) and a `ProbabilisticForecasterModule` whose validation samples an | ||
| ensemble and logs the RMSE of the ensemble mean. Move ownership of the | ||
| training objective from `ForecasterModule` onto the `Forecaster`: the | ||
| new abstract `Forecaster.compute_training_loss` returns a finished | ||
| `(loss, loss_components)` pair and `ForecasterModule.training_step` only | ||
| injects the configured scoring rule and interior mask and logs the | ||
| result. The deterministic `ARForecaster` training loss is unchanged in | ||
| value, only computed by the forecaster itself. | ||
| [\#685](https://github.com/mllam/neural-lam/issues/685) | ||
| @Sir-Sloth-The-Lazy |
There was a problem hiding this comment.
This entry is way too long. Should be a short 1-sentence summary.
There was a problem hiding this comment.
987fecd Done. Sorry for this, I went overboard with describing the changes
|
Thank you for taking out the time to review this @joeloskarsson means a lot. I will start working on the review now 😃 |
score_metric/per_var_std were injected into compute_training_loss by ForecasterModule and also used directly for val/test loss reporting, duplicating config the forecaster already needs for its own objective. ARForecaster/ProbabilisticARForecaster now own self.loss and self.per_var_std (computed from an optional config ctor arg), and ForecasterModule reads them off self.forecaster instead. Also trims the CHANGELOG entry for mllam#685 down to one sentence per review feedback.
A Forecaster built without config now silently has per_var_std=None when its predictor doesn't output its own std. Previously per_var_std was always computed by ForecasterModule itself, so this gap didn't exist; now that construction is split across two calls, catch it at ForecasterModule init instead of crashing at the first val/test step.
Each member's predicted std is its own, not a spread computed across the ensemble, so ensemble_std was a misleading name. Document on ProbabilisticForecaster that a per-member std makes the predictive distribution a mixture of Gaussians, and note in ProbabilisticARForecaster.compute_training_loss that averaging the per-member stds is a simplification of the true mixture variance (which also includes the spread between member means).
Co-authored-by: Joel Oskarsson <joel.oskarsson@outlook.com>
Scoring the ensemble mean with a pointwise metric only rewards the mean being right, giving the model no incentive to keep a calibrated spread, and risks training it to collapse to a point estimate. Redeclare compute_training_loss as abstract on ProbabilisticARForecaster instead of providing that as a default (it would otherwise silently fall back to ARForecaster's single-rollout objective via MRO, not even the ensemble mean). Concrete subclasses must define their own objective. Tests that only need an instantiable forecaster now use a local ConcreteProbabilisticARForecaster example (ensemble-mean scoring, moved out of the library code); a new test locks in that the base class itself cannot be instantiated.
3 participants
Describe your changes
This PR implements the abstract constructions agreed in the #685 discussion (RFC: a general probabilistic forecasting interface), deliberately without any concrete Graph-EFM instantiation, so the design can be reviewed on its own and this can server as the base for addition of any probabilistic model, not just
graph_efm.Move ownership of the training objective from
ForecasterModuleonto theForecaster. A new abstractForecaster.compute_training_lossreturns a finished(loss, loss_components)pair, so each forecaster owns its complete training objective (including assembling it from any internal terms).ForecasterModule.training_stepnow only injects the configured scoring rule (--loss), the interior mask andper_var_std, then logs the returned loss and components (component names prefixed with the phase). Per the discussion,this move is applied to the deterministic setup as well, so the same concept sits in the same place in both stacks: the deterministic
ARForecastertraining loss is unchanged in value (covered by an equality test), it is just computed by the forecaster itself. Note for reviewers: since the method is abstract onForecaster, any out-of-treeForecastersubclass now has to implement it; all in-repo forecasters go throughARForecaster.Add the probabilistic side, built on top of the deterministic one (no sample dimension leaks into deterministic components):
ProbabilisticForecaster(abstract): declaressample_ensemble, producing members stacked along a new dimension after batch,(B, S, pred_steps, num_grid_nodes, d_state). This encodes the module's only assumption "the forecaster can create ensemble forecasts of the correct shape" as a type contract, and leaves room for non-AR implementations (diffusion/flow) later.ProbabilisticARForecaster(ARForecaster, ProbabilisticForecaster): unrolls independent trajectories through a step predictor that samples its output (sequential loop as the safe first default), and by default trains on the configured scoring rule applied to the ensemble mean. Model-specific objectives (e.g. Graph-EFM's ELBO + CRPS) will live in subclasses that overridecompute_training_loss.ProbabilisticForecasterModule: training inherited unchanged; validation samples aneval_ensemble_sizeensemble and logs the RMSE of the ensemble mean (the example ensemble metric suggested in RFC: a general probabilistic forecasting interface #685, standing in until the ensemble-metrics PR), reusing theval_mean_loss/val_loss_unroll{i}log keys so existing checkpoint callbacks work unchanged.test_stepraisesNotImplementedErrorrather than silently evaluating a single member deterministically; ensemble test evaluation and plotting are a follow-up.Dependencies: none. In particular this PR does not depend on the planned ensemble metrics (
crps_ens,spread_squared). Follow-ups: ensemble metrics, the Graph-EFM forecaster (ELBO + CRPS, its ownkl_beta/crps_weightconfig), ensemble test evaluation + plotting, andtrain_model.pywiring forprobabilistic models.
Issue Link
Graph-EFM instantiation and ensemble evaluation follow in later PRs).
prob_model_lamonmain, see issue Merge Graph-EFM model fromprob_model_lambranch #62Type of change
Checklist before requesting a review
pullwith--rebaseoption if possible).Checklist for reviewers
Each PR comes with its own improvements and flaws. The reviewer should check the following:
Author checklist after completed review
reflecting type of change (add section where missing):
Checklist for assignee