[论文解读] Path-Gradient Estimators for Continuous Normalizing Flows

Lorenz Vaitl, Kim A. Nicoli+2 more|arXiv (Cornell University)|Jun 17, 2022

Model Reduction and Neural Networks被引用 7

一句话总结

本文提出了一种用于连续归一化流的高效路径梯度估计器，实现了变分推断中更低方差的训练。通过求解一个辅助ODE来计算对数似然的梯度，该方法在每次迭代中仅增加5–14%的运行时间，相比标准重参数化梯度，实现了更快的收敛速度和更优的性能。

ABSTRACT

Recent work has established a path-gradient estimator for simple variational Gaussian distributions and has argued that the path-gradient is particularly beneficial in the regime in which the variational distribution approaches the exact target distribution. In many applications, this regime can however not be reached by a simple Gaussian variational distribution. In this work, we overcome this crucial limitation by proposing a path-gradient estimator for the considerably more expressive variational family of continuous normalizing flows. We outline an efficient algorithm to calculate this estimator and establish its superior performance empirically.

研究动机与目标

• 解决在现代变分推断中广泛应用的高表达能力连续归一化流缺乏高效路径梯度估计器的问题。
• 克服先前路径梯度方法在归一化流中不适用或计算成本过高的局限性。
• 实现高维问题（如VAEs和格点场论）中连续归一化流的更快、更稳定的训练。
• 提供一种即插即用的替代方案，取代标准重参数化梯度，计算开销极小，且在实际性能上表现更优。

提出的方法

• 推导一个辅助ODE（公式11），将变分对数密度的梯度沿流反向传播，从而实现路径梯度的计算。
• 使用与主流相同ODE求解器向前求解该辅助ODE，确保兼容性并保持低内存开销。
• 利用辅助ODE的解来计算路径梯度，该梯度仅捕捉通过重参数化采样对ELBO的隐式依赖关系。
• 通过极少的代码修改，将路径梯度估计器集成到现有训练流程中，实现即插即用的梯度替换。
• 利用连续归一化流的结构高效计算梯度，避免对整个ODE积分过程进行反向模式微分。
• 通过复用主流ODE求解过程中的中间状态，确保每次迭代的内存使用恒定，从而保持计算效率。

实验结果

研究问题

• RQ1 是否能够为高度表达但路径梯度计算具有挑战性的连续归一化流高效实现路径梯度估计器？
• RQ2 在高维设置下，所提出的路径梯度估计器是否相比标准重参数化梯度能实现更快的收敛速度和更优的ELBO优化？
• RQ3 与标准训练相比，路径梯度估计器在运行时间和内存方面的计算开销如何？
• RQ4 在VAEs和格点场论中，路径梯度估计器在不同架构、数据集和ODE求解器下的表现如何？
• RQ5 路径梯度估计器是否可在大规模变分推断任务（如计算物理和生成建模）中实际部署？

主要发现

• 路径梯度估计器降低了训练方差并加速了收敛，在最高维的格点场论任务中，有效样本量（ESS）提升了30%以上。
• 在A100 GPU上，该方法每次迭代仅增加5–14%的运行时间，且随着格点尺寸增大而进一步降低，得益于架构的可扩展性。
• 在VAEs中，路径梯度估计器的性能与标准重参数化梯度相当或更优，Omniglot数据集上的实际运行时间在基线的111%以内，MNIST数据集上为98%以内。
• 与Agrawal等人（2020）的最先进方法相比，本方法每轮迭代速度显著更快，且内存使用仅为一半，后者内存成本翻倍。
• 路径梯度与标准梯度估计器之间在函数评估次数上未观察到明显差异，表明ODE求解器行为稳定。
• 实证结果表明，将标准梯度替换为路径梯度可在多种任务和架构中带来一致且显著的性能提升。