最優傳輸算法——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)，初始的密度分成了四片。