[论文解读] Randomized Automatic Differentiation

Deniz Oktay, Nick McGreivy+3 more|arXiv (Cornell University)|Jul 20, 2020

Gaussian Processes and Bayesian Inference参考文献 34被引用 8

一句话总结

本文提出了随机自动微分（RAD），一种通过在自动微分过程中引入随机性，以无偏、低内存的梯度估计器替代精确的反向模式自动微分的框架。通过以方差换取内存，RAD 在随机优化中实现了更快的收敛速度——在前馈和循环神经网络上得到验证，并应用于线性PDE的反应堆控制问题，显著降低了内存消耗。

ABSTRACT

The successes of deep learning, variational inference, and many other fields have been aided by specialized implementations of reverse-mode automatic differentiation (AD) to compute gradients of mega-dimensional objectives. The AD techniques underlying these tools were designed to compute exact gradients to numerical precision, but modern machine learning models are almost always trained with stochastic gradient descent. Why spend computation and memory on exact (minibatch) gradients only to use them for stochastic optimization? We develop a general framework and approach for randomized automatic differentiation (RAD), which can allow unbiased gradient estimates to be computed with reduced memory in return for variance. We examine limitations of the general approach, and argue that we must leverage problem specific structure to realize benefits. We develop RAD techniques for a variety of simple neural network architectures, and show that for a fixed memory budget, RAD converges in fewer iterations than using a small batch size for feedforward networks, and in a similar number for recurrent networks. We also show that RAD can be applied to scientific computing, and use it to develop a low-memory stochastic gradient method for optimizing the control parameters of a linear reaction-diffusion PDE representing a fission reactor.

研究动机与目标

• 解决传统自动微分产生的精确梯度与现代深度学习优化中的随机性之间的不匹配问题。
• 通过引入随机化估计器，减少反向模式自动微分中的内存使用，以方差换取更少的存储空间。
• 开发一个通用框架，实现无偏、低内存的梯度估计，适用于神经网络和物理模拟器。
• 证明在固定内存预算下，RAD 可实现比小批量训练更快的收敛速度。
• 将自动微分的适用范围扩展到高内存成本的大规模科学计算问题。

提出的方法

• 提出反向模式自动微分的随机变体，通过计算图路径的随机采样来计算无偏梯度估计。
• 对线性化计算图（LCG）应用随机化，其中边权重（局部导数）被随机子采样，以减少内存占用。
• 利用Bauer公式通过LCG中的路径枚举计算雅可比矩阵，并通过随机选择路径来减少中间导数的存储。
• 引入一种子采样方案，从LCG中选择部分路径，在保持无偏性的同时减少内存占用。
• 采用可伸缩级数估计器处理难以处理的图结构，使复杂模型中的梯度估计变得可行。
• 将RAD应用于前馈和循环神经网络，以及一个线性反应-扩散PDE，展示了其内存效率。

实验结果

研究问题

• RQ1 随机自动微分能否在保持无偏梯度估计的同时，减少反向微分中的内存使用？
• RQ2 在固定内存预算下，RAD 与小批量训练相比，收敛速度如何？
• RQ3 RAD 能否有效应用于高内存需求的科学计算问题，如PDE优化？
• RQ4 计算图的何种结构特性使得在自动微分中可有效进行随机化而不引入偏差？
• RQ5 RAD估计器的方差与标准自动微分相比如何，能否被有效控制？

主要发现

• 对于前馈网络，在相同内存预算下，RAD 的收敛迭代次数少于小批量训练。
• 对于循环网络，尽管内存使用减少，RAD 的收敛迭代次数与小批量训练相近。
• RAD 实现了对线性反应-扩散PDE的低内存优化，用于反应堆控制，展示了其在科学计算中的适用性。
• 该框架通过在线性化计算图中对路径进行子采样，在显著降低内存需求的同时保持了无偏梯度估计。
• 在内存受限环境下，RAD 通过利用问题特异性结构减少存储，同时保持无偏性，优于标准自动微分。
• 该方法具有通用性，可通过随机舍入和可伸缩估计器扩展至其他计算图，包括具有非线性项和循环结构的图。