最適輸送が遅すぎる（スライス法による解法）

Workshop OT 2023 http://webpark2072.sakura.ne.jp/otworkshop/ での講演スライドです。

最適輸送が重すぎて動かない問題に対する有効な解決策であるスライス法についてと "Fast Unbalanced Optimal Transport on a Tree" (NeurIPS 2020) を紹介します。
論文： https://arxiv.org/abs/2006.02703
コード: https://github.com/joisino/treegkr

『最適輸送の理論とアルゴリズム』好評発売中！ https://www.amazon.co.jp/dp/4065305144

Transcript

1. 1 KYOTO UNIVERSITY KYOTO UNIVERSITY 最適輸送が遅すぎる スライス法による解法 佐藤 竜馬
2. 2 / 39 KYOTO UNIVERSITY 点群の比較を行いたい  問題設定： 点群赤と青が与えられるので距離を求めたい
3. 3 / 39 KYOTO UNIVERSITY 点群の比較を行いたい  問題設定： 点群赤と青が与えられるので距離を求めたい → 最適輸送でできる！（もちろん！）
4. 4 / 39 KYOTO UNIVERSITY 最適輸送が遅すぎる  残念なお知らせ：最適輸送は遅い だいたい O(n3) 時間かかる n = 103 だと 1 秒 目安 n = 104 だと 20 分 n は点の数 n = 105 だと 300 時間
5. 5 / 39 KYOTO UNIVERSITY 補足：厳密に言うと三乗時間ではない  補足：「最適輸送は三乗時間」について 点群の場合は理論上 時間で解ける [Lee+ FOCS 2014] 論文に書く時の言い回しは注意が必要。 実用的な実装はないので現実的には「最適輸送は三乗時間」で 問題ないです。 Yin Tat Lee and Aaron Sidford. Path finding methods for linear programming: Solving linear programs in O˜(√rank) iterations and faster algorithms for maximum flow. FOCS 2014.
6. 6 / 39 KYOTO UNIVERSITY Sinkhorn アルゴリズムは早いとはいえ大規模だと苦しい  近似で良いなら Sinkhorn アルゴリズム O(n2) 時間で動く n = 103 だと 1 秒 目安 n = 104 だと 2 分 n は点の数 n = 105 だと 3 時間 n = 106 だと 300 時間 大規模データだとこれでも苦しい……
7. 7 / 39 KYOTO UNIVERSITY スライス法
8. 8 / 39 KYOTO UNIVERSITY スライス法は一次元に射影して最適輸送を計算する  スライス法の概要 オプション：様々な射影を試して平均する 入力点群を 適当な一次元に射影 一次元空間で最適輸送を計算
9. 9 / 39 KYOTO UNIVERSITY スライス法の擬似コード In: 点群 {x i }, {y i }, 試行回数 K Out: 点群の距離 For k in 1...K: 射影ベクトル w をランダムにサンプリング {x i }, {y i } を w 方向に射影して {x’ i }, {y’ i } を得る {x’ i } と {y’ i }（一次元点群）の最適輸送コストを計算 Return 最適輸送コストの平均値
10. 10 / 39 KYOTO UNIVERSITY 一次元の最適輸送は O(n log n) 時間で解ける  ポイント：一次元の最適輸送は簡単に解ける 一次元の点群は一番左の赤を一番左の青に 二番目に左の赤を二番目に左の青に……とマッチさせるのが最適 最適輸送ではマッチ相手の最適化に時間がかかるが、 一次元ではこれが簡単に定まる 一次元の最適輸送は O(n log n) 時間で解ける（ソートの時間）
11. 11 / 39 KYOTO UNIVERSITY スライス法は大規模でも計算できる  スライス法は O(n log n) 時間で動く n = 103 だと 0.00001 秒 n = 104 だと 0.0001 秒 目安 n = 105 だと 0.001 秒 n = 106 だと 0.01 秒 n = 107 だと 0.1 秒 n = 108 でも 1 秒 大規模データも余裕 最適輸送が重すぎて動かない問題に対してはスライス法は 簡単で有効な解決策
12. 12 / 39 KYOTO UNIVERSITY スライス法は距離の公理を満たす  スライス法の射影について期待値を取ったものは距離の公理を満たす → ちゃんと「点群の距離」になっている → 点群の比較を行いたいという当初の目的には適っている  一定の条件のもと、元の空間での最適輸送コストとスライス法は距離 として同値である（上下から定数倍で抑えられる）ことも示せる [Bonnotte 2013 Theorem 5.1.5.]。  現実的には射影についてモンテカルロ近似するので上記の性質は 厳密には満たさないことに注意 N. Bonnotte. Unidimensional and Evolution Methods for Optimal Transportation. Ph.D. Thesis 2013.
13. 13 / 39 KYOTO UNIVERSITY よりよいスライス法を求めて  スライス法の改善 1. 一次元点群の表現力は弱い。もっと良い表現方法はないか？ 二次元に射影したり、一般の低次元空間に射影したり……？ 2. 最適輸送の他の変種や最適輸送に基づいた問題も 同じアイデアで早く解けないか 不均衡最適輸送 ワッサースタイン重心など
14. 14 / 39 KYOTO UNIVERSITY もっと良い表現方法はないか
15. 15 / 39 KYOTO UNIVERSITY 二次元だと解けない。木だと解ける。  一次元点群の表現力は弱い。もっと良い表現方法はないか？  二次元の最適輸送は残念ながら効率よく解けない 一般の低次元空間への射影も当然無理  木上の点群だと効率よく解けることが知られている  木は枝分かれのある一次元の一般化  一次元（= パスグラフ）は 木の特殊ケース 50 60 10 40 70 20 30 50 のコストは 20 + 30 + 50 = 100 パスグラフ
16. 16 / 39 KYOTO UNIVERSITY 木の場合は線形時間で解ける  木の上の点群の場合も貪欲法で解ける  やり方: (1) DFS 帰りがけ順でノードを走査する (2) 現在のノードの子孫でまだ赤点と青点が残っていれば輸送する 子孫には一方の色の点しかなくなるまで繰り返し なくなれば次のノードへ  O(n) 時間で計算できる  一次元の場合は、ソートするのが木を構築するところに相当 ソートが終わればあとは線形時間なので、木の場合も一次元の場合も 同じ計算量
17. 17 / 39 KYOTO UNIVERSITY 高次元でも木に帰着して簡単にする  入力点群が高次元の場合、クラスタリング木などを構築して 木上の点群に変換（＝射影に対応）  さまざまな初期値のクラスタリング木を使って平均することも提案されて いる [Le+ NeurIPS 2019]。  良い木を使うと元の輸送距離の近似にもなる Tam Le, Makoto Yamada, Kenji Fukumizu, Marco Cuturi. Tree-Sliced Variants of Wasserstein Distances. NeurIPS 2019. http://www.sthda.com/english/articles/31-principal-component-methods-in-r-practical-guide/117-hcpc-hierarchical-clustering-on-principal-components-essentials/
18. 18 / 39 KYOTO UNIVERSITY 変種に適用できないか 不均衡最適輸送について
19. 19 / 39 KYOTO UNIVERSITY スライス法は不均衡最適輸送に適用できる？  最適輸送の他の変種や最適輸送に基づいた問題も同じアイデアで早 く解けないか  点群にノイズが多い場合には不均衡最適輸送がよく使われる 外れ値 コスト大 消去 通常の最適輸送： 外れ値（センサーエラーなど） の影響が大きい 不均衡最適輸送： 一定コストを払うと質量を生成・ 削除できる定式化 外れ値に頑健
20. 20 / 39 KYOTO UNIVERSITY スライス法は不均衡最適輸送に適用できる！  不均衡最適輸送は一次元の場合でもかなり難しい（後述） ……が実は効率よく解ける！[Sato+ NeurIPS 2020] 一次元の場合 O(n log2 n) 時間で解ける 木の場合でも O(n log2 n) 時間で解ける n = 106 でも 1 秒程度
21. 21 / 39 KYOTO UNIVERSITY スライス法でできることまとめ 　 　 　 　 　 　 　 空間 コスト・定式化＼ 一次元 超距離木 （木の特殊ケース） 一般の木 1-Wasserstein O(n log n) 時間 [folklore] O(n) 時間 [folklore, Charikar STOC 2002, Le+ NeurIPS 2019] O(n) 時間 [folklore, Charikar STOC 2002, Le+ NeurIPS 2019] p-Wasserstein O(n log n) 時間 [folklore] O(n) 時間 [Kloeckner Mathematika 2015] x 不均衡 1-Wasserstein O(n log2 n) 時間 [Sato+ NeurIPS 2020, Bonneel+ ToG 2019] O(n) 時間 [Heinemann+ AMO 2023] O(n log2 n) 時間 [Sato+ NeurIPS 2020] 不均衡 p-Wasserstein x O(n) 時間 [Heinemann+ AMO 2023] x 対応する論文名はスライド末尾の参考文献を参照
22. 22 / 39 KYOTO UNIVERSITY Fast Unbalanced Optimal Transport on a Tree. [Sato+ NeurIPS 2020]
23. 23 / 39 KYOTO UNIVERSITY 点群の比較を行いたい  問題設定： 一次元点群赤と青が与えられる。 不均衡最適輸送コストを求めたい。 コスト 1 で単位質量を一単位動かせる。 コスト λ で単位質量を消滅・生成できる。 最小コストで赤と青を一致させるコスト。 λ = ∞ なら普通の 最適輸送
24. 24 / 39 KYOTO UNIVERSITY 不均衡最適輸送は貪欲法で解けない  復習：一次元の通常の最適輸送は貪欲法で解ける 残念なお知らせ：不均衡最適輸送は貪欲法で解けない
25. 25 / 39 KYOTO UNIVERSITY 不均衡最適輸送の難しさ  点群のうち、左の方が明らかになったする  赤→青の輸送を発生させるべき？（通常の最適輸送なら常に yes) 未処理 x = 0 x = 5 x
26. 26 / 39 KYOTO UNIVERSITY 不均衡最適輸送の難しさ  この先がこうだったら  消すのが最適 x = 0 x = 5 x x = 6 x = 10 x = 11 x = 15 x = 0 x = 5 x x = 6 x = 10 x = 11 x = 15
27. 27 / 39 KYOTO UNIVERSITY 不均衡最適輸送の難しさ  この先がこうだったら  消さないのが最適 x = 0 x = 5 x x = 13 x = 14 x = 15 x = 16 x = 0 x = 5 x x = 13 x = 14 x = 15 x = 16
28. 28 / 39 KYOTO UNIVERSITY 不均衡最適輸送は途中まで見て解が定まらない  残念なお知らせ： 不均衡最適輸送では途中まで見て解が定まらない
29. 29 / 39 KYOTO UNIVERSITY 不均衡最適輸送は途中まででも少しは解が定まる  希望の光：ただし全く解が定まらないかというとそうではない 未処理 x 未処理 x
30. 30 / 39 KYOTO UNIVERSITY 不均衡最適輸送は途中まででも少しは解が定まる  希望の光：ただし全く解が定まらないかというとそうではない  上の解と下の解はどちらも同じ「未処理の右側」に繋がる どのような右側が来ても、上の解の方がコストが低い x x
31. 31 / 39 KYOTO UNIVERSITY 不均衡最適輸送は途中まででも少しは解が定まる  希望の光：ただし全く解が定まらないかというとそうではない  下の解はこの時点で最適解の可能性なしとして切り捨てることができる 未処理 x 未処理 x
32. 32 / 39 KYOTO UNIVERSITY 不均衡最適輸送は途中まででかなり解が絞れる  嬉しいお知らせ： 途中まで見た時点で解の候補は高々 O(n) 通りにまで絞れる 赤・青がどれだけ余っているかで分類する。 余っている個数が違う解どうしは比較ができない（一方に絞れない） 個数が同じ解どうしは、そこまでのコストを見ると一方に絞れる。  組合せ最適化の言葉でいうと、動的計画法の最適性の原理が 成り立つ  ということで、動的計画法で解ける
33. 33 / 39 KYOTO UNIVERSITY ナイーブな動的計画法では遅いが道具を駆使して高速に  残念なお知らせ： 状態数が O(n2) 個ある ナイーブに実装すると O(n3) 時間かかる 工夫すると O(n2) になるが状態数が多いのでこれ以上減らせない  嬉しいお知らせ：ここでは紹介しきれないが、キーワードだけ列挙すると  高速な凸 min-sum 畳み込み  平衡二分木  Weighted union heuristics を組合せると O(n log2 n) 時間で計算できる [Sato+ NeurIPS 2020] どれか一つでも欠けると O(n2) 時間になる
34. 34 / 39 KYOTO UNIVERSITY 提案法は 100 万点を 1 コアで 1 秒で処理できる  一般の木（一次元を特殊ケースとして含む一般的な設定）  不均衡最適輸送（通常の最適輸送を特殊ケースとして含む）  提案法の速度  106（100 万）サイズの点群を 1 コアノート PC で 1 秒で処理  論文: https://arxiv.org/abs/2006.02703  コード: https://github.com/joisino/treegkr
35. 35 / 39 KYOTO UNIVERSITY どこまで広い空間に一般化できるかは謎も多い  スライス法についての他のトピック： Series-Parallel graph というグラフクラス上の通常の最適輸送は O(n log2 n) 時間で解けることが知られている [Tseng+ J. Alg. 1996]  高さ 2 の二次元グリッドも効率よく解ける特殊ケース  最適輸送が効率よく解ける範囲 （とくに応用にとってできるだけ便利な空間） を広げていく・どこまでできるか見極める のが大事な未解決問題 s s s s t t t t + parallel composition series composition Paul Tseng. Zhi-Quan Luo. On Computing the Nested Sums and Infimal Convolutions of Convex Piecewise-Linear Functions. J. Alg. 1996. サイクルを含む 複雑な空間を 表現できる コレ
36. 36 / 39 KYOTO UNIVERSITY まとめ：スライス法は早くてすごい  非常に効率よく点群の距離を求められるスライス法を紹介した  スライス法は一次元だけでなく木に一般化できる  スライス法は最適輸送だけでなく不均衡最適輸送に一般化できる 最適輸送が重すぎて動かない問題に対しては スライス法は簡単で有効な解決策 Take Home Message
37. 37 / 39 KYOTO UNIVERSITY 宣伝：最適輸送の本を書きました  『最適輸送の理論とアルゴリズム』という本を書きました  スライス法にも 1 章割いて丁寧に解説しています（第 5 章） https://www.amazon.co.jp/dp/4065305144
38. 38 / 39 KYOTO UNIVERSITY Reference
39. 39 / 39 KYOTO UNIVERSITY 参考文献  Moses S. Charikar. Similarity estimation techniques from rounding algorithms. STOC 2002.  Tam Le, Makoto Yamada, Kenji Fukumizu, Marco Cuturi. Tree-Sliced Variants of Wasserstein Distances. NeurIPS 2019.  Benoît Kloeckner. A geometric study of Wasserstein spaces: ultrametrics. Mathematika 2015.  Ryoma Sato, Makoto Yamada, Hisashi Kashima. Fast Unbalanced Optimal Transport on a Tree. NeurIPS 2020.  Nicolas Bonneel and David Coeurjolly. SPOT: sliced partial optimal transport. ToG 2019.  Florian Heinemann, Marcel Klatt, Axel Munk. Kantorovich-Rubinstein distance and barycenter for finitely supported measures: Foundations and Algorithms. AMO 2023.