利用 JKO 格式驯服连续归一化流中的超参数调优

归一化流（NF）是将所选概率分布变换为标准正态分布的映射。此类流是机器学习与数据科学中用于数据生成和密度估计的常用技术。由 NF 获得的密度估计需要涉及 NF 变换的雅可比行列式计算的变量替换公式。为可处理地计算该行列式，连续归一化流（CNF）利用神经 ODE 估计映射及其雅可比行列式。最优传输（OT）理论已成功用于辅助寻找 CNF，将其表述为以软惩罚强制标准正态分布为目标测度的 OT 问题。基于 OT 的 CNF 的一个缺陷是增加了超参数 $\alpha$，其控制软惩罚的强度并需大量调优。我们提出 JKO-Flow，一种无需调优 $\alpha$ 的求解基于 OT 的 CNF 的算法。这是通过将 OT CNF 框架整合进 Wasserstein 梯度流框架（亦称 JKO 格式）实现的。我们不再调优 $\alpha$，而是对固定 $\alpha$ 重复求解优化问题，有效地执行带时间步 $\alpha$ 的 JKO 更新。因此，我们通过重复求解更简单的问题而非求解可能更难的大 $\alpha$ 问题，获得了一种“分而治之”的算法。