[論文レビュー] Scalable Optimal Transport in High Dimensions for Graph Distances, Embedding Alignment, and More

Johannes Klicpera, Marten Lienen+1 more|arXiv (Cornell University)|Jul 14, 2021

Caching and Content Delivery被引用数 3

ひとこと要約

本稿では、高次元空間におけるエントロピー正則化最適輸送（OT）を高速化するための2つの対数線形時間近似法——LSHに基づくスパースコスト行列と局所補正Nystrom（LCN）——を提案する。これらの手法により、深層学習応用におけるスケーラブルなOTが可能となり、単語埋め込みの整合性評価では3倍の高速化と3.1%の精度向上を達成し、グラフ距離回帰ではグラフ輸送ネットワーク（GTN）で48%の性能向上を達成した。

ABSTRACT

The current best practice for computing optimal transport (OT) is via entropy regularization and Sinkhorn iterations. This algorithm runs in quadratic time as it requires the full pairwise cost matrix, which is prohibitively expensive for large sets of objects. In this work we propose two effective log-linear time approximations of the cost matrix: First, a sparse approximation based on locality-sensitive hashing (LSH) and, second, a Nystrom approximation with LSH-based sparse corrections, which we call locally corrected Nystrom (LCN). These approximations enable general log-linear time algorithms for entropy-regularized OT that perform well even for the complex, high-dimensional spaces common in deep learning. We analyse these approximations theoretically and evaluate them experimentally both directly and end-to-end as a component for real-world applications. Using our approximations for unsupervised word embedding alignment enables us to speed up a state-of-the-art method by a factor of 3 while also improving the accuracy by 3.1 percentage points without any additional model changes. For graph distance regression we propose the graph transport network (GTN), which combines graph neural networks (GNNs) with enhanced Sinkhorn. GTN outcompetes previous models by 48% and still scales log-linearly in the number of nodes.

研究の動機と目的

• 高次元の深層学習応用における最適輸送（OT）の計算ボトル neck を解消すること。特に、完全なコスト行列の計算は計算的に非現実的である。
• 標準のSinkhornアルゴリズムの代替として、エントロピー正則化に依存するが、二次的コスト行列の計算を回避する効率的でスケーラブルな手法を開発すること。
• 単語埋め込みの整合性評価やグラフ距離回帰といった大規模な設定におけるOTの実用的応用を可能にすること。
• GNNと強化されたSinkhornを組み合わせた、新たなグラフ輸送ネットワーク（GTN）を導入し、性能を向上させつつ対数線形スケーラビリティを維持すること。
• 近似の品質と実世界のタスクにおけるエンドツーエンドの性能向上について、理論的および実験的検証を行うこと。

提案手法

• 局所性に敏感なハッシュ（LSH）を用いたスパースコスト行列近似を提案し、OTにおけるペairwise距離計算の回数を削減する。
• 局所補正Nystrom（LCN）法を導入し、コスト行列の低ランクNystrom近似に適したスパース補正項をLSHで特定する。
• LSHに基づく近似をエントロピー正則化OTに統合し、Sinkhorn反復を用いることで、収束保証を維持しつつ時間計算量を対数線形に削減する。
• GNNがノード特徴を生成するグラフ輸送ネットワーク（GTN）を設計し、LCN近似を用いた強化されたSinkhornを適用してグラフレベルの距離回帰を実行する。
• 理論的分析により、LCNが、元のデータ分布にやや弱い仮定のもとで収束性と精度を保持することを示した。
• 単語埋め込みの整合性評価やグラフ距離予測といった下流タスクにおいて、提案手法のエンドツーエンドの訓練と評価を実施した。

実験結果

リサーチクエスチョン

• RQ1 LSHに基づくコスト行列のスパース化は、対数線形時間計算量を達成しつつ、エントロピー正則化OTの精度を保持できるか？
• RQ2 局所補正Nystrom（LCN）法は、標準のNystrom法やLSH単体と比較して、近似品質と計算効率のトレードオフをより良く果たせるか？
• RQ3 提案された近似手法は、無教師学習による単語埋め込みの整合性評価といった実世界の深層学習応用において、効果を発揮するか？
• RQ4 LCN近似OTをグラフニューラルネットワークフレームワーク（GTN）に統合することで、グラフ距離回帰において最先端の性能を達成できるか？
• RQ5 既存のパイプラインにおいて標準のSinkhornを提案手法の近似に置き換えた場合、実験的にどれほど速度と精度が向上するか？

主な発見

• LSHに基づくスパース近似とLCN法により、エントロピー正則化OTの時間計算量が対数線形に削減され、大規模データセットへのスケーラビリティが実現された。
• 無教師学習による単語埋め込みの整合性評価では、提案手法により最先端手法の処理速度が3倍に向上し、精度も3.1ポイント向上した。
• GNNとLCN近似Sinkhornを組み合わせたグラフ輸送ネットワーク（GTN）は、グラフ距離回帰において、以前のモデルを48%上回る性能を発揮した。
• LCN近似は高い近似品質を維持しており、完全行列のSinkhornと比較してOT解の精度にほとんど劣化がなかった。
• 提案手法は、グラフタスクにおけるノード数に対して対数線形にスケーリング可能であり、大規模グラフにおいて実用的である。
• 実験的評価により、両近似手法が、深層学習表現を含む多様な高次元設定において有効であることが確認された。