Variational Grey-Box Dynamics Matching

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.

💡 Research Summary

The paper introduces Variational Grey‑Box Dynamics Matching (VGB‑DM), a novel framework that integrates incomplete physics models directly into flow‑matching based generative models while remaining simulation‑free. Traditional deep generative approaches such as Neural ODEs or flow‑matching excel at learning complex distributions but ignore known physical laws, whereas physics‑based simulators provide interpretability but often lack unknown terms or parameters. Existing grey‑box methods that combine the two typically rely on costly ODE solvers during training, leading to scalability and stability issues.

VGB‑DM addresses these problems by defining a structured variational distribution over the velocity field. Two latent variables are introduced: (1) θ, a latent encoding of interpretable physics parameters, equipped with a physics‑informed prior; and (2) z, a latent variable that captures stochasticity and multi‑modal behavior that arises when the physics model is incomplete. The combined velocity field is expressed as v_ϕ(t, x | θ, z) ∘ fₚ(x, θ), where fₚ is the known (but possibly incomplete) differential operator and v_ϕ is a neural network that learns the missing dynamics.

Training proceeds without any numerical integration. For each trajectory, consecutive states (x_k, x_{k+1}) are linearly interpolated to form an intermediate point x_t = (1‑t)x_k + t x_{k+1} and a target velocity \dot{x}t = x{k+1} – x_k. The loss matches the combined field to this target using a squared‑error term, yielding a flow‑matching objective that is fully differentiable and simulation‑free. A variational ELBO is derived, adding KL‑divergence penalties that regularize the posterior q_ψ(θ, z | segment) toward the priors p(θ) and p(z)=N(0,I). The posterior is factorized as q_ψ(z|segment) q_ψ(θ|segment, z) and parameterized by an encoder network, allowing analytic KL computation.

The framework is further extended to second‑order dynamics (e.g., pendulums, wave equations) by constructing a three‑point interpolant I(x_t) and matching both its first and second time derivatives (velocity and acceleration) with learned fields v_ϕ and a_ϕ. This enables the method to handle PDEs that involve second‑order time derivatives.

Experiments cover a range of representative ODE and PDE problems: an RLC circuit (first‑order ODE with unknown resistance), logistic growth, 1‑D heat diffusion, 1‑D wave propagation, and a weather‑forecasting time‑series task. VGB‑DM consistently outperforms pure data‑driven flow‑matching, Neural ODE, and recent grey‑box baselines such as Physics‑Integrated VAEs. It achieves lower negative log‑likelihood, smaller RMSE, and more accurate recovery of the hidden physics parameters θ. Moreover, because it avoids repeated ODE solves, training converges 2–3× faster and uses substantially less GPU memory than simulation‑based methods.

Key strengths of VGB‑DM include: (i) simulation‑free scalability to high‑dimensional systems, (ii) explicit separation of physics parameters and stochastic latent variables, preserving interpretability, and (iii) the ability to model multi‑modal velocity fields arising from missing physics. Limitations involve sensitivity to the dimensionality of the stochastic latent z and the need for a sufficiently expressive encoder to capture complex non‑linear physics. Future work is suggested on automatic latent dimension selection, richer physics‑informed priors, handling irregular time sampling or missing observations, and applying the method to large‑scale scientific simulators such as fluid dynamics or climate models.

In summary, VGB‑DM provides a principled, efficient, and interpretable approach to learning dynamical systems when only partial physical knowledge is available, bridging the gap between black‑box generative modeling and physics‑based simulation without incurring the computational burdens of traditional grey‑box techniques.