| Russian Mathematical Surveys |
|
|
|
|
|
Mathematical aspects of artificial intelligence On the inverse problem of flow matching in the one-dimensional and Gaussian cases A. A. Korotina, G. Pammerb a Skolkovo Institute of Science and Technology, Artificial Intelligence Research Institute b Graz University of Technology, Graz, Austria
Bibliographic databases:
    We study the inverse flow matching problem (FM [4]) between distributions $p_0, p_1 \in \mathcal{P}_{\exp}(\mathbb{R}^D)$ on $\mathbb{R}^D$ with finite exponential moment, that is, such that $\displaystyle\int \exp(\lambda \|x\|_2)\,d[p_0+p_1](x)<\infty$ for some $\lambda>0$. For this problem, we establish the uniqueness of its solution in two cases: the one-dimensional one ($D=1$) and the Gaussian one. We start by describing the forward problem of FM [4]. Let $\pi\in \Pi(p_0,p_1)$ be a transport plan between $p_0$ and $p_1$, that is, a joint distribution on $\mathbb{R}^{D}\times\mathbb{R}^{D}$ whose marginals are $p_0$ and $p_1$, respectively. Let $(X_0,X_1)\sim \pi$ be some random variable raking values in в $\mathbb{R}^{D}\times\mathbb{R}^{D}$ with distribution $\operatorname{Law}((X_0,X_1))=\pi$. For each $t\in(0,1)$ define the random variables $X_{t}=(1-t)X_0+tX_1$ and set $p_t^{\pi}=\operatorname{Law}(X_t)$, that is, $p_t^{\pi}$ is the distribution of $X_{t}$. The general goal of the flow matching problem is to find a velocity field $v\colon[0,1]\times\mathbb{R}^{D}\to \mathbb{R}^{D}$ such that transporting the mass of $p_0$ along it generates the sequence $\{p_{t}^{\pi}\}$, that is, $dp_t^{\pi}(x_t)/dt= {-}\mathrm{div}\bigl(p_t^{\pi}(x_t)v_t(x_t)\bigr)$ for $t\in [0,1]$. The idea of FM underlies state-of-the-art generative AI methods [1]. In general, the problem may not have a unique solution. Therefore, one typically seeks one specific solution which is relatively easy to find numerically in practice by solving a regression problem [4], Formula 1, and it satisfies the properties required above [4], § 2.2. Thus, the forward FM problem is associated precisely with finding this velocity field $v^{\pi}$ given the plan $\pi$.Recently the inverse problem of FM began to attract interest. Given $p_0$ and $,p_1$, as well as $\{p_t^{\pi}\}_{t\in (0,1)}$ and $v^{\pi}$ obtained via FM with some plan $\pi\in \Pi(p_0,p_1)$, the task is to find the plan $\pi$ itself. The reader interested in the practical motivation for solving the inverse problem are referred to [2], where numerically efficient generative models based on FM are developed by solving the inverse problem. Note that in practice, only $p_0$ and $v^{\pi}$ are considered as input data, since the remaining $\{p_{t}^{\pi}\}_{t\in (0,1]}$ are reconstructed using the continuity equation. It is still theoretically unknown whether the solution $\pi$ to the inverse problem is unique. Before studying this question, let us note one nuance. The velocity field $v^{\pi}_t$ must be considered only up to values outside the support of the distribution $p_t^{\pi}$, as they do not play a role in the continuity equation. Consider the case $D=1$. Theorem 1 (uniqueness of the solution to the inverse FM problem in the one-dimensional case). Consider two distributions $p_0,p_1\in\mathcal{P}_{\rm exp}(\mathbb{R}^{D})$. Let $\pi,\pi'\in \Pi(p_0,p_1)$ be two transport plans between $p_0$ and $p_1$. If $p_{t}^{\pi}=p_{t}^{\pi'}$ for all $t\in[0,1]$, then $\pi=\pi'$. Proof. Consider $\phi_{\pi}(\xi_0,\xi_1):=\mathsf{E}e^{i\xi_0 X_0+i\xi_1 X_1}$, which is the characteristic function of the distribution $\pi$. We will show that it is uniquely determined by the sequence $p_{t}^{\pi}$, $t\in [0,1]$. Indeed, knowing $p_{t}^{\pi}=\operatorname{Law}(X_t)$ we uniquely recover $\mathsf{E}e^{i\xi_{t}X_{t}}=\mathsf{E}e^{i\xi_{t}[(1-t)X_0+tX_1]}$ for any $t\in [0,1]$ and $\xi_{t}\in \mathbb{R}$. By varying $t$ and $\xi_{t}$ within the admissible range , we can obtain all expressions of the form $\mathsf{E}e^{i\xi_0 X_0+i\xi_1 X_1}$, where $\xi_0,\xi_1\in\mathbb{R}_{+}$, that is, we recover $\phi_{\pi}(\xi_0,\xi_1)$ on $\mathbb{R}_+ \times \mathbb{R}_+$. Since $p_0$ and $p_1$ have exponential moments, $\phi_{\pi}$is analytic in a neighborhood of $(0,0)$ and extends uniquely to $\mathbb{R}^{2}$ ([3], Corollary 2.3.8). Thus, from the sequence $\{p_t^{\pi}\}$ we uniquely recover $\phi_{\pi}$ and $\pi$. Since $p_{t}^{\pi}=p_{t}^{\pi'}$ for all $t$, we conclude that $\pi=\pi'$, as their characteristic functions are determined identically by the sequence $\{p_t^{\pi}\}$. $\Box$
In the proof we relied on information about the sequence of distributions $p_{t}^{\pi}$ and did not use the data about the field $v^{\pi}$ from FM. It might seem that knowing the vector field is redundant. However, when attempting to generalize the proof to the multivariate case ($D>1$), a limitation is revealed: the characteristic function for $\pi$ is recovered only at points $(\xi_0,\dots,\xi_0,\xi_1,\dots,\xi_1)\in \mathbb{R}^D\times\mathbb{R}^D$, where each $\xi_0,\xi_1\in \mathbb{R}_{+}$ is repeated $D$ times. This subset, being embedded in a two-dimensional plane in the space $\mathbb{R}^D\times\mathbb{R}^D$, is not open, which does not guarantee a unique extension of the characteristic function to the entire space. An illustrative example of such an ambiguity in recovering $\pi$ is the Gaussian case. Let $p_0=\mathcal{N}(\mu_0,\Sigma_0)$ and $p_1=\mathcal{N}(\mu_1,\Sigma_1)$ be Gaussian distributions with means $\mu_0,\mu_1\in \mathbb{R}^{D}$ and covariances $0\prec\Sigma_0,\Sigma_1\in\mathbb{R}^{D\times D}$, respectively, Consider the plan
$$
\begin{equation*}
\pi=\mathcal{N}\biggl(\begin{pmatrix}\mu_0 \\ \mu_1\end{pmatrix},\begin{pmatrix}\Sigma_0 & S \\ S^\top & \Sigma_1\end{pmatrix}\biggr)\in \Pi(p_0,p_1), \quad \text{where} \ \ S\in \mathbb{R}^{D\times D}.
\end{equation*}
\notag
$$
If $(X_0,X_1)\sim\pi$ and $X_{t}:=(1-t)X_0+tX_1$, then the triple $(X_0,X_t,X_1)$ has a Gaussian distribution. Furthermore,
$$
\begin{equation*}
\operatorname{Law}(X_{t})=p_{t}^{\pi}= \mathcal{N}(\mu_{t},\Sigma_t),
\end{equation*}
\notag
$$
where
$$
\begin{equation*}
\mu_{t}:=(1-t)\mu_0+t \mu_1 \quad \text{and}\quad \Sigma_{t}:=(1-t)^{2}\Sigma_0+t^{2}\Sigma_1+t(1-t)(S+S^{\top}).
\end{equation*}
\notag
$$
Note that the sequence $\{p_t^{\pi}\}$ is the same for all Gaussian plans $\pi$ that share the symmetric part of the matrix $S$, that is, $\frac{1}{2}(S + S^{\top})$. In the one-dimensional case ($D=1$ and $S \in \mathbb{R}^{1\times 1}$) we always have $S\!+\!S^{\top}\!=\!2S$, so from $p_{t}^{\pi}$ we uniquely obtain $S$ and the plan $\pi$. However, for $D>1$ we can construct Gaussian $\pi,\pi'\in \Pi(p_0,p_1)$ for which $S+S^{\top}=S'+(S')^{\top}$, but $S\ne S'$. From this observation it follows that information about the sequence $p_{t}^{\pi}$ alone is insufficient to precisely find the original $\pi$. Note that the additional assumption of having information about $v^{\pi}$ changes the situation. Theorem 2 (uniqueness of the Gaussian solution to the inverse FM problem in the multivariate Gaussian case). Let $p_0=\mathcal{N}(\mu_0,\Sigma_0)$ and $p_1=\mathcal{N}(\mu_1,\Sigma_1)$ be two Gaussian distributions on $\mathbb{R}^{D}$, and let $\pi,\pi'\in\Pi(p_0,p_1)$ be Gaussian plans with cross-covariance components $S$ and $S'$, respectively. If for $t=0$ we have $v_{0}^{\pi}=v_{0}^{\pi'}$, then also $\pi=\pi'$. Proof. Note that
$$
\begin{equation*}
v^{\pi}_{0}(x_0)=\mathsf{E}[X_1-X_0\mid X_0=x_0]= \mathsf{E}[X_{1}\mid X_{0}= x_0]-x_0= \mu_{1}+S^{\top}\Sigma_{0}^{-1}(x_{0}-\mu_{0})- x_0.
\end{equation*}
\notag
$$
Similarly, we obtain
$$
\begin{equation*}
v^{\pi'}_{0}(x_0)=\mu_{1}+(S')^{\top}\Sigma_{0}^{-1}(x_{0}-\mu_{0})- x_0.
\end{equation*}
\notag
$$
Equating these linear functions of $x_0$ we notice that $S^{\top}=(S')^{\top} \Rightarrow \pi=\pi'$. $\Box$ Theorem 2 shows that in the Gaussian case, to solve the inverse problem it is sufficient to know only the initial velocity $v_{0}^{\pi}$ — the intermediate values ${p_{t}^{\pi}}$ and $v_{t}^{\pi}$ for $t > 0$ are not needed. However, this theorem considers only solutions $\pi$ belonging to the Gaussian class. It is still unknown whether other solutions outside this class exist. Conclusion. We have considered important cases of the inverse flow matching problem, but general conditions for the uniqueness of the solution in the multivariate problem remain an open question. Its solution is important both for theory and for practice — for example, for a rigorous justification of modern FM-based generative AI methods.
Citation in format AMSBIB
\Bibitem{KorPam25} |
|
|||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
Contact us: |
Terms of Use
|
Registration to the website |
Logotypes |
|