[论文解读] 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

一句话总结

本文提出了一种新型期望最大化算法——Sinkhorn EM（sEM），用于高斯混合模型，该算法用基于熵正则化最优传输的责任分配替代标准责任分配，以整合已知的混合权重。通过利用Sinkhorn算法实现高效计算，sEM在全局收敛性、对局部最优解的敏感性以及在合成数据和秀丽隐杆线虫神经元分割任务中的性能方面，均优于标准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算法计算尊重已知混合权重的责任，确保数据点间的责任分配一致。
• 将sEM建模为对下界进行坐标上升的优化过程，类似于EM，但因熵正则化而使最优解附近的曲率更优。
• 保持EM的E步/M步结构，但将E步修改为求解熵正则化OT问题，而非独立计算后验概率。
• 将关于分量权重的先验知识直接整合到传输计划中，避免重加权或启发式调整的需要。
• 通过Sinkhorn算法实现高效优化，该算法具有良好的可扩展性，即使在高维数据上也能快速收敛。

实验结果

研究问题

• RQ1 熵正则化最优传输能否用于改善已知分量权重的混合模型中EM算法的收敛性与鲁棒性？
• RQ2 与标准EM和过参数化EM相比，sEM在收敛速度和全局优化性能方面表现如何？
• RQ3 通过OT整合关于混合权重的先验知识，是否能提升真实世界神经影像数据上的聚类准确率？
• RQ4 与标准EM相比，sEM在两分量高斯混合设置下的理论收敛保证是什么？
• RQ5 sEM能否在秀丽隐杆线虫神经元分割任务中优于现有方法，更准确地恢复细胞身份？

主要发现

• sEM在全局收敛性方面优于标准EM，显著降低了陷入次优局部最优解的可能性。
• 在模拟的两分量高斯混合模型中，sEM的收敛速度和可靠性均优于标准EM和过参数化EM。
• 在秀丽隐杆线虫神经元分割任务中，sEM的分割准确率显著更高，且均方误差更低，优于标准EM和过参数化EM。
• sEM在保持高测试对数似然的同时，提升了下游推理质量，表明其泛化能力更强。
• 可视化结果表明，sEM在多种随机初始化下均能产生更准确且一致的神经元分割结果。
• 熵正则化OT的引入使目标函数具有更优的曲率，尽管全局最优解相同，但在实际应用中收敛速度更快。