Gaussian Flow Dynamics: Simulation-Free Neural SDE Learning Beyond One-Time Marginals
Grigory Bartosh, Christian A. Naesseth
Abstract
Simulation-free training of latent Stochastic Differential Equations (SDEs) relies on a variational posterior process whose one-time marginals are tractable, typically Gaussian. Such marginals, however, do not determine the underlying dynamics: many processes share the same marginals while differing in their temporal structure, and existing parameterizations fix this structure implicitly, which restricts the posterior family and biases the learned model. We introduce Gaussian flow dynamics, which construct stochastic processes directly from smoothly evolving Gaussian marginals while making the marginal-preserving, or gauge, degrees of freedom explicit and parameterizable. The construction admits state-dependent diffusion coefficients and recovers every linear SDE with additive noise and a non-degenerate Gaussian initial distribution. Building on it, we propose Gauge Matching, a simulation-free method for latent SDE learning that combines Gaussian flow dynamics with the SDE Matching objective. Gauge Matching costs at most quadratically in the latent dimension per step, like SDE Matching, but learns the temporal structure of the posterior beyond its one-time marginals. It comes within a nat of Helmholtz-SDE, which computes the gauge from the prior Jacobian at cubic cost, on the linear benchmark where the exact posterior is known, matches it on nonlinear systems, and applies where Helmholtz-SDE does not, to state-dependent noise.