Variational Grey-Box Dynamics Matching

AI summary

Overview

• Research area: Scientific machine learning, specifically grey-box modelling of dynamical systems, combining partial differential equation (PDE) and ordinary differential equation (ODE) physics with deep generative models (flow matching and variational inference).
• Technical level: Advanced. The core intuition is accessible, but the method relies on flow matching, variational inference, the evidence lower bound, and second-order interpolation, and the paper assumes familiarity with neural ODEs and grey-box modelling.
• Scope: The paper proposes Variational Grey-Box Dynamics Matching (VGB-DM), a simulation-free framework that embeds incomplete physics equations into a generative model, infers unknown physics parameters from trajectory data alone, and matches or beats both black-box and prior grey-box baselines on four differential-equation benchmarks and real-world weather forecasting.

What This Paper Is About

Physics-based models of dynamical systems are interpretable but often incomplete, containing unknown parameters or missing terms that prevent them from fully describing real observations. Deep learning alternatives can fit almost any data, but they behave as black boxes, ignore physical structure, and degrade when data are scarce or when forecasting beyond the observed regime. This paper's goal is to combine the two: learn system dynamics from raw trajectories while retaining the interpretability of an explicit physics model, and do so without the computational cost and numerical instability of training through an ODE solver.

Key Contributions

1. Simulation-free grey-box modelling. The authors integrate an incomplete physics model directly into a flow-matching-style training objective. Instead of simulating the dynamics and backpropagating through a numerical integrator (as Neural ODE and PhysVAE-style methods do), they regress a target velocity field constructed from consecutive trajectory points, making training cheaper, more memory-efficient, and more numerically stable.
2. Structured variational inference with two latent variables. Rather than a single unstructured latent code, the method uses one latent variable to capture missing stochasticity and multimodal velocity fields, and a second to represent physics parameters drawn from a physics-informed prior. This separation lets the model represent multiple plausible futures from the same state while keeping the inferred physical parameters interpretable.
3. Extension to second-order dynamics. The framework is generalized to systems governed by second-order ODEs/PDEs (pendulums, wave equations) by introducing a non-linear interpolant over three consecutive trajectory points and regressing both velocity and acceleration fields, using a shared backbone with two output heads.
4. Empirical validation across synthetic and real systems. The method is evaluated on an RLC circuit, a damped pendulum, a reaction-diffusion system, a chaotic Lorenz attractor, and ERA5 weather data, with comparisons against PhysVAE, black-box Neural ODE, variational rectified flow matching, and trajectory flow matching.

Main Findings

• Best forecasting accuracy across all four differential-equation benchmarks. VGB-DM achieves the lowest mean squared error on the pendulum (0.283 × 10⁻²), RLC circuit (0.23 × 10⁻¹), reaction-diffusion (7.1 × 10⁻³), and Lorenz attractor (4.68 × 10⁻¹) tasks, outperforming PhysVAE, black-box Neural ODE, variational flow matching, and trajectory flow matching baselines.
• Largest gains on chaotic dynamics. On the Lorenz attractor, the error (4.68 × 10⁻¹) is roughly half that of the strongest baseline (9.85 × 10⁻¹) and nearly five times lower than black-box Neural ODE (22.15 × 10⁻¹), suggesting the physics prior is most valuable where trajectories are highly sensitive and multimodal.
• Faster and more stable convergence. Training on a single 12 GB RTX 3060, the method reaches lower test error in less wall-clock time than simulation-based approaches, which the authors attribute to avoiding backpropagation through an ODE solver, a known source of gradient instability.
• Stronger sample efficiency. Across training-set sizes from 10 to 1000 samples, grey-box models (VGB-DM and PhysVAE) consistently beat their black-box counterparts, with VGB-DM showing the lowest error and variance, particularly in low-data regimes.
• Competitive real-world weather forecasting. On ERA5 reanalysis data, the method matches or exceeds the physics-informed Neural ODE baseline ClimODE in latitude-weighted RMSE and anomaly correlation coefficient over 42-hour horizons, and outperforms it increasingly over a five-month horizon as compounding errors accumulate.
• Lower computational footprint with comparable or better accuracy. VGB-DM was trained on a single 12 GB GPU under a fixed six-hour budget with 2000 optimization steps, while ClimODE required at least 25 GB on an A5000.
• Physics parameters remain interpretable. Because physics parameters are explicit latent variables with a physics-informed prior, the model reports parameter estimation errors alongside forecasts, preserving mechanistic inspection of the learned system.

Methodology in Plain English

The starting point is an incomplete physics equation that describes part of a system's behaviour but is missing terms or has unknown constants. Standard grey-box approaches plug a neural network into that equation and then simulate the system forward, comparing predictions to observations and backpropagating through the solver. That simulation step is the bottleneck: it is slow, memory-hungry, and can produce unstable gradients.

The authors instead borrow the training trick from flow matching. Rather than simulating a full trajectory, they take pairs of consecutive observed states and build a smooth path between them, which gives a "target velocity" at every intermediate point. The model's job is to predict that velocity. The prediction is formed by combining (a) whatever the known physics equation says, evaluated with the current guess of the physics parameters, and (b) a neural network that supplies the missing dynamics.

The unknown physics parameters are not given in the data, so they must be inferred. The authors add two latent variables. The first, a stochastic latent, absorbs randomness and the multimodality that arises when one state could plausibly evolve in several directions. The second, the physics latent, represents the unknown parameters and is regularized toward a physically plausible prior. An encoder reads a short window of the trajectory and proposes distributions over both latents, with the physics latent conditioned on the stochastic one. Training maximizes a lower bound on the likelihood, which reduces to a velocity-matching loss plus two KL terms that keep the inferred distributions close to their priors.

For systems governed by second-order equations, the same idea is applied with a non-linear (Lagrange) interpolation across three consecutive states, which yields both a velocity and an acceleration target. A single backbone network with two heads predicts both, and a weighting term balances the acceleration loss.

At test time, forecasting works by integrating the combined physics-plus-network dynamics forward. Sampling different latent values produces an ensemble of plausible trajectories rather than one deterministic prediction.

Why This Matters

Impact on research. The work shows that grey-box modelling does not require differentiating through an ODE solver. Decoupling physics integration from training removes one of the main practical obstacles to scaling physics-informed learning, and it reframes physics parameters as objects of variational inference within a generative model. It also extends simulation-free generative modelling to second-order dynamics, which covers a much broader class of physical systems than first-order flow matching alone.

Real-world applications.

• Numerical weather and climate prediction, where physics-based formulations provide strong inductive biases and long-horizon stability matters more than short-term accuracy.
• Engineering and control of electrical, mechanical, and thermal systems, such as RLC circuits or damped oscillators, where some component behaviour is characterized but others are not.
• Geophysics and subsurface modelling, including geothermal exploration, which is one of the funding contexts acknowledged in the paper.
• Epidemiology and biological population dynamics, where compartmental models are structurally known but transmission and interaction parameters are uncertain and data are scarce.

Industry relevance. Simulation-based engineering and forecasting pipelines are expensive both in compute and in expert time. A method that trains on a single consumer-grade GPU while producing interpretable parameter estimates is attractive to industries building digital twins, energy load forecasting, or model-predictive control, where practitioners need to trust and inspect model components rather than accept a black-box output. The ability to learn useful dynamics from tens to hundreds of trajectories also matters in domains where collecting data is costly, dangerous, or slow.

Future Directions

• Stiff and discontinuous systems. The framework assumes differentiable, relatively smooth dynamics. How it behaves on stiff systems, shocks, or discontinuities remains untested, though the authors argue that avoiding solver backpropagation sidesteps some related instability.
• Transfer across related physical systems. The authors propose reusing a trained encoder as a pre-trained feature extractor for a related system (for instance, moving from a single pendulum to a double pendulum), adapting only output layers when physical parameters differ.
• Broader interpolation choices. The current second-order extension uses Lagrange interpolation; spline-based or Fourier interpolation are noted as untested alternatives that could affect accuracy on oscillatory or noisy systems.
• Scaling and scope of the variational structure. Whether the two-latent decomposition remains sufficient for high-dimensional, strongly coupled fields, and how the inferred physics parameters could be used for downstream scientific discovery rather than forecasting alone, are open questions.

Target Audience

Researchers and graduate students working at the intersection of machine learning and the physical sciences, particularly those interested in physics-informed learning, neural differential equations, and generative modelling of time series. It is also relevant to applied scientists and engineers who have partial mechanistic models and observational trajectories but not ground-truth parameters, and who need forecasts that remain interpretable and computationally tractable. Readers without a background in flow matching or variational inference will need to consult the cited background material to follow the derivations.

Authors’ abstract

Deep generative models such as flow matching and diffusion models have shown great potential in learning complex distributions and dynamical systems, but often act as black-boxes, neglecting underlying physics. In contrast, physics-based simulation models described by ODEs/PDEs remain interpretable, but may have missing or unknown terms, unable to fully describe real-world observations. We bridge this gap with a novel grey-box method that integrates incomplete physics models directly into generative models. Our approach learns dynamics from observational trajectories alone, without ground-truth physics parameters, in a simulation-free manner that avoids scalability and stability issues of Neural ODEs. The core of our method lies in modelling a structured variational distribution within the flow matching framework, by using two latent encodings: one to model the missing stochasticity and multi-modal velocity, and a second to encode physics parameters as a latent variable with a physics-informed prior. Furthermore, we present an adaptation of the framework to handle second-order dynamics. Our experiments on representative ODE/PDE problems show that our method performs on par with or superior to fully data-driven approaches and previous grey-box baselines, while preserving the interpretability of the physics model. Our code is available at https://github.com/DMML-Geneva/VGB-DM.