Title: [BUG] Probabilistic predictions suffer from variance underflow and ignore pred_std during ensemble generation
Description:
Description
While reviewing the probabilistic forecasting (--output_std) logic and verification pathways, I identified two critical bugs where the network's uncertainty predictions produce NaN mathematical collapses or are silently discarded during evaluation.
1. softplus Variance Underflow Leading to NaN Losses
In neural_lam/models/base_graph_model.py, the standard deviation is computed natively as:
pred_std = torch.nn.functional.softplus(pred_std_raw)
PyTorch's softplus evaluates to exactly 0.0 for sufficiently negative inputs (e.g., softplus(-100) == 0.0). When pred_std underflows to zero, it causes division-by-zero operations when calculating standard scores for nll or crps_gauss targets, producing cascading NaNs that irreversibly crash training.
2. Degenerate Ensembles Discarding True Uncertainty
In neural_lam/models/ar_model.py, the _sample_ensemble method hardcodes the generation of Gaussian noise to a static 0.01 standard deviation:
noise = torch.randn(...) * 0.01
Because this method ignores the model's dynamically predicted uncertainty map (pred_std), running an ensemble inference on a properly trained probabilistic model throws away the highly detailed uncertainty predictions and relies entirely on an arbitrary 0.01 fallback jitter.
Reproduction
- Manually inject a highly negative tensor (e.g.
[-100.0]) into pred_std_raw during a forward pass and evaluate mse/nll, resulting in a NaN crash.
- Run inference with
output_mode="ensemble" and --output_std, and observe that the variance of the resulting ensemble instances remains ~0.01 regardless of the model's actual confidence regarding the weather state.
Expected behavior
pred_std should be safely clamped to a small epsilon (e.g. 1e-6) to guarantee numerical stability during metrics computation._sample_ensemble should physically draw its noise distribution from the dynamically calculated pred_std matrix if it is available.
Title: [BUG] Probabilistic predictions suffer from variance underflow and ignore
pred_stdduring ensemble generationDescription:
Description
While reviewing the probabilistic forecasting (
--output_std) logic and verification pathways, I identified two critical bugs where the network's uncertainty predictions produceNaNmathematical collapses or are silently discarded during evaluation.1.
softplusVariance Underflow Leading toNaNLossesIn
neural_lam/models/base_graph_model.py, the standard deviation is computed natively as:PyTorch's
softplusevaluates to exactly0.0for sufficiently negative inputs (e.g.,softplus(-100) == 0.0). Whenpred_stdunderflows to zero, it causes division-by-zero operations when calculating standard scores fornllorcrps_gausstargets, producing cascadingNaNs that irreversibly crash training.2. Degenerate Ensembles Discarding True Uncertainty
In
neural_lam/models/ar_model.py, the_sample_ensemblemethod hardcodes the generation of Gaussian noise to a static0.01standard deviation:Because this method ignores the model's dynamically predicted uncertainty map (
pred_std), running an ensemble inference on a properly trained probabilistic model throws away the highly detailed uncertainty predictions and relies entirely on an arbitrary0.01fallback jitter.Reproduction
[-100.0]) intopred_std_rawduring a forward pass and evaluatemse/nll, resulting in aNaNcrash.output_mode="ensemble"and--output_std, and observe that the variance of the resulting ensemble instances remains~0.01regardless of the model's actual confidence regarding the weather state.Expected behavior
pred_stdshould be safely clamped to a small epsilon (e.g.1e-6) to guarantee numerical stability during metrics computation._sample_ensembleshould physically draw its noise distribution from the dynamically calculatedpred_stdmatrix if it is available.