[论文解读] Storchastic: A Framework for General Stochastic Automatic Differentiation

Emile van Krieken, Jakub M. Tomczak+1 more|arXiv (Cornell University)|Apr 1, 2021

Domain Adaptation and Few-Shot Learning参考文献 45被引用 8

一句话总结

Storchastic 是一种用于通用随机自动微分的新框架，使建模者能够在随机计算图的每个采样步骤中应用多种低方差梯度估计方法。通过将估计器分解为四个组成部分——提议分布、加权函数、梯度函数和控制变量——该框架在支持测度值导数和 SPSA 等高级技术的同时，为任意阶导数提供了可证明无偏的梯度估计。

ABSTRACT

Modelers use automatic differentiation (AD) of computation graphs to implement complex Deep Learning models without defining gradient computations. Stochastic AD extends AD to stochastic computation graphs with sampling steps, which arise when modelers handle the intractable expectations common in Reinforcement Learning and Variational Inference. However, current methods for stochastic AD are limited: They are either only applicable to continuous random variables and differentiable functions, or can only use simple but high variance score-function estimators. To overcome these limitations, we introduce Storchastic, a new framework for AD of stochastic computation graphs. Storchastic allows the modeler to choose from a wide variety of gradient estimation methods at each sampling step, to optimally reduce the variance of the gradient estimates. Furthermore, Storchastic is provably unbiased for estimation of any-order gradients, and generalizes variance reduction techniques to higher-order gradient estimates. Finally, we implement Storchastic as a PyTorch library at https://github.com/HEmile/storchastic.

研究动机与目标

• 解决现有随机自动微分框架的局限性，这些框架受限于连续变量，或因高方差的得分函数估计器而表现不佳。
• 使建模者能够为每个采样步骤选择最优的梯度估计方法，以最小化方差而不引入偏差。
• 将方差减少技术（包括控制变量）推广至随机计算图中的任意阶导数估计。
• 为 n 阶导数估计提供正式的数学基础，确保在估计器组件满足明确定义条件时，估计结果无偏。
• 实现一个实用的、开源的 PyTorch 库，以支持真实世界中的部署和多样化梯度估计器的实验。

提出的方法

• 将随机梯度估计器分解为四个核心组件：提议分布、加权函数、梯度函数和控制变量。
• 通过数学框架形式化前向模式自动微分的计算过程，以支持对梯度传播的严谨分析。
• 证明：若四个组件满足特定条件，则生成的代理损失可产生无偏的 n 阶导数估计。
• 通过组件分解，将高级梯度估计技术（如测度值导数 MVD 和 SPSA）集成到框架中。
• 通过将控制变量技术扩展至 n 阶导数估计，实现对高阶导数的方差减少。
• 将框架实现为 PyTorch 库，支持与深度学习工作流和动态计算图的无缝集成。

实验结果

研究问题

• RQ1 能否设计一种通用的随机自动微分框架，以支持广泛的梯度估计方法，同时确保任意阶导数估计的无偏性？
• RQ2 如何系统性地将控制变量等方差减少技术扩展至随机计算图中的高阶导数估计？
• RQ3 对梯度估计器四个组件的哪些正式条件可保证 n 阶导数估计的无偏性？
• RQ4 能否将测度值导数和 SPSA 等高级方法统一在一个可组合的单一框架中，用于随机自动微分？
• RQ5 建模者在不损害梯度估计正确性的前提下，能在多大程度上自由选择每个采样步骤的梯度估计器？

主要发现

• Storchastic 是首个正式证明：当四个估计器组件满足特定数学条件时，任意阶导数估计保持无偏的框架。
• 该框架支持广泛的梯度估计方法，包括重参数化、得分函数估计器、RELAX、MAPO、无序集合估计器以及测度值导数。
• Storchastic 实现了基于控制变量的方差减少技术向任意阶导数估计的扩展，这是该领域的一项新贡献。
• 在 PyTorch 中的实现使建模者能够通过极少的代码修改在不同估计器间切换，如在 VAE 实验中仅通过一行代码修改即可对比不同方法。
• 在 MNIST VAE 上的实验结果表明，该框架正确支持多种估计器，其预期性能排序与先前工作一致，尽管训练周期有限且未进行超参数调优。
• 该框架成功将 SPSA 和测度值导数集成到随机自动微分流程中，证明了其通用性和可扩展性。