私と理論

数学の話を主に書きます

この広告は、90日以上更新していないブログに表示しています。

2021-05-09

輸送問題を近似的に行列計算で解く（機械学習への応用つき）

輸送問題と呼ばれる問題があります．この問題は，普通は線形計画法やフローのアルゴリズムを使って解かれます．この記事では，この輸送問題を近似的に行列計算で解くアルゴリズム（エントロピー正則化 + Sinkhorn-Knoppアルゴリズム）を紹介します．

輸送問題とは

輸送問題とは以下のような問題です．

件の工場と件の店舗からなる，ある商品の流通圏があるとする．
各工場には $a_i \geq 0$ 個の在庫がある.．
各店舗では $b_i \geq 0$ 個の需要がある．
在庫の総和と需要の総和は等しいとする（すなわち $\sum_{i \in [n] } a_i = \sum_{i \in [n]} b_i$ ）．
工場から店舗に商品を一つ運ぶためには $C_{ij}$ の輸送コストがかかる．
各工場から各店舗への輸送量 $x_{ij}$ を適切に決めて，各店舗の需要を満たしつつ輸送コストの総和を最小化せよ．

輸送問題は最適化の歴史の中でもかなり初期から考えられている問題です．現実世界での応用が多く，さらに最大マッチングや最小重みマッチングを特殊ケースとして含む，重要な最適化問題です．(2023/11/13 註: 工場の数と店舗の数は違っていても以下と同様の議論が可能です．この記事では，説明の簡単のため工場の数と店舗の数が等しいとしています)

さて，この問題は数式で書くと以下のような最適化問題になります．

$\begin{align*}\begin{aligned}\min_{\boldsymbol{x} \in \mathbb{R_{\geq 0}^{n \times n}}}.& && \sum_{i, j \in [n]} C_{ij} \cdot x_{ij} \\\mathrm{s.t.}& && \sum_{j \in [n]} x_{ij} = a_i \quad \forall i \in [n], \\&&& \sum_{i \in [n]} x_{ij} = b_j \quad \forall j \in [n] \\\end{aligned}\end{align*} \tag{1}$

$\boldsymbol{x}$ は整数値であるという制約を加えることもありますが，今回は連続値だとしましょう．

この問題は線形計画問題と呼ばれる問題クラスに入りますので，線形計画問題のアルゴリズム（単体法，内点法）を用いて解くことができます．さらに言えば，この問題はその中でもさらに特殊な，最小費用流問題という問題に帰着することができ，専用のアルゴリズムを適用することでより高速に解くことができます．

しかし，線形計画問題や最小費用流問題を解くためのアルゴリズムは比較的複雑で，実装するのは一手間です．また，計算量が $\Theta (n^3)$ 程度かかるのであまりスケールしません．これらの問題を解決するアルゴリズムはないのでしょうか？

実は近似解で良ければそのようなアルゴリズムを構築することができます．このアルゴリズムは以下の特長を持ちます．

実装が簡単
速い
単純な行列計算のみからなるのでGPU などで高速化も可能

これが，本記事のテーマである「エントロピー正則化項を加えた輸送問題を Sinkhorn-Knoppアルゴリズムを用いて解く手法」です．

アルゴリズム

アルゴリズムは以下です．

ハイパーパラメータ $\lambda > 0$ を決める．
$K_{ij} = \exp(-\lambda \cdot C_{ij})$ で定義される行列を計算する．
ベクトル $\boldsymbol{u} \in \mathbb{R}^{n}$ を $\boldsymbol{u} \leftarrow [1, 1, \ldots, 1$ ] と初期化する．
「 $\boldsymbol{v} \leftarrow \boldsymbol{b} \ / \ (K^{\top} \boldsymbol{u} )$ の後に $\boldsymbol{u} \leftarrow \boldsymbol{a} \ / \ (K \boldsymbol{v} )$ 」という計算を， $\boldsymbol{u}$ が収束するまで繰り返す．ただし，はベクトルの要素ごとの除算を意味する．
をそれぞれ $\boldsymbol{u}, \boldsymbol{v}$ の要素を対角に並べた行列として，行列を出力する．

これだけです．とてもシンプルですね．20行くらいで実装できそうです．なんとこのシンプルなアルゴリズムがかなり良い解を出力してくれるのです．

ハイパーパラメータ $\lambda > 0$ は，アルゴリズムの性能をコントロールします． $\lambda$ が大きいほど出力が真の解に近づきますが，収束が遅くなります．また， $K_{ij} = \exp(-\lambda \cdot C_{ij})$ を計算する必要があるため， $\lambda$ が大きいとき数値的に扱いが難しくなります．

得られる解の例

ランダムに作った輸送問題のケースを，最小費用流及び Sinkhorn-Knoppアルゴリズムで解いた際の解 $\boldsymbol{x}$ の様子を図示したのが以下の図です．左上が最小費用流のアルゴリズムによる最適解であり，それ以外は $\lambda$ を色々変えた場合の Sinkhorn-Knoppアルゴリズムの出力です．

見ての通り， $\lambda$ を小さく設定した場合は Sinkhorn-Knoppアルゴリズムの出力は「ボヤッと」していて最適解からは程遠いですが， $\lambda$ が大きくなるにつれて最適解に近づいていることが分かります．それに伴い，目的関数値も最適値に近づいていきます．ちなみに全ての $\lambda$ について，制約を（ほぼ）満たす解が出力されています．

なぜこれで解けるのか？

気になるのはなぜこんなアルゴリズムで輸送問題が解けてしまうか？です．これについて，数学的に厳密な議論ではないですが直感的に説明します．

実は，このアルゴリズムは輸送問題を少し変形した以下の問題を解いています．

$\begin{align*}\begin{aligned}\min_{\boldsymbol{x} \in \mathbb{R_{\geq 0}^{n \times n}}}.& && \sum_{i, j \in [n]} C_{ij} \cdot x_{ij} \color{red}{+\frac{1}{\lambda} \sum_{i, j \in [n]} x_{ij} \log x_{ij}} \\\mathrm{s.t.}& && \sum_{j \in [n]} x_{ij} = a_i \quad \forall i \in [n], \\&&& \sum_{i \in [n]} x_{ij} = b_j \quad \forall j \in [n] \\\end{aligned}\end{align*} \tag{2}$

目的関数に何か余計な項が入っていますね． $\sum_{i, j \in [n] } x_{ij} \log x_{ij}$ という項は， $\boldsymbol{x}$ を確率分布と見做した場合のシャノンのエントロピー（にをかけたもの）として捉えることができます．そのため，この項はエントロピー正則化項と呼ばれます．

$\lambda$ はエントロピー正則化の強さをコントロールする係数です． $\lambda$ が大きくなるともとの輸送問題に近づき，小さくなるともとの輸送問題から遠ざかります．エントロピー正則化項を小さくするということは， $\boldsymbol{x}$ を全体的に「ボヤッとさせる」「曖昧にさせる」ことになります（ $\boldsymbol{x}$ の要素が全て等しい場合にエントロピーは最大になる）．

さて，エントロピー正則化項を足した新しい最適化問題について，ラグランジュ関数を書いてみると，

$\begin{align*}L(x, \boldsymbol{\alpha}, \boldsymbol{\beta}) = \sum_{i, j \in [n] } \left( \frac{1}{\lambda} x_{ij} \log x_{ij} + C_{ij} x_{ij} \right) + \boldsymbol{\alpha}^{\top} \left(x \boldsymbol{1}_{n} - \boldsymbol{a} \right) + \boldsymbol{\beta}^{\top} \left(x^{\top} \boldsymbol{1}_{n} - \boldsymbol{b} \right)\end{align*}$

となります．ただし， $\boldsymbol{1}_{n}$ は全ての要素がの次元縦ベクトル，は $x_{ij}$ を並べた $n \times n$ 行列です．

このとき簡単な計算によって，全ての $i, j \in [n$ ] について

$\begin{align*}\frac{\partial L}{\partial x_{ij}} = 0 \Rightarrow x_{ij} = u_i \cdot K_{ij} \cdot v_j\end{align*} \tag{3}$

となります．ただし， $u_i := e^{-\frac{1}{2} -\lambda \alpha_i}, v_j := e^{-\frac{1}{2} -\lambda \beta_j}$ とおきました．

数式 (3) の言っていることは，「行列は，行列の行目に u_i を，列目に v_j を掛けた形になっている」ということです．

さらに，行列はもとの問題の制約

$\begin{align*}\sum_{j \in [n]} x_{ij} &= a_i \quad \forall i \in [n] \tag{4} \\\sum_{i \in [n]} x_{ij} &= b_j \quad \forall j \in [n] \tag{5}\end{align*}$

を満たす必要があります．

(3)(4)(5) を満たすような行列（もしくは $\boldsymbol{u}$ 及び $\boldsymbol{v}$ ）を見つけるため，以下のような簡単なアルゴリズムを考えてみます．

$\boldsymbol{u}$ を適当に初期化する．
$\boldsymbol{u}$ を固定し，(4) の条件は一旦無視して，(3)(5) の条件を満たすようにうまく $\boldsymbol{v}$ を決める．これは， $\boldsymbol{v} \leftarrow \boldsymbol{b} \ / \ (K^{\top} \boldsymbol{u} )$ とすれば OK．
$\boldsymbol{v}$ を固定し，(5) の条件は一旦無視して，(3)(4) の条件を満たすようにうまく $\boldsymbol{u}$ を決める．これは， $\boldsymbol{u} \leftarrow \boldsymbol{a} \ / \ (K \boldsymbol{v} )$ とすれば OK．
2, 3 を収束するまで繰り返す．

行の定数倍と列の定数倍を，制約を満たすように交互に繰り返すイメージです．このアルゴリズムが，まさに Sinkhorn-Knoppアルゴリズムと呼ばれるものです．直観的で単純なアルゴリズムですが，実はこのアルゴリズムがうまくいくことを理論的に示すことができます．特に，

問題 (2) にはユニークな最適解が存在すること
上記のアルゴリズムの出力がユニークな最適解に収束すること

が示されています．厳密な数学的な議論に関してはこの論文の Chapter 4 などを参照してください．

競プロの問題を解いてみる

紹介したアルゴリズムは近似解法であり，厳密な最適解を求めさせることの多い競プロの問題（特にマラソンではなくアルゴリズムのコンテスト）では適用は難しそうです．しかしどこまで解けるのかも気になるということで， opt さん (@opt_cp) がいくつかの問題（これとこれ）に挑戦してくれました！

方法は，輸送問題として定式化→エントロピー正則化 + Sinkhorn-Knoppアルゴリズムで解く→結果を適当に丸めることで整数実行可能解を得る，というものです．

結果としては，

ACを取るのは難しそう
ハイパーパラメータ $\lambda$ の調整が難しい．小さいと最適解から離れてしまうし，大きいと収束が遅い＆オーバーフローする．オーバーフローの問題は logsumexp を用いたテクニック（詳細はこれを参照）で回避できるが，丸めて厳密解が得られるくらいの精度に辿り着くにはやはり反復回数が足りない．
ただ，それほど大きくないケースなら通せる．また，行列のサイズが $10000 \times 3$ くらいなら（この記事ではコスト行列が正方行列の場合を扱ったが，そうでない場合にも簡単に一般化可能），最適値が $10^{13}$ くらい (整数) になる問題も 2sec 以内に厳密に解けることがあった．

という感じのようです．ただ，今回適用したのはかなりシンプルな手法であり，色々と発展手法もあるようなので，それらを用いればまた状況も変わるかもしれません．また，厳密解が必要ではない場面（マラソンなど）では使える可能性はあると思います．

機械学習界隈における流行

さて，話は変わるのですが，輸送問題を Sinkhorn-Knoppアルゴリズムで解く周辺の話は機械学習界隈で大流行しています．このことについて少し説明します．

冒頭の問題では $\boldsymbol{a}$ や $\boldsymbol{b}$ は商品の量として説明しましたが，これを次元空間における確率分布（確率ベクトル）として捉えてみましょう．例えば $\boldsymbol{a} = [0.3, 0.1, 0.6$ ] などは次元確率ベクトルです．

このとき，問題 (1) の最適値は，「確率分布 $\boldsymbol{a}$ から確率分布 $\boldsymbol{b}$ に確率を輸送する際に，必要な最小コスト」になります．これを $L(\boldsymbol{a}, \boldsymbol{b})$ と書くことにしましょう．実は， $L(\boldsymbol{a}, \boldsymbol{b})$ は，確率分布 $\boldsymbol{a}$ と確率分布 $\boldsymbol{b}$ の近さ，つまり確率分布間の距離として用いることができます．実際，コスト行列 $C_{ij}$ が距離の公理を満たすとき， $L(\boldsymbol{a}, \boldsymbol{b})$ も距離の公理を満たす距離関数になることが知られています． $L(\boldsymbol{a}, \boldsymbol{b})$ は最適輸送距離と呼ばれます．

さて，機械学習においては確率分布間の近さを測ることはひじょーーーーーーに重要です．最も基本的な操作と言っても良いくらいだと思います．一般にこの用途でよく用いられるのは Kullback-Leibler Divergence や Jensen–Shannon Divergence などの尺度です．しかし，最適輸送距離はこれらにはない特長をいくつも持つため，重要視されています（「もとの空間の距離を反映した確率分布間の距離になっている」「二つの確率分布のサポートが一致しない場合にも定義できる」など）．

このような重要な最適輸送距離ですが，一つ大きな欠点を抱えていました．計算量です．というのも，最適輸送距離を求めるためには上で述べたように輸送問題を解く必要があり，大きな計算量がかかってしまいます．機械学習ではかなり高次元の確率分布も扱うので，この計算量は重すぎます．

この問題を解決し，最適輸送距離の応用範囲を大きく広げたのが，今回紹介した Sinkhorn-Knoppアルゴリズムによる近似的な最適輸送距離の計算です．この手法は NIPS 2013 の"Sinkhorn Distances: Lightspeed Computation of Optimal Transport" という論文で提案されました．現在ではこの論文は 1400 件近く引用されている，この分野では記念碑的な論文になっています．流行は続いており，最適輸送に関する論文が機械学習の国際会議である NeurIPS や ICML，ICLR などでも現在でもたくさん発表されているという状況です．（ただ，最適輸送が流行っているのは，最適輸送を使った GAN が有名になった等の別の理由もあったりします）

ちなみに，この最適輸送問題は Monge-Kantorovich の問題と呼ばれたりもして，競プロでお馴染み Monge さんがこんなところにも現れます．また，コスト行列が Monge 性を満たす場合，輸送問題は効率的に解けるという話もあったりします．