Gibbs Variational Last Layers for Long-Horizon Dynamics Learning

AI summary

Overview

Research area: Machine learning for dynamical systems — probabilistic (Bayesian) dynamics modeling and long-horizon / multi-step prediction, with a focus on variational Bayesian last layers (VBLLs) and variational inference for decision-making under uncertainty.

Technical level: Advanced. The paper relies on variational inference, Gibbs posteriors, proper scoring rules (log score, CRPS), and backpropagation through time. The core idea is explainable to an intermediate reader, but the derivations assume familiarity with Bayesian last-layer models.

Scope: The paper reinterprets variational Bayesian last layers through Gibbs variational inference so that the training loss becomes a free design choice, then derives and empirically compares a hierarchy of single-step, multi-step, CRPS-based, and fully autoregressive training objectives for learned dynamics models rolled out over long horizons.

What This Paper Is About

Learned probabilistic dynamics models are almost always trained by scoring one-step-ahead predictions against ground-truth transitions, even though at deployment they are rolled out autoregressively over many steps, where small errors compound. The authors show that the standard training objective of variational Bayesian last layers (VBLLs) is only one arbitrary choice within a much larger family, because the Gaussian posterior over the last layer is mathematically independent of the loss used to fit it. Their goal is to design better training losses — multi-step, robust to outliers, and exposed to the model's own rollout errors — while keeping the VBLL architecture and its cheap, tractable predictive inference completely unchanged.

Key Contributions

1. Gibbs variational last layers (C1). The authors formalize a framework in which VBLLs are trained by minimizing an arbitrary expected loss plus a KL penalty toward the Gaussian prior, rather than a Gaussian log-likelihood. They show this Gibbs VI objective strictly generalizes the standard VBLL/ELBO objective, which is recovered exactly when the loss is the negative log-likelihood with horizon H = 1 and λ = 1/N.

2. A hierarchy of multi-step objectives (C2). They derive a ladder of losses that progressively adds long-horizon structure, robustness via the continuous ranked probability score (CRPS), and exposure to autoregressive errors: single-step NLL (s-NLL), teacher-forced multi-step NLL (m-NLL), teacher-forced multi-step CRPS (m-CRPS), and autoregressive NLL/CRPS (m-arNLL, m-arCRPS).

3. An equivalence result and its limits (Corollary 1). For a fixed feature extractor, m-NLL equals H times s-NLL when λ_M = H·λ_S, so the two share minimizers; differences arise only because the feature extractor is optimized jointly, which reshapes the learned representation.

4. Broad empirical validation (C3). Evaluation spans 18 dynamical systems (including a suite of chaotic systems), illustrative synthetic environments, and data from a real-world unstable robotic system, plus an isolated demonstration of a sampling-bias failure mode of single-step likelihood training.

Main Findings

• Multi-step CRPS objectives outperform single-step likelihood training. The abstract states that multi-step, CRPS-based objectives "substantially improve long-horizon accuracy and distributional fit, consistently outperforming single-step likelihood training" across synthetic environments, chaotic systems, and real-world data. Specific per-system numeric scores are in appendices not included in the excerpt reviewed.

• The CRPS gives bounded gradients; the log score does not. The derivative of the CRPS with respect to the predictive mean is 1 − 2Φ(z), which lies in (−1, 1), and with respect to σ lies in [−1/√π, (√2 − 1)/√π]. The NLL gradients −z/σ and −z²/σ + 1/σ diverge as |z| → ∞. Empirically, CRPS attenuates gradient growth enough to keep autoregressive training finite across all tested systems and horizons, while log-score training produces gradient norms orders of magnitude larger.

• Autoregressive NLL is unstable; autoregressive CRPS is stable. Under backpropagation through time, the unbounded log-score gradients compound, and m-arNLL destabilizes on all chaotic systems tested, whereas m-arCRPS trains stably.

• Teacher-forced multi-step training alone yields limited gains. Because m-NLL is just a sum of H one-step losses (Corollary 1), the authors explicitly note the gains from moving from single-step to teacher-forced multi-step log-likelihood are "inherently limited"; the scoring rule is the second, orthogonal lever.

• A sampling-bias failure mode in the variable-speed limit cycle (VSLC). In this 2D system with ω(θ) = ω₀(1 + ε sin(n_modes θ)), ω₀ = 2.0, ε = 0.85, n_modes = 2, data sampled at fixed Δt concentrate in slow zones (ω(θ) < ω₀), which contain roughly 82% of samples (F_slow ≈ 0.823). Single-step training therefore focuses its gradient signal there, and the learned noise departs from the ground-truth σ₁, σ₂ = 0.01; multi-step objectives mitigate this.

• A tractable Gaussian surrogate justifies autoregressive training. The true multi-step predictive is non-Gaussian for nonlinear features, so the authors use P̃ = N(x̂_h(W), h Σ_ν), i.e. the deterministic rollout mean with covariance growing linearly in h. Proposition 1 shows this is exact in the linear case when A Σ_ν Aᵀ = Σ_ν, and that the covariance error is O(h²‖E‖‖Σ_ν‖) for A = I + E.

• Cost stays comparable in the forward pass. With S posterior samples and horizon H the forward cost is O(SH), as in teacher forcing, but autoregressive losses are sequential and require backpropagation through time, making them less parallelizable and more memory-intensive.

• One-step predictives remain Gaussian for every loss. The one-step predictive distribution stays Gaussian under all objectives considered, so inference tractability is preserved regardless of the training loss chosen.

Methodology in Plain English

The model is a residual dynamics model: the next state is the current state plus a learned feature extractor passed through a last linear layer, plus Gaussian process noise. Only the last layer is treated as uncertain — it gets a Gaussian prior and a mean-field Gaussian variational posterior factorized across output columns — so predictions are cheap and closed-form, close to the cost of an ordinary neural network.

The key move is conceptual. Standard training for this model maximizes a single-step Gaussian log-likelihood. The authors instead invoke Gibbs variational inference, which swaps the log-likelihood for any loss while keeping the KL penalty to the prior. That turns the loss into a knob. They then turn the knob along two axes:

• Scoring rule: log score (likelihood, quadratic tails, unbounded gradients) versus CRPS (linear tails, bounded gradients, closed form for Gaussians).
• Evaluation mode: teacher forcing (always condition on ground-truth states) versus autoregressive (condition on the model's own predictions).

Turning both knobs produces the four new losses. For the autoregressive ones they avoid the intractable exact multi-step distribution by scoring against a Gaussian surrogate whose mean follows the deterministic rollout and whose variance grows linearly with the horizon, then still deploying full stochastic rollouts at test time.

Experiments hold the architecture and initialization fixed and vary only the loss, to isolate the effect of the objective. Baselines are the standard VBLL ELBO (s-NLL) and an autoregressive kernel-density baseline (m-KDE, from Ivanovic et al.). Rollouts are generated by sampling last-layer weights from the posterior and simulating forward with learned noise; λ is set to 1/|D|, and s-NLL uses batch size B × H so every method sees the same number of transitions per gradient step. Evaluation covers a controlled synthetic system, 18 dynamical systems (including chaotic ones), and a real-world unstable robotic platform.

Why This Matters

Impact on research. The paper reframes a widely used model class: it shows that the objective, not the architecture, is what limits VBLLs in autoregressive deployment, and that the likelihood-based Bayesian interpretation can be traded for a PAC-Bayesian one in exchange for task-aligned training. This connects Bayesian last-layer models to the Gibbs-posterior and scoring-rule literatures, and gives a principled route to using non-likelihood losses (CRPS) that classical Bayesian inference cannot accommodate. It also documents a concrete failure mode — sampling bias induced by non-uniform data coverage — that single-step training cannot correct.

Real-world applications:

• Model-based control and sampling-based MPC, where the model is rolled out many steps inside the optimizer.
• Uncertainty-aware policy search and reinforcement learning, where miscalibrated long-horizon uncertainty corrupts the policy update.
• Robotics on unstable platforms (the paper's own real-world demonstration), where compounding prediction error is the dominant failure mode.
• Any safety-critical forecasting setting that needs calibrated distributional predictions under occasional large deviations, where the CRPS's linear tails are preferable to quadratic log-score tails.

Industry relevance. VBLLs give near-deterministic inference cost with uncertainty quantification, which makes them attractive for embedded and real-time control. The paper's message — that a different loss, not a bigger model, buys better long-horizon behavior — is directly actionable for practitioners who have already deployed a last-layer Bayesian model and want to improve rollouts without changing their inference stack or latency budget.

Future Directions

• Better multi-step distribution surrogates. The linearly growing Gaussian covariance is exact only in the linear/near-identity case; relaxing this to handle strongly amplifying Jacobians, correlated state dimensions, or skewed/multimodal rollout distributions is left open.
• Beyond mean-field, column-factorized posteriors. The paper keeps the original VBLL posterior family throughout; whether richer posteriors interact differently with non-likelihood losses is untested.
• Controlled and higher-dimensional systems. The paper notes the framework extends to controlled systems by augmenting the state with a control input, but the experiments focus on autonomous dynamics.
• Tuning and theory for λ. With non-likelihood losses, λ becomes a tunable regularization parameter rather than a fixed 1/N (as in tempered posteriors); principled selection and sharper PAC-Bayesian guarantees for the multi-step and autoregressive losses are open questions.
• Reducing the cost of autoregressive training. Backpropagation through time makes these losses less parallelizable and more memory-intensive than teacher forcing, motivating cheaper rollout-based objectives.

Target Audience

Researchers and graduate students working on probabilistic dynamics models, Bayesian deep learning, or model-based control and reinforcement learning. It is most useful to practitioners who already use variational Bayesian last layers (or similar last-layer Bayesian models) in an autoregressive deployment setting and want to improve long-horizon accuracy and calibration without changing the model architecture. Readers chiefly interested in benchmark leaderboards should note that detailed numeric results sit in appendices not included in the excerpt reviewed.

Authors’ 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.