Hilbert Space Setting
Let \(\mathcal{H}\) be a real separable Hilbert space with inner product \(\langle \cdot, \cdot \rangle_{\mathcal{H}}\) and countable ONB \(\{e_k\}_{k \in \mathbb{N}}\).
For \(K \in \mathbb{N}\), define the canonical projection \(\pi_K : \mathcal{H} \to \mathbb{R}^K\) by: \[\pi_K(x) = \bigl(\langle x, e_1 \rangle, \ldots, \langle x, e_K \rangle\bigr)\]
Consider two independent \(\mathcal{H}\)-valued random variables:
- \(X_0 \sim \mu_0 = \mathcal{N}(0, C)\) — Gaussian reference with trace-class covariance \(C\)
- \(X_1 \sim \mu_1\) — the data distribution
The Cameron–Martin space is represented as \(\mathcal{H}_C\).
Functional Flow Matching (FFM)
Construct the linear interpolation between \(X_0\) and \(X_1\): \[X_t = (1-t) X_0 + t X_1, \quad t \in [0,1], \qquad \mu_t = \mathrm{Law}(X_t)\]
A velocity field \(v : \mathcal{H} \times [0,1] \to \mathcal{H}\) generates the path \((\mu_t)\) via the continuity equation: \[\int_{\mathcal{H}} \varphi(x,1)\,d\mu_1 - \int_{\mathcal{H}} \varphi(x,0)\,d\mu_0 = \int_0^1 \int_{\mathcal{H}} \bigl(\partial_t \varphi + \langle v_t(x), \nabla_x \varphi \rangle\bigr)\,d\mu_t\,dt\]
Training: minimize the tractable conditional flow-matching loss
Generation: simulate the learned ODE \(\partial_t \phi_t(x) = v_\theta(\phi_t(x), t)\) from \(x \sim \mu_0\)
Absolute Continuity and KL Divergence
Under Assumptions 2.1–2.2, for \(t \in [0,1)\), we assume absolute continuity \(\mu^A_t \ll \mu^B_t\) with density ratio \(r_t(x) := \frac{d\mu^A_t}{d\mu^B_t}(x)\).
Lemma 3.4: Under mild regularity, the KL divergence between path measures is: \[\mathrm{KL}(\nu^A \| \nu^B) = \int_0^1 \int_{\mathcal{H}} \langle v_t^A(x) - v_t^B(x), \nabla_x \log r_t(x) \rangle_{\mathcal{H}} \, d\mu_t^A(x) \, dt\]
The score difference \(\nabla_x \log r_t(x)\) is computed via a reference measure \(\rho_t\): \[\nabla_x \log r_t(x) = \nabla_x \log \frac{d\mu_t^A}{d\rho_t}(x) - \nabla_x \log \frac{d\mu_t^B}{d\rho_t}(x)\]
Step 1: Basis Projection
For \(K \in \mathbb{N}\), define projected measures \(\hat{\mu}^{A,B}_{t,K} = (\hat{\pi}_K)_\# \mu^{A,B}_t\).
The projected density ratio is \(r^K_t = \frac{d\hat{\mu}^A_{t,K}}{d\hat{\mu}^B_{t,K}} = \mathbb{E}_{P^B}[r_t \mid \mathcal{G}_K]\), with \(\mathcal{G}_K = \sigma(\hat{\pi}^K(X_t))\).
From the weak continuity equation with \(r^K_0 = 1\) (\(P^B\)-a.s.): \[\mathrm{KL}(\nu^A_K \| \nu^B_K) = \int_0^1 \int_{\mathcal{H}} \left(\partial_t \log r^K_t(x) + \langle v^A_t(x), \nabla_x \log r^K_t(x) \rangle \right) d\mu^A_t(x)\,dt\]
Step 2: Combining the Continuity Equations
From the continuity equation: \[\int_0^1 \int_{\mathcal{H}} \partial_t \log r^K_t(x)\,d\mu^A_t(x)\,dt = -\int_0^1 \int_{\mathcal{H}} \langle v^B_t(x), \nabla_x \log r^K_t(x) \rangle\,d\mu^A_t(x)\,dt\]
Substituting back: \[\mathrm{KL}(\nu^A_K \| \nu^B_K) = \int_0^1 \int_{\mathcal{H}} \langle v^A_t(x) - v^B_t(x),\, \nabla_x \log r^K_t(x) \rangle\,d\mu^A_t(x)\,dt\]
Note that \(\mathrm{KL}(\nu^A_K \| \nu^B_K) \to \mathrm{KL}(\nu^A \| \nu^B)\) as \(K \to \infty\).
Step 3: Expressing the Log-Gradient via Velocity
The logarithmic gradients can be expressed as: \[\nabla_x \log \frac{d\mu^{A,B}_t}{d\rho_t}(x) = \frac{t}{(1-t)^2} C^{-1} \mathbb{E}_{P^{A,B}}[X_1 \mid X_t = x]\]
For the linear interpolation used in FFM: \[v^{A,B}_t(x) = \frac{1}{1-t}\mathbb{E}_{P^{A,B}}[X_1 \mid X_t = x] - \frac{x}{1-t}\]
\(\Rightarrow\) The log-gradient of \(r_t = \frac{d\mu^A_t}{d\mu^B_t}\) can be expressed entirely in terms of \(v^A_t - v^B_t\).
Implementation Details
Architecture: Mesh-Informed Neural Operator (MINO) — transformer version
Projection: all computations use \(K\)-mode truncation of the ONB \(\{e_k\}\)
Noise: trace-class noise (Matérn or rougher data covariance) rather than white noise
- Better generation quality and more robust KL estimation
- Resolution invariant
Asymmetry: both \(\mathrm{KL}(\nu^A \| \nu^B)\) and \(\mathrm{KL}(\nu^B \| \nu^A)\) are reported:
- Forward KL: penalizes missing coverage of the reference
- Reverse KL: penalizes hallucinated mass outside the reference support
Special Case 1: Gaussian Measures
Consider \(\nu^A = \mathcal{N}(m^A, R)\) and \(\nu^B = \mathcal{N}(m^B, R)\) on \(\mathcal{H} = L^2(\mathbb{T}; \mathbb{R})\) with Matérn covariance \(R\).
Closed-form KL: \[\mathrm{KL}(\nu^A \| \nu^B) = \frac{1}{2}\|m^A - m^B\|^2_{\mathcal{H}_R}\]
Setting \(m^A = \cos(2\pi x)\), \(m^B = 0\); ONB \(\{e_k(x) = e^{2\pi i k \cdot x}\}\) with \(\mathcal{K}e_k = \lambda_k e_k\), \(\lambda_k = \sigma^2(4\pi^2 k^2 + \tau^2)^{-\alpha}\).
Validation: FKL estimator closely matches the closed-form value, and is resolution invariant.
Special Case 1: Results
| 1 |
1 |
1 |
0.5 |
3.64 |
3.95 |
3.64 |
3.88 |
| 2 |
1 |
1 |
1.5 |
32.79 |
32.61 |
32.79 |
33.31 |
| 3 |
1 |
3 |
1.5 |
50.00 |
49.58 |
50.00 |
49.00 |
| 4 |
1 |
5 |
1.5 |
103.74 |
104.46 |
103.74 |
100.58 |
| 5 |
2 |
1 |
0.5 |
7.29 |
6.77 |
7.29 |
6.72 |
| 6 |
3 |
1 |
0.5 |
10.93 |
10.48 |
10.93 |
10.49 |
| 7 |
5 |
1 |
0.5 |
18.22 |
22.32 |
18.22 |
22.28 |
| 8 |
10 |
1 |
0.5 |
36.43 |
45.12 |
36.43 |
45.73 |
Gaussian Measure: analytic vs estimated functional KL.
Special Case 2: Linear SDEs (Girsanov)
Two SDEs with same diffusion, different drifts: \[A: dY_t = c^A Y_t\,dt + g\,dW_t, \qquad B: dY_t = c^B Y_t\,dt + g\,dW_t\]
Girsanov’s theorem + Itô isometry give the closed-form KL: \[\mathrm{KL}(\nu^A \| \nu^B) = \frac{(c^A - c^B)^2}{2g^2} \left[ (M_0 + S_0)\,\frac{e^{2c^A} - 1}{2c^A} + \frac{g^2 D}{2c^A}\!\left(\frac{e^{2c^A} - 1}{2c^A} - 1\right) \right]\]
where \(M_0 = \|Y_0\|^2\), \(S_0 = \text{Tr}(\Sigma_0)\).
Our FKL estimates match this ground truth closely across all tested settings.
Special Case 2: Results
| 1 |
1 |
0.01 |
1.50 |
0.75 |
8.93 |
8.87 |
54.71 |
53.77 |
| 2 |
1 |
0.10 |
2.00 |
0.75 |
15.89 |
13.89 |
186.19 |
200.64 |
| 3 |
2 |
0.01 |
1.50 |
0.75 |
17.86 |
18.05 |
109.43 |
127.67 |
| 4 |
3 |
0.01 |
1.50 |
0.75 |
26.79 |
34.58 |
164.14 |
145.06 |
| 5 |
5 |
0.01 |
1.50 |
1.00 |
26.34 |
32.82 |
158.22 |
135.28 |
Linear SDE: analytic (Girsanov) vs estimated FKL.
FKL vs Upsampling Ratio
- Trace-class noise: stable KL estimates at all resolutions
- White noise: degrades at super-resolution — trace-class noise covariance is essential
TI Methods Compared
| SBIRR |
Schrödinger Bridge |
Gradient fields + stochasticity; iterative ref. refinement |
| vSB |
Schrödinger Bridge |
Vanilla IPF-based; snapshot matching |
| MSBM |
Multi-marginal SB |
Iterative Markovian Fitting; multi-snapshot constraints |
| MFL |
Mean-Field Langevin |
Path-space optimization; min-entropy w.r.t. Wiener measure |
| AM |
Action Matching |
Minimizes kinetic energy; optional entropic variant |
| TIGON |
ODE-based |
Neural ODE; dynamic unbalanced OT + growth |
Reference: VAL = resampled ground-truth trajectories (oracle upper bound)
Marginal Metrics Evaluated
| Earth Mover’s Distance |
EMD (\(W_1\)) |
OT |
Geometry-aware |
| 2-Wasserstein Distance |
\(W_2\) |
OT |
Quadratic cost |
| Sliced Wasserstein Distance |
SWD |
OT |
Random projections |
| Max-sliced Wasserstein |
MWD |
OT |
Max-sliced variant |
| Maximum Mean Discrepancy |
MMD |
Kernel |
RBF kernel; bandwidth-dependent |
\[W_p(\mu,\nu) = \left(\inf_{\gamma\in\Pi(\mu,\nu)}\int \|x-y\|^p\,d\gamma\right)^{1/p}\] \[\mathrm{MMD}^2(\mu,\nu) = \mathbb{E}_{x,x'\sim\mu}[k(x,x')] - 2\,\mathbb{E}_{x\sim\mu,y\sim\nu}[k(x,y)] + \mathbb{E}_{y,y'\sim\nu}[k(y,y')]\]
Limitation: All these metrics operate on marginals — blind to temporal coupling structure.
Synthetic Datasets
Lotka–Volterra
- Predator-prey SDE dynamics
- 9 snapshots (5 train, 4 val)
- 100 samples/snapshot
- Non-linear oscillatory behavior
Repressilator
- 3-gene regulatory oscillator
- Cyclic inhibition loop
- 11 snapshots (6 train, 5 val)
- 100 samples/snapshot
Petal
- Bifurcations & merges
- Mimics cell differentiation
- 5 snapshots (all training)
- 500 samples/snapshot
Ground-truth trajectories available \(\Rightarrow\) FKL computed against true path measure
Real Single-Cell RNA-seq Datasets
Embryoid Body (EB)
- Moon et al., 2019
- 3 training + 2 validation timepoints
- 300 cells per snapshot
- 5 dimensions (PCA-reduced)
- Stem cell differentiation
Human Embryonic Stem Cell (hESC)
- Chu et al., 2016
- 3 training + 2 validation snapshots
- 92 / 66 / 138 cells (train)
- 102 / 172 cells (validation)
- 5 dimensions (PCA-reduced)
No ground-truth trajectories on real data. We use SBIRR as reference: it models differentiation via gradient fields with stochasticity, matching biological intuition of Waddington’s landscape.
