On suppose que l'on possède une distribution source $p_0$ et une distribution de données $p_{data}$. On peut supposer que $p_0 = \mathcal{N}(0, I_d)$ sans perte de généralité. On considère également $N$ données d'entraînement $x^{(1)}, \dots, x^{(n)} \in \mathbb{R}^d$ tirées aléatoirement suivant $p_{data}$. On considère alors $\hat{p}_{data} = \frac{1}{N} \sum_{i=1}^N \delta_{X_i}$, la distribution empirique des données d'entraînement.
L'objectif du flow matching est de trouver un champ de vitesse $u: \mathbb{R}^d \times [0,1] \rightarrow \mathbb{R}^d$ qui transporte $p_0$ vers $\hat{p}_{data}$ de manière à ce que si $u$ est solution de l'équation différentielle $$ \begin{cases} x(0) \in \mathbb{R}^d \\ \dot{x}(t) = u(x(t), t) \end{cases} $$ Sur $[0, 1]$ alors si $x(0) \sim p_0$ alors $x(1) \sim p_{\text{data}}$. Une fois ce champ déterminé, il est possible de générer $x_1 \sim p_{data}$ en tirant aléatoirement $x_0 \sim p_0$ et en lui faisant suivre ce champ, par exemple par une méthode de Runge-Kutta.
On commence par définir une variable $z$ indépendante de $t$ ; par exemple, pour le flow matching par interpolation linéaire, $z \sim p_0 \times p_{data}$. Ensuite, on choisit un chemin $p(\cdot | t, z=(x_0, x_1))$. On peut par exemple choisir une gaussienne centrée sur la droite reliant $x_0$ à $x_1$ : $$ p(x | t, z) = \mathcal{N}((1-t)x_0 + t x_1, \sigma^2 I_d) $$ pour retrouver $p(x, t=0) = \int p(x | t, z) p(z) dz = p_0$, on doit avoir $\sigma=0$. On obtient donc finalement: $$ p(x | t, z) = \delta_{(1-t)x_0 + t x_1}(x) $$ On peut maintenant introduire l'équation de continuité, en notant $p_t = p(\cdot | t)$ et $u_t = u(\cdot | t)$ : $$ \partial_t p_t + \operatorname{div}(u_t\; p_t) = 0 $$ Cette équation traduit le fait que, si l'on a trouvé $u$, on est capable de pousser la matière sans rien perdre, de sorte que la densité se transforme en ce que l'on cherche. On va montrer que $u^{c}(x, t, z) = x_1 - x_0$ respecte cette équation, avec $p_t$ défini plus haut.
On pose $\mu(t) = (1-t)x_0 + t x_1$, donc $\dot{\mu}(t) = x_1 - x_0$. Alors : $$ \begin{align*} \partial_t p_t + \operatorname{div}(u^c \; p_t) &= \partial_t [\delta(x - \mu(t))] + p_t(\operatorname{div} u^{c}) + \nabla p_t \cdot u^{c} \\ &= \partial_t (x - \mu(t)) \cdot \nabla \delta (x - \mu(t)) + \langle u^c, \nabla \delta_{\mu(t)}(x) \rangle \\ &= - \dot{\mu}(t) \cdot \nabla \delta(x - \mu(t)) + \langle u^c, \nabla \delta_{\mu(t)}(x) \rangle \\ &= \langle u^c - (x_1 - x_0), \nabla \delta_{\mu(t)} (x) \rangle \\ &= 0 \end{align*} $$ D'où le résultat énoncé précédemment : c'est le théorème de marginalisation qui rend le flow matching possible. $$ \forall t, x \quad u^\star(x, t) = \mathbb{E}_{z | x, t} [u^c(x, t, z)] = \mathbb{E}_{z | x, t}[x_1 - x_0] $$ Ce théorème est très puissant puisqu'il relie le champ $u^\star(x, t)$ que l'on cherche au champ conditionnel $u^c$, qui est très simple à calculer.
La flèche noire est $u(x, t)$, les traits rouges sont $x_1 - x_0$.
Attention, on pourrait croire que l'on a gagné, que l'on sait définir $u(x, t)$ en chaque point et donc le champ de vecteurs mais ceci est faux. En effet pour calculer $u(x, t)$ il nous faut trouver tous les couples $(x_0, x_1)$ qui passent par $(x, t)$ donc itérer sur tout le dataset, à chaque itération de Runge-Kutta. C'est impraticable.
La solution est donc d'approcher $u$ par un réseau de neurones $u_\theta$ en minimisant la perte $$ \mathcal{L}_{\text{FM}}(\theta) = \mathbb{E}_{\substack{t \sim \mathcal{U}([0,1]) \\ x_t \sim p(\cdot | t)}} \| u_\theta(x_t, t) - u^\star(x_t, t) \|^2 $$ Le calcul de cette perte sous-entend le calcul de $u^\star$ qui est coûteux. Pourtant il est possible de montrer que minimiser $\mathcal{L}_{\text{FM}}(\theta)$ équivaut à une constante près à minimiser $$ \mathcal{L}_{\text{CFM}}(\theta) = \mathbb{E}_{\substack{x_0 \sim p_0 \\x_1 \sim p_{\text{data}} \\ t \sim \cal{U}([0,1])}} \| u_\theta(x_t, t) - u^c(x_t, z, t) \|^2 $$ où $x_t = (1- t) x_0 + t x_1$. Le calcul de $\mathcal{L}_{\text{CFM}}(\theta)$ est facile en approximant le choix de $x_1$ dans $p_{\text{data}}$ par le choix de $x_1$ dans $\hat{p}_{\text{data}}$