1.
2021/01/10
最適輸送の計算アルゴリズム
の研究動向
1
@ohken322

2.
目次
1. 背景：最適輸送理論と応用
2. 事前準備：Kantorovichの問題とWasserstein距離
3. エントロピー正則化とSinkhornアルゴリズム
4. トピック1 ： エントロピー正則化のDebiasing
5. トピック2 ： ロバストな最適輸送
2

3.
重要な参考文献
1. Peyré, Cuturi “Computational Optimal Transport”,
Foundations and Trends in Machine Learning 2019
• 最適輸送のアルゴリズムに関する非常に明快なサーベイ論文
(引用元を書いていない図はこの論文から)
• 今日の目標は、この論文後の理論の発展をざっくり追うこと
2. Cuturi “Sinkhorn Distances: Lightspeed Computation of Optimal Transport”, NIPS 2013
• 正則化付きOTを考案し近似解の軽量・高速な計算を実現した記念碑的論文
3

4.
記号の準備
• 𝒳, 𝑑 : 距離空間 (ユークリッド空間 𝒳 = ℝ 𝑑
, 𝑑 𝑥, 𝑦 = | 𝑥 − 𝑦 | と思ってよい)
• 𝒫 𝒳 : 𝒳上の確率測度全体の集合
• 行列 𝐴 ∈ ℝ 𝑛×𝑛, ベクトル 𝒂 ∈ ℝ 𝑛
• diag 𝒂 ∈ ℝ 𝑛×𝑛
は 𝑖, 𝑖 -成分が𝑎𝑖の対角行列
• 𝑒 𝐴,
𝒂
𝒃
などは要素ごとの演算を表す
• 𝒂 ⊙ 𝒃 ≔ 𝑎𝑖 𝑏𝑖 𝑖 ∈ ℝ 𝑛, 𝐴, 𝐵 ≔ tr 𝐴 𝑇 𝐵 = σ𝑖,𝑗 𝐴𝑖𝑗 𝐵𝑖𝑗
• すべての成分が1の𝑛次元ベクトルを𝟏 𝑛と書く
4

5.
最適輸送(Optimal Transport)理論とは?
𝜇, 𝜈 ∈ 𝒫 𝒳 ∶距離空間 (𝒳, 𝑑) 上の確率測度
最適輸送 = 「𝜇を𝜈に最も効率よく動かす」
𝑇: 𝒳 → 𝒳 s.t. 𝑇# 𝜇 = 𝜈 を輸送写像と呼ぶ
単位量を 𝑥 ∈ 𝒳から𝑦 ∈ 𝒳に運ぶために
コスト 𝑐 𝑥, 𝑦 = 𝑑(𝑥, 𝑦) がかかる
Mongeの最適輸送問題 (Monge, 1781)
輸送コスト 𝐶 𝑇 ≔ ‫׬‬𝒳
𝑐(𝑥, 𝑇 𝑥 )d𝜇 を最小化せよ
一般にはコスト関数 𝑐 ∶ 𝒳 × 𝒳 → ℝ は距離に限らない
5

6.
応用分野
 画像処理
 自然言語処理
 機械学習
その他いろいろ…
3D形状マッチング
(Eisenberger, et al. 2020)
色の変換、脳の活性部位検知、形状変換、自然言語処理、データ分布比較
(workshop[1], 2017)
クロスドメインの概念マッチング
(Chen, et al. 2020)
データ生成/データ補完
(Muzellec, et al. 2020)
個体群のモデリング
(Tong, et al. 2020)
ICML/NeurIPS 2020 での応用研究例
[1] Optimal Transport meets Probability, Statistics and Machine Learning 6

7.
Kantorovichの問題
Mongeの問題は不良設定問題 (#Supp(𝜇) < #Supp 𝜈 のときは輸送写像が存在しない)
Kantorovichの問題 (Kantorovich, 1942)
𝜋 ∈ 𝒫(𝒳2
) s.t. proj1 # 𝜋 = 𝜇, proj2 # 𝜋 = 𝜈 を輸送計画という.
輸送コスト 𝐶 𝜋 ≔ ‫׬‬𝒳×𝒳
𝑐 𝑥, 𝑦 d𝜋(𝑥, 𝑦) を最小化せよ.
Rem: 𝒳:コンパクト、𝑐:連続ならば解が存在する
7

8.
Wasserstein距離
 最適輸送のコストは、確率測度間の距離を定める
𝑝 ≥ 1, 𝑐 𝑥, 𝑦 = 𝑑 𝑝 𝑥, 𝑦 (𝑑 ∶ 𝒳上の距離)に対し
𝑊𝑝 𝜇, 𝜈 ≔ min
𝜋
න
𝒳×𝒳
𝑑 𝑝 𝑥, 𝑦 d𝜋 𝑥, 𝑦
1/𝑝
命題
𝑊𝑝 は 𝒳 上の確率測度全体の集合𝒫 𝒳 の距離を定める.
8

9.
離散測度の場合
𝒫 𝒳
𝜇 = ෍
𝑖
𝑁
𝑎𝑖 𝛿 𝑥 𝑖
, 𝜈 = ෍
𝑗=1
𝑀
𝑏𝑗 𝛿 𝑦 𝑗
Π 𝜇, 𝜈 ≔
{𝜋 ∈ 𝒫 𝒳 × 𝒳 ∣ p1 # 𝜋 = 𝜇, (p2)𝜋 = 𝜈 }
𝑐: 𝒳 × 𝒳 → ℝ
OT 𝜇, 𝜈 ≔ min
𝜋
න
𝒳×𝒳
𝑐(𝑥, 𝑦)d𝜋
ベクトル/行列表現
𝒂 = 𝑎𝑖 ∈ ℝ 𝑛, 𝒃 = (𝑏𝑗) ∈ ℝ 𝑚
𝒰 𝒂, 𝒃 ≔
{𝑃 ∈ ℝ≥0
𝑛×𝑚
∣ 𝑃𝟏 𝑚 = 𝒂, 𝑃 𝑇
𝟏 𝑛 = 𝒃}
𝐶 = 𝑐 𝑥𝑖, 𝑦𝑗
𝑖𝑗
∈ ℝ 𝑛×𝑚
OT 𝒂, 𝒃 ≔ min
𝑃∈𝒰(𝒂,𝒃)
〈𝐶, 𝑃〉
 離散測度：
 計算理論：ベクトルとして表現可能
 統計：未知の分布からのサンプリング
Remark.
統計的には 𝑁, 𝑀 ≫ 1, 𝑎𝑖 =
1
𝑁
, 𝑏𝑗 =
1
𝑀
の場合に興味がある
線形計画
問題(LP)
9

10.
線形計画問題を解くのは計算コスト大
二つの輸送写像はともに最適
(輸送計画もすべて最適)
線形計画問題を直接解くと 𝑂 𝑁3 の時間計算量を免れない
これは解が退化しているため(右図)
エントロピー正則化(Cuturi 2013)
𝐻 𝑃 ≔ − σ𝑖,𝑗 𝑃𝑖𝑗 log 𝑃𝑖𝑗 ≥ 0 : エントロピー
𝐶𝜖 𝑃 ≔ 𝐶, 𝑃 − 𝜖𝐻 𝑃 を最小化する問題を考える
OT𝜖 𝒂, 𝒃 ≔ min
𝑃
𝐶, 𝑃 − 𝜖𝐻(𝑃)
命題
上の問題の最適解は一意で、
𝑃𝑖𝑗 = 𝑢𝑖 𝑒−𝐶 𝑖𝑗/𝜖
𝑣𝑗, ∃ 𝑢, 𝑣 ∈ ℝ 𝑛
× ℝ 𝑚
と書ける
10

11.
Sinkhornのアルゴリズム
 ヒューリスティックに最適解に近づく繰り返しアルゴリズムを考える
最適解 𝑃 = diag 𝒖 𝐾diag 𝒗 , 𝐾 ≔ 𝑒−𝐶/𝜖は
𝒂 = 𝑃𝟏 𝑚 = 𝒖 ⊙ 𝐾𝒗
𝒃 = 𝑃 𝑇 𝟏 𝑚 = 𝒗 ⊙ 𝐾 𝑇 𝒖
を満たす。
Sinkhornのアルゴリズム：初期値 𝒗0 ∈ ℝ>0
𝑚
、以下を繰り返す：
𝒖(𝑙+1) =
𝒂
𝐾𝒗 𝑙
𝒗(𝑙+1) =
𝒃
𝐾 𝑇 𝒖 𝑙+1
 Sinkhornは線形収束する
 元のOTに対する任意の近似誤差𝑒を
達成するための 𝜖 を明示的にとれる
総計算量 ෨𝒪(𝑛2/𝑒3) (Altschuler 2017)
11

12.
最適輸送の計算理論の研究動向
12

13.
トピック1：
エントロピーのDebiasing
エントロピー正則化の問題点
解のスパースさが失われ「ぼやけて」しまう
→ 最適輸送計画自体に興味がある場合
(画像変換や中間状態の計算など)に問題
問題
 エントロピーの質的な影響を定量化できるか?
cf. 正規分布に対する解析解(Janati 2020b)
 効率よくエントロピーバイアスを解消できるか?
例：Wasserstein重心のぼやけ
(Janati 2020a)
正則化
有り
無し
Wasserstein重心
測度 𝜇 𝑘 ∈ 𝒫 𝒳 ,重み𝑤 𝑘 ∈ 0,1 , σ 𝑘 𝑤 𝑘 = 1
𝜇∗ ≔ arg min
𝜇
෍
𝑘
𝑤 𝑘 𝑊𝑝 𝜇, 𝜇 𝑘
13

14.
アプローチ1：
Debiased Sinkhorn Divergence
定義 (Genevey 2018, Feydy 2018)
S 𝜖 𝜇, 𝜈 ≔ OT𝜖 𝜇, 𝜈 −
1
2
OT𝜖 𝜇, 𝜇 + OT𝜖 𝜈, 𝜈
良い性質
 𝜖 → ∞ でも非自明な距離(MMD)を与える
 Sinkhornをそのまま使えて計算量を抑制できる
課題
 理論がまだ浅い
(バイアス除去の効果(Janati 2020a,後述)、統計的推定の効率(Xu 2020))
14

15.
論文紹介：Debiased Sinkhorn Barycenter
Janati, Cuturi, Gramfort (ICML 2020)
エントロピー正則化
= 一様分布に近づける
OT𝜖(𝜇, 𝜈) = min
𝑃
𝐶, 𝑃 + 𝜖KL 𝑃 𝒰
𝒰 : supp𝜇 × supp𝜈 上一様分布
目的
 重心 𝜇∗ ≔ arg min σ 𝑘 𝑤 𝑘 𝑊𝑝 𝜇, 𝜈 の計算を
正則化付きで行いたい
 しかしぼやけてしまうことが実験で知られていた
KL 𝛼 𝛽 ≔ න log
d𝛼
d𝛽
d𝛼 (= ∞ if supp𝛼 ⊂ supp𝛽)
−𝐻 𝑃 = ෍
𝑖𝑗
𝑃𝑖𝑗 log 𝑃𝑖𝑗 = KL 𝑃 𝒰 + const.
15

16.
正規分布の重心へのエントロピーの影響
重心のぼかされ方を初めて解析的に計算
正則化付きの重心
連続化
OT𝜖
ℒ
≔ min
𝜋
න
𝒳×𝒳
𝑐 𝑥, 𝑦 𝜋 + 𝜖KL 𝜋 ℒ
Lebesgue測度
分散が正則化に依存
16

17.
Sinkhorn divergenceを使った重心
積測度に近づける正則化
OT𝜖
𝜇⊗𝜈
≔ min
𝜋
𝐶, 𝜋 + 𝜖KL 𝜋 𝜇 ⊗ 𝜈
今度はとがりすぎてしまう…
Sinkhorn Divergence
S 𝜖 𝜇, 𝜈
≔ OT𝜖 𝜇, 𝜈 −
1
2
OT𝜖 𝜇, 𝜇 + OT𝜖 𝜈, 𝜈
= OT𝜖
⊗
𝜇, 𝜈 −
1
2
OT𝜖
⊗
𝜇, 𝜇 + OT𝜖
⊗
𝜈, 𝜈
定理
等分散正規分布 𝒩(𝑚 𝑘, 𝜎2) の
𝑆𝜖による重心は𝜖によらず𝒩( ഥ𝑚, 𝜎2)
17

18.
アプローチ2：
エントロピー以外の正則化を使う
最適輸送の正則化をシステマチックに一般化する枠組み (Blondel 2018)
OTΩ(𝒂, 𝒃) ≔ min
𝑃∈𝒰(𝒂,𝒃)
𝐶, 𝑃 + Ω 𝑃
 双対性を用いて理論保証のある凸最適化に帰着できる
左図：適切な正則化と最適化法を選ぶことで
定性的に良い解が効率的に得られる
OT理論の王道 × 最適化理論の王道
という感じ
UOT(後述)を含む強力で一般的な枠組みだが
(まだ?)あまり流行っていない印象
強凸関数
Ω 𝑃 = |𝑃| 𝐹, −𝐻(𝑃)など
18
(Blondel 2018)

19.
トピック2：
異常値に頑健な最適輸送
Source 𝜇 やTarget 𝜈 に少量含まれるノイズが
最適輸送に大きな影響を与える
→ 実用的には無視してほしい
 特に高次元の場合に深刻な問題
(𝑊𝑝の統計的推定における次元の呪い)
- 例えば…
(Niles-Weed 2019)
(Balaji 2020)
19

20.
アプローチ1：
Unbalancedな最適輸送(UOT)
質量の増加・減少を許して 𝜇, 𝜈 を一般の正測度に拡張
UOT 𝒂, 𝒃 ≔ min
𝑃∈ℝ≥0
𝑛×𝑚
〈 𝐶, 𝑃〉 + 𝜏KL(𝑃𝟏 𝑚, 𝒂) + 𝜏KL(𝑃 𝑇
𝟏 𝑛, 𝒃)
近年の発展
 理論的な整備 e.g. Liero 2018
 効率的に計算する定式化
Blondel 2018, Pham 2020, Balaji 2020, Sato 2020
UOT(右)では山の分裂を回避できている
↓
サンプル数のランダムさに対して頑健
増減に対するペナルティ
UOT版Sinkhornの解析 木構造でのUOTの計算
凸最適化 マージン制約付きで定式化
20

21.
アプローチ2：
Projection Robust Wasserstein
最適輸送を低次元に射影して考える
 1次元の場合 ： Sliced Wasserstein (Rabin 2011, Bonneel 2015, Deshpande 2019)
SW 𝜇, 𝜈 2 ≔ න
𝕊 𝑑−1
𝑊2 𝑃 𝜃 # 𝜇, 𝑃 𝜃 # 𝜈
2
d𝜃
maxSW 𝜇, 𝜈 ≔ max
𝜃∈𝕊 𝑑−1
𝑊2 𝑃 𝜃 # 𝜇, 𝑃 𝜃 # 𝜈
𝑃 𝜃 ∶ 𝒳 = ℝ 𝑑 → ℝ は𝜃 ∈ 𝕊 𝑑−1 = {𝑥 ∈ ℝ 𝑑 ∣ 𝑥 = 1} 方向の直線への射影
 1次元でのOTが陽に解けるので計算が楽
 SWも距離(しかもWasserstein距離と同値)
 生成モデル(GAN)への応用 (Deshpande 2018, Deshpande 2019, Kolouri 2019, Chen 2020)
21

22.
アプローチ2：
Projection Robust Wasserstein
最適輸送を低次元に射影して考える
 k次元の場合 ： Projection Robust Wasserstein (Paty 2019, Niles-Weed 2019)
PRW𝑘 𝜇, 𝜈 2 ≔ max
𝐸:ℝ 𝑑→ℝ 𝑘
orthonormal
𝑊2 𝐸# 𝜇, 𝐸# 𝜈 2
 Sliced Wasserstein よりも柔軟な射影を扱いながら
統計的にロバスト (Lin 2020a)
 実際に計算するのが困難
• 凸緩和 (Paty 2019) をしてもまだ計算量が大きい
(SVDで𝑂 𝑑3
)
22
通常のWasserstein PRWの凸緩和
(Paty 2019)

23.
論文紹介： Projection Robust Wasserstein distance and
Riemannian Optimization
Lin, Fan, Ho, Cuturi, Jordan (NeurIPS 2020)
 Stiefel 多様体
St = St 𝑑,𝑘 = {𝑈 ∈ ℝ 𝑑×𝑘 ∣ 𝑈 𝑇 𝑈 = 𝐼 𝑘}
上の最適化として効率的にPRWを計算するアルゴリズムを提案
 𝑓𝜂 𝑈 に関して多様体上の勾配法で最大化。任意の近似誤差を達成する計算量を導出
エントロピー正則化
23
𝜋𝟏 𝒏 = 𝒂, 𝜋𝟏 𝒏 = 𝒃
Sinkhornで計算

24.
多様体上の勾配法
 勾配法：
max
𝑈∈ℝ 𝑑×𝑘
𝑓(𝑈) iterate 𝑈(𝑙+1) = 𝑈(𝑙) + 𝛻 𝑈 𝑓(𝑈(𝑙))
 行列多様体 ℳ ⊂ ℝ 𝑑×𝑘 上の勾配法 (cf. Absil 2009)
1. 勾配 𝛻 𝑈 𝑓 𝑈 𝑙 ∈ T 𝑈 𝑙 ℝ 𝑑×𝑘 ≃ ℝ 𝑑×𝑘 を計算
2. 勾配を接空間に射影 pr 𝑈(𝑙): ℝ 𝑑×𝑘
→ T 𝑈(𝑙)ℳ
3. レトラクションを使って更新： ෩𝑈(𝑙+1)
= Retr 𝑈 𝑙 pr 𝑈 𝑙 𝛻 𝑈 𝑓 𝑈 𝑙
※ レトラクション Retr 𝑈: T 𝑈ℳ → ℳ s.t. (i) Retr 𝑈 0 = 𝑈, (ii) Retr 𝑈 𝜉 − 𝑈 + 𝜉 → 0 as 𝜉 → 0
精密な解析を必要とするが、通常の
勾配法と同様の収束解析ができる
24

25.
参考文献(Intro)
1. Monge, “Mémoire sur la théorie des déblais et des remblais”, De l’Imprimerie Royale, 1781.
2. Tong, Huang, Wolf, Dijk, Krichnaswamy, “TrajectoryNet: A Dynamic Optimal Transport Network for
Modeling Cellular Dynamics”, ICML 2020.
3. Muzellec, Josse, Boyer, Cuturi, “Missing Data Imputation using Optimal Transport”, ICML 2020.
4. Chen, Gan, Cheng, Li, Carin, Liu, “Graph Optimal Transport for Cross-Domain Alignment”, ICML 2020.
5. Eisenberger, Toker, Taixe, Cremers, “Deep Shells: Unsupervised Shape Correspondence with Optimal
Transport”, NeurIPS 2020.
6. Kantorovich. “On translation of mass”, Proceedings of the USSR Academy of Sciences, 37:199–201,
1942.
7. Cuturi “Sinkhorn Distances: Lightspeed Computation of Optimal Transport”, NIPS 2013.
8. Peyré, Cuturi “Computational Optimal Transport”, Foundations and Trends in Machine Learning 2019.
25

26.
参考文献 (Entropy Debiasing)
1. Genevay, Peyre, Cuturi “Learning Generative Models with Sinkhorn Divergences”, AISTATS
2018.
2. Feydy, Sejourne, Vialard, Amari, Trouve, Peyre, “Interpolating between Optimal Transport
and MMD using Sinkhorn Divergence, AISTATS 2019.
3. Janati, Muzellec, Peyre, Cuturi, “Entropic Optimal Transport between Unbalanced Gaussian
Measures has a Closed Form”, NeurIPS 2020b.
4. Janati, Cuturi, Gramfort, “Debiased Sinkhorn Barycenter”, ICML 2020a.
5. Xu, Wenliang, Munn, Acciaio, “COT-GAN: Generating Sequential Data via Causal Optimal
Transport”, NeurIPS 2020.
6. Blondel, Seguy, Rolet, “Smooth and Sparse Optimal Transport”, AISTATS 2018
26

27.
参考文献 (Unbalanced Optimal Transport)
1. Liero, Mielkem, Savare, “Optimal Entropy-Transport problems and a new Hellinger-
Kantorovich distance between positive measures”, Invent. math. 211:969-1117, 2018.
2. Pham, Le, Ho, Pham, Bui, “On Unbalanced Optimal Transport: An Analysis of Sinkhorn
Algorithm”, ICML 2020.
3. Balaji, Chelleppa, Feizi, “Robust Optimal Transport with Applications in Generative Modeling
and Domain Adaptation”, NeurIPS 2020.
4. Sato, Yamada, Kashima, “Fast Unbalanced Optimal Transport on a Tree”, NeurIPS 2020.
27

28.
参考文献 (Sliced Wasserstein Distance)
1. Rabin, Peyre, Delon, Marc, “Wasserstein Barycenter and its Application to Texture Mixing”,
SSVM’11, 435-446, 2011.
2. Bonneel, Rabin, Peyre, Pfister, “Sliced and Radon Wasserstein Barycenters of Measures”,
Journal of Mathematical Imaging and Vision, Springer Verlag, 1 (51), 22-45, 2015.
3. Deshpande, Zhang, Schwing, “Generative Modeling Using the Sliced Wasserstein Distance”,
CVPR 2018.
4. Deshpande, Hu, Sun, Pyrros, Siddiqui, Koyejo, Zhao, Forsyth, Schwing, “Max-Sliced
Wasserstein distance and its use for gans”, CVPR 2019.
5. Kolouri, Nadjahi, Simsekli, Badeau, Rohde “Generalized Sliced Wasserstein Distances”,
NeurIPS 2019.
6. Chen, Yang, Li, “Augmented Sliced Wasserstein Distances”, arXiv:2006.08812, 2020.
28

29.
参考文献 (Projection Robust Wasserstein)
1. Niles-Weed, Rigollet, “Estimation of Wasserstein distances in the Spiked Transport Model”,
arXiv:1909.07513, 2019.
2. Paty, Cuturi, “Subspace Robust Wasserstein Distances”, ICML 2019.
3. Lin, Zheng, Chen, Cuturi, Jordan, “On Projection Robust Optimal Transport: Sample
Complexity and Model Misspecification”, arXiv:2006.12301, 2020a.
4. Lin, Fan, Ho, Cuturi, Jordan, “Projection Robust Wasserstein Distance and Riemannian
Optimization”, NeurIPS 2020b.
5. Absil, Mahony, Sepulchre, “Optimization Algorithms on Matrix Manifolds”, Princeton
University Press, 2009.
29