[論文レビュー] Importance Sparsification for Sinkhorn Algorithm

Mengyu Li, Jun Yu+2 more|arXiv (Cornell University)|Jun 11, 2023

Bayesian Modeling and Causal Inference被引用数 5

ひとこと要約

本稿では、未知の最適輸送計画に対する自然な上界を用いてサンプリング確率を導出することで、スパースなカーネル行列を構築する新しい重要度スパース化手法Spar-Sinkを提案する。この手法により、エントロピー正則化最適輸送（OT）およびアンバランスOT（UOT）のSinkhornアルゴリズムの計算コストを、1反復あたり$O(n^2)$から$\tildewide{O}(n)$に削減し、弱い条件下でも一貫性を持つ推定器を達成する。合成データおよび心エコー画像データ（心拍周期解析や生成モデル構築を含む）において、速度と精度の両面で先行手法を上回る性能を示した。

ABSTRACT

Sinkhorn algorithm has been used pervasively to approximate the solution to optimal transport (OT) and unbalanced optimal transport (UOT) problems. However, its practical application is limited due to the high computational complexity. To alleviate the computational burden, we propose a novel importance sparsification method, called Spar-Sink, to efficiently approximate entropy-regularized OT and UOT solutions. Specifically, our method employs natural upper bounds for unknown optimal transport plans to establish effective sampling probabilities, and constructs a sparse kernel matrix to accelerate Sinkhorn iterations, reducing the computational cost of each iteration from $O(n^2)$ to $\widetilde{O}(n)$ for a sample of size $n$. Theoretically, we show the proposed estimators for the regularized OT and UOT problems are consistent under mild regularity conditions. Experiments on various synthetic data demonstrate Spar-Sink outperforms mainstream competitors in terms of both estimation error and speed. A real-world echocardiogram data analysis shows Spar-Sink can effectively estimate and visualize cardiac cycles, from which one can identify heart failure and arrhythmia. To evaluate the numerical accuracy of cardiac cycle prediction, we consider the task of predicting the end-systole time point using the end-diastole one. Results show Spar-Sink performs as well as the classical Sinkhorn algorithm, requiring significantly less computational time.

研究の動機と目的

• 大規模な最適輸送（OT）およびアンバランスOT（UOT）問題におけるSinkhornアルゴリズムの高い計算コストを軽減すること。
• 推定精度を維持しつつ、計算複雑性を著しく低減するスパース近似手法を開発すること。
• 心エコー画像解析や生成モデル構築を含む、実世界の応用に向けた効率的かつ正確なOTおよびUOT計算を可能にすること。
• 弱い正則性条件の下で、提案手法の理論的一貫性を確立すること。
• 合成データおよび実データの両方において、既存の最先端手法を速度と推定誤差の両面で上回ること。

提案手法

• 未知の最適輸送計画に対する自然な上界を用いて、カーネル行列のスパース化に適した有効なサンプリング確率を導出するSpar-Sinkを提案する。
• これらの重要度確率に基づく要素ごとのサンプリングによりスパースなカーネル行列を構築し、Sinkhorn反復ごとの計算コストを$O(n^2)$から$\tildewide{O}(n)$に削減する。
• スパース化されたカーネル行列をSinkhornアルゴリズムに適用し、エントロピー正則化OTおよびUOTの解を効率的に近似する。
• 生成モデルでは$s = 10s_0(n)$、カラーテンプレート移行では$s = 8s_0(n)$のサブサンプリングパラメータを用い、精度と速度のバランスを図る。
• OTフレームワークにおける比較のためのロバスト版であるRobust-NysSinkを採用し、Spar-Sinkを新たなオートエンコーダー変種、Spar-Sinkオートエンコーダー（SSAE）に統合する。
• 理論的分析により、弱い正則性条件の下で、正則化OTおよびUOTの両方における推定器の一貫性を証明する。

実験結果

リサーチクエスチョン

• RQ1 自然な上界に基づく重要度スパース化により、推定精度を損なわずにSinkhornアルゴリズムの計算コストを著しく削減できるか？
• RQ2 Nys-Sink や Robust-NysSink と比較して、Spar-Sink はOTおよびUOT問題において速度と精度の両面で優れているか？
• RQ3 Spar-Sink は心エコー画像データにおける心拍周期解析を効果的に高速化でき、収縮期終了時刻を収縮期開始時刻から正確に予測できるか？
• RQ4 Spar-Sink を生成モデル（SSAE）に統合した場合、元のSinkhornオートエンコーダー（SAE）と比較して性能が維持または向上するか？
• RQ5 弱い正則性条件の下で、Spar-Sink推定器の正則化OTおよびUOTにおける理論的一貫性はどのように保証されるか？

主な発見

• カラーテンプレート移行の輸送計画計算において、CPU時間はSinkhornの60.45秒からSpar-Sinkでは3.15秒に削減され、視覚的結果はほぼ同一であった。
• 心エコー画像データにおいて、Spar-Sinkは古典的手法のSinkhornと同等の精度で収縮期終了時刻を予測したが、はるかに短時間で処理を完了した。
• 生成モデルにおいて、SSAEはFIDスコア23.65 ± 0.06を達成し、SAEの24.72 ± 0.13を下回り、学習時間も半分の64.81秒（SAEは125.87秒）で完了した。
• 本手法は1反復あたり$\tildewide{O}(n)$の計算コストを達成し、標準的なSinkhornの$O(n^2)$に比べて顕著な改善を示した。
• 合成データを用いた実験により、Spar-Sinkが主流の競合手法を推定誤差および計算速度の両面で上回ることが確認された。
• 理論的分析により、弱い正則性条件の下で、Spar-Sink推定器が正則化OTおよびUOTにおいて一貫性を持つことが確認された。