Randomized Automatic Differentiation

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.

arXiv commentICLR 2021

References in corpus (11)

• A Simple Way to Initialize Recurrent Networks of Rectified Linear Units
• Training Deep Nets with Sublinear Memory Cost
• The Reversible Residual Network: Backpropagation Without Storing Activations
• Residual Flows for Invertible Generative Modeling
• meProp: Sparsified Back Propagation for Accelerated Deep Learning with Reduced Overfitting
• Unbiasing Truncated Backpropagation Through Time
• Minimal Effort Back Propagation for Convolutional Neural Networks
• Tangent: Automatic differentiation using source-code transformation for dynamically typed array programming
• Efficient Optimization of Loops and Limits with Randomized Telescoping Sums
• SUMO: Unbiased Estimation of Log Marginal Probability for Latent Variable Models
• Faster Neural Network Training with Approximate Tensor Operations

Cited by in corpus (2)

• Distributed Training and Optimization Of Neural Networks
• Bias-Free Scalable Gaussian Processes via Randomized Truncations