Rectified Flow
Overview
数据分布为 \(p(x_0)\) \(p(x_1)\) ） ，我们需要在这两个分布之间构建映射。简而言之，Rectified Flow 通过
\[
\min_{\theta} \int_0^1 \mathbb{E}_{x_0, x_1}\left[\left\|(x_1 - x_0) - v_{\theta}(x_t, t)\right\|^2\right]\mathrm{d}t,\quad x_t = (1-t)x_0 + tx_1
\]
学习一个向量场 \(v_{\theta}(x_t, t)\) Rectified Flow (well-defined) ODE:
\[
\frac{\mathrm{d}z_t}{\mathrm{d}t} = v(z_t, t)
\]
其中 \(p(z_0)=p(x_0)\) \(p(z_1)=p(x_1)\) 。给定 \(z_0=x_0\) ODE \(z_1\)
直观理解
直观来看，\(x_0\) \(x_1\) \(z_0\) \(z_1\) \(x_0\) \(x_1\) \(\{z_t\}_{[0, 1]}\) \(\{z_t\}_{[0, 1]}\) \(\{x_t\}_{[0, 1]}\) Rectified Flow 。
Marginal Preserving
正如 DDIM Markov non-Markov \(p(x_t)\) Rectified Flow 希望“走直线”，也需要保证边缘分布不变：
\[
p(z_t) = p(x_t), \quad \forall t \in [0, 1]
\]
从希望高质量地跨步去噪的角度来看，Rectified Flow 和 DDIM
probability density function v.s. probability law
原论文 使用 Law 而不是概率密度函数 进行描述，对于两者之间的区别，我截取了 wikipedia Random variable
In measure-theoretic terms, we use the random variable \(X\) to "push-forward" the measure \(P\) on \(\Omega\) to a measure \(p_{X}\) on \(\mathbb {R}\) . The measure \(p_{X}\) is called the "(probability) distribution of \(X\) " or the "law of \(X\) ". The density \(f_{X}=\mathrm{d}p_{X}/\mathrm{d}\mu\) , the Radon–Nikodym derivative of \(p_{X}\) with respect to some reference measure \(\mu\) on \(\mathbb {R}\) (often, this reference measure is the Lebesgue measure in the case of continuous random variables, or the counting measure in the case of discrete random variables).
也就是说，Law 只是将概率测度 \(P: \Omega\to \mathbb{R}\) \(p_X: \mathcal{B}(\mathbb{R})\to \mathbb{R}\)
\(\mathcal{B}(\mathbb{R})\) \(\mathbb{R}\) Borel \(\mathbb{R}\) \(\sigma\) -algebra
probability law of a random variable
对概率空间 \((\Omega, \mathcal{F}, P)\) \(X: \Omega \to \mathbb{R}\) probability law \(p_{X}\)
\[
p_{X}(B) = P(X \in B), \quad \forall B \in \mathcal{B}(\mathbb{R})
\]
为了方便，我们假定讨论连续随机变量，直接使用概率密度函数描述概率分布
先引入必要的定义：
Expected Velocity
对沿路径连续可微 (path-wise continuously differentiable) \(\{x_t\}_{[0, 1]}\) (expected velocity)
\[
v^x(x, t) =
\begin{cases}
\mathbb{E}[\dot{x}_t | x_t = x], & x \in \mathrm{supp}(x_t) \\
0, & \text{otherwise}
\end{cases}
\]
随机变量的支撑集 (support) \(\{x\in \mathbb {R} :P(X=x)>0\}\) \(\{x\in \mathbb {R} :f_{X}(x)>0\}\) Support (mathematics) - Wikipedia
Rectifiable
当满足以下条件时，我们把随机过程 \(\{x_t\}_{[0, 1]}\) Rectifiable 的
\(v^x\) (locally bounded)通过如下积分式，可以解得唯一的 \(z_t\)
\[
z_t = x_0 + \int_0^t v^x(z_s, s)\mathrm{d}s, \quad \forall t \in [0, 1], \quad z_0 = x_0
\]
由此，\(\{z_t\}_{[0, 1]}\) \(\{x_t\}_{[0, 1]}\) Rectified Flow 。
接下来正式地描述 Rectified Flow
Marginal Preserving
假设 \(\{x_t\}_{[0, 1]}\) Rectifiable \(\{z_t\}_{[0, 1]}\) Rectified Flow
\[
p(z_t) = p(x_t), \quad \forall t \in [0, 1]
\]
原论文中是 \(\mathrm{Law}(z_t) = \mathrm{Law}(x_t)\)
Proof
将会使用测试函数法证明，首先简单介绍测试函数法。
测试函数法
测试函数法推导连续性方程和 Fokker-Planck - 中简单地给出了测试函数法的基本原理：如果对于任意函数 \(\phi(x)\)
\[
\int f(x)\phi(x)\mathrm{d}x
= \int g(x)\phi(x)\mathrm{d}x
\]
那么就成立 \(f(x)=g(x)\) \(\phi(x)\)
然而还存在一些细节，\(\phi(x)\) / / ） ？参考 Weak solution - Wikipedia Rectified Flow
注意，在 SGM Fokker-Planck Fokker-Planck 测试函数法推导连续性方程和 Fokker-Planck - ） ，因此在基本证明路径上其实是相似的。
测试函数法证明连续性方程 (continuity equation)
连续性方程作为 F-P ODE SGM F-P Rectified Flow ） 。
在这里将测试函数法简单应用于连续性方程的证明（来自测试函数法推导连续性方程和 Fokker-Planck - ） ，首先阐明 ODE:
\[
\frac{\mathrm{d}x_t}{\mathrm{d}t} = f_t(x_t)
\]
泰勒展开、舍去高阶项：
\[
\phi(x_{t + \Delta t}) = \phi(x_t + f_t(x_t)\Delta t) \approx \phi(x_t) + \Delta t\cdot f_t(x_t) \nabla_{x_t}\phi(x_t)
\]
两边取期望有
\[
{ \color{red}
\int \phi(x_{t + \Delta t}) p_{t + \Delta t}(x_{t + \Delta t})\mathrm{d}x_{t + \Delta t}
} \approx
\int \left[
{ \color{orange} \phi(x_t) } +
{ \color{skyblue}
\Delta t\cdot f_t(x_t) \nabla_{x_t}\phi(x_t)
}
\right] p_t(x_t)\mathrm{d}x_t
\]
注意到有
\[
\int \phi(x_{t + \Delta t}) p_{t + \Delta t}(x_{t + \Delta t})\mathrm{d}x_{t + \Delta t}
= { \color{red}
\int \phi(x_t) p_{t + \Delta t}(x_t)\mathrm{d}x_t
}
\]
另一方面，根据分部积分有
\[
\int p_t(x_t) f_t(x_t) \nabla_{x_t}\phi(x_t) \mathrm{d}x_t
= - {\color{skyblue}
\int \nabla_{x_t}\left[p_t(x_t) f_t(x_t) \right] \phi(x_t)\mathrm{d}x_t
}
\]
测试函数 \(\phi\)
高维空间分部积分
在高维空间中，分部积分的公式为
\[
\int _{\Omega} v\nabla u \mathrm{d}x = \int _{\partial\Omega} u v\cdot \hat{n} \mathrm{d}S - \int _{\Omega} u \nabla\cdot v \mathrm{d}x
\]
其中 \(u\) \(x\) \(v\) \(x\)
\(\Omega\) : , \(\partial\Omega\) \(\hat{n}\) : \(\mathrm{d}S\)
Proof
基于
\[
\nabla(u v) = u \nabla v + v \nabla u
\]
我们有
\[
\int _{\Omega} v\nabla u \mathrm{d}x = \int _{\Omega} \nabla(u v) \mathrm{d}x - \int _{\Omega} u \nabla\cdot v \mathrm{d}x
\]
根据高斯散度定理 (divergence theorem) ，有
\[
\int _{\Omega} \nabla(u v) \mathrm{d}x = \int _{\partial\Omega} u v\cdot \hat{n} \mathrm{d}S
\]
特别地，我们考虑概率密度函数 \(p(x)\) \(u=p(x)\) \(v=\nabla p(x)\) \(\Omega\) \(p(x)\to 0\) \(\nabla p(x)\to 0\) 0
\[
\begin{aligned}
\int v \nabla p \mathrm{d}x
&= -\int p \nabla\cdot v \mathrm{d}x \\
\int u \nabla\cdot \nabla p \mathrm{d}x
&= -\int \nabla p \cdot \nabla u \mathrm{d}x
\end{aligned}
\]
上面给了一个比较模糊的“无穷远处”“趋向于 0 \(p\) ） ，这样才能让面积积分为 0
于是
\[
\begin{aligned}
&\int [
{\color{red} p_{t + \Delta t}(x_t)}
- {\color{orange} p_t(x_t)}
]\phi(x_t)\mathrm{d}x_t
= - {\color{skyblue}
\Delta t
\int \nabla_{x_t}\left[p_t(x_t) f_t(x_t) \right] \phi(x_t)\mathrm{d}x_t
} \\
\Rightarrow & \int \frac{\partial p_t(x_t)}{\partial t} \phi(x_t)\mathrm{d}x_t = -\int \nabla_{x_t}\left[p_t(x_t) f_t(x_t)\right] \phi(x_t)\mathrm{d}x_t \\
\Rightarrow & \frac{\partial p_t(x_t)}{\partial t} = -\nabla_{x_t}\left[p_t(x_t) f_t(x_t)\right]
\end{aligned}
\]
先让 \(\Delta t\to 0\)
Rectified Flow 最关键的设计就是 \(x_t=\varphi_t(x_0, x_1)\) 可逆 变换，可以与 normalizing flow \(\varphi_t(x_0, x_1)\) \(\varphi_t\)
在 Rectified Flow \(\varphi_t(x_0, x_1)=(1-t)x_0 + tx_1\)
在 \(\varphi_t\) \(x_1\) \(x_1 = \psi_t(x_0, x_t)\) \(\varphi_t(x_0, x_1)\) \(\varphi_t(x_0, x_t) = \varphi_t(x_0, \psi_t(x_0, x_t))\)
\[
\begin{aligned}
{\color{red} v^x(x_t, t)}
&= \mathbb{E}[\dot{x_t}\mid x_t] \\
&= \mathbb{E}_{x_0, x_1\mid x_t}[\dot{\varphi_t}(x_0, x_1)] \\
&= \mathbb{E}_{x_0\mid x_t}[\dot{\varphi_t}(x_0, x_t)]
\xrightarrow{\text{simplify symbol}}
{\color{red}
\mathbb{E}_{x_0\mid x_t}[\dot{\varphi_t}]
} \\
\end{aligned}
\]
注意到，在测试函数法证明连续性方程时，实际上是利用了
\[
\mathbb{E}_{x_t + \Delta t}\left[\phi(x_{t + \Delta t})\right]
= \mathbb{E}_{x_t}\left[\phi(x_t + f_t(x_t)\Delta t)\right]
\]
也就是说，我们也只需证到这一步即可证明连续性方程 ，也就证明边缘分布一致了 。对于任意紧支撑 ( compact support ) \(\phi(x): \mathbb{R}^n\to \mathbb{R}\)
\[
\begin{aligned}
\mathbb{E}_{x_{t + \Delta t}}\left[
\phi({\color{orange}
x_{t + \Delta t}
})
\right]
&= \mathbb{E}_{x_0, x_1}\left[
\phi({\color{orange}
\varphi_{t + \Delta t}
})
\right] & {\color{orange}
x_{t+\Delta t} = \varphi_{t + \Delta t}(x_0, x_1)
} \\
&\approx \mathbb{E}_{x_0, x_1}\left[
\phi({
\color{green} \varphi_t
}) +
\Delta t
{\color{green}
\frac{\partial \varphi_t}{\partial t}
} \cdot
\nabla_{\varphi_t} \phi({\color{green}
\varphi_t
})
\right] & \text{Taylor Expansion to }O(\Delta t)\\
&= \mathbb{E}_{\color{skyblue} x_0, x_1}\left[
\phi({\color{green}
x_t
}) +
\Delta t
{\color{green}
\dot{\varphi_t}
} \cdot
\nabla_{\varphi_t} \phi({\color{green}
x_t
})
\right] & { \color{green}
x_t = \varphi_t(x_0, x_1)
}\\
&= \mathbb{E}_{\color{skyblue} x_t}[\phi(x_t)]
+ \Delta t
\mathbb{E}_{\color{skyblue} x_0, x_t}\left[
\dot{\varphi_t}
\nabla_{\varphi_t} \phi(x_t)
\right] & { \color{skyblue}
\varphi_t(x_0, x_1) \to \varphi_t(x_0, x_t)
}\\
&= \mathbb{E}_{x_t}[\phi(x_t)]
+ \Delta t
\mathbb{E}_{x_t}\left[
{ \color{red}
\mathbb{E}_{x_0\mid x_t}\left[
\dot{\varphi_t}
\right]
}
\nabla_{\varphi_t} \phi(x_t)
\right] \\
&= \mathbb{E}_{x_t}[\phi(x_t)]
+ \Delta t
\mathbb{E}_{x_t}\left[
{ \color{red}
v^x(x_t, t)
}
\nabla_{\varphi_t} \phi(x_t)
\right] & { \color{red}
v^x(x_t, t) = \mathbb{E}_{x_0\mid x_t}\left[
\dot{\varphi_t}
\right]
}\\
&= \mathbb{E}_{x_t}\left[
\phi(x_t)
+ \Delta t
v^x(x_t, t)
\nabla_{\varphi_t} \phi(x_t)
\right] \\
&\approx \mathbb{E}_{x_t}\left[
\phi\left(
x_t
+ \Delta t
v^x(x_t, t)
\right)
\right] & \text{Reverse Taylor Expansion}
\end{aligned}
\]
对任意测试函数都成立。于是我们证明了连续性方程，有
\[
x_{t + \Delta t} = x_t + \Delta t v^x(x_t, t) \Rightarrow
\frac{\mathrm{d}x_t}{\mathrm{d}t} = v^x(x_t, t)
\]
根据 \(v^x\) \(v^x\) \(v_{\theta}\)
\[
\min_{\theta} \int_0^1 \mathbb{E}_{x_0, x_1}\left[\left\|\dot{\varphi_t} - v_{\theta}\right\|^2\right]\mathrm{d}t,\quad \dot{\varphi_t}(x_0, x_1) = x_1 - x_0
\]
注意到我们这里的证明全程没有提及 \(z_t\) \(z_t\) Rectifiable ODE \(f_t(x_t)\) \(v^x(x_t, t)\)
\[
\frac{\partial p_t(x_t)}{\partial t} = -\nabla_{x_t}\left[p_t(x_t) v^x(x_t, t)\right]
\]
也就是说，我们训练拟合的最优解 \(v^x\)
从更直接的角度理解，我们实际上证明的是从任意的 \(p(z_t)=p(x_t)\) \(p(z_{t + \Delta t}) = p(x_{t + \Delta t})\) \(p(z_0)=p(x_0)\) \(p(z_t)=p(x_t)\)
这个保持边缘分布一致的性质非常重要，实际上保证了从 \(z_0\) \(v_{\theta}(x_t, t)\) ODE \(z_1\) \(p(z_1)=p(x_1)\) Rectified Flow