Probabilistic Programming with Programmable Variational Inference

与现代概率编程语言（PPL）支持的各种高级蒙特卡罗方法相比，PPL对变分推断（VI）的支持较少：用户通常只能使用预定义的一些变分目标和梯度估计器，这些估计器在PPL后端中是作为整体实现的（且没有形式上的正确性论证）。在本文中，我们提出了一种更模块化的方法来支持PPL中的变分推断，该方法基于组合程序转换。在我们的方法中，变分目标被表示为程序，可以使用一流的构造来计算用户定义的模型和变分族下的密度和期望值。然后，我们将这些程序系统地转换为无偏梯度估计器，以优化它们定义的目标。我们的设计使得可以对许多相互作用的问题进行模块化推理，包括自动微分、密度累积、跟踪和无偏梯度估计策略的应用。此外，相对于现有的PPL对VI的支持，我们的设计在三个方面增加了表达能力：（1）它支持开放式的用户定义变分目标，而不是固定的选项菜单；（2）它支持梯度估计策略的组合空间，其中许多策略目前的PPL并不能自动化；（3）它支持更广泛的模型和变分族，因为它支持用于近似边缘化和归一化的构造（之前仅用于蒙特卡罗推断）。我们在Gen概率编程系统的扩展中实现了我们的方法（genjax.vi，在JAX中实现），并在几个深度生成建模任务上进行了评估，表明与手工编码实现相比，几乎没有性能开销，并且性能与已经建立的开源PPL相当。

本论文旨在提出一种更加模块化的方法来支持概率编程语言中的变分推断，以解决当前PPL支持变分推断的局限性。具体来说，论文试图解决PPL对于变分目标和梯度估计器的限制，以及缺乏形式正确性证明的问题。

论文的关键思路是将变分目标表达为程序，并使用程序转换技术将其系统地转换为无偏梯度估计器。这种设计使得自动微分、密度积累、追踪和无偏梯度估计策略等多个交互关注点可以进行模块化推理。相比于现有的PPL支持，该设计增加了三个方面的表达能力。

论文的亮点包括：1. 支持开放式的用户定义变分目标；2. 支持多种梯度估计策略，并且不仅限于当前自动化的PPL；3. 支持更广泛的模型和变分族，因为它支持近似边缘化和归一化构造。论文在Gen概率编程系统中实现了该方法，并在几个深度生成建模任务上进行了评估，结果表明性能与已建立的开源PPL相当。

在最近的相关研究中，也有一些关于变分推断的工作，例如：Variational Inference: A Review for Statisticians (Blei et al., 2017)、Black Box Variational Inference (Ranganath et al., 2014)、Automatic Differentiation Variational Inference (Kucukelbir et al., 2017)等。