[论文解读] Pathwise Derivatives Beyond the Reparameterization Trick

Martin Jankowiak, Fritz Obermeyer|arXiv (Cornell University)|Jun 5, 2018

Gaussian Processes and Bayesian Inference参考文献 33被引用 26

一句话总结

该论文将最优质量传输（OMT）梯度作为重参数化技巧的推广，实现了对伽马分布、贝塔分布和狄利克雷分布等无法应用标准重参数化技巧的分布的低方差路径梯度。研究表明，基于Cholesky分解的多变量正态分布的重参数化技巧在最优传输意义下并非最优，并推导出改进的梯度，显著降低了方差，提升了高斯过程回归和变分推断任务中的性能。

ABSTRACT

We observe that gradients computed via the reparameterization trick are in direct correspondence with solutions of the transport equation in the formalism of optimal transport. We use this perspective to compute (approximate) pathwise gradients for probability distributions not directly amenable to the reparameterization trick: Gamma, Beta, and Dirichlet. We further observe that when the reparameterization trick is applied to the Cholesky-factorized multivariate Normal distribution, the resulting gradients are suboptimal in the sense of optimal transport. We derive the optimal gradients and show that they have reduced variance in a Gaussian Process regression task. We demonstrate with a variety of synthetic experiments and stochastic variational inference tasks that our pathwise gradients are competitive with other methods.

研究动机与目标

• 将路径梯度估计方法从重参数化技巧推广至无法应用标准重参数化技巧的分布，如伽马、贝塔和狄利克雷分布。
• 基于最优传输理论，识别并推导出最优路径梯度，表明标准重参数化梯度在某些情况下存在非最优性。
• 在随机变分推断中降低梯度方差，特别是针对通过Cholesky分解参数化的多变量正态分布。
• 在实际应用中展示性能提升，包括高斯过程回归和合成变分推断任务。

提出的方法

• 建立重参数化技巧与最优传输理论中传输方程解之间的对应关系。
• 将最优路径梯度（OMT梯度）作为最小化方差的最优传输意义下的传输方程解进行推导。
• 将OMT框架应用于通过Cholesky因子参数化的多变量正态分布，推导出闭式解梯度，并证明标准重参数化梯度的非最优性。
• 利用传输方程的数值解，为伽马、贝塔和狄利克雷等单变量分布开发高精度近似路径梯度。
• 通过速度场可视化展示多变量设置下重参数化梯度与OMT梯度之间的几何差异。
• 通过合成实验和真实世界任务（包括高斯过程回归和稀疏伽马缺陷模型）验证性能，采用单样本梯度估计器。

实验结果

研究问题

• RQ1 能否系统地为缺乏标准重参数化形式的分布（如伽马、贝塔和狄利克雷分布）推导出路径梯度？
• RQ2 在最优传输意义下，基于Cholesky分解的多变量正态分布的重参数化技巧是否非最优？
• RQ3 OMT梯度能否降低随机变分推断中的梯度方差并提升收敛性？
• RQ4 在ELBO优化和训练稳定性方面，OMT梯度与得分函数估计器和重参数化估计器相比表现如何？

主要发现

• 对于通过Cholesky分解参数化的多变量正态分布，重参数化技巧产生的速度场中存在显著的旋转分量，表明其在最优传输意义下存在非最优性。
• 多变量正态分布的OMT梯度已通过解析方法推导，并在高斯过程回归任务中显示其方差低于标准重参数化梯度。
• 对于伽马、贝塔和狄利克雷分布，成功利用最优传输框架计算出高精度近似路径梯度。
• 在合成实验中，OMT梯度的方差更低，收敛速度更快，优于重参数化技巧和得分函数估计器。
• 在稀疏伽马缺陷模型中，OMT梯度估计器优于重参数化技巧，实现了更高的ELBO值并更快收敛。
• OMT梯度估计器在多个变分推断任务中表现出一致的性能提升，包括具有复杂、不可重参数化似然函数的任务。