[论文解读] Unbalanced Sobolev Descent.

Youssef Mroueh, Mattia Rigotti|arXiv (Cornell University)|Jan 1, 2020

Statistical Methods and Inference被引用 1

一句话总结

Unbalanced Sobolev Descent (USD) 是一种基于粒子的算法，无需质量守恒即可将源分布传输至目标分布，通过在再生核希尔伯特空间（RKHS）中使用 Sobolev-Fisher 见证函数的梯度流实现。该方法结合了对流（粒子运动）和反应（通过出生-死亡过程的质量重加权），在最大均值差异（MMD）意义下实现对目标分布的渐近收敛，并在单细胞 RNA 测序分析中展现出更快的传输速度和更优的性能。

ABSTRACT

We introduce Unbalanced Sobolev Descent (USD), a particle descent algorithm for transporting a high dimensional source distribution to a target distribution that does not necessarily have the same mass. We define the Sobolev-Fisher discrepancy between distributions and show that it relates to advection-reaction transport equations and the Wasserstein-Fisher-Rao metric between distributions. USD transports particles along gradient flows of the witness function of the Sobolev-Fisher discrepancy (advection step) and reweighs the mass of particles with respect to this witness function (reaction step). The reaction step can be thought of as a birth-death process of the particles with rate of growth proportional to the witness function. When the Sobolev-Fisher witness function is estimated in a Reproducing Kernel Hilbert Space (RKHS), under mild assumptions we show that USD converges asymptotically (in the limit of infinite particles) to the target distribution in the Maximum Mean Discrepancy (MMD) sense. We then give two methods to estimate the Sobolev-Fisher witness with neural networks, resulting in two Neural USD algorithms. The first one implements the reaction step with mirror descent on the weights, while the second implements it through a birth-death process of particles. We show on synthetic examples that USD transports distributions with or without conservation of mass faster than previous particle descent algorithms, and finally demonstrate its use for molecular biology analyses where our method is naturally suited to match developmental stages of populations of differentiating cells based on their single-cell RNA sequencing profile. Code is available at this https URL .

研究动机与目标

• 开发一种粒子下降算法，可在不需质量守恒的情况下实现分布传输。
• 定义一种新的差异度量——Sobolev-Fisher 差异，使其与对流-反应动力学及 Wasserstein-Fisher-Rao 度量相联系。
• 在见证函数于 RKHS 中估计的温和假设下，实现 MMD 意义下对目标分布的收敛。
• 设计基于神经网络的 USD 实现，以实现可扩展且灵活的分布匹配。
• 在合成数据和真实世界单细胞 RNA 测序数据上展示该方法在发育阶段分析中的有效性。

提出的方法

• USD 采用两步过程：通过 Sobolev-Fisher 见证函数的梯度流实现对流，通过与见证函数成比例的质量重加权实现反应。
• 反应步骤通过与见证函数成比例的生长速率建模粒子的出生与死亡，实现质量再分配。
• Sobolev-Fisher 差异通过见证函数上的希尔伯特空间范数定义，使其与 Wasserstein-Fisher-Rao 度量相联系。
• 当见证函数在 RKHS 中估计时，在温和假设下，USD 可在 MMD 意义下渐近收敛至目标分布。
• 提出了两种基于神经网络的变体：一种使用粒子权重上的镜像下降，另一种使用随机出生-死亡过程进行重加权。
• 该方法利用由见证函数梯度流导出的对流-反应 PDE 所支配的粒子动力学。

实验结果

研究问题

• RQ1 当源分布与目标分布不守恒质量时，粒子下降算法能否有效实现分布传输？
• RQ2 如何定义 Sobolev-Fisher 差异，以支持分布传输中的对流与反应动力学？
• RQ3 当见证函数在 RKHS 中估计时，USD 是否能在 MMD 意义下收敛至目标分布？
• RQ4 神经网络能否有效估计 Sobolev-Fisher 见证函数，以实现可扩展且可微的分布匹配？
• RQ5 在合成数据和真实生物数据上，USD 与现有粒子下降方法相比，在收敛速度和精度方面表现如何？

主要发现

• 当 Sobolev-Fisher 见证函数在再生核希尔伯特空间（RKHS）中估计时，在温和假设下，USD 可在最大均值差异（MMD）意义下实现对目标分布的渐近收敛。
• 在合成示例中，与以往的粒子下降方法相比，USD 展现出更快的分布传输速度，尤其在非平衡设置下表现更优。
• USD 成功建模了单细胞 RNA 测序数据中的发育过渡，实现了不同分化阶段细胞群之间的精确匹配。
• 基于神经网络的 USD 变体——采用镜像下降和随机出生-死亡过程——实现了高效且可扩展的分布匹配。
• USD 中的反应步骤通过出生-死亡过程建模，实现了动态质量重加权，从而提升了传输效率和分布保真度。
• 实证结果表明，USD 在匹配质量不等的分布时，收敛速度和精度均优于现有方法。