技術概要：FlashSinkhorn — GPU 上の IO 意識的エントロピー正則化最適輸送

問題提起

エントロピー正則化最適輸送（EOT）は、Sinkhorn 反復法によって解かれ、生成モデルから分布整合に至るまで、現代の機械学習において確率測度を比較するための基本的なツールである。しかし、GPU 上で大規模な点群（例：$n \approx 70,000$）に対して EOT をスケーリングすることは、依然として持続的なボトルネックとなっている。

既存の GPU ソルバーは、主に 2 つの非効率性に直面している：

1. テンソル化実装：これらは密な $n \times m$ コスト行列を具体化（マテリアライズ）するため、高帯域幅メモリ（HBM）のトラフィックが二次関数的に増加し（$O(nm)$の読み書き）、$n$ が大きくなるにつれてメモリ容量の問題が発生し、Out-of-Memory（OOM）エラーを引き起こす。
2. オンラインバックエンド：これらは KeOps や OTT-JAX などを介して相互作用をオンザフライで評価することで密な行列の保存を回避するが、汎用的なタイル化されたマップ・リデュース還元カーネルに依存している。これらのカーネルは融合が限られており、テンソルコアの GEMM 構造を活用できず、メモリ使用量が削減されているにもかかわらず、計算利用率が低く、実行時間が長くなる結果となっている。

核心的なシステム上のボトルネックは、現在の手法が過剰な HBM トラフィックに起因するメモリ束縛、あるいは低い命令効率を伴う計算束縛のいずれかに陥っており、下流パイプラインで必要とされる大規模な反復解法にスケーリングできない点にある。

手法

著者らは、二乗ユークリッド距離コストを持つ EOT 向けの IO 意識的 GPU ソルバー「FlashSinkhorn」を提案する。この手法は、安定化されたログドメインの Sinkhorn 更新とトランスフォーマーの注意機構（アテンション機構）との間に直接的な類似性を引き合いに出し、FlashAttentionの原則を援用してメモリアクセスパターンを最適化する。

1. 注意形式への再定式化

鍵となる理論的洞察（命題 1）は、二乗ユークリッド距離コストの場合、安定化されたログドメインの Sinkhorn 更新は、バイアス付きドットプロダクトスコア行列の行ごとの LogSumExp（LSE）還元として書き換えられるという点である。

• 具体的には、更新 $f_i \leftarrow -\varepsilon \text{LSE}_j [(g_j - C_{ij})/\varepsilon + \log b_j]$ は、行列 $S_X = (QK^\top + \mathbf{1}(\hat{g} + \delta))/\varepsilon$ の LSE を計算することと等価である。ここで、$Q=\sqrt{2}X$、$K=\sqrt{2}Y$ であり、$\delta, \gamma$ は事前に計算されたログ重みバイアスである。
• この再定式化は、Sinkhorn 反復をスケーリングされたドットプロダクト注意で使用される正規化統計量と完全に一致させ、IO 意識的なストリーミング技術の適用を可能にする。

2. ストリーミングカーネルとタイル化

FlashSinkhorn は、この再定式化を、2 段階のメモリモデル（低速な HBM と高速なオンチップ SRAM）を採用した融合されたTritonカーネルを用いて実装する：

• 具体化の回避：$n \times m$ コスト行列は HBM 上で決して具体化されない。代わりに、$Q$ と $K$ のタイルがストリーミングされる。
• オンライン LSE：カーネルは、スコア行列のタイルを処理するにつれて、SRAM 上で実行中の行ごとの最大値と指数和の統計値を維持する。これにより、HBM アクセスを大幅に減らしながら正確な LSE 還元を計算できる。
• 融合：スコア計算、バイアス追加、LSE 還元、および潜在的な更新を含む全体更新が、単一のカーネルパスに融合される。これにより、中間的な HBM への書き込みが排除され、カーネル起動のオーバーヘッドが削減される。
• 行固定トラバーサル：多くの場合、キー/バリューを外部ループで反復する FlashAttention とは異なり、FlashSinkhorn は Sinkhorn 更新の行ごとの還元性質により適合するよう、行固定トラバーサル（SRAM 内のクエリ行ブロックを固定し、すべてのキーブロックをストリーミングする）を採用する。

3. 微分と二次最適化

本論文は、この IO 意識的戦略を微分および二次的手法へ拡張する：

• 輸送行列の適用：重心射影と勾配計算は、注意に似た操作（加重平均）として表現され、輸送行列 $P^*$ を具体化することなく、同じストリーミングカーネルを通じて計算可能である。
• ヘッセ行列 - ベクトル積（HVP）：著者らは、ヘッセ行列 - ベクトル積を輸送ベクトル積と輸送行列積の系列に分解するストリーミング HVP オラクルを導出した。これにより、$O(n^2 d^2)$ のヘッセテンソルや $O(nm)$の輸送行列を具体化することを回避し、作業メモリを$O((n+m)d)$ に削減する。

主要な貢献

1. 注意形式 Sinkhorn 更新：安定化された Sinkhorn 更新をバイアス付きドットプロダクト LSE 還元として表現し、EOT に対する正確な FlashAttention 様ストリーミングを可能にする。
2. IO 意識的 EOT ソルバー：ストリーミングされた Sinkhorn 反復用の融合 Triton カーネルと、微分のための輸送行列不要のストリーミング演算子を提供し、HBM IO ボリュームを大幅に削減する。
3. スケーラブルな二次最適化：ストリーミング HVP を導入し、EOT ベースの損失関数に対するスケーラブルなニュートン CG 法および信頼領域法を可能にする。これらは以前、メモリ制約により大規模スケールでは実行不可能であった。
4. 大規模評価：合成ベンチマークおよび最適輸送データセット距離（OTDD）やシャッフル回帰などの下流タスクにおいて、A100 GPU 上で大幅な高速化とメモリ効率を実証する。

実験結果

NVIDIA A100-80GB GPU 上で評価した結果、FlashSinkhorn は最先端のベースライン（KeOps/テンソル化バックエンドおよび OTT-JAX を用いた GeomLoss）に対して顕著な性能向上を示した：

• 速度：
  • フォワードパス：KeOps に対して最大32 倍、OTT-JAX に対して5.1 倍の高速化を達成。
  • エンドツーエンド（フォワード＋バックワード）：高次元（$d=512$）において、KeOps に対して最大161 倍の高速化を達成。
  • HVP：次元性に応じて3〜25 倍の高速化を達成し、他の手法がタイムアウトするスケールでもニュートン CG 法を実用的にする。
• メモリ効率：
  • FlashSinkhorn は線形メモリ複雑度 $O((n+m)d)$ を維持するのに対し、テンソル化ベースラインは $O(nm)$または$O(n^2)$ としてスケーリングする。
  • テンソル化手法がメモリ不足となり、オンラインベースライン（KeOps）が時間制限を超えるような、$n=50,000$ の点群と最大 $d=1024$ の次元まで正常に処理可能。
• 下流タスク：
  • OTDD：FlashSinkhorn は、ピークメモリが 1GB 未満で、$n=60,000$（MNIST $\leftrightarrow$ Fashion-MNIST）まで OTDD 距離計算と勾配流をスケーリングする。一方、テンソル化ベースラインは $n > 20,000$ で OOM となる。
  • 鞍点脱出：シャッフル線形回帰において、ストリーミング HVP により最小ヘッセ固有値を監視して鞍点を検出可能とし、ハイブリッド Adam-ニュートン最適化器が鞍点から脱出して迅速に収束することを可能にする。

意義と主張

本論文は、FlashSinkhorn がメモリ束縛領域において大規模 EOT を実用的にすると主張する。Sinkhorn 更新を FlashAttention の IO 意識的構造に一致するように書き換えることで、著者らは最適輸送と現代のトランスフォーマー加速技術の間の溝を埋めた。

その意義は以下の点にある：

• システムレベルの最適化：この gains は、Sinkhorn 反復そのものへのアルゴリズム的変更ではなく、カーネルレベルの専門化（融合とタイル化）に起因する。
• 新たなワークロードの実現：メモリフットプリントと実行時間の削減により、以前は計算的に実行不可能とみなされていたスケールで、データセット距離や暗黙的微分などの下流タスクにおける EOT の利用が可能になる。
• 再現性：著者らは、スケーラブルな最適輸送に関するさらなる研究を促進するため、オープンソース実装を公開している。

著者らは謙虚に、このアプローチは現在、特定の構造（二乗ユークリッド距離および正規化後のコサイン距離）を持つコストに適用可能であると指摘している。この構造を持たないコスト（例、平方根を含む生ユークリッド距離や学習されたニューラルコスト）は、将来の課題として特定されている。