cs.LGNov 2, 2026

Gibbs Variational Last Layers for Long-Horizon Dynamics Learning

Authors: Paul Brunzema, Henrik Hose, Sebastian Trimpe

Organizations: RWTH Aachen University, Aachen, Germany

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Accurate long-horizon probabilistic predictions of dynamical systems are a prerequisite for model-based decision-making. Variational Bayesian last layers (VBLLs) are an attractive model class for this setting, offering tractable uncertainty at near-deterministic cost, yet they are trained exclusively with single-step likelihood objectives that may be insufficient for reliable multi-step rollouts. We revisit VBLLs through the lens of Gibbs variational inference, which decouples posterior inference from the choice of training loss. Building on this perspective, we introduce a family of multi-step training losses that progressively incorporate long-horizon structure, robustness via continuous ranked probability scores (CRPS), and training on full autoregressive rollouts. Our approach retains the simplicity of VBLLs while directly optimizing for multi-step predictive quality. Across illustrative synthetic environments, a diverse suite of chaotic dynamical systems, and real-world data, we show that multi-step, CRPS-based objectives substantially improve long-horizon accuracy and distributional fit, consistently outperforming single-step likelihood training.

Figures & tables

Figure 1: Overview of the developed loss functions for VBLLs under the Gibbs VI framework. We extend existing single-step ( s-NLL ) to a multi-step ( m-NLL ), robust ( m-CRPS ), autoregressive loss ( m-arCRPS ).
Figure 2: Results on the VSLC. Top: Rollouts of all baselines against rollouts (color) from the true dynamics (black). Bottom row: Performance across various metrics computed from rollouts from multiple initial conditions.
Figure 3: Learned noise for the VSLC. GT noise is σ1,σ2=0.01\sigma_{1},\sigma_{2}=0.01 .
Figure 4: Rank summary ( ↓\downarrow ) across benchmark systems. Full results in Tab. 3 in Appendix I .
Figure 5: Qualitative results on data from the Mini Wheelbot.
 One-step metricsMulti-step metricsTraining 1 1 1 All baselines were trained on the same HPC cluster, each on a single NVIDIA H100 GPU.
MethodMSE ↓\downarrowNLL ↓\downarrowCRPS ↓\downarrowFinal L2 ↓\downarrowCumul. L2 ↓\downarrowDTW ↓\downarrowEnergy dist. ↓\downarrowTime [min]
s-NLL [ 31 , 7 ]$$0.005^{0.019}_{0.002}$$-8.273^{-6.084}_{-8.932}$$0.328^{0.509}_{0.273}$$0.371^{0.759}_{0.223}$$12.690^{20.954}_{8.285}$$2.194^{3.919}_{1.161}$$1.909^{2.314}_{1.561}1.3661.366
m-NLL$$0.005^{0.019}_{0.002}$$-8.304^{-6.071}_{-8.982}$$0.340^{0.515}_{0.286}$$0.403^{0.722}_{0.217}$$13.350^{22.801}_{7.982}$$2.228^{3.848}_{1.089}$$5.181^{8.139}_{2.879}1.8611.861
m-CRPS0.0050.0010.017\mathbf{0.005}^{\mathbf{0.017}}_{\mathbf{0.001}}−10.676−12.260−5.489\mathbf{-10.676}^{\mathbf{-5.489}}_{\mathbf{-12.260}}0.2480.1690.479\mathbf{0.248}^{\mathbf{0.479}}_{\mathbf{0.169}}$$0.200^{0.259}_{0.131}$$7.792^{11.838}_{5.540}$$1.117^{1.582}_{0.767}$$1.456^{1.931}_{1.148}2.0482.048
m-KDE [ 34 ]$$0.009^{0.034}_{0.003}$$-8.249^{1.104}_{-11.956}$$0.352^{0.707}_{0.197}$$0.026^{0.044}_{0.017}$$1.076^{1.776}_{0.678}$$0.087^{0.123}_{0.043}$$0.236^{0.285}_{0.190}3.7253.725
m-arNLL$$0.010^{0.039}_{0.003}$$-8.011^{5.626}_{-11.677}$$0.363^{0.787}_{0.197}$$0.031^{0.045}_{0.020}$$1.111^{1.728}_{0.777}$$0.089^{0.141}_{0.053}$$0.223^{0.285}_{0.190}3.9333.933
Table 1: Quantitative results on Wheelbot data. Multi-step metrics are reported at 100 steps. Each cell reports medianq25q75\mathrm{median}^{\mathrm{q75}}_{\mathrm{q25}} . Bold, shaded values are significantly best (paired Wilcoxon signed-rank test, α=0.05\alpha=0.05 ).
Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Model rollouts on Swimmer (flow field in gray). We test whether the different losses allow the model to capture complex nonlinear dynamics resulting in multi-modal rollouts. We can see that all losses allow the model demonstrate such complex behavior. Still, multi-step losses allow for a better approximation of the distribution.
Figure 7: The Mini Wheelbot [ 32 ] . Banana for scale.
Figure 8: Viz. of m-KDE .
MethodMeanMed.Std1st2nd3rd4th5th6th
m-CRPS2.1221.2685493516112
m-NLL2.8321.4846562333391
m-arCRPS3.2431.51285027473214
s-NLL3.7841.3291766502927
m-KDE3.8341.43202031407710
m-arNLL5.2161.47106161210144
Appendix
Table 2: Statistical evaluation across 198 (environment, metric) blocks. Left: Rank distribution (lower mean rank is better). Right: Post-hoc pairwise comparisons (Nemenyi test, critical difference Δr=0.536\Delta r=0.536 , and Holm-corrected Wilcoxon signed-rank test); pairs below the line are not significant at α=0.05\alpha=0.05 under the Nemenyi test. Friedman test: χ2=314.1\chi^{2}=314.1 , p<10−65p<10^{-65} . Diverged runs (143 blocks of m-arNLL ) are ranked last.
Figure 9: Feature sensitivity (top), mean prediction sensitivity (middle), and Fisher sensitivity (bottom) for the learned model on the VSLC (Fig. 2 ). All models are evaluated on the same data yet there are distinct differences in behavior, most prominently in feature and Fisher sensitivity.
Figure 10: Ablation of model size (top) and the sample size SS (bottom) on Mini Wheelbot data.
Figure 11: Sensitivity to different horizon HH . A longer horizon during training increases performance. Training time for small HH increases as the effective data set size increase.
Figure 12: Influence of changing the KL weight λ\lambda for the CRPS baselines.
Figure 13: Backbone gradient norms during training. 95th percentile of the pre-clipping backbone gradient norm per epoch for all objectives across all benchmark environments; the dashed line marks the clipping threshold. m-arNLL develops non-finite gradients on all chaotic systems and is therefore only shown for the non-chaotic environments. CRPS attenuates but does not prevent the gradient growth induced by autoregressive training.
Figure 14: Relative update of the backbone parameters per epoch. m-arNLL develops non-finite gradients on all chaotic systems and is therefore only shown for the non-chaotic environments.
Figure 15: Relative update of the posterior mean of the last-layer weights per epoch. The posterior mean keeps updating on a similar relative scale as the backbone ( cf . Fig. 14 ).
Figure 16: Relative update of the posterior covariance of the last-layer weights per epoch. As for the posterior mean, the covariance does not stall while the backbone continues to update ( cf . Fig. 14 ).
Figure 17: Rollouts on Mini Wheelbot data. Figure shows the complete set of states to complement the rollouts shown in Fig. 5 in the main body of the paper.
EnvironmentMetrics-NLLm-NLLm-CRPSm-arNLLm-arCRPSm-KDE
SpiralFinal L2 ↓\downarrow$$0.325^{+0.236}_{-0.159}$$0.056^{+0.061}_{-0.031}0.033−0.016+0.016\mathbf{0.033}^{+\mathbf{0.016}}_{-\mathbf{0.016}}$$0.228^{+0.041}_{-0.079}$$0.122^{+0.020}_{-0.031}$$0.063^{+0.081}_{-0.047}
Cumul. L2 ↓\downarrow$$28.269^{+23.817}_{-19.836}$$5.732^{+6.806}_{-2.495}2.605−0.480+1.451\mathbf{2.605}^{+\mathbf{1.451}}_{-\mathbf{0.480}}$$13.834^{+3.687}_{-6.170}$$7.788^{+1.644}_{-2.293}$$6.926^{+11.042}_{-4.906}
DTW ↓\downarrow$$0.982^{+0.479}_{-0.273}$$0.580^{+0.345}_{-0.183}0.351−0.067+0.191\mathbf{0.351}^{+\mathbf{0.191}}_{-\mathbf{0.067}}$$0.732^{+0.166}_{-0.221}$$0.631^{+0.111}_{-0.076}$$0.546^{+0.387}_{-0.263}
MMD 2 ↓\downarrow$$0.020^{+0.031}_{-0.018}$$0.001^{+0.003}_{-0.001}0.000−0.000+0.000\mathbf{0.000}^{+\mathbf{0.000}}_{-\mathbf{0.000}}$$0.005^{+0.003}_{-0.003}$$0.002^{+0.001}_{-0.001}$$0.003^{+0.010}_{-0.003}
Energy Dist. ↓\downarrow$$2.867^{+3.429}_{-1.998}$$0.339^{+0.616}_{-0.199}0.082−0.042+0.089\mathbf{0.082}^{+\mathbf{0.089}}_{-\mathbf{0.042}}$$1.123^{+0.478}_{-0.488}$$0.440^{+0.240}_{-0.220}((4.519^{+4.034}_{-4.519})\times 10^{8}
Term. MMD 2 ↓\downarrow$$0.041^{+0.063}_{-0.032}$$0.001^{+0.003}_{-0.001}0.000−0.000+0.001\mathbf{0.000}^{+\mathbf{0.001}}_{-\mathbf{0.000}}$$0.018^{+0.008}_{-0.010}$$0.005^{+0.003}_{-0.003}$$0.003^{+0.009}_{-0.003}
Appendix
Table 3: Results on the benchmarking of Sec. 4.2 . (Part 1/4)
EnvironmentMetrics-NLLm-NLLm-CRPSm-arNLLm-arCRPSm-KDE
LorenzFinal L2 ↓\downarrow$$11.642^{+20.187}_{-7.280}$$11.676^{+20.290}_{-7.289}$$12.378^{+21.322}_{-8.270}div.10.755−6.936+16.145\mathbf{10.755}^{+\mathbf{16.145}}_{-\mathbf{6.936}}$$11.388^{+15.737}_{-7.501}
Cumul. L2 ↓\downarrow$$866.827^{+1,085.481}_{-577.921}$$865.478^{+1,100.729}_{-592.923}$$888.891^{+1,070.867}_{-607.088}div.805.358−605.800+1090.680\mathbf{805.358}^{+\mathbf{1090.680}}_{-\mathbf{605.800}}$$939.147^{+993.302}_{-664.319}
DTW ↓\downarrow$$101.314^{+176.614}_{-74.273}$$101.067^{+154.610}_{-75.645}$$106.475^{+163.943}_{-78.152}div.91.795−70.442+158.475\mathbf{91.795}^{+\mathbf{158.475}}_{-\mathbf{70.442}}$$96.209^{+183.364}_{-76.473}
MMD 2 ↓\downarrow$$0.018^{+0.034}_{-0.017}$$0.016^{+0.049}_{-0.016}0.002−0.002+0.001\mathbf{0.002}^{+\mathbf{0.001}}_{-\mathbf{0.002}}div.$$0.032^{+0.128}_{-0.031}$$0.060^{+0.163}_{-0.052}
Energy Dist. ↓\downarrow$$32.190^{+41.168}_{-21.357}$$27.615^{+43.562}_{-19.143}6.585−4.436+10.295\mathbf{6.585}^{+\mathbf{10.295}}_{-\mathbf{4.436}}div.$$40.672^{+91.551}_{-28.389}$$66.410^{+87.538}_{-36.004}
Term. MMD 2 ↓\downarrow$$0.052^{+0.181}_{-0.050}$$0.047^{+0.214}_{-0.046}0.007−0.007+0.036\mathbf{0.007}^{+\mathbf{0.036}}_{-\mathbf{0.007}}div.$$0.089^{+0.164}_{-0.087}$$0.138^{+0.335}_{-0.118}
Appendix
Table 4: Results on the benchmarking of Sec. 4.2 . (Part 2/4)
EnvironmentMetrics-NLLm-NLLm-CRPSm-arNLLm-arCRPSm-KDE
DadrasFinal L2 ↓\downarrow0.625−0.504+1.436\mathbf{0.625}^{+\mathbf{1.436}}_{-\mathbf{0.504}}$$0.971^{+1.933}_{-0.808}$$0.747^{+1.564}_{-0.628}div.$$1.658^{+2.602}_{-1.316}$$1.623^{+2.702}_{-1.391}
Cumul. L2 ↓\downarrow$$64.177^{+142.690}_{-48.394}$$86.223^{+163.618}_{-72.316}63.810−53.438+143.738\mathbf{63.810}^{+\mathbf{143.738}}_{-\mathbf{53.438}}div.$$165.292^{+242.559}_{-123.316}$$155.975^{+262.996}_{-132.129}
DTW ↓\downarrow5.327−4.150+13.106\mathbf{5.327}^{+\mathbf{13.106}}_{-\mathbf{4.150}}$$7.868^{+16.345}_{-6.798}$$6.245^{+15.722}_{-5.639}div.$$14.775^{+25.840}_{-12.719}$$14.076^{+27.889}_{-12.537}
MMD 2 ↓\downarrow$$0.003^{+0.006}_{-0.003}$$0.004^{+0.006}_{-0.004}0.001−0.001+0.001\mathbf{0.001}^{+\mathbf{0.001}}_{-\mathbf{0.001}}div.$$0.034^{+0.013}_{-0.032}$$0.021^{+0.128}_{-0.021}
Energy Dist. ↓\downarrow$$5.273^{+8.197}_{-3.674}$$4.221^{+7.851}_{-3.194}1.905−1.287+2.392\mathbf{1.905}^{+\mathbf{2.392}}_{-\mathbf{1.287}}div.$$19.186^{+14.491}_{-15.053}$$14.222^{+40.953}_{-10.325}
Term. MMD 2 ↓\downarrow$$0.012^{+0.030}_{-0.011}$$0.010^{+0.022}_{-0.009}0.003−0.003+0.007\mathbf{0.003}^{+\mathbf{0.007}}_{-\mathbf{0.003}}div.$$0.080^{+0.154}_{-0.077}$$0.050^{+0.181}_{-0.048}
Appendix
Table 5: Results on the benchmarking of Sec. 4.2 . (Part 3/4)
EnvironmentMetrics-NLLm-NLLm-CRPSm-arNLLm-arCRPSm-KDE
Nose HooverFinal L2 ↓\downarrow$$0.402^{+0.207}_{-0.321}$$0.192^{+0.211}_{-0.158}0.191−0.144+0.148\mathbf{0.191}^{+\mathbf{0.148}}_{-\mathbf{0.144}}div.$$0.193^{+0.149}_{-0.148}$$0.318^{+0.245}_{-0.284}
Cumul. L2 ↓\downarrow$$21.467^{+17.987}_{-16.129}$$14.180^{+17.436}_{-11.045}13.527−9.536+13.745\mathbf{13.527}^{+\mathbf{13.745}}_{-\mathbf{9.536}}div.$$14.617^{+18.193}_{-8.045}$$21.230^{+38.851}_{-15.677}
DTW ↓\downarrow$$2.316^{+1.011}_{-1.992}1.063−0.877+1.892\mathbf{1.063}^{+\mathbf{1.892}}_{-\mathbf{0.877}}$$1.172^{+0.818}_{-0.907}div.$$1.205^{+1.479}_{-0.898}$$1.903^{+2.600}_{-1.626}
MMD 2 ↓\downarrow$$0.005^{+0.004}_{-0.005}$$0.002^{+0.002}_{-0.002}0.001−0.001+0.002\mathbf{0.001}^{+\mathbf{0.002}}_{-\mathbf{0.001}}div.$$0.004^{+0.006}_{-0.004}$$0.007^{+0.016}_{-0.007}
Energy Dist. ↓\downarrow$$1.090^{+1.581}_{-0.890}0.465−0.383+0.555\mathbf{0.465}^{+\mathbf{0.555}}_{-\mathbf{0.383}}$$0.488^{+0.702}_{-0.390}div.$$1.148^{+1.264}_{-0.883}$$1.484^{+3.402}_{-1.184}
Term. MMD 2 ↓\downarrow$$0.017^{+0.025}_{-0.016}0.004−0.004+0.002\mathbf{0.004}^{+\mathbf{0.002}}_{-\mathbf{0.004}}$$0.004^{+0.008}_{-0.004}div.$$0.011^{+0.013}_{-0.011}$$0.017^{+0.044}_{-0.017}
Appendix
Table 6: Results on the benchmarking of Sec. 4.2 . (Part 4/4)
HyperparameterValue
Hidden width64 (128 for Glycolytic oscillator)
Hidden layers3
Feature dimension32 (64 for Glycolytic oscillator)
Activationmish
Prior scale11
Appendix
Table 7: Training hyperparameters for the cluster sweep: model configuration.
LossEpochsLRBatchWeight decayGrad. clipλ\lambdaSS
s-NLL3000.0010.00164 ×\times H1.0\times10−041.0\text{\times}{10}^{-04}111/∣Dsingle∣1/|{\mathcal{D}}_{\text{single}}|–
m-NLL3000.0010.001641.0\times10−041.0\text{\times}{10}^{-04}111/∣Dtraj∣1/|{\mathcal{D}}_{\text{traj}}|–
m-CRPS3000.0010.001641.0\times10−041.0\text{\times}{10}^{-04}111/∣Dtraj∣1/|{\mathcal{D}}_{\text{traj}}|10
m-KDE3000.0010.001641.0\times10−041.0\text{\times}{10}^{-04}111/∣Dtraj∣1/|{\mathcal{D}}_{\text{traj}}|10
m-arNLL3000.0010.001641.0\times10−041.0\text{\times}{10}^{-04}111/∣Dtraj∣1/|{\mathcal{D}}_{\text{traj}}|10
m-arCRPS3000.0010.001641.0\times10−041.0\text{\times}{10}^{-04}111/∣Dtraj∣1/|{\mathcal{D}}_{\text{traj}}|10
Appendix
Table 8: Training hyperparameters for the cluster sweep: loss-specific training settings.
NameDim.Norm.Δt\Delta tσ\sigmaTraj.Len.Stride ssHHEval stepsStatesSamples SS for eval
Spiral2no0.10.10.010.013009014012080100
Swimmer2no0.10.10.010.013009014012080100
Variable-speed limit cycle2no0.10.10.010.013009014012080100
Bifurcating limit cycle2no0.050.050.030.033009014012080100
Glycolytic oscillator7yes0.010.010.010.013009014012080100
Lorenz3yes0.010.010.010.0140090204012080100
Appendix
Table 9: Training hyperparameters for the cluster sweep: environment, data, and post-evaluation settings.
HyperparameterValue
Hidden width128
Hidden layers3
Feature dimension64
Activationmish
Prior scale11
Appendix
Table 10: Training hyperparameters for the cluster sweep: Wheelbot model configuration.
LossEpochsLRBatchWeight decayGrad. clipHHStride ssSSSvalS_{\mathrm{val}}
s-NLL4005.0\times10−045.0\text{\times}{10}^{-04}2561.0\times10−041.0\text{\times}{10}^{-04}11––––
m-NLL4005.0\times10−045.0\text{\times}{10}^{-04}2561.0\times10−041.0\text{\times}{10}^{-04}11402010100
m-CRPS4005.0\times10−045.0\text{\times}{10}^{-04}2561.0\times10−041.0\text{\times}{10}^{-04}11402010100
m-KDE4005.0\times10−045.0\text{\times}{10}^{-04}2561.0\times10−041.0\text{\times}{10}^{-04}11402010100
m-arNLL4005.0\times10−045.0\text{\times}{10}^{-04}2561.0\times10−041.0\text{\times}{10}^{-04}11402010100
m-arCRPS4005.0\times10−045.0\text{\times}{10}^{-04}2561.0\times10−041.0\text{\times}{10}^{-04}11402010100
Appendix
Table 11: Training hyperparameters for the cluster sweep: Wheelbot loss-specific training settings.
More in this topic

Probabilistic Time Series Forecasting

Probabilistic time series forecasting estimates distributions of future values or trajectories to quantify uncertainty and support risk-aware decisions.

See all 174 papers

Explore similar work

CardsList

  1. Forecasting Multiple Observables with SCROLL: Score-Trained Uncertainty for Stochastic Dynamics

    Aug 26, 2026Pavel ProchazkaProbabilistic Time Series ForecastingUncertainty Quantification

    Cisco Inc.

  2. Stable by Construction: Variational Latent Markov Operators for Long-Horizon PDE Prediction

    Sep 15, 2026Junyi Liao, Johann Guilleminot, Vahid TarokhVariational AutoencodersLatent Dynamics Modeling

    †Department of Electrical and Computer Engineering, Duke University · ‡Department of Mechanical Engineering and Materials Science, Duke University

  3. Variational Augmented Invertible Koopman Autoencoder for probabilistic time series forecasting

    Sep 29, 2026Anthony Frion, Lucas Drumetz, Guillaume Tochon +3Probabilistic Time Series ForecastingLong-Term Time Series Forecasting

    Institute of coastal modelling Helmholtz-Zentrum Hereon · IMT Atlantique Lab-STICC, UMR CNRS 6285, Brest, France · LRE EPITA, Le Kremlin-Bicêtre, France +1