Conditional flow-matching training and exact-flow sampling

Let $d,q,B\ge 1$, $N\ge 0$, and $\mathrm{\Theta }\subseteq {\mathbb{R}}^q$ be open. Let $\lambda$ be a sigma-finite Borel measure on ${\mathbb{R}}^d$, let $\pi$ be a Borel probability measure on ${\mathbb{R}}^d$, and let $\rho :[0,1]\times {\mathbb{R}}^d\times {\mathbb{R}}^d\to [0,\mathrm{\infty }]$ and $u^c:[0,1]\times {\mathbb{R}}^d\times {\mathbb{R}}^d\to {\mathbb{R}}^d$ be jointly measurable. Assume $\int {\rho }_t(x\mid x_1)\lambda (dx)=1$ for every $t$ and for $\pi$-almost every $x_1$, and that for $\pi$-almost every $x_1$, ${\rho }_t(\cdot \mid x_1)$ and every component of ${\rho }_t{u}_t^c(\cdot \mid x_1)$ are continuously differentiable in $(t,x)$ on $(0,1)\times {\mathbb{R}}^d$ and satisfy ${\mathrm{\partial }}_t{\rho }_t(x\mid x_1)+{\sum }_{j=1}^d{\mathrm{\partial }}_{x_j}({\rho }_t(x\mid x_1)(u_t^c{)}^j(x\mid x_1))=0.$ For every compact $K\subset (0,1)\times {\mathbb{R}}^d$, suppose there is a nonnegative $G_K\in L^1(\pi )$ bounding on $K$, for $\pi$-almost every $x_1$, the absolute values of ${\rho }_t$, ${\mathrm{\partial }}_t{\rho }_t$, each component of ${\rho }_t{u}_t^c$, and each corresponding spatial flux derivative.

Define $p_t(x):=\int {\rho }_t(x\mid x_1)\pi (d{x}_1)$ and $J_t(x):=\int {\rho }_t(x\mid x_1)u_t^c(x\mid x_1)\pi (d{x}_1)$. Assume $0<p_t(x)<\mathrm{\infty }$ and $J_t(x)$ is finite on $[0,1]\times {\mathbb{R}}^d$, and put $u_t^m(x):=J_t(x)\mathrm{/}{p}_t(x)$. Let $v:[0,1]\times {\mathbb{R}}^d\times \mathrm{\Theta }\to {\mathbb{R}}^d$ be jointly measurable and continuously differentiable in $\theta$, with jointly measurable parameter derivative. Define ${\mathcal{L}}_C(\theta ):={\int }_0^1\text{\hspace{0.17em}⁣}\int \text{\hspace{0.17em}⁣}\int {\rho }_t(x\mid x_1)\mathrm{\parallel }v(t,x;\theta )-u_t^c(x\mid x_1){\mathrm{\parallel }}_2^2\pi (d{x}_1)\lambda (dx)dt$ and ${\mathcal{L}}_M(\theta ):={\int }_0^1\text{\hspace{0.17em}⁣}\int p_t(x)\mathrm{\parallel }v(t,x;\theta )-u_t^m(x){\mathrm{\parallel }}_2^2\lambda (dx)dt\mathrm{.}$ Assume, almost surely, that ${\mathcal{L}}_C({\theta }_k)$ and ${\mathcal{L}}_M({\theta }_k)$ are finite for every $k=0,\dots ,N$, and that $C:={\int }_0^1\text{\hspace{0.17em}⁣}\int \text{\hspace{0.17em}⁣}\int {\rho }_t(x\mid x_1)\mathrm{\parallel }{u}_t^c(x\mid x_1)-u_t^m(x){\mathrm{\parallel }}_2^2\pi (d{x}_1)\lambda (dx)dt$ is finite. Write $D_{\theta }v$ for the matrix with entries $(D_{\theta }v{)}_{ij}={\mathrm{\partial }}_{{\theta }_j}{v}^i$. For every compact $H\subset \mathrm{\Theta }$ and coordinate $j$, suppose nonnegative integrable majorants $G_{H,j}^C$ on $dt\lambda (dx)\pi (d{x}_1)$ and $G_{H,j}^M$ on $dt\lambda (dx)$ satisfy, for every $\theta \in H$ outside fixed null sets, $2{\rho }_t(x\mid x_1)\mid ⟨{\mathrm{\partial }}_{{\theta }_j}v(t,x;\theta ),v(t,x;\theta )-u_t^c(x\mid x_1)⟩\mid \le G_{H,j}^C(t,x,x_1)$ and $2p_t(x)\mid ⟨{\mathrm{\partial }}_{{\theta }_j}v(t,x;\theta ),v(t,x;\theta )-u_t^m(x)⟩\mid \le G_{H,j}^M(t,x)\mathrm{.}$ Let $Q(dt,d{x}_1,dx):=dt\pi (d{x}_1){\rho }_t(x\mid x_1)\lambda (dx)$ be the probability law of a training triple, with $dt$ Lebesgue probability measure on $[0,1]$. Take ${\theta }_0$ as fixed and let ${\mathcal{F}}_0$ be the trivial history. At iteration $k$, let ${\mathcal{F}}_k$ be the history before the batch, so ${\theta }_k$ is ${\mathcal{F}}_k$-measurable; draw $B$ independent triples with law $Q$, independently of ${\mathcal{F}}_k$, and set ${\mathcal{F}}_{k+1}$ to be the sigma-algebra generated by ${\mathcal{F}}_k$ and that batch. For a triple $w=(t,x_1,x)$, define $\zeta (\theta ,w):=2D_{\theta }v(t,x;\theta {)}^T(v(t,x;\theta )-u_t^c(x\mid x_1))\mathrm{.}$ For each $k,b$, set ${\zeta }_{k,b}:=\zeta ({\theta }_k,(T_{k,b},X_{1,k,b},X_{k,b}))$ and require it to be integrable on the joint training probability space.

Given the fixed initial parameter ${\theta }_0\in \mathrm{\Theta }$ and deterministic step sizes ${\eta }_0,\dots ,{\eta }_{N-1}\ge 0$, for each $k=0,\dots ,N-1$ set $g_k:=\frac{1}{B}{\sum }_{b=1}^B{\zeta }_{k,b},{\theta }_{k+1}:={\theta }_k-{\eta }_k{g}_k\mathrm{.}$ Assume all iterates ${\theta }_0,\dots ,{\theta }_N$ remain in $\mathrm{\Theta }$ almost surely. Training uses only conditional samples and does not solve the parameterized ODE during these updates. After training, condition on the history ${\mathcal{F}}_N$, draw $X_0$ from a supplied initial probability law ${\mu }_0$ independently of ${\mathcal{F}}_N$, and return $X_1={\mathrm{\Phi }}_{1,0}^{{\theta }_N}(X_0)$, where ${\mathrm{\Phi }}_{t,s}^{{\theta }_N}$ solves $\dot{z}(t)=v(t,z(t);{\theta }_N)$ with $z(s)=X_0$. This last step is well-defined if the field extends to an open time interval containing $[0,1]$, is jointly $C^{\mathrm{\infty }}$ in time, state, and parameter, and on every compact parameter set $H\subset \mathrm{\Theta }$ has uniform constants $L_H,M_H\ge 0$ satisfying $\mathrm{\parallel }v(t,z;\theta )-v(t,w;\theta ){\mathrm{\parallel }}_2\le L_H\mathrm{\parallel }z-w{\mathrm{\parallel }}_2,\mathrm{\parallel }v(t,0;\theta ){\mathrm{\parallel }}_2\le M_H$ for $t\in [0,1]$, $\theta \in H$, and $z,w\in {\mathbb{R}}^d$. Conditional on ${\mathcal{F}}_N$, the output law assigns mass ${\mu }_0(\{x:{\mathrm{\Phi }}_{1,0}^{{\theta }_N}(x)\in A\})$ to each Borel set $A\subseteq {\mathbb{R}}^d$. The procedure gives no convergence guarantee for the parameter updates or equality of this output law with a desired data law.

Lipman et al., §3.2 and §6.2

Holderrieth and Erives, §3.3 Algorithm 3

Conditional batch-gradient unbiasedness and exact-flow output well-definedness under the full stated hypotheses. ↗

The cited papers specify the conditional sample, squared regression loss, parameter update, and ODE sampling steps. This algorithm makes the sampler law, batch gradient, exact global-flow conditions, and output law explicit. Its parameter updates carry no optimization convergence guarantee, and the exact flow is a mathematical solution rather than a numerical solver with an error bound.