Table of Contents
Fetching ...

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.

Flow Matching for Averaged Systems

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 and 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.
Paper Structure (8 sections, 6 theorems, 61 equations, 6 figures, 2 algorithms)

This paper contains 8 sections, 6 theorems, 61 equations, 6 figures, 2 algorithms.

Key Result

Theorem 2.1

ZE:14 The ensemble of systems eq:stochastic ensemble of systems, where $\epsilon=0$, is said to be averaged controllable if and only if the vector space spanned by the columns of $\left\{\int_{0}^{1}A(\theta)^kB(\theta)d\theta \right\}_{k= 0}^{\infty}$ is of rank $d$.

Figures (6)

  • Figure 1: Interpolation of fixed endpoints $x_0=[1,0]$ and $x_f=[1,1]$ by the averaged of linear systems in \ref{['eq:stochastic ensemble of systems']} governed by the Ornstein-Uhlenbeck dynamics parameterized as \ref{['eq: 2D-Ornstein-Uhelbeck_parameters']}. Each color corresponds to a different noise intensity $\epsilon \in \{0,0.5,1\}$.
  • Figure 2: Interpolation of fixed endpoints $x_0=[1,0]$ and $x_f=[1,1]$ by averaged of linear systems in \ref{['eq:stochastic ensemble of systems']} characterized by the parameters in \ref{['eq: 2D_anti-damped_parameters']}. Compared with the Ornstein-Uhlenbeck case, trajectories exhibit outward spiralling and non-reverting behavior.
  • Figure 3: Comparison of initial and final distributions for the ensemble of Ornstein-Uhlenbeck and anti-damped processes. Both subplots are scaled to equal size for visual comparison.
  • Figure 4: Comparison of controlled trajectories using LSTM control for two systems: (A) the Ornstein-Uhlenbeck process and (B) the anti-damped process. Both subplots visualize how the learned control steers the stochastic ensemble from $\mu_0$ toward $\mu_f$ under different system dynamics.
  • Figure 5: Comparison between OT-paired samples and controlled trajectories learned from the OT map. (A) OT coupling between initial and target distributions for a bimodal Gaussian target. (B) Trajectories generated by the learned open-loop control $u(x_0,t)$ that dynamically transport samples from the initial Gaussian $\mu_0$ to the bimodal target $\mu_f$.
  • ...and 1 more figures

Theorems & Definitions (13)

  • Theorem 2.1
  • Proposition 2.1
  • Remark 2.1
  • proof
  • Proposition 2.2
  • Remark 2.2
  • Theorem 3.1
  • proof
  • Theorem 3.2
  • Remark 3.1
  • ...and 3 more