[論文レビュー] Sinkhorn EM: An Expectation-Maximization algorithm based on entropic optimal transport

Gonzalo E. Mena, Amin Nejatbakhsh+2 more|arXiv (Cornell University)|Jun 30, 2020

Bayesian Modeling and Causal Inference参考文献 26被引用数 7

ひとこと要約

本稿では、既知の混合重みを組み込むために標準的な責任割り当てをエントロピー最適輸送に置き換えることで、ガウス混合モデル向けの新たな期待最大化（EM）アルゴリズム、Sinkhorn EM（sEM）を提案する。Sinkhornのアルゴリズムを用いることで、計算を効率化し、より良いグローバル収束性、局所最適解への感受性の低減、および合成データとC. elegans神経細胞セグメンテーションの両方で、通常のEMおよび過パラメータ化されたEMよりも優れた性能を達成する。

ABSTRACT

We study Sinkhorn EM (sEM), a variant of the expectation maximization (EM) algorithm for mixtures based on entropic optimal transport. sEM differs from the classic EM algorithm in the way responsibilities are computed during the expectation step: rather than assign data points to clusters independently, sEM uses optimal transport to compute responsibilities by incorporating prior information about mixing weights. Like EM, sEM has a natural interpretation as a coordinate ascent procedure, which iteratively constructs and optimizes a lower bound on the log-likelihood. However, we show theoretically and empirically that sEM has better behavior than EM: it possesses better global convergence guarantees and is less prone to getting stuck in bad local optima. We complement these findings with experiments on simulated data as well as in an inference task involving C. elegans neurons and show that sEM learns cell labels significantly better than other approaches.

研究の動機と目的

• 混合重みに関する事前知識が利用可能な状況で、標準的なEMの収束性の悪さを是正すること。
• 混合重みを無視し、収束が遅い過パラメータ化されたEMの限界を克服すること。
• 理論的裏付けが強く、計算的にも効率的なEMの変種を開発し、事前知識としての混合重みを自然に統合すること。
• 合成混合データおよび実際の神経画像データの両方で、細胞ラベルの回復を含め、クラスタリング性能の向上を図ること。

提案手法

• EMにおける標準的な対数尤度関数の代わりに、既知の成分重みを保持するエントロピー最適輸送（OT）に基づく損失関数を導入する。
• Sinkhornのアルゴリズムを用いて、既知の混合重みを尊重する責任を計算し、データポイント間で一貫した割り当てを保証する。
• EMと同様に、下界における座標昇下手続きとしてsEMを定式化するが、エントロピー正則化のおかげで最適解周辺の曲率が向上し、収束性が向上する。
• EMのEステップ/Mステップ構造を維持するが、Eステップを独立した事後確率の計算からエントロピーOT問題の解法に変更する。
• 重みの再調整やヒューリスティックな調整を回避するため、成分重みに関する事前知識を輸送計画に直接組み込む。
• Sinkhornアルゴリズムを用いることで、高次元データに対してもスケーラブルかつ高速に収束する効率的な最適化を実現する。

実験結果

リサーチクエスチョン

• RQ1 エントロピー最適輸送を用いることで、既知の成分重みを持つ混合モデルにおけるEMアルゴリズムの収束性とロバスト性を向上させることができるか？
• RQ2 sEMは、通常のEMおよび過パラメータ化されたEMと比較して、収束速度およびグローバル最適化性能において優れているか？
• RQ3 OTを用いて混合重みに関する事前知識を統合することで、実世界の神経画像データにおけるクラスタリング精度が向上するか？
• RQ4 特に2成分ガウス混合モデルの設定において、sEMは標準的なEMと比較して理論的な収束保証をどのように持っているか？
• RQ5 C. elegans神経細胞セグメンテーションタスクにおいて、sEMは既存の手法を上回る細胞識別回復性能を示せるか？

主な発見

• sEMは、通常のEMと比較して、より優れたグローバル収束性を示し、悪質な局所最適解に陥る可能性が著しく低減される。
• 2成分ガウス混合モデルのシミュレーションにおいて、sEMは通常のEMおよび過パラメータ化されたEMと比較して、より速くかつ信頼性の高い収束を達成する。
• C. elegans神経細胞セグメンテーションにおいて、sEMは通常のEMおよび過パラメータ化されたEMと比較して、顕著に高いセグメンテーション精度と低い平均二乗誤差を達成する。
• sEMは、推論品質の向上を図りながらも、高いテスト対数尤度を維持しており、一般化性能の向上を示唆している。
• 可視化結果から、sEMは複数の異なる初期化条件において一貫性があり、より正確な神経細胞セグメンテーションを生成することが確認された。
• エントロピーOTの使用により、目的関数の曲率が向上し、最適解が同一であっても、実用的にはより速い収束を実現している。