[论文解读] Automatic Differentiation of Programs with Discrete Randomness

Gaurav Arya, Moritz Schauer+2 more|arXiv (Cornell University)|Oct 16, 2022

Data Stream Mining Techniques被引用 12

一句话总结

本文提出了一种基于重参数化的自动微分方法，适用于具有离散随机性的程序，通过在随机路径上传播无穷小扰动，实现无偏、低方差的梯度估计。该方法支持马尔可夫链、基于代理的模型和粒子滤波等随机模拟中的前向和反向模式自动微分，并通过剪枝和光滑化技术实现计算开销为 O(1) 的开源 Julia 实现。

ABSTRACT

Automatic differentiation (AD), a technique for constructing new programs which compute the derivative of an original program, has become ubiquitous throughout scientific computing and deep learning due to the improved performance afforded by gradient-based optimization. However, AD systems have been restricted to the subset of programs that have a continuous dependence on parameters. Programs that have discrete stochastic behaviors governed by distribution parameters, such as flipping a coin with probability $p$ of being heads, pose a challenge to these systems because the connection between the result (heads vs tails) and the parameters ($p$) is fundamentally discrete. In this paper we develop a new reparameterization-based methodology that allows for generating programs whose expectation is the derivative of the expectation of the original program. We showcase how this method gives an unbiased and low-variance estimator which is as automated as traditional AD mechanisms. We demonstrate unbiased forward-mode AD of discrete-time Markov chains, agent-based models such as Conway's Game of Life, and unbiased reverse-mode AD of a particle filter. Our code package is available at https://github.com/gaurav-arya/StochasticAD.jl.

研究动机与目标

• 解决具有离散随机行为（如伯努利试验）的程序的微分问题，其中经典自动微分因非可微、不连续的结果而失效。
• 开发一种方法，在保持随机程序离散结构的同时，实现统计期望的梯度计算。
• 实现无偏且低方差的梯度估计器，其性能可与连续路径式自动微分相媲美，且不依赖连续松弛或可调参数。
• 在复杂随机系统（包括马尔可夫链和粒子滤波）中支持前向和反向模式自动微分，且计算开销最小化。
• 提供一种可组合的、端到端可微分的框架，可无缝集成现有自动微分工具，支持概率建模中的基于梯度的优化。

提出的方法

• 引入了‘随机导数’的概念——一种路径式估计器，用于追踪由于无穷小参数扰动导致的结果概率的相对变化。
• 使用重参数化将离散随机变量表示为均匀随机变量和参数的确定性函数，从而实现梯度在随机路径上的传播。
• 应用‘平滑随机导数’技术处理离散事件，通过相关采样路径实现无偏梯度估计。
• 采用在线剪枝策略，通过动态消除不活跃或低影响路径，将前向模式自动微分的计算开销保持在 O(1)。
• 结合平滑化与路径式估计，推导出反向模式自动微分机制，将直通梯度估计器作为特例恢复。
• 在 StochasticAD.jl 中实现该方法，该 Julia 包可与 ForwardDiff.jl、Zygote.jl 和 ChainRulesCore.jl 集成，实现无缝的前向和反向模式微分。

实验结果

研究问题

• RQ1 自动微分能否扩展到具有离散随机组件（如伯努利或泊松抽样）的程序，其中经典导数未定义？
• RQ2 如何为具有离散随机性的随机程序的期望构造一个无偏且低方差的梯度估计器？
• RQ3 能否在不产生高计算成本的前提下，高效地将前向模式自动微分应用于复杂随机模拟（如基于代理的模型或马尔可夫链）？
• RQ4 是否可能为离散随机程序推导出一种反向模式自动微分机制，其性能可与连续路径式方法相媲美？
• RQ5 如何系统性地生成相关采样路径，以最小化离散随机程序梯度估计器的方差？

主要发现

• 所提方法实现了对离散随机程序期望的无偏梯度估计，且当扰动大小趋近于零时，方差保持有界。
• 该框架通过剪枝技术实现了对离散时间马尔可夫链和基于代理的模型（如康威生命游戏）的前向模式自动微分，计算开销为 O(1)。
• 该方法实现了粒子滤波的端到端反向模式自动微分，将先前工作中已知的技术作为特例恢复。
• 平滑随机导数技术使得前向和反向模式微分成为可能，克服了传统得分函数估计器的局限性。
• 实验评估表明，该方法在方差方面优于有限差分法，且与连续松弛方法具有竞争力，同时无需可调或学习参数。
• 开源的 StochasticAD.jl 包在多种应用场景中（包括随机模拟和概率推理）展示了该方法的实用性。