Flow Matching for Averaged Systems
Daniel Owusu Adu, Yongxin Chen
TL;DR
This work extends Flow Matching to ensembles of linear systems by averaging over the parameter index θ, revealing memory and non-Markovian interpolation in the stochastic case. It derives explicit Dirac- and general-distribution controllability results for both deterministic (ε=0) and stochastic (ε>0) averaged dynamics, with conditional-expectation formulations $u^z(t)$ and $u_ε^z(t)$ that realize distributional transport; for Gaussian and Gaussian-mixture targets, closed-form expressions for conditional expectations are provided to plug into the control law. To handle intractable conditional expectations in practice, the authors introduce a two-stage Flow Matching procedure: memory-aware learning (LSTM/Transformer) for ε>0 and memoryless, OT-guided learning (FNN) for ε=0, followed by a prediction stage using appropriate numerical schemes. The paper also emphasizes a practical numerical split between the stochastic and deterministic regimes, showing how OT coupling or SB-inspired approaches can realize full target distributions, and demonstrates the method on Ornstein–Uhlenbeck and anti-damped systems. Overall, the framework connects non-Markovian averaged control with Flow Matching, enabling scalable, trainable interpolation of high-dimensional distributions under uncertainty for generative modeling and distributed control.
Abstract
We extend flow matching to ensembles of linear systems in both deterministic and stochastic settings. Averaging over system parameters induces memory leading to a non-Markovian interpolation problem for the stochastic case. In this setting, a control law that achieves the distributional controllability is characterized as the conditional expectation of a Volterra-type control. This conditional expectation in the stochastic settings motivates an open-loop characterization in the noiseless-deterministic setting. Explicit forms of the conditional expectations are derived for special cases of the given distributions and a practical numerical procedure is presented to approximate the control inputs. A by-product of our analysis is a numerical split between the two regimes. For the stochastic case, history dependence is essential and we implement the conditional expectation with a recurrent network trained using independent sampling. For the deterministic case, the flow is memoryless and a feedforward network learns a time-varying gain that transports the ensemble. We show that to realize the full target distribution in this deterministic setting, one must first establish a deterministic sample pairing (e.g., optimal-transport or Schrodinger-bridge coupling), after which learning reduces to a low-dimensional regression in time.
