Density Controlによる最適輸送(2)

前回の記事「Density Controlによる最適輸送(1)」において、最適輸送問題を流体を用いて定式化することで最適制御問題として捉える方法であるDensity Controlについて概要を述べました。この記事ではDensity Controlを行うための具体的なアルゴリズムの導出を行います。このアルゴリズムはJ.D. Benamou, Y. Brenier, "A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem"において提案されたものです。

初めに前回記事の復習をして、その後アルゴリズムの導出を行います。アルゴリズムの導出のために比較的多くの理論的な計算を行いますので、アルゴリズムだけ知りたい場合は最後のまとめの章のみ読んでください。数値計算による実験はこの記事では行わず次の記事で行う予定です。

前回記事の復習

最適輸送問題はいかに輸送コストを小さくできるかという古くから知られた問題です。最適輸送問題はいくつかのバリエーションがありますが、ここでは$L^2$-MKP (Monge-Kantorovich Problem)と呼ばれるものを扱います。$L^2$-MKPは以下です。

$\mathbb{R}^n$において2つの密度分布$\rho_0,\rho_T$で
$\int_{\mathbb{R}^n}\rho_0(x)dx^n=\int_{\mathbb{R}^n}\rho_T(x)dx^n=1$
を満たすものが与えられているとき、滑らかな最適写像$M:\mathbb{R}^n\to \mathbb{R}^n$で以下の２つを満たすものを求める問題です。
(i) 任意の可測集合$A\subset\mathbb{R}^n$に対して
$\int_{x\in A}\rho_T(x)dx^n=\int_{M(x)\in A}\rho_0(x)dx^n$
を満たす。
(ii)
$\int_{\mathbb{R}^d}|M(x)-x|^2\rho_0(x)dx^n$
を最小にする。

最適制御問題とは、何らかの時間変化するシステムにおいて、コストなどの性能指標を最小化（または最大化）するために、外部からそのシステムにどのように影響（制御入力と呼ばれる）を与えていけばよいかを考える問題です。
Density Controlとは最適輸送問題に時間変数を導入し流体の流れの最適制御問題として扱うことで最適な流速場を求め、各時刻において各位置にある質量をその最適な流速に沿って一定時間動かすことで最終的に最適写像 $M$ を実現することを目指すというアイデアです。具体的には以下のような定式化になります。

$\mathbb{R}^n$上の流速場を $v(t,x)$ 、密度を$\rho(t,x)$ とする流体に対して、任意の時刻 $t$ において $\rho(t,x)$ に関して (i) の条件が成り立つ（流体と共に動く有界領域内の質量が一定）とすると、レイノルズの輸送定理より

$$0=\frac{d}{dt}\int_{\Omega_t}\rho(t,x)dx^n=\int_{\Omega_t}\left(\frac{\partial\rho}{\partial t}+\textrm{div}(\rho v)\right)dx^n$$

となるので、連続の式

$$\partial_t\rho+{\rm div}(\rho v)=0$$

と境界条件

$$\rho(0,x)=\rho_0,\rho(T,x)=\rho_T$$

が満たされる流体を考えればよいことになります。さらに、$v(t,x)$ の流れに沿って点を $T$ 秒間移動させる写像が最適写像 $M$ になっているためには条件 (ii) を満たすようにしなければなりません。ポイントになるのが次の不等式です（証明はAppendix）。

$$T\int_{\mathbb{R}^n}\int^T_0\rho(t,x)|v(t,x)|^2dx^ddt\geq \int_{\mathbb{R}^n}|M(x)-x|^2\rho_0(x)dx^d$$

このことから左辺の積分を連続の式という制約条件の下で最小化する問題を解けば、求めたい $v,\rho$ が分かります。従ってラグランジュ関数を

$$L(\phi,\rho,v):=\int_0^T\int_{\mathbb{R}^n} \{\frac{1}{2}\rho(t,x)|v(t,x)|^2 + \phi(t,x)(\partial_t\rho+\textrm{div}(\rho v))\}dtdx$$

と置くとき、

$$\inf_{\rho, v} \sup_{\phi} L(\phi, \rho, v)$$

を考えればよいことになります。また $\frac{1}{2}\rho(t,x)|v(t,x)|^2$ は流体の運動エネルギーのようなものだと考えることができます。前回記事では話をコンパクトにまとめるためにこのラグラジアンを変分して場の方程式を出して定性的な話のみしましたが、この極値問題を解くにはより込み入った議論を必要とします。

拡張ラグランジュ関数の導入

以下では登場する場は全て $\mathbb{R}^n$ においてコンパクト台を持つことを仮定し、$D$ は関連する全ての場の台を含む領域とします。

新たに運動量 $m=\rho v$ を定義すると、$L(\phi,\rho,v)$ は

$$L(\phi, \rho, m) = \int_0^T \int_D \left( \frac{|m|^2}{2\rho} - \partial_t \phi \rho - \nabla_x \phi \cdot m \right)dxdt- \int_D \left( \phi(0, \cdot) \rho_0 - \phi(T, \cdot) \rho_T \right)dx$$

となるから

$$\inf_{\rho, m} \sup_{\phi} L(\phi, \rho, m) \tag{1}$$

を考えればよいです。変分すると

$$\partial_t \phi + \frac{|m|^2}{2\rho^2} = 0,\ \frac{m}{\rho} = \nabla_x \phi,\ \partial_t \rho + \textrm{div}m = 0$$

$$\rho(0, \cdot) = \rho_0,\ \rho(T, \cdot) = \rho_T$$

が得られます。

これまで空間 $\mathbb{R}^n\ (n=2,3)$ を考えてきましたが、時間方向の次元を加えて全体として $M:=\mathbb{R}\times \mathbb{R}^n$ を考え、これを拡張空間と呼ぶことにします（数学における時空とは別の概念なのでこのように呼び分けることにします）。拡張空間 $M$ のベクトル $\xi$ は時間成分が $\tau$ で空間成分($\mathbb{R}^n$ のベクトル)が $s$ のとき、

$$\xi = (\tau, s)$$

と書くことにします。

このとき

$$\mu= (\rho, m)\\ q = (a,b)$$

と置きます（$a, b$ は $\rho, m$ の双対変数として導入します）。また拡張空間のベクトル同士の内積 $\cdot$ はEuclid内積（成分ごとの積の和）とします。さらに $L^2$ 内積を

$$\langle\mu,q\rangle:=\int_0^T\int_D\mu\cdot q\ dtdx$$

と定義します。

ここで

$$K =\{ q: \mathbb{R} \times \mathbb{R}^3 \rightarrow \mathbb{R} \times \mathbb{R}^3;\ \ a + \frac{|b|^2}{2} \leq 0 \}$$

と定義すると、

$$\frac{|m(t, x)|^2}{2\rho(t, x)} = \sup_{(a, b) \in K} [a(t, x) \rho(t, x) + b(t, x) \cdot m(t, x)]$$

となります。実際、$(a, b) \in K$ に対して、

$$a(t, x) \rho(t, x) + b(t, x) \cdot m(t, x) \le \left(-\frac{|b|^2}{2}\rho(t, x) + b(t, x) \cdot m(t, x) \right)\\ \le -\frac{\rho}{2}\left| b - \frac{m}{\rho} \right|^2 + \frac{|m|^2}{2\rho} \le \frac{|m|^2}{2\rho}$$

であることから成り立つことが分かります。

$$F(q) = \begin{cases}0 \ \ (q\in K)\\ \infty \ \ (q\notin K)\end{cases}\\ G(\phi) = \int_D \left[ \phi(0, \cdot) \rho_0 - \phi(T, \cdot) \rho_T \right] \, dx$$

と定義すると、(1)は

$$\sup_\mu \inf_{\phi, q} \left[ F(q) + G(\phi) + \langle \mu, \nabla_{t,x} \phi - q \rangle \right] \tag{2}$$

と同値です。ここで拡張空間のgradientを

$$\nabla_{t,x} \phi(t,x) := (\partial_t\phi, \textrm{grad}\phi)$$

と書きました。

さらに（２）に対して、$\mu$ を制限 $\nabla_{t,x} \phi(t,x)-q=0$ に関するLagrange乗数と読み替えることにより、拡張ラグランジュ関数を

$$L_r(\phi, q, \mu) := F(q) + G(\phi) + \langle \mu, \nabla_{t,x} \phi - q \rangle + \frac{r}{2} \langle \nabla_{t,x} \phi - q, \nabla_{t,x} \phi - q \rangle$$

と定義すると、

$$\sup_{\mu} \inf_{\phi, q} L_r(\phi, q, \mu) \tag{3}$$

を考えてもよいことになります。

極値問題（３）を解くには以下のアルゴリズムを利用します。

$(\phi^{n-1}, q^{n-1}, \mu^n)$ が与えられたとします。
Step A : 以下を満たす $\phi^n$ を見つけます。

$$L_r(\phi^n, q^{n-1}, \mu^n) \leq L_r(\phi, q^{n-1}, \mu^n), \quad \forall \phi.$$

Step B : 以下を満たす $q^n$ を見つけます。

$$L_r(\phi^n, q^n, \mu^n) \leq L_r(\phi^n, q, \mu^n), \quad \forall q.$$

Step C : 以下のように $\mu^n$ を更新します。

$$\mu^{n+1} = \mu^n + r (\nabla_{t,x} \phi^n - q^n)$$

Step A に戻ります。
以下では各stepについて詳細に論じます。

Step A

Step Aの必要条件として $\phi_n$ が満たす条件を求めるために変分します。ここの計算は物理などでラグラジアンから変分法で運動方程式を導出するテクニックと同様です（例えばスカラー場の理論など）。

$$\phi_t ^n= \phi_n + t\varphi$$

とすると、

$$\frac{d}{dt}L_r(\phi_t^n, q^{n-1}, \mu^n)=G(\varphi) + \langle \mu^n, \nabla_{t,x} \varphi \rangle + r \langle \nabla_{t,x} \phi^n - q^{n-1}, \nabla_{t,x} \varphi \rangle = 0 \tag{A1}$$

が得られます。

$$r \langle \nabla_{t,x} \phi^n - q^{n-1}, \nabla_{t,x} \varphi \rangle = r \langle (\partial_t \phi^n - a^{n-1}, \textrm{grad} \phi^n - b^{n-1}), (\partial_t\varphi, \textrm{grad} \varphi) \rangle \\ = r\int_0^Tdt\int_Ddx\left( \partial_t\varphi(\partial_t\phi^n-a^{n-1})+(\textrm{grad}\phi^n-b^{n-1})\cdot\textrm{grad}\varphi \right)\\ \langle \mu^n, \nabla_{t,x} \varphi \rangle = \langle (\rho^n, m^n), (\partial_t \varphi, \textrm{grad} \varphi) \rangle =\int_0^Tdf\int_Ddx( \partial_t\varphi\rho^n + m^n\cdot \textrm{grad}\varphi )$$

であるから、

$$\langle \mu^n, \nabla_{t,x} \varphi \rangle + r \langle \nabla_{t,x} \phi^n - q^{n-1}, \nabla_{t,x} \varphi \rangle\\ =\int_0^Tdt\int_Ddx\left( \partial_t\varphi (\rho^n + r(\partial_t\phi^n-a^{n-1})) + (m^n + r(\textrm{grad}\phi^n-b^{n-1}))\cdot\textrm{grad}\varphi \right)\\ =\int_0^Tdt\int_Ddx \frac{d}{dt}(\varphi (\rho^n + r(\partial_t\phi^n-a^{n-1}))) - \int_0^Tdt\int_Ddx( \varphi (\partial_t\rho^n + r\partial_t(\partial_t\phi^n-a^{n-1}))\\ + \int_0^T\int_Ddx(m^n + r(\textrm{grad}\phi^n-b^{n-1}))\cdot\textrm{grad}\varphi )\\ =\int_Ddx \{\varphi(T,x) (\rho^n(T,x) + r(\partial_t\phi^n(T,x)-a^{n-1}(T,x)))\} - \int_Ddx \{\varphi(0,x) (\rho^n(0,x) + r(\partial_t\phi^n(0,x)-a^{n-1}(0,x)))\}\\ - \int_0^Tdt\int_Ddx \varphi \{\partial_t\rho^n + r\partial_t(\partial_t\phi^n-a^{n-1}) + \textrm{div}(m^n + r(\textrm{grad}\phi^n-b^{n-1}))\}\\ + \int_0^Tdt\int_Ddx \textrm{div}(\varphi(m^n + r(\textrm{grad}\phi^n-b^{n-1})))$$

ここで、場がコンパクト台を持つと仮定しているから、

$$\langle \mu^n, \nabla_{t,x} \varphi \rangle + r \langle \nabla_{t,x} \phi^n - q^{n-1}, \nabla_{t,x} \varphi \rangle\\ =\int_Ddx \{\varphi(T,x) (\rho^n(T,x) + r(\partial_t\phi^n(T,x)-a^{n-1}(T,x)))\} - \int_Ddx \{\varphi(0,x) (\rho^n(0,x) + r(\partial_t\phi^n(0,x)-a^{n-1}(0,x)))\}\\ - \int_0^Tdt\int_Ddx \varphi \{\partial_t\rho^n + \textrm{div}(m^n) + r(\partial^2_t\phi^n + \Delta\phi^n-\partial_ta^{n-1} - \textrm{div} b^{n-1})\}$$

となります。従って、(A1)は

$$G(\varphi) + \langle \mu^n, \nabla_{t,x} \varphi \rangle + r \langle \nabla_{t,x} \phi^n - q^{n-1}, \nabla_{t,x} \varphi \rangle\\ =\int_Ddx \{\varphi(T,x) (\rho^n(T,x) -\rho(T,x) + r(\partial_t\phi^n(T,x)-a^{n-1}(T,x)))\}\\ - \int_Ddx \{\varphi(0,x) (\rho^n(0,x) -\rho(0,x) + r(\partial_t\phi^n(0,x)-a^{n-1}(0,x)))\}\\ - \int_0^Tdt\int_Ddx \varphi \{\partial_t\rho^n + \textrm{div}(m^n) + r(\partial^2_t\phi^n + \Delta\phi^n-\partial_ta^{n-1} - \textrm{div} b^{n-1})\}=0$$

となるから、結局

$$-r\Delta_{t,x}\phi^n=\nabla_{t,x}\cdot(\mu^n-rq^{n-1}) \tag{A2}\\ r\partial_t\phi^n(T,x)=\rho_T -\rho^n(T,x) + ra^{n-1}(T,x)\\ r\partial_t\phi^n(0,x)=\rho_0 -\rho^n(0,x) + ra^{n-1}(0,x)$$

が得られます。
$\mu^n, q^{n-1}$ が与えられているとすると、(A2)は $\phi^n$ に関するPoisson方程式です。

Step B

$q$ に関する変分はできないですが、$q$ の処理はより単純で、$\phi^n,\mu^n$ が与えられたとき、

$$\inf_q\{F(q)+ \frac{r}{2} \langle \nabla_{t,x} \phi^n - q, \nabla_{t,x} \phi^n - q \rangle + \langle \mu^n, \nabla_{t,x} \phi^n - q \rangle\}\\ =\frac{r}{2}\inf_{q\in K}\{ \langle \nabla_{t,x} \phi^n - q, \nabla_{t,x} \phi^n - q \rangle + \frac{2}{r}\langle \mu^n, \nabla_{t,x} \phi^n - q \rangle\}\\ =\frac{r}{2}\inf_{q\in K}\{ \langle \nabla_{t,x} \phi^n - q + \frac{1}{r}\mu^n, \nabla_{t,x} \phi^n - q + \frac{1}{r}\mu^n\rangle - \frac{1}{r^2}||\mu^n||^2\}$$

を解けばよいです(最後の式は平方完成した)。
よって

$$\inf_{q\in K}\{ \langle \nabla_{t,x} \phi^n - q + \frac{1}{r}\mu^n, \nabla_{t,x} \phi^n - q + \frac{1}{r}\mu^n\rangle\}$$

を考えればよいです。

$$\alpha^n(t,x):=\partial_t\phi^n(t,x) + \frac{\rho^n(t,x)}{r}\\ \beta^n(t,x):=\textrm{grad}_x\phi^n(t,x) + \frac{m^n(t,x)}{r}$$

と置くと、$q=(a,b)$ は

$$\inf_{(a,b)}\{(a-\alpha^n)^2+|b-\beta^n|^2;\ a+\frac{|b|}{2}\le0\} \tag{B1}$$

を求めればよいです。これは各 $(t,x)$ 毎に求めることができます。

step C

最後に $\mu$ を $L_r$ が大きくなるように更新するためには、

$$\mu^{n+1} = \mu^n + r (\nabla_{t,x} \phi^n - q^n) \tag{C1}$$

とすればよいです。

アルゴリズムの停止判定

StepA, B, Cを繰り返すときにいつ手続きを停止させるかを考えましょう。極値問題 (1) の解 $(m,\rho,\phi)$ は

$$L(\phi, \rho, m) = \int_0^T \int_D \left( \frac{|m|^2}{2\rho} - \partial_t \phi \rho - \nabla_x \phi \cdot m \right)dxdt- \int_D \left( \phi(0, \cdot) \rho_0 - \phi(T, \cdot) \rho_T \right)dx$$

を変分した方程式

$$\partial_t\phi+\frac{|m|^2}{2\rho^2}=0,\\ \frac{m}{\rho}=\textrm{grad}_x\phi,\\ \partial_t\rho+\textrm{div}_xm=0$$

を満たします。特に $\phi$ は

$$\partial_t\phi+\frac{1}{2}|\textrm{grad}_x\phi|^2=0$$

を満たします。従って、アルゴリズムが極値問題 (1) の解への近似列を生成するとき、

$$\textrm{res}^n:=\partial_t\phi^n+\frac{1}{2}|\textrm{grad}_x\phi^n|^2$$

は０に収束するはずです。この量を時間と空間で積分して元のラグラジアンで規格化した量

$$\textrm{crit}^n:=\sqrt{\frac{\int_0^T\int_D\rho^n\textrm{res}^n}{\int_0^T\int_D\rho^n|v^n|^2}} =\sqrt{\frac{\int_0^T\int_D\rho^n\textrm{res}^n}{\int_0^T\int_D\rho^n|\textrm{grad}_x\phi^n|^2}}$$

が０に近づくことを見ればよいです。

最適輸送を実現する流れを求めるアルゴリズムまとめ

これまでの結果をまとめるとアルゴリズムは以下のようになります。

$\mathbb{R}^n$ の有界領域を $D$ とし、全ての場は $D$ 内にsupportを持つとします。
始めに２つの $D$ 上の密度 $\rho_0(x),\rho_T(x)$ が与えられているとします。
時刻を $t\in [0,T]$ とします。

$k=1,2,\cdots$ に対して、$\phi^k(t,x)$ を 時間に依存する $D$ 上のスカラー関数、$\rho^k(t,x),\ a^k(t,x)$ を 時間に依存する $D$ 上の非負のスカラー関数、$m^k(t,x),\ b^k(t,x):[0,T]\times D\to \mathbb{R}^n$ を $D$ 上の時間に依存するベクトル場とし、$\mu^k=(\rho^k,m^k),\ q^k=(a^k,b^k)$ を $[0,T]\times D$ 上のベクトル場とします。

$(\phi^{n-1}, q^{n-1}, \mu^n)$ が与えられたとします。
Step A : 以下を満たす $\phi^n$ を見つけます。

$$-r\Delta_{t,x}\phi^n=\textrm{div}_{t,x}(\mu^n-rq^{n-1})\\ r\frac{\partial}{\partial t}\phi^n(T,x)=\rho_T(x) -\rho^n(T,x) + ra^{n-1}(T,x)\\ r\frac{\partial}{\partial t}\phi^n(0,x)=\rho_0(x) -\rho^n(0,x) + ra^{n-1}(0,x)$$

Step B : 以下を満たす $q^n=(a^n,b^n)$ を見つけます。

$$\alpha^n(t,x):=\partial_t\phi^n(t,x) + \frac{\rho^n(t,x)}{r}\\ \beta^n(t,x):=\textrm{grad}_x\phi^n(t,x) + \frac{m^n(t,x)}{r}\\ \inf_{(a,b)}\{(a^n-\alpha^n)^2+|b^n-\beta^n|^2;\ a+\frac{|b|}{2}\le0\}$$

Step C : 以下のように $\mu^n$ を更新します。

$$\mu^{n+1} = \mu^n + r (\nabla_{t,x} \phi^n - q^n)$$

Step A に戻ります。

StepA, B, Cを繰り返し、

$$\textrm{crit}^n =\sqrt{\frac{\int_0^T\int_D\rho^n\textrm{res}^n}{\int_0^T\int_D\rho^n|\textrm{grad}_x\phi^n|^2}}$$

が０に収束したときに停止し、得られた解を $(m^\ast, \rho^\ast, \phi^\ast)$ とします。

時間に依存するベクトル場

$$v^\ast(t,x)=\frac{m^\ast(t,x)}{\rho^\ast(t,x)}$$

に対する非自励系

$$\partial_t\Phi(t,x)=v^\ast(t,x)$$

の解 $\Phi:[0,T]\times D\to D$ が$\rho_0(x)$ を $\rho_T(x)$ に写す最適写像を与える１パラメータ変換です（非自励系なので一般には１パラメータ変換群にはならないことに注意）。

このアルゴリズムの利点として著者は以下のようなことが挙げています。

• 一番計算が重いStepがPoisson方程式を解くStep Aだが、Poisson方程式の解き方は開発されており高速なsolverが利用できる
• コーディングが比較的簡単である
• 計算コストが比較的低い

実装は次の記事で行う予定です。

Appendix

ここでは復習の節で出てきた不等式

$$T\int_{\mathbb{R}^n}\int^T_0\rho(t,x)|v(t,x)|^2dx^ddt\geq \int_{\mathbb{R}^n}|M(x)-x|^2\rho_0(x)dx^d$$

の証明をします。

$\Phi:[0,T]\times D\to D$ を $v(t,x)$ が作る１パラメータ変換とします。任意の関数 $f(t,x)$ に対して、

$$f(t,x)\rho(t,x)dx^n=f(t,\Phi(t,x'))\rho(t,\Phi(t,x'))|d\Phi(t,x')|(dx')^n\\ =f(t,\Phi(t,x'))\rho_0(x')(dx')^n$$

が成り立つので、

$$T\int_{\mathbb{R}^n}\int^T_0\rho(t,x)|v(t,x)|^2dx^ddt =T\int_{\mathbb{R}^n}\int^T_0\rho_0(x)|v(t,\Phi(t,x))|^2dx^ddt\\ =T\int_{\mathbb{R}^n}\int^T_0\rho_0(x)|\partial_t\Phi(t,x)|^2dx^ddt =\int_{\mathbb{R}^n}\rho_0(x)\sum_{i=1}^nT\int^T_0|\partial_t\Phi_i(t,x)|^2dx^ddt\\ \geq \int_{\mathbb{R}^n}\rho_0(x)\sum_{i=1}^n|\Phi_i(T,x)-\Phi_i(0,x)|^2dx^d = \int_{\mathbb{R}^n}\rho_0(x)|\Phi(T,x)-\Phi(0,x)|^2dx^d\\ = \int_{\mathbb{R}^n}|M(x)-x|^2\rho_0(x)dx^d$$

となります。3行目の不等式はCauchy–Schwarzの不等式を使うと任意の関数 $f(t)$ に対して

$$T\int_0^T|\partial_t f(t)|^2dt\geq |f(T)-f(0)|^2$$

となることを使いました。

さらに古典的な結果として、最適写像 $M$ は適当な関数 $\Psi(x)$ があり、

$$M(x)=\textrm{grad}\Psi(x)$$

で与えられることが知られています。今、

$$\Phi(t,x) = x + \frac{t}{T}(\textrm{grad}\Psi - x)$$

とおくと、$t$ に関して一次式であるから上の不等式は等号が成り立ち、かつ右辺の最小値でもああります。よって左辺の最小値を与える $(\rho, v)$ は左辺を最小にします。