最適輸送の基礎から最近の動向まで 2024.6.13 佐藤竜馬（国立情報学研究所）

2 National Institute of Informatics 自己紹介 • 名前：佐藤 竜馬（さとう りょうま） • 国立情報学研究所 助教 好評発売中！

3 National Institute of Informatics 最適輸送は距離、対応、変換ができる • 最適輸送にできること： [距離] 点群（確率分布）どうしの距離を求める [対応] 点の対応関係を得る [変換] 点群から点群への変換方法を得る • 本日はこれら 3 つの「できること」を基に、 最適輸送が使える場面の感覚を身につける ことを目標とします

4 National Institute of Informatics 最適輸送は確率分布の「山」を輸送 • 確率分布の「山」を移動させて一致させるのにかかるコスト が最適輸送の直感的な定義 • 輸送という単語はここからきている

5 National Institute of Informatics 点群比較の場合の例 • 点群（離散分布）比較の場合 • 各点に質量 1/n の砂山があると考えその輸送距離を考え る

6 National Institute of Informatics 最適な輸送を用いる • 輸送の仕方は複数通りあるが、最もコストが少ない輸送方 法を選ぶ • ヒストグラムの場合も同様 • 「最適」という言葉はここからきている 最適ではない輸送の例

7 National Institute of Informatics もう少し規模の大きい点群の最適輸送 • [距離] [対応] もう少し規模の大きい点群の最適輸送

8 National Institute of Informatics 実際に輸送する際はまっすぐ動かす • [変換] 対応だけでなく実際に輸送するには対応する点ま でまっすぐ移動させればよい 変位補間という • 寄り道せずにまっすぐ動くのが最適（省エネルギー）である https://github.com/nbonneel/network_simplex

9 National Institute of Informatics POT の使用例: 簡単に使えます • POT ライブラリを使えば簡単に最適輸送を求められる • 数値例: 二つの正規分布からの点群の比較 import numpy as np import matplotlib.pyplot as plt import ot # POT ライブラリ n = 100 # 点群サイズ mu = np.random.randn(n, 2) # 入力分布 1 （青） nu = np.random.randn(n, 2) + 1 # 入力分布 2 (オレンジ） a = np.ones(n) / n # 点重み (1/n, ..., 1/n) b = np.ones(n) / n # 点重み (1/n, ..., 1/n) C = np.linalg.norm(nu[np.newaxis] - mu[:, np.newaxis], axis=2) # コスト行列 P = ot.emd(a, b, C) # 最適輸送行列の計算（POT ライブラリを使用） plt.scatter(mu[:, 0], mu[:, 1]) # mu の散布図描写 plt.scatter(nu[:, 0], nu[:, 1]) # nu の散布図描写 for i in range(n): j = P[i].argmax() # i の対応相手: 最もたくさん輸送している先 plt.plot([mu[i, 0], nu[j, 0]], [mu[i, 1], nu[j, 1]], c='grey', zorder=-1)

10 National Institute of Informatics 最適輸送の計算は三乗時間かかる • 注意点： ot.emd の計算量は三乗時間（おおよそ） • 数千点くらいのサイズが限界 • より大きなデータを扱うには シンクホーンアルゴリズム（二乗時間） スライス法（線形時間） 敵対的学習（償却可能） などを用いる • シンクホーンアルゴリズムは単純・高速・微分可なので人気

11 National Institute of Informatics シンクホーンアルゴリズムは超シンプル • シンクホーンアルゴルズムの実装 • 導出は 『最適輸送の理論とアルゴリズム』 第三章や 『最適輸送の解き方』 p.198- を参照してください。 https://speakerdeck.com/joisino/zui-shi-shu-song-nojie-kifang?slide=198 超シンプル！ K = np.exp(- C / eps) u = np.ones(n) for i in range(100): v = b / (K.T @ u) u = a / (K @ v) P = u.reshape(n, 1) * K * v.reshape(1, m) d = (C * P).sum()

12 National Institute of Informatics シンクホーンアルゴリズムは自動微分可 • 四則計算と exp だけからなるので自動微分が可能 a.requires_grad = True K = torch.exp(- C / eps) u = torch.ones(n) for i in range(100): v = b / (K.T @ u) u = a / (K @ v) P = u.reshape(n, 1) * K * v.reshape(1, m) d = (C * P).sum() d.backward() print(a.grad)

13 National Institute of Informatics ニューラルネットワークと相性が良い • 四則計算と exp だけからなるので自動微分が可能 C = net1(z) K = torch.exp(- C / eps) u = torch.ones(n) for i in range(100): v = b / (K.T @ u) u = a / (K @ v) P = u.reshape(n, 1) * K * v.reshape(1, m) d = (C * P).sum() loss = net2(P, d) loss.backward() 何かしらのニューラルネットワーク

14 National Institute of Informatics 重みと位置 in → 距離とマッチ out 𝑎 𝑏 𝑋 𝑌 𝐶 最適輸送問題 𝑃 𝑑 [距離] [対応] 点重み 点の位置

15 National Institute of Informatics 入力についての勾配を求めて最適化 𝑎 𝑏 𝑋 𝑌 𝐶 最適輸送問題 𝑃 𝑑 損失 誤差逆伝播 近似誤差が最小となる サンプル重みづけを求めたい 輸送コストが最小となるような 点の配置を求めたい

16 National Institute of Informatics 自動微分を使って配置を最適化する例 動画 ソースコード https://drive.google.com/file/d/19XNtttaSr-Kc8yfv1VKRz0O8dUpcxSZM/view?usp=sharing https://colab.research.google.com/drive/1u8lu0I7GwzR48BQqoGqOp2A_7mTHxzrk?usp=sharing • 数値例：点群 A を点群 B に近づける パラメータは位置 X Adam で最適化

17 National Institute of Informatics 距離、対応、変換を最適化できる • 最適輸送にできること： [距離] 点群（確率分布）どうしの距離を求める [対応] 点の対応関係を得る [変換] 点群から点群への変換方法を得る が求まるだけでなく、これらを勾配法で最適化できる

18 National Institute of Informatics 応用例：文書比較・文書検索 2 つの文書は 似ている？ 文書検索など • ナイーブな方法は bag-of-words ベクトルの L2 距離 = 単語ごとの出現数で比較 • 同じ意味でも語彙が違えば似ていないと判定されてしまう

19 National Institute of Informatics 応用例：文書比較・文書検索 • [距離] 最適輸送距離を使って文書を比較 → 語彙が被っていなくても意味が近ければ似ていると判定 パーツに分かれるモノを比較するのに最適輸送は使える 文書のパーツは単語、分子 → 原子、画像 → パッチ 赤: 文書 1 の単語分布 青: 文書 2 の単語分布 各粒が単語に対応 単語の位置は word2vec Kusner et al. From Word Embeddings To Document Distances. ICML 2015

20 National Institute of Informatics 応用例：転移学習 • シミュレーションデータで訓練して実環境で運用・テストする 転移学習を考える ニューラルネットワーク 入力 （人工 or 実環境） 予測ヘッド 予測 埋め込みベクトル

21 National Institute of Informatics 応用例：転移学習 • シミュレーションデータで訓練して実環境で運用・テストする 転移学習を考える • [距離 最適化 ] 予測誤差+赤と青の最適輸送距離を最小化 • シミュレーションデータと本番環境の差を中間層で吸収する ニューラルネットワーク 入力 予測ヘッド 予測 1600 サンプルの埋め込み 赤：シミュレーションデータについての埋め込み 青：本番環境データについての埋め込み

22 National Institute of Informatics 損失を最適輸送距離で抑えられる • 最適輸送を用いた転移学習の理論 特定の条件※の下 （実でのモデル h の損失） ≤（シでのモデル h の損失） + (シと実の最適輸送距離） が成り立つ 実：実環境での入力分布 シ：シミュレーションでの入力分布 ※ モデルが L1 ノルムについて ½ リプシッツ連続かつ損失が L1 ノルムについて 1 リプシッツ連続かつリスク 0 を達成可能 など参照

23 National Institute of Informatics 損失を最適輸送距離で抑えられる • 最適輸送を用いた転移学習の理論 （実でのモデル h の損失） ≤（シでのモデル h の損失） + (シと実の最適輸送距離） 右辺は計算可能 右辺を最小化すると、左辺も（具体値は分からないが） 小さいことが保証される 実環境でのラベルが無いので計算不可能 シミュレーション環境では ラベルが作れるので計算可能 ラベルなしデータだけなら実環境 でも手に入るので計算可能

24 National Institute of Informatics 応用例：転移学習 • [距離 最適化 ] 赤と青の最適輸送距離を最小化、のほか • [対応] このシデータは実データで言えばこれ、を知る • [変換] シミュレーションデータを実データに変換する、なども ニューラルネットワーク 入力 予測ヘッド 予測 1600 サンプルの埋め込み 赤：シミュレーションデータについての埋め込み 青：本番環境データについての埋め込み

25 National Institute of Informatics 応用例：モデルマージ • 二つの訓練済みモデルからよりよいモデルを得たい • 問題点：モデル A の i 次元目はモデル B の i 次元目と 対応しているとは限らない • [対応] 最適輸送で次元の対応関係を得る → 次元が揃ったあとはパラメータ平均を取れば OK Singh et al. Model Fusion via Optimal Transport. NeurIPS 2020. Imfeld et al. Transformer Fusion with Optimal Transport. ICLR 2024.

26 National Institute of Informatics 応用例：ピクセル色変換 • [変換] 最適輸送を使って内容を保ちつつ色味を変換 • 応用例：合成画像の色味や光加減の補正など https://drive.google.com/ file/d/1_X- 6pbBTb5GLYclwW9FcgYjat guDn94T/view?usp=shari ng

27 National Institute of Informatics 応用例：生成モデル • 最適輸送は生成モデルとも密接に関係している • 特に敵対的生成ネットワーク (GAN) と拡散モデル と関係が深い • [距離 最適化 ] GAN では分布の距離を測るために使う • [対応][変換] 拡散モデルでは分布の変換に使う

28 National Institute of Informatics 生成モデルの訓練と評価に使える • [距離] 生成した画像の集合と訓練データセット内の画像 の集合の最適輸送コストで生成モデルの品質を評価 = Fréchet Inception Distance (FID) • [距離 最適化 ] 生成画像の集合と訓練データセット内の画像の 集合が最適輸送の意味で近くなるように訓練する = Wasserstein GAN 赤: モデルが生成した画像の集合 青: 訓練データセット内の画像の集合 各粒が画像に対応 画像の位置は特徴ベクトル

29 National Institute of Informatics 拡散モデルはノイズ分布をデータに変換 • 拡散モデルはノイズ入りの画像を受け取り、ノイズを 除去することを目指す • 生成時には、完全なノイズ画像からはじめて、モデル によりノイズを除去することと、微小なノイズを加える ことを繰り返して洗練させていき、自然画像を得る https://joisino.hatenablog.com/entry/2024/03/08/174412

30 National Institute of Informatics 最適輸送と拡散モデルのつながり • [変換] 最適輸送は分布と分布をまっすぐ結ぶ → 拡散モデルよりも安定している 寄り道が無いので早い

31 National Institute of Informatics 移動方向をモデリングする • 生成モデルでは訓練時に見ていない点も輸送する 必要がある • 基本方針：点 x を受け取りどちらの方向に動くべ きかを返す関数 f を訓練する モデル f 右上に動け 時刻 t

32 National Institute of Informatics シミュレーションベースの訓練法 • 訓練方法にはいくつかの方針がある • 方針 1:[変換]モデル f に従って点を動かして画像 を生成してみて、そのクオリティを目的関数とする Entropic Neural Optimal Transport via Diffusion Processes など メリット：直観的に分かりやすい ちゃんと最適輸送の原理に基づいている デメリット：毎回生成してみるのが時間がかかる 安定しない（特に最初は生成結果が めちゃくちゃ）

33 National Institute of Informatics フローマッチングによる簡単化 • 方針 2：フローマッチング Flow Matching for Generative Modeling Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport • 最適輸送ではデータ分布 P とノイズ分布 Q 全体 の対応関係を得ようとしていた → 難しい • [変換] フローマッチングではノイズ分布とデータ x （を中心とした分布 P x ）の移動方法を得る → これは閉じた式で得られる ノイズ分布から x にまっすぐ移動 → 移動方法の ground truth が 明示的に求まり回帰で解ける

34 National Institute of Informatics フローマッチングによる簡単化 • 方針 2：フローマッチング • 一点 x だけを生成できても意味なくない？ 実は x について期待値を取ると、データ分布 P を 生成する正しい動かし方モデル f が得られる • メリット：訓練が簡単かつ安定的 • デメリット：正確には最適輸送とはいえない x を固定したとき最適でも 期待値を取ると一般に最適でなくなる • [対応] ミニバッチ {x i } とノイズ分布の対応を使う方法もある。 左のデータ x i は左のノイズ z j に対応させよう、というようなイメージ。 従来のフローマッチングは各データ x j に全ノイズを対応させる Tong et al. Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport

35 National Institute of Informatics フローマッチングの例 • [変換] https://github.com/atong01/conditional-flow-matching

36 National Institute of Informatics 最適輸送は距離、対応、変換ができる • 最適輸送にできること： [距離] 点群（確率分布）どうしの距離を求める 例：文書の比較、転移学習、生成データの評価 [対応] 点の対応関係を得る 例：モデルマージ、ノイズとデータの対応 [変換] 点群から点群への変換方法を得る 例： ピクセル色変換、画像生成 最適輸送は距離測定、対応、変換ができる take home message 本もよろしく お願いします！