[论文解读] Matching Normalizing Flows and Probability Paths on Manifolds

Heli Ben-Hamu, Samuel M. Cohen+6 more|arXiv (Cornell University)|Jul 11, 2022

Generative Adversarial Networks and Image Synthesis被引用 5

一句话总结

该论文提出了CNF匹配（CNFM），一种新颖的框架，通过最小化概率路径分歧（PPD）在流形上训练连续归一化流（CNF），其中PPD是一种在训练过程中避免求解ODE的分歧度量。PPD利用对数质量守恒PDE，将CNF的流与目标概率路径匹配，从而在低维和中等维流形（包括球面和建模机器人状态空间的乘积流形）上实现了最先进水平的似然度与样本质量。

ABSTRACT

Continuous Normalizing Flows (CNFs) are a class of generative models that transform a prior distribution to a model distribution by solving an ordinary differential equation (ODE). We propose to train CNFs on manifolds by minimizing probability path divergence (PPD), a novel family of divergences between the probability density path generated by the CNF and a target probability density path. PPD is formulated using a logarithmic mass conservation formula which is a linear first order partial differential equation relating the log target probabilities and the CNF's defining vector field. PPD has several key benefits over existing methods: it sidesteps the need to solve an ODE per iteration, readily applies to manifold data, scales to high dimensions, and is compatible with a large family of target paths interpolating pure noise and data in finite time. Theoretically, PPD is shown to bound classical probability divergences. Empirically, we show that CNFs learned by minimizing PPD achieve state-of-the-art results in likelihoods and sample quality on existing low-dimensional manifold benchmarks, and is the first example of a generative model to scale to moderately high dimensional manifolds.

研究动机与目标

• 解决使用传统基于似然的方法在流形上训练CNF时计算成本高且可扩展性有限的问题。
• 开发一种分歧度量，实现在黎曼流形上高效、无需ODE的CNF训练，同时保持模型的表达能力。
• 实现在中等高维流形（例如S^15）上的训练，而先前的方法因计算或数值限制而失效。
• 提供一个兼容多种目标路径的灵活框架，这些路径在噪声分布与数据分布之间插值。

提出的方法

• 提出概率路径分歧（PPD），一种在CNF生成的概率路径与目标路径之间的新型分歧度量，通过对数质量守恒（LMC）PDE定义。
• LMC PDE将模型路径的对数密度与CNF的向量场联系起来，使得PPD可直接从向量场及其导数计算，而无需评估模型密度。
• PPD以单个参数ℓ ≥ 1的形式表述，且上界经典分歧度量：TV（ℓ=1）、α-分歧度（1<ℓ<∞）和反向KL（ℓ=∞）。
• 通过最小化PPD来训练CNF，避免了训练过程中的ODE求解，使时间和内存成本降低1至2个数量级。
• 通过使用内在几何（包括球面的立体投影和机器人状态空间的乘积流形）支持嵌入流形的数据。
• 采用灵活的目标路径设计，其中p_t从纯噪声过渡到数据分布，且log p_t及其导数可通过解析或数值方法计算。

实验结果

研究问题

• RQ1 能否定义一种分歧度量，实现在流形上无需ODE的CNF训练，同时保持高模型表达能力？
• RQ2 PPD是否为非欧几里得数据上的CNF提供一种理论基础扎实且可扩展的替代于基于似然的训练方法？
• RQ3 CNFM能否在低维和中等维流形上实现似然度与样本质量的最先进性能？
• RQ4 基于PPD的训练在更高维流形（如S^15或建模机器人运动学的乘积流形）上如何扩展？

主要发现

• CNFM在S^1、S^2和S^15上实现了最先进水平的对数似然度，优于vMF混合模型和S-FFJORD，在S^15上达到0.83的NLL得分。
• 在S^15上，CNFM将NLL降低至0.83，显著优于最佳vMF-MM基线（1.05）和S-FFJORD（1.15）。
• CNFM的训练时间比基于ODE的似然训练快1至2个数量级，4K秒的训练运行显示出显著的速度提升。
• CNFM在建模机器人状态空间的乘积流形（如Cheetah、Walker、Humanoid）上成功生成了定性上逼真的样本，具有合理的从噪声到数据的轨迹。
• 尽管在高维空间（如MNIST，d=784）中存在梯度估计偏差，CNFM仍保持稳定并生成合理样本，表明对近似误差具有鲁棒性。
• PPD被证明可上界经典分歧度量（TV、α-分歧度、反向KL），为其作为训练目标提供了理论依据。