RFC: a general probabilistic forecasting interface

[Sir-Sloth-The-Lazy]

Status: Draft for review
Audience: @joeloskarsson , @leifdenby , @sadamov , @observingClouds and anyone interested ! 🤩
Scope: how to support any probabilistic forecasting model in neural-lam

---

1. Summary

The deterministic stack has a property worth preserving: adding a new model (GraphLAM, HiLAM, …) means implementing one abstraction — StepPredictor and registering it. The core loop (StepPredictor → ARForecaster → ForecasterModule) is never touched. It is closed for modification, open for
extension.

This proposal extends that same property to probabilistic models. The goal:

A new probabilistic model is added by implementing a single predictor
abstraction and registering it. ProbabilisticARForecaster and
ProbabilisticForecasterModule are generic and never change per model, and
contain no mention of ELBO, KL, CRPS, priors, posteriors, or β-weights.

ELBO becomes one example of a probabilistic objective, expressed entirely
inside the Graph-EFM model class — not baked into the core.

---

2. Why the deterministic core stays fixed today

Every deterministic model conforms to one contract:

StepPredictor.forward(prev_state, prev_prev_state, forcing) -> (pred, std | None)

and they all share one objective shape, assembled in
ForecasterModule.training_step:

loss = mean( self.loss(prediction, target, pred_std, mask) )

The only thing that varies between models is forward; the only thing that varies between runs is the configured scoring rule (--loss wmse|mae|nll|…). Because both the contract and the objective shape are fixed, the core never changes. That is exactly why [ARForecaster] and [`ForecasterModule'] are model-agnostic.

---

3. What probabilistic models break

Probabilistic models keep the forward contract (they can still produce a predictive step) but they break the shared objective shape. In prob_model_lam, GraphEFM.training_step assembles:

loss = -mean_likelihood + kl_beta * mean_kl   (+ crps_weight * crps_ens)

— reconstruction likelihood + KL(posterior‖prior) + an optional ensemble-CRPS term over sampled trajectories. A pure-ensemble model would instead train on CRPS alone; a diffusion model would train on a denoising loss. There is no single objective formula that covers these.

So if we want the core to stay fixed, the objective, the part that varies, must be owned by the model and reached through a uniform contract, exactly the way forward already is.

---

4. The generalization

The deterministic objective is

loss = reduce( scoring_rule(prediction, target, std) )

Two ingredients generalize:

1. Predictive object — deterministic: a point (+ std). General: a Gaussian (mean, std), an ensemble of S trajectories, a per-step reconstruction sampled from a prior/posterior, a diffusion denoiser. Produced by the model.

2. Objective = a proper scoring rule on the predictive object, plus model-internal penalties, then combined and reduced:

   Model	Named terms
   deterministic / Gaussian	{score: wMSE or NLL}
   ensemble (CRPS-trained)	{score: crps_ens}
   Graph-EFM (ELBO)	{nll: NLL(reconstruction), kl: KL(post‖prior)} (+ optional {crps})
   diffusion	{score: denoising-MSE}

The universal shape is

loss = reduce( Σ_k  w_k · term_k )

The core knows only this shape: collect named per-sample terms, multiply each by a configured weight, sum, average over the batch, and log each. It never sees the strings "kl"/"crps". KL/likelihood/CRPS are named only inside the concrete model.

The contract that keeps the core truly generic: each term arrives already
reduced to shape (B,) (the model alone knows whether a term sums over
grid+vars or over latent dims). The core does only
loss = Σ_k w_k · mean_B(term_k) and per-term logging. That is, literally, "the
module reduces and combines."

---

5. The uniform contract (the entire extension surface)

A new probabilistic model implements only ProbabilisticStepPredictor:

class ProbabilisticStepPredictor(StepPredictor):           # core, abstract

    # 1. eval / sampling step — fresh randomness each call.
    #    (Already what GraphEFM.forward does: sample prior, decode.)
    def forward(prev, prev_prev, forcing) -> (pred, std | None): ...

    # 2. per-step training contributions, each ALREADY reduced to (B,),
    #    plus the state to AR-condition on.
    def training_step_terms(
        prev, prev_prev, forcing, target, score_fn, interior_mask
    ) -> (terms: dict[str, Tensor[B]], next_state: Tensor): ...

    # 3. (optional) terms that need full sampled trajectories, not a step.
    def trajectory_terms(trajectories, target, mask) -> dict[str, Tensor[B]]:
        return {}                                          # default: none

    # 4. model-owned weights + how many trajectories training needs.
    loss_term_weights: dict[str, float]   # keys match the dicts above
    train_ensemble_size: int = 0          # 0 = no trajectory sampling

What makes this satisfy the Open/Closed requirement:

• The model names its own terms. ELBO → {"nll", "kl"}; CRPS-trained → {"crps"}; deterministic → {"score"}. The core never sees the names.
• The model owns its weights (kl_beta, crps_weight, …), read from its own config in its own constructor never plumbed into the core module.
• The scoring rule and mask are injected (score_fn, interior_mask) by the module, consistent with the already-settled boundary that loss_fn/mask are parameters rather than predictor-owned state. The pointwise scoring rule stays config-driven and shared (--loss); an ensemble scoring rule (CRPS) is chosen internally by a model that trains on ensembles.

---

6. The core, now fixed

# ProbabilisticARForecaster.training_rollout — generic, never edited per model
def training_rollout(init, forcing, target, score_fn, mask):
    terms = defaultdict(lambda: 0)
    prev_prev, prev = init[:, 0], init[:, 1]
    for i in range(pred_steps):
        step_terms, next_state = predictor.training_step_terms(
            prev, prev_prev, forcing[:, i], target[:, i], score_fn, mask
        )
        for k, v in step_terms.items():
            terms[k] = terms[k] + v                 # sum over the rollout
        new_state = boundary_mask * target[:, i] + interior_mask * next_state
        prev_prev, prev = prev, new_state

    if predictor.train_ensemble_size > 0:
        traj = self.sample_trajectories(init, forcing, target,
                                        predictor.train_ensemble_size)
        terms.update(predictor.trajectory_terms(traj, target, mask))
    return terms                                    # dict[str, (B,)]


# ProbabilisticForecasterModule.training_step — generic, never edited per model
def training_step(self, batch):
    terms = self.forecaster.training_rollout(*batch, self.loss, self.interior_mask_bool)
    weights = self.forecaster.predictor.loss_term_weights
    loss = sum(weights[k] * terms[k].mean() for k in terms)
    self.log_dict({f"train_{k}": terms[k].mean() for k in terms} | {"train_loss": loss})
    return loss

No ELBO, KL, CRPS, prior, posterior, or β anywhere in these classes. train_ensemble_size = 0 with one-member evaluation recovers exactly today's deterministic behavior evidence the generalization is real rather than a
special case bolted on.

---

7. Worked examples (core untouched in every case)

Graph-EFM (ELBO). In its training_step_terms: embed including the target →
encoder posterior → sample z → decode → return
{"nll": -score_fn(recon, target, std, mask).sum(grid, vars), "kl": KL(posterior‖prior).sum(latent)}, next_state = recon_mean.
trajectory_terms → {"crps": crps_ens(traj, target, mask)}.
loss_term_weights = {"nll": 1, "kl": kl_beta, "crps": crps_weight},
train_ensemble_size = 2. Every ELBO/prior/posterior detail lives here.

A pure ensemble / CRPS model. training_step_terms returns ({}, sample);
train_ensemble_size = K; trajectory_terms → {"crps": …};
loss_term_weights = {"crps": 1}.

Deterministic, re-expressed. training_step_terms → ({"score": score_fn(pred, target, std)}, pred), loss_term_weights = {"score": 1},
train_ensemble_size = 0.

In each case the only thing written is a predictor subclass + a MODELS
registration.

---

8. Evaluation generalizes the same way

Deterministic eval scores a point. Probabilistic eval scores a distribution,
so the generic eval step:

1. Samples an ensemble via forecaster.sample_trajectories(...)
   (sample_trajectories is a generic forecaster capability: loop forward).
2. Runs the configured ensemble metrics (CRPS, spread/skill, ens-MSE/MAE) —
   the module runs "the configured metrics," it does not name them.
3. Plots ensemble examples.

With ensemble_size = 1 this collapses to the current deterministic eval, so the
same code path serves both.

Note: crps_ens and spread_squared are not yet in metrics.py (it has
wmse/mse/wmae/mae/nll/crps_gauss). They are the planned ensemble-metrics PR and
are a prerequisite.

---

9. What changes vs. the current WIP branch draft PR #678

File	Change
probabilistic_module.py	Remove ELBO/kl_beta/elbo_*. Becomes generic: weighted sum of named terms + per-term logging + ensemble eval.
forecasters/probabilistic.py	training_rollout stops computing KL; it accumulates named terms the predictor returns. Keeps sample_trajectories (generic).
step_predictors/probabilistic.py	Reframed as the concrete latent-variable family (may keep prior/posterior naming), or split so the generic base only declares the contract in §5.
graph/graph_efm.py	Implement training_step_terms / trajectory_terms / loss_term_weights here — this is where ELBO moves to.

---

10. Open questions for review

1. Where does the objective live? This proposal puts the term-assembly on the predictor (training_step_terms), which is the closest analogue to the deterministic pattern (one class per model). The alternative is a tiny model-owned Objective object the predictor declares, which keeps the predictor a pure distribution producer. The latter directly addresses your earlier comment that a step predictor shouldn't compute likelihood/KL, at the cost of a second small class per model. Which do you prefer?

2. This reverses the boundary we earlier agreed (predictor = distribution producer, module = ELBO assembler). That agreement assumed the module was allowed to be ELBO-specific. The new requirement (generic core, untouched per model) forces ELBO assembly off the generic module and onto the model side. Confirming you're happy with that reversal is the crux of this review.

3. Generality ceiling. The contract is AR-structured (per-step terms + optional trajectory terms), which covers ELBO, ensembles, and CRPS. A genuinely non-AR objective (e.g. diffusion denoising over sampled noise levels) would need an escape hatch where the model overrides the training rollout entirely. Acceptable to design for AR now and leave that hatch, or do you want the more general "model owns its whole objective given tools" contract from the start?

4. Config. Model-specific hyperparameters (kl_beta, crps_weight,ensemble_size, …) would live in a probabilistic: config section and be read by the model into its own loss_term_weights, keeping them out of the core module. Agreed?

5. Ensemble batch-folding. Eval/training trajectory sampling can either loop S times or fold S into the batch dim (one B·S rollout). The fold trades memory for speed; OK as the default?

---

11. Proposed PR sequencing

1. Ensemble metrics — crps_ens, spread_squared in metrics.py. No deps.
2. Ensemble batch-fold utility — expand/fold S↔B helpers.
3. Generic probabilistic core — de-ELBO ProbabilisticForecasterModule (named-terms weighted sum + per-term logging), generic ProbabilisticARForecaster.training_rollout (term accumulation) + sample_trajectories, and the abstract ProbabilisticStepPredictor contract from §5. Depends on 1–2.
4. Graph-EFM as the first model — implement training_step_terms / trajectory_terms / loss_term_weights on the EFM predictor; register in MODELS; config-aware assembly in train_model.py. Depends on 3.
5. Ensemble eval + plotting wired into the generic module. Depends on 3–4.
6. Eval scripts. Optional / defer.

Each step keeps the core closed: only step 4 adds a model, and it touches no core
class.

[joeloskarsson]

Had a read through this, I think this is going in a good direction, and contains lots of good ideas! The code snippets under points 5 and 6 are very good, I think we can discuss and refine this a lot by just looking at (and updating) those.

Starting with a some high-level thoughts:

• While extension to new models only requiring creating a single new StepPredictor sounds nice in one way, I am not sure if this will be flexible enough. I don’t think we should be strict about this if it leads to cumbersome design.
• We don’t need to view the probabilistic extension to only land on top of the deterministic path. We can very well also change the structure of the deterministic path to make that more general or similar to the probabilistic. If we here generalize the idea of how loss is computed, then we should probably fit the deterministic models into this generalized framework as well.

Some specific things

• It seems conceptually wrong that the StepPredictor should deal with full trajectories. This seems to be the concern of the Forecaster.
  • Maybe this is a consequence of trying to make extension = implement just a new StepPredictor, and we can get rid of it if we allow for also making a new Forecast for a new modeling setup.
  • Note that the CRPS loss term can be computed per-step as well, not just for full trajectories. The challenge is more how to roll out 3 separate trajectories (1 for ELBO, 2 for CRPS) at the same time.
  • On this note, in general it would be best to compute as many loss terms as possible for full trajectories. This is computationally more efficient (there are some consideration for gradient checkpointing, but we don't need to worry about that here for now). Can be done both for ELBO and CRPS.
  • Trying to remember why we did not put most of the loss calculation in the Forecaster when refactoring to these classes 🤔 Might be something smart we thought of then that I am missing.
• What I describe above could look like a base AR probabilistic forecaster that just rolls out an ensemble (instead of a single forecast, which the corresponding deterministic AR forecaster does) and computes some loss over this. This would work for an ensemble model that only relies on a CRPS loss. Then, for something like Graph-EFM that has a more involved loss, it will have to extend this to a Forecaster with a loss calculation that does both the ELBO rollout and optionally the CRPS ones. That can then be more restrictive in which StepPredictor it works with (only Graph-EFM) and that way rely on specific methods in the StepPredictor to get all the parts needed for its loss. And something like a diffusion or flow-matching model could also change this loss calculation to denoising or score-mathcing and perhaps disallow rollout steps during training.
• The module should probably work with different types of Forecasters, and not assume that it uses (or how it uses) a StepPredictor. So e.g. weights = self.forecaster.predictor.loss_term_weights would not work. This relates to the generality vs restriction to AR, point 3 below.

Open questions for review

1. I am actually leaning towards putting more of the objective on the Forecaster, due to the points above. The problem with putting most of it on the StepPredictor is that we conceptually can only compute things one step at a time.
2. I don't really agree that this was something agreed upon, so probably some misunderstandings around this. So no need to stick to this boundary.
3. This is a key point I think. I would like to make this general enough and not AR-specific. However we can implement the AR-version as a specific setup within this more general framework. This should be the only one needed for Graph-EFM. I am thinking if we could somehow build things so that one typically (in an AR setup) only needs to extend the StepPredictor, and for some more advanced losses (the ELBO + CRPS of Graph-EFM could fall in this category) need to also extend Forecaster. Best would be if one never had to extend the ForecasterModule, and that just assumes that the Forecaster can at least create ensemble forecasts of the correct shape, and works with that minimal assumption.
4. Ensemble size is kind of inherent to the whole probabilistic setup, so should always exist. Others would sit in specific classes. But exactly how to specify config params is a later question. None of these should be in the generic base class though I agree.
5. Optimally we would have an option for this and allow for both. Ok to just choose one as a first default (sequentially is safest).

[Sir-Sloth-The-Lazy]

Thanks for the thorough read @joeloskarsson, this is really helpful, and I'm convinced by the pivot. Let me restate the new direction to confirm we're aligned, answer the "why wasn't the loss in the Forecaster" aside (I dug it up), and then condense what's left to decide.

Agreed: the objective moves onto the Forecaster
I'm dropping the RFC's proposal of putting term-assembly on the StepPredictor. You're right that a step predictor can only ever reason one step at a time, so anything trajectory-shaped (the ELBO rollout, CRPS over members) belongs on the Forecaster. Concretely I'll aim for:

ForecasterModule: generic, never extended. Its only assumption is that the forecaster can produce an ensemble forecast of the right shape and compute its own training loss. No ELBO/KL/CRPS/beta strings, and crucially no reaching through the forecaster into a predictor.
A base AR probabilistic forecaster, rolls out an ensemble and computes a CRPS loss. Enough on its own for a pure-ensemble model, with no predictor-specific calls.
A Graph-EFM forecaster extending it owns the posterior-conditioned ELBO rollout (likelihood + KL) plus optional CRPS, is allowed to be tied to the latent StepPredictor and call specific methods on it (step_distributions) to assemble its loss.
Diffusion/flow as a sibling forecaster that swaps the loss to denoising/score-matching and may disallow rollout steps during training.
And per your second high-level point, I'll treat deterministic AR as a special case of this same framework (ensemble size 1 + a pointwise score) rather than bolting probabilistic on top though I'd like to agree on the timing of that (see Q5).

Re: "why we didn't put most of the loss in the Forecaster"
I went back through the history. It was the #208 refactor (the composable-layers one). Before it, ARModel held everything; after it, Forecaster/ARForecaster got only forward (the rollout → (prediction, pred_std)) and predicts_std deliberately a pure trajectory producer while ForecasterModule kept self.loss = metrics.get_metric(loss) and applied it identically in train/val/test.

So there was no clever hidden reason. The loss stayed on the module simply because for deterministic models the objective is uniform and config-driven: the same self.loss(prediction, target, pred_std) over the forecaster's output across every split and every model. The only thing that varied per model was the rollout strategy, which is what went onto the forecaster; scoring/normalization/masking/logging were shared module concerns, so the loss sat with them. That reasoning only holds while the objective shape is shared and trivial, which is exactly what ELBO/CRPS break, so there's nothing to be careful of in moving it down to the forecaster now.

Sketch of the new contract

# generic module never edited per model
def training_step(self, batch):
    loss, logs = self.forecaster.compute_training_loss(batch, score_fn=self.loss)
    self.log_dict({"train_loss": loss, **logs})
    return loss

def eval_step(self, batch):
    ensemble = self.forecaster.predict_ensemble(...)   # (B, S, T, N, d_state)
    # run the configured ensemble metrics; S=1 recovers deterministic eval

The pointwise scoring rule stays config-driven (--loss) and is injected into the forecaster; the ensemble score (CRPS) is chosen internally by the specific forecaster. Model weights (kl_beta, crps_weight, …) live on the specific forecaster, read from its own config never on the generic module.

Open questions to converge on

• Module↔Forecaster boundary. Forecaster returns a finished (loss, log_dict), or returns named per-sample terms + its own loss_term_weights and the module does the weighted sum + logging? I lean toward the former (forecaster fully owns assembly, module stays term-agnostic), agree?

• Model → Forecaster selection. MODELS currently maps --model to a StepPredictor with Forecaster/Module fixed. Should a registry entry become a (StepPredictor, Forecaster) pair, or should a StepPredictor advertise its compatible Forecaster?

• The 3 trajectories. Interleave a per-step CRPS into the posterior ELBO rollout, or sample a separate prior ensemble (full-trajectory CRPS)? Share the rollout, or keep separate passes for clarity first?

• CRPS estimator. Confirm (almost-)fair CRPS with m = 2 training members, prior-sampled and distinct from the posterior-conditioned rollout.

• Deterministic unification now or follow-up? I'd lean toward proving the probabilistic shape first, then unifying.

• Eval contract. Is ensemble (B, S, T, N, d_state) (+ optional std), deterministic = S=1, the single shared eval output?

Revised PR sequence

• Ensemble metrics: crps_ens, spread_squared in metrics.py. No deps.

• The pivot : move likelihood + ELBO assembly out of the module into the forecaster; make ProbabilisticForecasterModule generic.

• Forecaster hierarchy, base ensemble/CRPS forecaster + Graph-EFM forecaster; resolves Q2.
  Ensemble eval + plotting in the generic module.

• Deterministic unification, gated on Q5.
  I'd nail down Q1, Q2 and Q5 before writing PR 2/3 (above). Does this match what you had in mind?

[joeloskarsson]

I think this sounds pretty good, and something we can move forward with now, looking at the concrete code solutions.

What I am still a bit unsure about (i.e. I don't really know what's the best solution) is the integration with the deterministic branch. Treating deterministic as a special case of 1 ensemble member would make the codebase ensemble-first. While this might minimize duplicate code, I can really see a point in differentiating between a 1 member ensemble and a deterministic prediction. These are conceptually different (one represents a single sample probabilistically, while the deterministic prediction represents a conditional mean). So I am leaning towards building the probabilistic framework largely on top of the deterministic one, trying to never mix in any sample S tensor dimension in the deterministic components. But we should make sure that the same concepts sit at the same place for both approaches, in particular if we now move loss to the Forecaster we should do so also for the deterministic setup (which has the benefit of improved extendability also there). This does probably mean that we need separate ForecasterModules and Forecasters for the probabilistic case, but I don't see that as a bad thing.

For the specific questions:

1. I also think that the whole responsibility for the loss should sit in one place (now, Forecaster), including assembling it. It is however very useful to be able to log separate loss components, so we still need a way to provide these separately for logging purposes. Maybe sufficient to just return both the complete loss and a dict to log with loss-related quantities.
2. Yea this is the big headache atm, how we actually assemble models in train_model.py. Relates to Fix graph LAM GNN kwarg handling #688 and [RFC]: Configuration Architecture - Requirements for Explicit APIs and Checkpoint Stability #672. As a first solution here I think we can just have some specialized code for instantiating each --model choice, and then we can see if this can be improved by tackling this configuration design problem.
3. (Note that this does not relate to the probabilistic interface defined here now. With the discussed setup this is a pure graph-efm problem). Still: No strong opinions. If you find a way to do the three rollouts at the same time that would be optimal. But it might get very messy, e.g. if you assemble the three batches of states into one batch, but then some parts of that batch should have a latent sampled from prior and some should have it from the VI-approximation. Doing first the ELBO rollout, then (optionally) CRPS is also ok. I think that's how the original code worked.
4. (Similar to point above, not actually related to this) The CRPS-loss should always be computed with trajectories created by sampling from the prior. The CRPS estimator and number of members would optimally be configurable. But defaults should be unbiased/fair CRPS and 2 members (matches original paper).
5. Relates to my thoughts above about collapsing det/prob parts. Not sure if there is a need for a "shared eval contract". One should always be careful with the difference in evaluating an ensemble (even with 1 member) and a deterministic prediction, so there is a point in keeping this separate. Metrics in metrics.py accepts arbitrary batch dimensions, so should work fine, and it is responsibility of the caller to handle batch/sample dimension. Also for plotting we should treat plotting of an ensemble different to plotting a det prediction.

Regarding PRs
I think the best next step now would be to move this discussion into a draft PR, to look at this more in detail (can either change #678 or open a new one). Best would probably be to just look at the abstract constructions, and wait before integrating the concrete Graph-EFM instantiation of this. No need to depend on the probabilistic metrics for this. If you need an ensemble-specific metric as an example for how the metric calculation should sit in all this you can use RMSE of ensemble mean.