PaperlayerView exact revision on arXiv

Source-linked AI summary

Variational Streaming Flow: Probabilistic Forecasting in Physical Time

Hans Hao-Hsun Hsu, Minseon Gwak, Soon Hoe Lim, Pan Li, N. Benjamin Erichson

arXiv:2610.00976v1cs.LG

TL;DR

Stochastic dynamics and incomplete observations can make a single observed history consistent with multiple futures, while deterministic physical-time flow models produce only one. VSF adds a latent variable to physical-time forecasting to represent multiple possible futures. It improves long-horizon forecasting and planning outcomes while retaining one velocity-field evaluation per imagined transition.

  • Problem

    Stochastic dynamics and incomplete observations can make the same observed history consistent with multiple future trajectories, while standard deterministic physical-time models produce only one future.

    Source: §1 · p. 2

  • Method

    VSF uses variational inference and an energy-based reconstruction to train a physical-time velocity field conditioned on one latent variable sampled for an entire rollout.

    Source: §1 · pp. 2–3, §5 · p. 10

  • Results

    VSF improves long-horizon forecasting on deterministic and partially observed stochastic systems, captures future distributions more accurately, and improves planning success while retaining one velocity-field evaluation per imagined transition.

    Source: §5 · p. 11

  • Takeaways & Limitations

    As a predictor in an action-conditioned world model, VSF improves planning success while preserving the physical-time generation efficiency of Streaming Flow.

    Source: §1 · p. 2, §5 · p. 11

    VSF assumes adjacent states provide a meaningful local direction of motion, an assumption less reliable with abrupt impacts, discontinuities, or strong observation noise.

    Source: §5 · p. 11

Abstract

from arXiv · show

Probabilistic forecasting is important for predicting complex dynamical systems because intrinsic randomness and incomplete observations can cause the same observed state to evolve into multiple plausible futures. While flow matching is a flexible approach for probabilistic forecasting, it is computationally expensive. Streaming flow (SF) reformulates this approach to model temporal evolution efficiently by learning a continuous velocity field directly in physical time. However, SF learns a deterministic velocity field. Thus, it provides only a single future trajectory for a given fixed initial state and observation history. To overcome this limitation, we introduce Variational Streaming Flow (VSF). Our approach learns a latent distribution that is conditioned on the dynamics of interest. In turn, this enables probabilistic forecasting. Importantly, we retain the computational efficiency of SF by generating in physical time. Across deterministic and stochastic dynamical systems, VSF demonstrates superior predictive accuracy and distributional fidelity. We demonstrate the advantage for both long-horizon rollouts exceeding 1,000 steps, and settings with bifurcating dynamics. Moreover, VSF can be integrated into existing Joint-Embedding Predictive Architecture (JEPA)-based world models as a plug-and-play predictor to improve temporal dynamics and goal-directed success rate in navigation, motion planning, and manipulation.

1 Introduction

Long-horizon forecasting is difficult, and stochastic or partially observed systems can have multiple plausible futures, while deterministic physical-time models produce only one. VSF adds a latent variable to physical-time forecasting; evaluations report gains in forecasting and planning while retaining efficient generation.

1 Introduction

  • Probabilistic forecasts matter because stochastic dynamics and incomplete observations can make the same observed history consistent with multiple futures, making a single trajectory insufficient for forecasting or planning.Source: §1 · p. 2
  • VSF adds a latent variable sampled once per forecast and held fixed throughout the rollout, allowing different trajectories from the same observed state while preserving physical-time generation.During training, an encoder infers the latent from the observed future transition and regularizes its distribution toward a fixed prior; at inference, the latent is sampled from that prior.Source: §1 · pp. 2–3
  • VSF improves long-horizon accuracy on deterministic systems and recovers long-horizon future distributions more accurately on stochastic Rössler and Lorenz systems.The deterministic comparison includes DeepAR, diffusion- and flow-based forecasters, SF, and Chronos-2.Source: §1 · p. 2
  • Across four robotics tasks, VSF improves average goal-directed planning success by 7.34% over LeWM’s deterministic predictor while retaining one velocity-field evaluation per imagined transition.The evaluation uses VSF as the predictor in a JEPA-based world model.Source: §1 · p. 3
  • The variational framework combines a conditional-velocity ELBO with energy-based reconstruction over multiple latent samples; under posterior–prior alignment, the reconstruction controls energy distance between predicted and true velocity distributions.Source: §1 · p. 3

2 Problem Formulation and Preliminaries

The section formulates forecasting as learning a conditional distribution over future trajectories from an observed history. It presents Streaming Flow as a physical-time predictor and describes the trajectory-centered velocity field and its discrete-observation training setup.

2 Problem Formulation and Preliminaries

  • Forecasting aims to learn a conditional distribution over future trajectories given the observed history; fully observed deterministic dynamics yield one future, while partial observations or stochastic dynamics can permit several.Source: §2 · p. 3
  • Streaming Flow starts from the latest observation and advances the forecast state directly through physical time using a learned velocity field conditioned on state, time, and history.Training uses observed future states and physical velocities to construct a velocity target; the correction term pulls nearby states toward the trajectory and vanishes on it.Source: §2 · pp. 3–4
  • The SF target contracts a cloud of nearby initial states toward the observed future trajectory, forming a Gaussian tube whose mean follows the trajectory and whose width contracts over time.Source: §2 · pp. 3–4
  • With discrete observations, the trajectory is interpolated piecewise linearly between observation times, giving each transition a local state and constant physical-time velocity for VSF training.Source: §2 · p. 4

3 Variational Streaming Flow

Variational Streaming Flow (VSF) addresses the single-future limitation of non-latent physical-time flows by conditioning the velocity field on a latent variable. Its training combines variational learning with a distributional velocity reconstruction for representing multiple possible futures.

3 Variational Streaming Flow

  • Near bifurcations, non-latent streaming flow must distinguish similar states to select a future branch, while VSF conditions its physical-time velocity field on a latent variable.Changing the latent can produce a different future from the same observed history.Source: §3 · pp. 4–5

3.1 Training and Inference

  • During training, an encoder infers latent values from local future motion and KL regularization keeps its posterior close to the prior; inference samples one prior latent and retains it throughout the rollout.The generated history may be refreshed as the forecast advances, while the same latent continues conditioning it.Source: §3.1 · p. 5

3.2 Variational Training Objective

  • The local variational objective pairs a latent-conditioned reconstruction of observed physical-time velocity with KL regularization toward the inference-time prior.For the Gaussian likelihood, the ELBO reconstruction reduces to latent-conditioned MSE.Source: §3.2 · p. 6

3.3 Distributional Velocity Reconstruction

  • With MSE reconstruction, the learned velocity is the conditional mean of velocities compatible with the same context and latent, so overlapping future continuations are averaged.Such overlap can arise from imperfect inference, overlapping posterior distributions, or posterior collapse.Source: §3.3 · p. 7
  • The distributional reconstruction compares multiple latent-conditioned velocity predictions with the observed velocity, accounting for prediction spread rather than forcing all samples to match independently.Under aggregate posterior–prior alignment, its expected population reconstruction upper-bounds half the energy distance between prior-predictive and true conditional velocity distributions; the resulting objective is not an ELBO.Source: §3.3 · p. 7

4 Experimental Results

Experiments assess long-horizon forecasting on deterministic and partially observed stochastic systems, then test VSF as a predictor in a JEPA-based world model. VSF improves forecasting, distribution matching, and planning outcomes while retaining efficient physical-time prediction.

4 Experimental Results

4.1 Deterministic Forecasting

  • On deterministic systems, VSF maintains lower long-horizon error than SF, whose forecasts drift even when initial-state noise is added.Source: §4.1 · p. 9

4.2 Stochastic Forecasting

  • On stochastic Rössler and Lorenz systems, VSF best matches conditional future distributions; Rössler forecasts contrast its variability with narrower Block diffusion and drift-prone AR diffusion.Source: Figure 4 · p. 8, §4.2 · p. 9
  • For the 2,500-step stochastic Lorenz rollout, VSF has the lowest CME, CSE, and ED, and its samples preserve both attractor lobes.Source: Table 2 · p. 9, §4.2 · p. 9, Figure 5 · p. 9

4.3 World Model Predictor

  • Across four world-model tasks, VSF achieves the highest planning success on each, improves average success by 7.34%, and retains one velocity-field evaluation per imagined transition.Source: Table 3 · p. 10, §4.3 · p. 10
  • The Cube example shows a predicted latent trajectory closely following ground truth and leading to a successful plan.Source: Figure 6 · p. 10, §4.3 · p. 10
  • The world-model comparison holds the JEPA encoder, CEM planner, and planning budget fixed while comparing AR, SF, and VSF predictors.Source: §4.3 · p. 10

5 Conclusion

The conclusion presents VSF as a probabilistic physical-time extension of streaming flow, using one latent variable to condition a rollout. It reports forecasting and planning gains with a population-level distribution guarantee, while noting assumptions and numerical limitations.

5 Conclusion

  • VSF extends streaming flow by conditioning an entire physical-time rollout on one latent variable, enabling different possible futures while retaining direct physical-time generation.Source: §5 · p. 10
  • Empirically, VSF improves long-horizon forecasting on deterministic and partially observed stochastic systems, future-distribution quality, and planning success while retaining one velocity-field evaluation per imagined transition.Source: §5 · p. 11
  • The variational training objective has a population-level guarantee connecting it to how well prior sampling captures the distribution of possible local dynamics.Source: §5 · p. 11
  • The local-velocity assumption is most reliable for smoothly varying dynamics; impacts, discontinuities, strong observation noise, and coarse Euler steps can impair resolution or accumulate error.Source: §5 · p. 11

A Related Work

  • Related work spans probabilistic time-series and flow-matching forecasters, physical-time flows, and stochastic or feature-space world models for planning.Source: A Related Work · pp. 15–16
  • VSF combines physical-time flow with variational latent inference for probabilistic forecasting; the paper distinguishes this formulation from being the first latent streaming-flow model.Source: Table 4 · p. 15, A Related Work · p. 16

B Motivation for VSF

Standard non-latent Streaming Flow faces a fidelity–sensitivity tradeoff when futures share a physical state before bifurcating. VSF moves branch information into a trajectory-level latent variable, while a variational objective trains the latent-conditioned velocity model.

B Motivation for VSF

  • Near a bifurcation, standard SF must encode branch identity through small physical-state differences, creating a fidelity–sensitivity tradeoff even at the population optimum.Keeping sensitivity below a threshold imposes a strictly positive lower bound on fidelity error in this model.Source: B Motivation for VSF · pp. 16, 17, 19
  • As ετ→0, standard-SF sensitivity grows as JSFτ,t∼exp(C/ε²τ), while the ideal latent construction retains sensitivity e−k(t−τ).In the two-branch example, the futures share a physical trajectory up to τ and separate afterward.Source: Figure 7 · p. 16, B Motivation for VSF · pp. 16–17
  • In the idealized construction, latent branch information lets physical-state perturbations contract at rate e−k(t−τ), unlike standard SF’s amplification near the branch separator.This is a demonstrated possibility for an ideal latent construction, not a claim about the learned model itself.Source: B Motivation for VSF · pp. 19–20
  • VSF samples a latent variable once per forecast and keeps it fixed, allowing the same physical state to generate different future trajectories.The latent branch code distinguishes continuations while the physical state remains on the common stem.Source: B Motivation for VSF · p. 20

C A variational objective for clean streaming trajectories

  • The conditional-velocity variational bound uses a posterior over z conditioned on the observed local velocity and a likelihood conditioned on state, time, and history.The target velocity is supplied to the training-time posterior, not used as a conditioning variable of the velocity likelihood.Source: C A variational objective for clean… · p. 21
  • The bound is tight exactly when the variational posterior agrees almost everywhere with the exact model posterior.For the piecewise-linear path, local endpoints provide the same conditioning information to the posterior as (us, xs, s), without changing likelihood conditioning.Source: C A variational objective for clean… · pp. 21–22

D Population Energy Reconstruction and Energy-Distance Control

The energy-reconstruction analysis links the training objective to a distribution-level discrepancy, under aggregate posterior–prior alignment. That alignment assumption can hold even when the encoder collapses to the prior.

D Population Energy Reconstruction and Energy-Distance Control

  • The population reconstruction objective is the counterpart of the finite-ensemble V-statistic LES, obtained by replacing its expectations with M posterior samples.The analysis defines energy distance between prior-predictive and true conditional velocities.Source: D Population Energy Reconstruction and… · p. 23
  • Under aggregate posterior–prior alignment, expected energy reconstruction upper-bounds half the energy distance between the prior-predictive velocity law and the true conditional velocity law.Thus, reducing reconstruction controls a distribution-level discrepancy, not only conditional-mean error.Source: D Population Energy Reconstruction and… · p. 23
  • Aggregate posterior–prior alignment does not guarantee informative latent codes: posterior collapse, with qϕ(z|u,c)=p(z) for every u, also satisfies it.Finite-capacity aggregate-matching training does not guarantee exact alignment.Source: D Population Energy Reconstruction and… · p. 24

E.1 Forecasting training and inference

Forecasting training constructs local physical-time training tuples and optimizes the VSF objective using posterior samples. At inference, a single prior latent is held fixed while the history window is refreshed during rollout.

E.1 Forecasting training and inference

  • The history encoding stays fixed within each K-step segment and is recomputed from the generated-state window at segment boundaries.The window is a fixed-length queue that enqueues generated states and dequeues the oldest when full.Source: §E.1 · p. 25
  • Training samples a window and forecast time, constructs (xs, us), evaluates the VSF loss with posterior samples, and updates the velocity field and posterior.The training procedure uses M posterior samples and the specified α and βKL settings.Source: §E.1 · p. 25
  • Inference draws z from the prior once for the complete rollout and holds it fixed as the state advances through physical forecast time.Each update advances the state by ∆t and emits the resulting forecast state.Source: §E.1 · p. 25

F World-model VSF predictor: training and inference

The world-model predictor trains on teacher-forced transitions in a frozen JEPA representation and performs open-loop latent rollouts with a prior sample. The evaluation metrics distinguish pointwise marginal accuracy from trajectory-level distributional fidelity.

F World-model VSF predictor: training and inference

  • The world-model VSF trains only the predictor and trajectory posterior over frozen JEPA representations, using causal context refreshed after each imagined transition.Training uses teacher forcing with causal masking; inference feeds generated latents back into the rolling context.Source: F World-model VSF predictor: training… · p. 26
  • For each training clip, a shared interpolation fraction forms teacher-forced states and transition velocities from adjacent encoded observations.The stated construction is xs=(1−ρ)xk+ρxk+1 and us=N(xk+1−xk).Source: F World-model VSF predictor: training… · p. 26
  • Inference draws one prior latent for the rollout and predicts the next latent state from the rolling causal context and aligned actions.The predictor emits each generated state, which is then appended to the context window.Source: F World-model VSF predictor: training… · p. 26

G Evaluation metrics

  • Deterministic evaluation uses time-local CRPS, RMSE, and MAE, while stochastic evaluation includes marginal-moment errors and trajectory-level W1 and energy distance.W1 and energy distance retain complete trajectories and therefore assess cross-time dependence; lower reported metric values are better.Source: §G.1 · p. 27, §G.2 · p. 28

H Dataset Statistics

The datasets span deterministic physical systems, parameter-mixture futures, and visual-control tasks. Parameter-mixture trajectories share an initial state and remain aligned through the conditioning history before their futures can diverge.

H Dataset Statistics

  • Forecasting datasets were generated by integrating each system with an adaptive DOP853 solver in 64-bit floating-point precision, using relative and absolute error tolerances of 10^-9.The resulting trajectories were stored in 32-bit floating-point precision for training.Source: H Dataset Statistics · p. 30
  • Across parameter-mixture trajectories, conditioning histories remain closely aligned, with mean pairwise RMSE of 0.008 for Rössler and 0.047 for Lorenz, while later futures may diverge.Each initial condition produces one deterministic trajectory or 100 parameter-mixture trajectories.Source: Table 5 · p. 30

H.1 Deterministic Datasets

  • The deterministic benchmarks include double pendulum, Rössler, and Lorenz systems, with initial conditions and system parameters specified for dataset generation.The double-pendulum state is represented in Cartesian coordinates; Rössler and Lorenz use their system dynamics.Source: §H.1 · p. 30

H.2 Stochastic Datasets

  • Parameter-mixture data pair each initial condition with 100 futures generated using independently sampled, trajectory-level dynamical parameters withheld from the forecaster.Splits are made over the 40 initial conditions, and evaluation uses all 100 futures for each held-out initial condition.Source: §H.2 · p. 31

H.3 World-Model Datasets

  • The four visual-control datasets cover pushing, cube placement, navigation between rooms, and reaching a target joint configuration.Their datasets use 20,000 PushT episodes and 10,000 episodes each for Cube, TwoRoom, and Reacher.Source: §H.3 · p. 31

I Experimental Setup

Forecasting experiments compare VSF with diffusion, flow, autoregressive, and pretrained time-series baselines under long-rollout evaluation. World-model experiments compare predictors while holding the visual representation and planning setup fixed.

I.1 Forecasting

  • Forecasting comparisons include eight baselines, spanning block and autoregressive diffusion, rectified flow, DeepAR, SF variants, and frozen or calibrated Chronos-2.The calibrated Chronos-2 is a post-hoc affine calibration of predictive quantiles.Source: §I.1 · p. 32
  • VSF uses SF’s conditional U-Net with a four-dimensional trajectory-level latent fused at the bottleneck using 90% latent and 10% bottleneck features.The latent is mapped to bottleneck width by a two-layer MLP.Source: §I.1 · p. 32
  • Forecasts are evaluated over 1,000-step rollouts, with an additional 2,500-step stochastic-Lorenz experiment; results average three seeds.Deterministic data are split by trajectory and parameter-mixture data by initial condition.Source: §I.1 · p. 33

I.2 World Models

  • World-model comparisons freeze LeWM’s visual and action encoders and use the same 192-dimensional latent space, CEM planner, tasks, and planning budget.The compared predictors are LeWM’s pretrained autoregressive predictor, SF, and VSF.Source: §I.2 · p. 33
  • The world-model VSF predictor samples one prior latent per imagined rollout and, like SF, uses one velocity-field evaluation per transition.Evaluation uses the same planning setup, including 30 CEM updates of 300 action-sequence candidates each.Source: §I.2 · p. 33

J Full ablation results

The ablations support generalized energy reconstruction, prior alignment, latent conditioning, and keeping one latent fixed through a rollout. The reported comparisons indicate that these design choices contribute in complementary ways.

Energy-score exponent α and MSE

  • Performance improves consistently as the energy-score exponent α increases; all tested settings outperform the baseline, and near-quadratic α also surpasses MSE.For 1 < α < 2, the residual-term gradient scales as αr^(α−1), emphasizing large local errors; the energy loss also has a pairwise term.Source: J Full ablation results · p. 34, Figure 8 · p. 34

KL regularization

  • Performance remains better overall than SF across tested KL weights, with an intermediate weight giving the strongest aggregate result.Too-weak regularization (βKL = 1e−3) leaves a posterior–prior gap that makes prior samples less reliable at inference.Source: J Full ablation results · p. 34, Figure 9 · p. 34

Latent contribution

  • Every positive latent contribution improves overall performance relative to SF, with an intermediate contribution performing best.Maximal latent contribution weakens performance but retains state and physical-time conditioning.Source: J Full ablation results · p. 35, Figure 10 · p. 35

Latent resampling

  • A fixed latent performs similarly to per-step resampling early, but the gap widens with horizon in favor of the fixed latent.The comparison covers deterministic double pendulum and stochastic Rössler, using three-seed means.Source: J Full ablation results · p. 35, Figure 11 · p. 35

Latent-source and prior-alignment variants

  • Full VSF has lower error than Prior-latent SF on all five metrics, while Prior-latent SF improves over SF across all metrics.The comparison separates latent-conditioned multi-sample prediction from the additional contribution of future-conditioned variational training.Source: J Full ablation results · p. 36, Figure 12 · p. 36

Why prior alignment matters

  • Removing posterior–prior alignment degrades every metric, especially the Rössler distributional scores.The proposed explanation is train–test mismatch: prior samples may miss or misweight continuations represented by posterior codes.Source: J Full ablation results · p. 36, Figure 12 · p. 36

K Rollout Error Plots

The supplementary figures visualize per-timestep rollout errors, sampled trajectories, and world-model latent trajectories. They show evaluation views across deterministic and stochastic systems and qualitative examples of world-model rollouts.

K Rollout Error Plots

  • Per-timestep plots show deterministic 1,000-step forecasts across double pendulum, Rössler, and Lorenz, and marginal distributional scores for stochastic Rössler and Lorenz rollouts.The deterministic plots report CRPS, RMSE, and MAE; stochastic plots report CME, CSE, marginal W1, and marginal energy distance.Source: Figure 13 · p. 37, Figure 14 · p. 38

L Trajectory Forecast Visualizations

  • Trajectory visualizations compare ground truth with VSF samples and ensemble mean, SF, and block diffusion on deterministic systems, plus ground truth and generated stochastic-Lorenz trajectories.The deterministic panels include ten VSF trajectories and magnified regions of chaotic forecasts.Source: Figure 15 · p. 39, Figure 16 · p. 40

M World-Model Latent-Trajectory Visualizations

  • World-model latent-trajectory plots show an AR-predictor rollout diverging and failing, while the VSF-predictor rollout follows ground truth and succeeds; PushT and TwoRoom examples also show successful VSF rollouts.The latent trajectories are compared after joint-PCA projections.Source: Figure 17 · p. 41, Figure 18 · p. 42, Figure 19 · p. 42
Loading 2610.00976v1…