梯度流

欧式空间

如果我们想要求解 $f(x)$ 的最小值，使用梯度下降法可知： $$x_{t+1}=x_t-\alpha\nabla_{x_t} f(x_t)\tag{3.1}$$
其中 $\alpha$ 为学习率。如果使学习率趋向于0并记为 $\mathrm{\Delta }t$ ， $x_{t+1}$ 记为 $x_{t+\mathrm{\Delta }t}$ ，那么有： $$\frac{dx}{dt}=-\nabla_{x}f(x)\tag{3.2}$$
则（3.2）式称为梯度流。将（3.2）写成逆向Euler形式（ODE前向和逆向是相等的），有： $$\begin{array}{c} x_{t+1}=x_t-\alpha\nabla_{x_{t+1}} f(x_{t+1})\ \end{array}\tag{3.3}$$
那么 $$\nabla_{x} \left(\frac{|x-x_t|^2}{2\alpha}+ f(x)\right)\bigg|{x=x{t+1}} =0\tag{3.4}$$
而（3.4）式可以解释为 $x_{t+1}$ 是目标函数： $$\frac{|x-x_t|^2}{2\alpha}+f(x)\tag{3.5}$$
的最优值。即： $$x_{t+1}=\operatorname*{argmin}_{x}\left{\frac{|x-x_t|^2}{2\alpha}+f(x)\right}\tag{3.6}$$

这告诉我们，梯度下降法每一步的优化，相当于在给定以之前参数 $x_t$ 为圆心的欧式距离约束下，找到下降最快的方向。显然，我们可以把（3.6）推广到更一般的情况： $$x_{t+1}=\operatorname*{argmin}_{x}\left{\frac{d^2(x,x_t)}{2\alpha}+f(x)\right}\tag{3.7}$$
其中， $d$ 是任意的距离度量。那么（3.7）式就表示欧式空间中在 $d$ 度量下对 $f(x)$ 优化的梯度流。

概率空间

更一般的，我们还可以推广到概率空间中在某个概率度量下对泛函优化的梯度流，即把 $x$ 换成概率分布 $p(x)$ ，把 $d$ 换成概率距离（如KL散度，W距离等）， $f(x)$ 换成概率的泛函 $\mathcal{F}(p(x))$ 。特别的，如果我们想研究W-2度量下的梯度流，那么就是需要求解： $$p_{t+1}=\operatorname*{argmin}_{p}\left{\frac{W_2^2(p,p_t)}{2\alpha}+\mathcal{F}(p)\right}\tag{3.8}$$
求解右边的泛函的极值，需要用到变分导数和Euler-Lagrange方程的概念，这里简单介绍一下：

假设 $y(x)$ 在 $x=a$ 和 $x=b$ 处为定点，考虑如下形式的泛函： $$I=\int_a^bf(y,y')dx\tag{3.9}$$
我们的目的是找到某个 $y(x)$ 使得这个积分最大，假设 $y$ 有微小的改变 $\delta y$ ，将 $\delta y$ 叫做 $y$ 的变分，那么 $f$ 的变分为： $$\delta f=\frac{\partial f}{\partial y}\delta y + \frac{\partial f}{\partial y'}\delta y' \tag{3.10}$$
严格的证明（3.10）需要用到变分关于测试函数的定义和线性的性质，但它与微分的运算性质相近，很容易推广和理解，所以为了简便就不详细展开了。那么 $I$ 的变分为： $$\delta I=\int_a^b \left(\frac{\partial f}{\partial y}\delta y + \frac{\partial f}{\partial y'}\delta y'\right)dx \tag{3.11}$$
对于积分里第二项： $$\begin{align} \int_a^b \frac{\partial f}{\partial y'}\delta y' dx&=\int_a^b \frac{\partial f}{\partial y'}\delta \left(\frac{dy}{dx}\right) dx\ &=\int_a^b \frac{\partial f}{\partial y'} \frac{d(\delta y)}{dx} dx & \text{变分与微分互换顺序}\ &=\frac{\partial f}{\partial y'}\delta y\bigg|_a^b-\int_a^b \frac{d}{dx}\left(\frac{\partial f}{\partial y'} \right)\delta y & \text{分部积分}\ &=-\int_a^b \frac{d}{dx}\left(\frac{\partial f}{\partial y'} \right)\delta y & \text{端点为定点} \end{align}\tag{3.12}$$
带入（3.11）有： $$\delta I=\int_a^b \left[\frac{\partial f}{\partial y} -\frac{d}{dx}\left(\frac{\partial f}{\partial y'} \right)\right]\delta ydx \tag{3.13}$$
因此，如果 $I$ 有极值，那么 $\delta I=0$ ，则： $$\frac{\partial f}{\partial y} -\frac{d}{dx}\left(\frac{\partial f}{\partial y'} \right)=0\tag{3.14}$$
（3.14）就叫做Euler-Lagrange方程，是变分法里面最重要的方程之一。并且因为（3.14）的解与 $I$ 的极值能对应上，所以从某种程度上可以把 $I$ 的变分导数看做是（3.14）的左边，记作： $$\frac{\delta I}{\delta y} = \frac{\partial f}{\partial y} -\frac{d}{dx}\left(\frac{\partial f}{\partial y'} \right)\tag{3.15}$$
可以发现，对变分的导数只需要对积分里面的函数做运算即可，不用考虑积分号。

为了求解（3.8）这样的泛函，我们同样可以对其求导应用Euler-Lagrange方程： $$\frac{1}{2\alpha}\frac{\delta W_2^2(p,p_t)}{\delta p}+\frac{\delta \mathcal{F}(p)}{\delta p}=0\tag{3.16}$$
根据之前提到的W距离的对偶形式，有： $$\begin{align} W_2^2(p,p_t)&=\max_{f} \left{\int [p(x)f(x) -p_t(x)f(x)] dx \mid f(x) - f(y) \leq |x-y|^2\right}\ &= \int [p(x)f^(x) -p_t(x)f^(x)] dx \end{align} \tag{3.17}$$
假设 $f^{\ast }$ 是满足约束的最优解，那么（3.16）式就转化为了： $$\frac{f^}{2\alpha}+\frac{\delta \mathcal{F}(p)}{\delta p}=0\tag{3.18}$$
为了继续化简，需要引入Monge和Kantorovich问题的概念。事实上，之前介绍的推土机距离问题（1.8）为Kantorovich问题，而如果要求从 $x$ 搬走的土都填到 $T(x)$ 上，那么就变为了Monge问题： $$\begin{align} W[\gamma^]&=\inf_{T,T(p)=q} \int \gamma(x,T(x))c(x,T(x))dx\ &=\inf_{T,T(p)=q} \int p(x)c(x,T(x))dx\ \end{align} \tag{3.19}$$
$T$ 也叫做 $p(x)$ 到 $q(x)$ 的传输映射。可以看到，Kantorovich问题是Monge问题更松弛的情况，显然它的最优解也一定要小于等于Monge问题的最优解。不过在历史上，Monge问题先被提出，但是难以给出解的存在性和唯一性。直到1940年Kantorovich才将这个问题推广到更加松弛的情况下，得到了解决Monge问题的许多定理。Monge问题在狄拉克分布到任意一连续分布的情况下是无解的，但是在Kantorovich问题中有解；而如果是连续分布到连续分布，这种情况下Monge和Kantorovich问题是等价的（实际上有一定条件，但我们实际应用考虑的情况都是符合的），也就是说通过求解Kantorovich问题，可以得到一个传输映射 $T^{\ast }(x)$ 。如何通过Kantorovich问题的解来构造出Monge问题的解呢？其实关键在于（2.9）式的操作，这里换个记号重写一下： $$f^{}(x)=\min_{y} {c(x,y)-g(y) } \tag{3.20}$$
在OT的研究中，这样的操作也称为c-transform。对于新解 $f^{\ast }(x)$ 和 $g(y)$ ， $y$ 是由 $x$ 唯一确定的，那么这样的新的传输方案 $\gamma (x,y)$ 就一定属于 $\gamma (x,T(x))$ ，也就确定了Monge问题的解。事实上，在动力学中，Monge问题的解才是我们更应该关心的，传输映射 $T$ 能够将 $p(x)$ 映射到 $q(x)$ ，这对我们研究运动轨迹对概率的变化，即概率流，具有很重要的意义。由于我们现在考虑的是W-2距离，假设对应Monge问题的最优解为 $y^{\ast }=T^{\ast }(x)$ ，令（3.20）两边对 $x$ 求梯度，则： $$\begin{align} \nabla_x f^{}(x)&=\nabla_x {c(x,T^(x))-g(T^(x)) }\ &=\nabla_x c(x,T^(x))\ &=\nabla_x |x-T^(x)|^2\ &=2(x-T^*(x)) \end{align}\tag{3.21}$$

关于对最优化求导，主要应用了信封定理（Envelope Theorem），这里简单介绍一下：如果我们需要优化 $f(x)=\underset{y}{min}g(x,y)$ ，假设最优值为 $y^{\ast }(x)$ ，那么 $\frac{df(x)}{dx}=\frac{\partial f(x,y^{\ast }(x))}{\partial x}+\frac{\partial f(x,y^{\ast }(x))}{\partial y}\frac{d{y}^{\ast }(x)}{dx}$ ，由于 $y^{\ast }$ 是最优值，那么关于它的偏导应该为0（假设 $y$ 是连续的），那么 $\frac{df(x)}{dx}=\frac{\partial f(x,y^{\ast }(x))}{\partial x}$ 。也就是说对于优化问题 $\underset{y}{min}g(x,y)$ ，只需要考虑 $x$ 的偏导即可。

（3.21）式揭示了Monge和Kantorovich问题最优解之间的联系，其中更一般的第二个等号： ${\mathrm{\nabla }}_x{f}^{\ast }(x)={\mathrm{\nabla }}_{x}c(x,T^{\ast }(x))$ 也叫做Brenier定理。

回到（3.18）式，两边同时对 $x$ 求梯度： $$\frac{\nabla_xf^}{2\alpha}+\nabla_x\frac{\delta \mathcal{F}(p)}{\delta p}=0\tag{3.22}$$
应用（3.21）进行代换并整理： $$\frac{T^(x)-x}{\alpha}=\nabla_x\frac{\delta \mathcal{F}(p)}{\delta p}\tag{3.23}$$
类似（3.1），当 $\alpha$ 趋向于0时， $\frac{T^{\ast }(x)-x}{\alpha }$ 可以看做 $dx/dt$ ，所以（3.32）可以看成概率流ODE。但要注意的是，（3.23）的概率流向是最优值 $p_{t+1}$ 流向 $p_t$ ，因为（3.16）的定义是 $p$ 到 $p_t$ 。因此 $p_t$ 到 $p_{t+1}$ 的概率流ODE为： $$\frac{dx}{dt}=-\nabla_x\frac{\delta \mathcal{F}(p)}{\delta p}\tag{3.24}$$
那么（3.24）的FP方程为： $$\frac{\partial}{\partial t}p_t=\nabla_x\cdot\left[p_t\nabla_x\frac{\delta \mathcal{F}(p_t)}{\delta p_t}\right]\tag{3.25}$$
（3.25）就为（3.8）式的解，也叫作概率空间中在W-2度量下对泛函 $\mathcal{F}(p)$ 优化的梯度流。写成离散形式就有： $$p_{t+1}=p_t+\nabla_x\cdot\left[p_t\nabla_x\frac{\delta \mathcal{F}(p_t)}{\delta p_t}\right]\tag{3.26}$$
举个例子，如果要求KL散度在W-2度量下的梯度流，则令 $\mathcal{F}(p)=KL(p\parallel q)$ ， $q$ 为某个给定的分布，带入到（3.25）式，再应用变分导数运算，则： $$\begin{align} \frac{\partial}{\partial t}p_t&=\nabla_x\cdot\left[p_t\nabla_x\frac{\delta \mathcal{F}(p_t)}{\delta p_t}\right]\ &=\nabla_x\cdot\left[p_t\nabla_x\left(\log p_t +1-\log q \right)\right]\ &=\nabla^2_xp_t -\nabla_x\cdot[p_t\nabla_x \log q] \end{align}\tag{3.27}$$
（3.27）刚好就是： $$dx=\nabla_x\log q(x)dt+\sqrt{2}dw\tag{3.28}$$
对应的FP方程，也就是说如果能知道 ${\mathrm{\nabla }}_x\log q(x)$ ，那么就能从（3.28）采样得到 $q(x)$ ，因为最终是在优化 $KL(p\parallel q)$ 。把（3.28）离散化，其实也是之前提到的Langevin方程，也是多步从 ${\mathrm{\nabla }}_x\log q(x)$ 采样出 $q(x)$ 。因此梯度流把之前的一些方法全部联系了起来，是个非常有意思的理论。