最优传输算法——Benamou Brenier算法

Benamou Brenier算法

Brief

是一种连续数值方法，将最优传输问题转化为一个容易处理的\(d+1\)维凸变分问题。我们将会用Wasserstein测地线的理论描述它（相比于找到映射，这个方法是找到测地曲线\(\mu_t\)）。

另外两个经典的连续方法是：

• Angenent-Hacker-Tannenbaum：基于最优传输映射应该是一个梯度的事实，移除非梯度项来减少能量；

• Loeper-Rapetti：基于Monger-Ampere方程的解析(resolution)。

这两种方法都需要光滑和非消逝(non-vanishing)的密度，以及特殊的域来处理边界数据(一个长方形，或者更好的，一个环面)。Monge -Ampère方程在更一般的域上的解决问题更加微妙，最近已经被解决了。

Benamou Brenier构想和它的数值应用

5.3和5.4节的结果允许我们将对应于代价为\(|x-y|^P\)的优化问题用一种聪明的方式重写，将它转化成一种凸优化问题。实际上就是：

• 寻找代价为\(|x-y|^P\)的OT等价于寻找\(\mathbb W_P\)空间中的常速测地线，因为从最优计划中，我们可以重建测地线；从测地线中(通过它们的速度场)，重构OT也是有可能的；

• 常速测地线可以通过最小化\(\int_0^1|\mu'|(t)^pdt\)来找到；

• 在Wasserstein空间的情况下，我们有\(|\mu'|(t)^p = \int _\Omega|\mathbf v_t|^pd\mu_t\)，其中\(\mathbf v\)是一个和\(\mu\)一起求解连续方程的速度场。这个场不是唯一的，但是度量导数\(|\mu'|(t)\)等于所有可能的场的\(L^p\)范数的最小值。

  作为这些考虑的结果，对于\(p>1\)，解决动能最小问题

  \[\min \{\int_0^1\int_\Omega|\mathbf v_t|^pd\varrho_t dt:\partial_t\varrho_t + \nabla\cdot(\mathbf v_t\varrho_t) = 0,\varrho_0=\mu,\varrho_t=\nu\} \]

选择了连接\(\mu\)到\(\nu\)的常速测地线，因此允许找到两个测度间的最优传输。

​ 另一方面，变量\((\varrho_t,\mathbf v_t)\)的这个最小化问题有非线性约束（由于乘积\(\mathbf v_t\varrho_t\)），而且泛函是非凸的（由于\((t,x)\mapsto t|x|^p\)是非凸的）。但是，我们可以知道，结合5.3.1节的工具，将它转化为凸问题是可能的。

​ 为此，转换变量就足够了，从\((\varrho_t,\mathbf v_t)\)到\((\varrho_t,E_t)\)，其中\(E_t=\mathbf v_t\varrho_t\)，在时空(space-time)中使用泛函\(\mathscr B_p\)。回忆\(\mathscr B_p(\varrho,E):=\int f_p(\varrho,E)\)，其中\(f_p:\mathbb R \times \mathbb R^d \to \mathbb R \cup \{+\infty\}\)，定义在引理5.17中。

\[f_{p}(t, x):=\sup _{(a, b) \in K_{q}}(a t+b \cdot x)=\left\{\begin{array}{ll} \frac{1}{p} \frac{|x|^p}{t^{p-1}} & \text { if } t>0, \\ 0 & \text { if } t=0, x=0 \\ +\infty & \text { if } t=0, x \neq 0, \text { or } t<0, \end{array}\right. \]

其中 \(K_{q}:=\left\{(a, b) \in \mathbb{R} \times \mathbb{R}^{d}: a+\frac{1}{q}|b|^{q} \leq 0\right\}\)。问题变为

问题6.1 求解

\[\left(\mathrm{B}_{p} \mathrm{P}\right) \quad \min \left\{\mathscr{B}_{p}(\varrho, E): \partial_{t} \varrho_{t}+\nabla \cdot E_{t}=0, \varrho_{0}=\mu, \varrho_{1}=v\right\}. \]

我们可以这样写： \(\mathscr{B}_{p}(\varrho, E)=\int_{0}^{1} \mathscr{B}_{p}\left(\varrho_{t}, E_{t}\right) \mathrm{d} t=\int_{0}^{1} \int_{\Omega} f_{p}\left(\varrho_{t}(x), E_{t}(x)\right) \mathrm{d} x \mathrm{~d} t\),
该函数的第三个表达式隐式假定 \(\varrho_{t}, E_{t} \ll \mathscr{L}^{d}\)。 确实，正如我们在命题 5.18 中看到的那样，泛函 \(\mathscr{B}_{p}\) 具有此形式的完整表示，只要 \(\varrho_{t}\) and \(E_{t}\) 关于耦合相同的正测度是绝对连续的。

​ 我们还想强调，由 \(\partial_{t} \varrho_{t}+\nabla \cdot E_{t}=0, \varrho_{0}=\) \(\mu, \varrho_{1}=v\) 给出的约束实际上是时空(space-time)中的发散约束 (考虑向量 \(\left.(\varrho, E):[0,1] \times \Omega \rightarrow \mathbb{R}^{d+1}\right)\)。空间边界约束已经是无通量类型(no-flux type)， \(\varrho\) 的初始值和终值提供了边界 \(\{0\} \times \Omega\) 和\(\{1\} \times \Omega\)上的非齐次Neumann数据。确实，整个约束可以理解为 \(\nabla_{t, x} \cdot(\varrho, E)=\delta_{0} \otimes \mu-\delta_{1} \otimes v\) （注意，使用\((t, x)\)的导数的观点会在本节再次出现 ）。我们将最小化的泛函是一个1维齐次泛函，这样\(\left(\mathrm{B}_{p} \mathrm{P}\right)\) 变成了我们在4.2节所见的Beckmannn问题的动态版本（在时空上）。 这就是如果我们应用[196]中提出的成本\(|x-y|^{p}\)降低而得到的问题，这将其转化为时空上的1齐次运输成本。

​ 现在，约束是线性的，而泛函是凸的。 然而，泛函（以及函数\(f_P\)）是凸的，但不是很多，因为正如我们所说的，是1-齐次的。 特别的，它不是严格凸且是不可微的。这降低了任一梯度下降算法解决问题的效率，但可以使用一些改进的方法。

​ 在[34]中，作者在此基础上提出了一种数值方法，基于变量的凸性变化，对偶性以及所谓的“增强拉格朗日(argumented Lagrangian)”。

注 鞍点(saddle point)，Uzawa和增广Lagrangian

假设我们需要最小化一个凸函数\(f: \mathbb{R}^{N} \rightarrow \mathbb{R}\)，满足\(k\)个线性等式约束\(A x=b\) (其中 \(b \in \mathbb{R}^{k}\)，以及 \(A \in \mathrm{M}^{k \times N}\) )。 这个问题等价于：

\[\min _{x \in \mathbb{R}^{N}} f(x)+\sup _{\lambda \in \mathbb{R}^{k}} \lambda \cdot(A x-b). \]

这给出了一个具有Lagrangian函数\(L(x, \lambda):=f(x)+\lambda \cdot(A x-b)\)的min-max问题。如果我们相信对偶，对\(f\)在约束下找到一个最小化子等价于找到一个\(g(\lambda):=\inf _{x \in \mathbb{R}^{N}} f(x)+\lambda \cdot(A x-b)\)的最大化子。这也等价于为\(L\)找到鞍点 (一个点 \((\bar{x}, \bar{\lambda})\)，满足 \(L(x, \bar{\lambda}) \geq L(\bar{x}, \bar{\lambda}) \geq L(\bar{x}, \lambda)\) ，对于每个\(x\) 和每个 \(\lambda\),，即一个点，在 \(x\) 处最小，在 \(\lambda\) 处最大)。

​ 在 \(\lambda\) 处的最大化是更容易的，因为问题是非约束的：可以用一个梯度算法(见Box 6.8)。我们只需要计算 \(\nabla g(\lambda)\)，通过 \(\nabla g(\lambda)=A x(\lambda)-b\)，这也是比较容易得到的，其中 \(x(\lambda)\)是(希望唯一) \(x \mapsto f(x)+\lambda \cdot(A x-b)\)的最小化子。
​ 我们从这个想法中获得的算法，叫做Uzawa algorithm，如下：给定 \(\left(x_{k}, \lambda_{k}\right)\)，设\(x_{k+1}:=x\left(\lambda_{k}\right)\) 以及\(\lambda_{k+1}=\lambda_{k}+\tau \nabla g\left(\lambda_{k}\right)=\lambda_{k}+\tau\left(A x_{k+1}-b\right)\)，对于一个给定的小的\(\tau\)，点 \(x_{k}\) 的序列收敛，在合理的假设下，收敛到原始问题的最小化子。

​ 一个使得点\(x(\lambda)\)的计算更容易以及加速收敛的替代的想法如下：相比于使用Lagrangian\(L(x, \lambda)\)，使用以下变体\(-\widetilde{L}(x, \lambda):=L(x, \lambda)+\frac{\tau}{2}|A x-b|^{2}\)，对于一个\(\tilde{\tau}\)的给定值。成为\(\tilde{L}\)和\(L\)的一个鞍点的条件分别是：

\[\text { for } \tilde{L}:\left\{\begin{array}{l} \nabla f(x)+A^{t} \lambda+\tilde{\tau}(A x-b)=0, \\ A x-b=0, \end{array} \quad \text { for } L:\left\{\begin{array}{l} \nabla f(x)+A^{t} \lambda=0 \\ A x-b=0 \end{array}\right.\right. \]

因此，\(\tilde{L}\) 和 \(L\) 的鞍点是相同的，但是使用\(\tilde{L}\)使得问题更凸，在\(x\)点更容易处理。然后，在\(\tilde{g}(\lambda):=\inf _{x}f(x)+\lambda \cdot(A x-b)+\frac{\tilde{\tau}}{2}|A x-b|^{2}\)上使用相同的梯度算法，使用步长\(\tau\)。取\(\tau=\tilde{\tau}\)，迭代以下算法：

\[\left\{\begin{array}{l} x_{k+1}=\operatorname{argmin} f(x)+\lambda_{k} \cdot(A x-b)+\frac{\tau}{2}|A x-b|^{2} \\ \lambda_{k+1}=\lambda_{k}+\tau\left(A x_{k+1}-b\right) \end{array}\right. \]

​ 下面是构思这个算法的主要步骤。

​ 首先，我们用弱形式来写出约束（根据（4.3））。这意味着我们想求解

\[\min _{\varrho, E} \quad \mathscr{B}_{p}(\varrho, E)+\sup _{\phi}\left(-\int_{0}^{1} \int_{\Omega}\left(\left(\partial_{t} \phi\right) \varrho_{t}+\nabla \phi \cdot E_{t}\right)+G(\phi)\right), \]

其中我们设

\[G(\phi):=\int_{\Omega} \phi(1, x) \mathrm{d} v(x)-\int_{\Omega} \phi(0, x) \mathrm{d} \mu(x), \]

这个sup是通过遍历定义在\([0,1]\times \Omega\)上的所有函数计算出来的。（我们在这里不关系它们的正则性，因为它们将由\([0,1] \times \mathbb R^d\)中在一个网格的点上定义的函数表示）

注6.2 我们来观察上述问题和Hamilton-Jacobi方程的关系。我们将考虑最简单的情形，\(p=2\)。在这种情况下，我们可以将问题写做：

\[\min _{(E, \varrho): \varrho \geq 0} \int_{0}^{1} \int_{\Omega} \frac{|E|^{2}}{2 \varrho}+\sup _{\phi}-\int_{0}^{1} \int_{\Omega}\left(\left(\partial_{t} \phi\right) \varrho+\nabla \phi \cdot E\right)+G(\phi), \]

其中，我们用 \(\mathscr{B}_{2}\) 的积分形式来表达它， 在绝对连续测度下是有效的，其中 \(\varrho \geq 0\) (如果\(\varrho=0\),，我们必须使得 \(E=0\) ，为了得到有限能量)。

​ 如果我们交换inf和sup，我们将得到下面的问题：

\[\sup _{\phi} \quad G(\phi)+\inf _{(E, \varrho): \varrho \geq 0} \int_{0}^{1} \int_{\Omega}\left(\frac{|E|^{2}}{2 \varrho}-\left(\partial_{t} \phi\right) \varrho-\nabla \phi \cdot E\right) . \]

我们可以首先计算出最优\(E\), 对于固定的 \(\varrho\) 和 \(\phi\), 因此得到 \(E=\varrho \nabla \phi\)。问题变为

\[\sup _{\phi} \quad G(\phi)+\inf _{\varrho \geq 0} \int_{0}^{1} \int_{\Omega}\left(-\left(\partial_{t} \phi\right) \varrho-\frac{1}{2}|\nabla \phi|^{2} \varrho\right) . \]

下确界有限 (因此消逝) 的条件是 \(\partial_{t} \phi+\frac{1}{2}|\nabla \phi|^{2} \leq 0\)，在最优条件下，我们必须有相等，在 \(\{\varrho>0\}\)上。 这就给出了Hamilton-Jacobi等式：

\[\partial_{t} \phi+\frac{1}{2}|\nabla \phi|^{2}=0 \quad \varrho \text { -a.e. } \]

而且，从最优\(\phi\)中，我么可以使用 \(\psi(x):=\phi(1, x)\) 和\(\varphi(x):=-\phi(0, x)\)来恢复Kantorovich势。
在这个变分问题中， \(\mathscr{B}_{p}\) 可能会表示成sup，因此我们得到：

\[\min _{\varrho, E} \sup _{(a, b) \in K_{q}, \phi} \int_{0}^{1} \int_{\Omega}\left(a(t, x) \mathrm{d} \varrho+b(t, x) \cdot \mathrm{d} E-\partial_{t} \phi \mathrm{d} \varrho-\nabla \phi \cdot \mathrm{d} E\right)+G(\phi). \]

​ 令\(m=(\varrho, E)\)。 这里\(m:[0,1] \times \Omega \rightarrow \mathbb{R}^{d+1}\)是一个定义在\(d+1\)维空间的\(d+1\)维向量场。在这里，我们不考虑\(m\)是一个测度还是一个实函数，因为不管怎样我们都会在一个离散设定中考量，\(m\)会是一个在\([0,1] \times \mathbb{R}^{d}\)中的一个网格的每个点上定义的函数。类似的，我们设 $\xi=(a, b) $ 。我们也用\(\nabla_{t, x} \phi\)表示\(\phi\)的时空梯度，即\(\nabla_{t, x} \phi=\left(\partial_{t} \phi, \nabla \phi\right)\)。

​ 问题被重写为：

\[\min _{m} \sup _{\xi, \phi: \xi \in K_{q}}\left\langle\xi-\nabla_{t, x} \phi, m\right\rangle+G(\phi). \]

​ 这就是使用增强拉格朗日方法的想法。实际上，上面的问题让人想起了拉格朗日，但是是以相反的方式。我们必须认为对偶变量是\(m\)，原始变量是\((\xi, \phi)\)。函数\(f(\xi, \phi)\)包括项\(G(\phi)\)和约束\(\xi \in K_{q}\)，还有一个等式约束\(\xi=\nabla_{t, x} \phi\)。实际上我们不关心产生这个Lagrangian的原始约束问题，而是打算为优化加入一项\(\frac{\tau}{2}\left|\xi-\nabla_{t, x} \phi\right|^{2}\)（对于一个小步长\(\tau\)）。

​ 因此，我们寻找以下问题的解：

\[\min _{m} \sup _{\xi, \phi: \xi \in K_{q}}\left\langle\xi-\nabla_{t, x} \phi, m\right\rangle+G(\phi)-\frac{\tau}{2}\left|\xi-\nabla_{t, x} \phi\right|^{2} \]

​ 找寻最优的算法需要做以下工作：产生一个序列\(m_{k}\)，对于每个\(k\)，寻找最优\(\left(\xi_{k}, \phi_{k}\right)\)。因此，相比于精确找到最优\(\left(\xi_{k}, \phi_{k}\right)\)，我们将会用两步来优化（首先固定\(\xi\)，优化\(\phi\)，对于这个\(\phi\)，再优化\(\xi\)）。

​ 算法包括以下三个迭代步骤。假设我们有一个三元组\(\left(m_{k}, \xi_{k}, \phi_{k}\right)\)：

• 在给定 \(m_{k}\) 和 \(\xi_{k}\) 的情况下，通过求解下式，找到最优的 \(\varphi_{k+1}\),
  \[ \max _{\varphi}-\left\langle\nabla_{t, x} \varphi, m_{k}\right\rangle+G(\varphi)-\frac{\tau}{2}\left\|\xi_{k}-\nabla_{t, x} \varphi\right\|^{2}, \]

​ 这归结为最小化关于 \(\nabla_{t, x} \varphi\) 的二次问题。解可以通过求解 Laplace方程\( \tau \Delta_{t, x} \varphi=\nabla \cdot\left(\tau \xi_{k}-m_{k}\right)\)来找到，以及一个时空Laplacian算子和Neumann边界条件。这些条件关于空间是齐次的，由于\(G\)的存在，在\(t=0\)和\(t=1\)时刻是非齐次的。大多数Laplace求解器可以在 \(O(n \log n)\)的倍数复杂度内找到解，这里\(n\)是离散化的点的个数。

• 给定 \(m_{k}\) 和 \(\varphi_{k+1}\) 的情况下，通过求解下式，找到最优的 \(\xi_{k+1}\),

  \[\max _{\xi \in K_{q}}\left\langle\xi, m_{k}\right\rangle-\frac{\tau}{2}\left|\xi-\nabla_{t, x} \varphi_{k+1}\right|^{2}, \]

  通过展开平方项，我们看到这个问题等价于计算 \(\nabla_{t, x} \varphi_{k+1}+\frac{1}{\tau} m_{k}\)的投影，并且没有梯度出现在这一优化步骤里。优化算法逐点执行，通过对每个\((t, x)\)选取点\(\xi=(a, b)\)，这个点是在凸集\(K_{q}\) 中最接近 \(\nabla_{t, x} \varphi_{k+1}(t, x)+\frac{1}{\tau} m_{k}(t, x)\)的点。如果我们有一个在\(\mathbb{R}^{n+1}\)中的投影算法，并且此投影算法需要一个常数的操作，则这一逐点的步骤的代价为\(O(N)\)。

• 最后，我们通过设置 \(m_{k+1} \leftarrow m_{k}-\tau\left(\xi_{k+1}-\nabla_{t, x} \varphi_{k+1}\right)\) 来更新 \(m\)。

  ​ 这个算法在每个迭代过程中共需要\(O(N \log N)\)次操作（N是由时空中的离散化给定的）。它可以被证明是收敛的。与其他算法相比，Benamou-Brenier方法有一些重要的优点：

  1. 到目前为止，它几乎是唯一一个毫不费力地处理密度消失的方法，而且对它们的支集也不需要特殊的假设;

  2. 不需要特定的二次代价：我们用幂次代价\(c(x,y)=|x-y|^p\)提出它，但是它也可以适用于其他凸代价函数\(h(x-y)\)，以及所有拉格朗日作用产生的代价函数；它可以包含取决于位置和时间的成本，适用于黎曼流形;

  3. 它可能适合考虑密度\(\varrho\)上的凸约束(例如，下界或上界);

  4. 它可以处理不同的“动态”问题，在密度或速度上添加惩罚，就像在平均场博弈论中发生的那样(见章节8.4.4)，并且已经被利用，例如，在[35]中;

  5. 它可以处理多个种群，具有可能的相互作用行为。

​ 我们参考[96,248]来对上述2和3点的变体进行数值处理，并参考[37]来实现多种群情况。这里我们给出一个简单的图(图6.1)，在一个环面上通过Benamou-Brenier方法得到的最简单的情况。

2D环面的一个高斯分布和加入一个向量（1/2,1/2）位移的相同的高斯分布间的测地线。注意到，为了避免割迹(cut locus)，初始的密度分成了四片。