Gibbs variational last layers improve long horizon dynamics prediction
Gibbs Variational Last Layers for Long-Horizon Dynamics Learning
Machine LearningRobotics
Summary
Predicting how systems behave over a long time is important but hard. The paper looks at a method called variational Bayesian last layers (VBLLs) that is fast and can show uncertainty, but usually trains on one-step predictions. The authors rethink this using Gibbs variational inference, allowing training focused on multi-step, longer-term predictions. They test this on various systems and find that training with longer steps and special scoring leads to better accuracy and more reliable long-term forecasts.
What this means in practice
- •For robotics engineers: Improve long-term robot motion and control predictions by using multi-step trained VBLL models for better uncertainty estimation.
- •For energy grid operators: Enhance forecasting of complex electrical grid dynamics over extended periods with improved probabilistic models using multi-step training.
Authors
Paul Brunzema, Henrik Hose, Sebastian Trimpe
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.