Randomized Automatic Differentiation
Authors: Deniz Oktay
Nick McGreivy
Joshua Aduol
Alex Beatson
Ryan P Adams
Venue: ICLR 2021 Oral (2021)
Keywords: automatic differentiation
autodiff
backprop
deep learning
pdes
stochastic optimization

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.

Published as a conference paper at ICLR 2021
RANDOMIZED AUTOMATIC DIFFERENTIATION
Deniz Oktay, Nick McGreivy, Joshua Aduol, Alex Beatson, Ryan P. A.damss Princeton University
Princeton, NJ
Megstay,mgrrivy. jaduol, abeatson, spw.pprincetoo..ccccmmmmmm edu
ABSTRACT
The successes of deep learning, variational inference, and many other fields have been aided by specialized implementations of reverse-mode automatic differenti- ation (AD) to compute gradients of mega-dimensional objectives. The AD tech- niques underlying these tools were designed compute exact gradients nu- merical 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 low unbiased gradient estimates to be computed with reduced memory in return for variance. We examine limitations of the general approach, and argue that must leverage problem specific structure to realize benefits. We develop RAD techniques for a variety of simple neural network architectures, and show that for fixed memory budget, RAD converges in fewer iterations than using small batch size for feedforward networks, and in a similar number for recurrent networks. We also show that RAD can be applied scientific computing, and use it to develop low-memory stochastic gradient method for optimizing the control parameters of a linear reaction-diffusion PDE representing fission reactor.
INTRODUCTION
Deep neural networks have taken center stage as a powerful way to construct and train massively- parametric machine learning (ML) models for supervised, unsupervised, and reinforcement learning tasks. There are many reasons for the resurgence of neural networks- large data sets, GPU numerical computing, technical insights into overyaraaeeterization, and more-but one major factor has been the development of tools for automatic differentiation (AD) of deep architectures. Tools like PyTorch and TensorFlow provide a computational substrate for rapidly exploring wide variety of differentiable architectures without performing tedious and error-prone gradient derivations. The flexibility of these tools has enabled a revolution in AI research, but the underlying ideas for reverse- mode AD go back decades. While tools like Py PyTch and TensorFlow have received huge dividends from half-century of AD research, they are also burdened by the baggage of design decisions made in a different computational landscape. The research on AD that led to these ubiquitous deep learning frameworks is focused on the computation of Jacobians that are exact up to numerical precision. However, in modern workflows these Jacobians used for stochastic optimization. We ask:
Why spend resources on exact gradients when we' 'e going to use stochastic optimization?
This question is motivated by the surprising realization over the past decade that deep neural network training can be performed almost entirely with first-order stochastic optimization. fact, empirical evidence supports the hypothesis that the regularizing effect of gradient noise assists model gen- eralization (Keskar et al., 2017: Smith Le, 2018: Hochreiter & Schmidhuber, 1997). Stochastic gradient descent variants such as AdaGrad (Duchi et al., 2011) and Adam (Kingma & Ba, 2015) form the core of almost all successful optimization techniques for these models, using small subsets of the data to form the noisy gradient estimates.
Department of Computer Science
2 DDpartment of Astrophysical Sciences

---

Published as a conference paper at ICLR 2021
from math import sin, exp
def f(xl, x2):
a = exp(x1)
eap bol
b (x2) c = b * x2
exp
d =
return a * d
costef1
1923
(a) Differentiable Python function
(b) Primal graph
(c) Linearized graph
(d) Bauer paths
Figure I:
Illustration of the basic concepts of the linearized computational graph and Baucr's 's formula. (a) a simple Python function with intermediate variables; (b) the primal computational graph, a DAG with variables as vertices and flow moving upwards to the output; (c) the linearized computational graph (LCG) in which the edges labeled with the values of the local derivatives; (d) illustration of the four paths that must be evaluated to compute the Jacobian. (Example from Paul D. Hovland.)
The goals and assumptions of automatic differentiation as performed in classical and modern systems are mismatched with those required by stochastic optimization. Traditional AD computes the derivative or Jacobian of function accurately to numerical precision. This accuracy is required for many problems in applied mathematics which AD has served, e.g., solving systems of differential make do with inaccurate gradients,
equations. But in stochastic optimization our estimator is unbiased and has reasonable variance. We ask the same question that motivates mini-batch SGD: why compute exact gradient if we can get noisy estimates cheaply? By thinking of this question in the context of AD, we can go beyond mini- batch SGD to more general schemes for developing cheap gradient estimators: in this paper, we focus on developing gradient estimators with low memory cost. Although previous research has investigated approximations in the forward or reverse pass of neural networks reduce computational requirements, here we replace deterministic AD with randomized automatic differentiation (RAD), trading off of computation for variance inside AD routines when imprecise gradient estimates are tolerable, while retaining unbiasedness.
long as
2 AUTOMATIC DIFFERENTIATION
Automatic (or algorithmicc differentiation is a family of techniques for taking program that computes a differentiable function f : IR" + R", and producing another program that computes the associated derivatives; most often the Jacobian: [[]]
f' R" + R= (For comprehensive treatment of AD, see Griewank & Walther (2008); for an ML-focused review see Baydin et al. (2018).) In most machine learning applications, f is loss function that produces scalar output, i.e., m 1, for which the gradient with respect to parameters is desired. AD techniques are contrasted with the method of finite differences, which approximates derivatives numerically using a small but non-zero step size, and also distinguished from symbolic differentiation in which mathematical expression is processed using standard rules to produce another mathematical expression, although Elliott (2018) argues that the distinction is simply whether or not it is the compiler that manipulates the symbols.
There are a variety of approaches to AD: source-code transformation (c.g., Bischof et al. (1992);
Hascoet & Pascual (2013); van Merrienboer et al. (2018)), execution tracing (c.g., Walther & Griewank (2009); Maclaurin et al.), manipulation of explicit computational graphs (e.g., Abadi et (2016); Bergstra et al. (2010)), and category-theored exist for many different host languages, although they vary in the extent to which they take advantage of native programming patterns, control flow, and language features. Regardless of whether it is constructed at compile-time time, run-time, or via an embedded domain-speciff language, all AD approaches can be understood manipulating the linearized computational graph (LCG) to collapse out intermediate variables. Figure shows the LCG for a simple example. These computational graphs are always directed acyelic graphs (DAGs) with vertices variables.
transformations (Elliott, 2018). AD implementations
2

---

Published as a conference paper at ICLR 2021
Let the outputs of f be yj, the inputs 0,, and the intermediates 21. AD can be framed as the computation of a partial derivative as a sum over all paths through the LCG DAG (Bauer, 1974):
dy - Jollis Σ ;0:
dz
(1)
Zzk
[-0 (i-1)
where li + indexes paths from vertex to vertex j and (k, E li + -] denotes the set of edges in that path. See Figure ld for an illustration. Although general, this naïve sum over paths does not take advantage of the structure of the problem and so, as in other kinds of graph computations, dynamic programming (DP) provides better approach. DP collapses substructures of the graph until it becomes bipartite and the remaining edges from inputs to outputs represent exactly the entries of the Jacobian matrix. This is referred to the Jacobian accumulation problem (Naumann, 2004) and there are a variety of ways to manipulate the graph, including vertex, edge, and face elimination (Griewank & Naumann, 2002). Forward-mode AD and reverse-mode AD (backpropagation) are special cases of more general dynamic programming strategies to perform this summation; determination of the optimal accumulation schedule is unfortunately NP-complete (Naumann, 2008).
While the above formulation in which each variable is a scalar can represent any computational graph, it can lead structures that are difficult to reason about. Often we prefer to manipulate vectors and matrices, and we can instead let each intermediate represent di dimensional vector. In this case, 82//22, E Tedx xd. represents the intermediate Jacobian of the operation zạ + zt. Note that Equation now expresses the Jacobian of f as a sum over chained matrix products.
3 RANDOMIZING AUTOMATIC DIFFERENTIATION
We introduce techniques that could be used to decrease the resource requirements of AD when used for stochastic optimization. We focus on functions with a scalar output where we are interested in the gradient of the output with respect to some parameters, [o] Reverse-mode AD efficiently calculates [fl.) but requires the full linearized computational graph to either be stored during the forward pass, or be recomputed during the backward pass using intermediate variables recorded during the forward pass. For large computational graphs this could provide large memory burden.
The most common technique for reducing the memory requirements of AD is gradient checkpoint- ing (Griewank& & Walther, 2000; Chen et al., 2016), which saves memory by adding extra forward pass computations. Checkpointing is effective when the number of "layers" " a computation graph is much larger than the memory required at each layer. We take different approach; we instead aim to save memory by increasing gradient variance, without extra forward computation.
Our main idea to consider an unbiased estimator Jo [f] such that EEA = Jolf] which allows us to memory required for reverse mode AD. Our approach is to determine sparse (but random) linearized computational graph during the forward pass such that reverse-mode AD applied on the sparse graph yields an unbiased estimate of the true gradient. Note that the original computational graph is used for the forward pass, and randomization is used to determine a LCG to use for the backward pass in place of the original computation graph. We may then decrease memory costs by storing the sparse LCG directly or storing intermediate variables required to compute the sparse LCG.
In this section we provide general recipes for randomizing AD by sparsifying the LCG. In sections 4 and 5 we apply these recipes to develop specific algorithms for neural networks and linear PDEs which achieve concrete memory savings.
3.1 PATH SAMPLING
Observe that in Bauer's s formula each Jacobian entry is expressed as sum over paths in the LCG. simple strategy is to sample paths uniformly at random from the computation graph, and form Monte Carlo estimate of Equation| 1. Naively this could take multiple passes through the graph. However, multiple paths can be sampled without significant computation overhead by performing a topological sort of the vertices and iterating through vertices, sampling multiple outgoing edges for each. We provide a proof and detailed algorithm in the appendix. Dynamic programming methods such reverse-mode automatic differentiation can then be applied to the sparsified L.C..
3

---

Published as a conference paper at ICLR 2021
3.2 RANDOM MATRIX INJECTION
In computation graphs consisting of vector operations, the vectorized computation graph is a more compact representation. We introduce an alternative view on sampling paths in this case. single path in the vectorized computation graph represents many paths in the underlying scalar computation graph. an example, Figure 2c a vector representation for Figure 2b. For this example,
By By 8C dB DA
(2)
00 aC ²B 0A 00 where A, B, C vectors with entries ag b, C22 8C/0B, BB/aA are 3 x 3 Jacobian matrices for the intermediate operations, 0y/ac is x 3, and 0A/a0 is 1.
We now note that the contribution of the path p = 0 a1 b2 C2 + y to the gradient is,
dy aC 0B
0A
ac ac²a
OB
P2
0A
Pi
00
(3)
where P: = e,ee' (outer product of standard basis vectors). Sampling from {P1, P2, 2, P and right multiplying a Jacobian is equivalent to sampling the paths passing through vertex in the scalar graph.
(a) Independent Paths Graph
In general, if we have transition B + C in a vectorized com- putational graph, where E C Rm can insert a ran- dom matrix = d/L Σ E Ps where each P, sampled uniformly from {PI. P2..... Pa.. With this construction, EP = Id, so
aC
ac AB
(4)
(b) Fully Interlcaved Graph
a
If we have a matrix chain product, we can use the fact that the expectation of product of independent random variables is equal to the product of their expectations, so drawing independent random
(c) Vector graph for (b).
matrices PB, Pe would give
Common compu- tational graph patterns. The graphs may be arbitrarily deep and wide. (a) A small number independent paths. Path sam- pling has constant variance with depth. (b) The number of paths in- creases exponentially with depth; path sampling gives high variance. Independent paths are common when a loss decomposes over data. Fully interlcaved graphs are com- mon with vector operations.
Figure 2:
ay a0 aB
ac
dy ac
OC E P
Ty ac ac 0B
Pc
PB
E
Pcl
(5)
JB
Right multiplication by P may be achieved by sampling the in- termediate Jacobian: one does not need to actually assemble and multiply the two matrices. For clarity tion Sp [0C/8л] = aC/anP. This is sampling (with replacement) & out of the d vertices represented by B, and only considering paths that pass from those vertices.
adopt the nota-
The important properties of P that enable memory savings with an
unbiased approximation are
EP Ia
and
P RR R <
(6)
We could therefore consider other matrices with the same properties. In our additional experiments in the appendix, we also let be a random projection matrix of independent Rademacher random variables, a construction common in compressed sensing and randomized dimensionality reduction.
In vectorized computational graphs, we can imagine two-level sampling scheme. We can both sample paths from the computational graph where each vertex on the path corresponds to vector. We can also sample within each vector path, with sampling performed via matrix injection as above.
In many situations the full intermediate Jacobian for a vector operation is unreasonable to store. Consider the operation -> where B, E RRd The Jacobian d x d. Thankfully many common operations are element-wise, leading to diagonal Jacobian that can be stored as a d-vector. Another common operation is matrix-vecto products. Consider Ab = c, 0c/at = A. Although A has many more entries than or 6, in many applications A is either parameter to be optimized or is easily recomputed. Therefore in our implementations, we do not directly construct and sparsify the Jacobians. We instead sparsify the input vectors the compact version of the Jacobian in a way that has the same effect. Unfortunately, there are some practical operations such as softmax that do not have compactly-reprrsentabbe Jacobian and for which this is not possible.
4

---

Published as a conference paper at ICLR 2021
3.3 VARIANCE
The variance incurred by path sampling and random matrix injection will depend on the structure of the LCG. We present two extremes Figure 2. In Figure 2a, each path is independent and there are small number of paths. If we sample fixed fraction of all paths, variance will be constant in the depth of the graph. In contrast, in Figure 2b, the paths overlap, and the number of paths increases exponentially with depth. Sampling fixed fraction of all paths would require almost all edges in the graph, and sampling fixed fraction vertices at each layer (using random matrix injection, as an example) would lead exponentially increasing variance with depth.
is thus difficult to apply sampling schemes without knowledge of the underlying graph. Indeed, our initial efforts to apply random matrix injection schemes to neural network graphs resulted in variance exponential with depth of the network, which prevented stochastic optimization from converging. We develop tailored sampling strategies for computation graphs corresponding problems of common interest, exploiting properties of these graphs to avoid the exploding variance problem.
CASE STUDY: NEURAL NETWORKS
We consider neural networks composed of fully connected layers, convolution layers, ReLU nonlinear- ities, and pooling layers. We take advantage of the important property that many of the intermediate Jacobians can be compactly stored, and the memory required during reverse-mode is often bottle- necked by a few operations. We draw a vectorized computational graph for typical simple neural network in figure 3. Although the diagram depicts a dataset of size of 3, mini- bithhatch size of size 1, and hidden layers, we assume the dataset size N. Our analysis is valid for any number of hidden layers, and also recurrent networks. We are interested in the gradients 8y/aw, and 0u/aws.
4.1 MINIBATCH SGD AS RANDOMIZED AD
At first look, the diagram has a very similar pattern to that 2a, so that path sampling would a good fit. Indeed, we could sample B < N paths from W, y, and also B paths from W2 y. Each path corresponds to processing a different mini- batch element, and the computations are independent.
In empirical risk minimization, the final loss function is an average of the loss Therefore, the intermediate partials 00/ah2, - for each data point æ will be independent of the other data points. As result, if the same paths are chosen in path sampling for Wi and W2, and if are only interested in the stochastic gradient (and not the full function evaluation), the computation graph only needs to be evaluated for the data points corresponding the sampled paths. This exactly corresponds to mini-batching. The paths are visually depicted in Figure 3b.
data points.
4.2 ALTERNATIVE SGD SCHEMES WITH RANDOMIZED AD
gPAsi
We wish to use our principles derive randomization scheme that can be used on top of mini-batch SGD. We ensure our estimator is unbiased as we randomize by ap- plying random matrix injection independently to various intermediate Jacobians. Consider path corresponding to data point 1. The contribution to the gradient dy/ow, is
Oy 0h2, 001,1 0hi,1 Thg,1 da,,1 Th1,1 aW1
(7)
(a) Neural network computational graph
pras.
Using random matrix injection sample every Jacobian would lead to exploding variance. Instead, we analyze each term to see which are memory bottlenecks.
0y/Oha, is the Jacobian with respect to (typically) the loss. Memory requirements for this Jacobian are independent of depth of the network. The dimension of the classifier is
usually smaller (10 1000) than the other layers (which have dimension 10, 000 more in convolutional net- works). Therefore, the Jacobian at the output layer is not a memory bottleneck.
(b) Computational Graph with Mini-batching
Figure 3: NN computation graphs.
5

---

Published as a conference paper at ICLR 2021
Figure Convnet activation sampling one mini- batch element. X is the image, H is the pre-activation, and A is the activation. A is the output of ReLU, so can store the Jacobian 0Ai/am, with bit per entry. For X and H sample spatial elements and compute the Jacobians @Hts/ww and oHaa/ows with the sparse tensors.
0/2./001, is the Jacobian of the hidden layer with respect to the previous layer activation. This can be constructed from 22, which must be stored in memory, with memory cost independent of mini-batch size. In convnets, due weight sharing, the effective dimensionality is much smaller than H1 x H2- In recurrent networks, it is shared across timesteps. Therefore, these are not memory bottleneck.
001,/001.1 contains the Jacobian of the ReLU activation function. This using 1-bit per entry, as the gradient can only be 1 or O. Note that this is true for ReLU activations in particular, and not true for general activation functions, although ReLU is widely used in deep learning. For ReLU activations, these partials are not a memory bottleneck.
be compactly stored
a/a,, contains the memory bottleneck for typical ReLU neural networks. This is the Jacobian of the hidden layer output with respect to W1, which, in a multi-layer perceptron, is equal to T1. For B data points, this B D dimensional matrix.
Accordingly, we
choose to
sample 8h1.//w,, replacing the matrix chain
with By 812,1 001.1 SPw
aha.a aw,
For an arbitrarily deep NN, this can be generalized:
1h2.1 d 1b.,1
dy Tha 00s-1,1
001,1
0h1,1
dy Iha,1 00s-1.1
0a2,1
@h2,1 aw2
Spw;
Spws 0h2.1
Ōha,1 0a.-1,1 0ha-1.1
0h1,1
OW,
0ha,1 das-1,1 Tha-1,1
This can be interpreted as sampling activations on the backward pass. This is our proposed alternative SGD scheme for neural networks: along with sampling data points, can also sample activations, while maintaining an unbiased approximation to the gradient. This does not lead to exploding variance, as along any path from a given neural network parameter the loss, the sampling operation is only applied to single Jacobian. Sampling for convolutional networks is visualized in Figure 4.
4.3 NEURAL NETWORK EXPERIMENTS
We evaluate our proposed RAD method on two feedforward architectures: a small fully connected network trained on MNIST, and a small convolutional network trained on CIFAR-10. We also evaluate method calculations for the associated memory savings from our method are available in the appendix. In Figure we include empirical analysis of gradient noise caused by RAD mini-batching.
RNN trained on Sequential- - The exact architectures and the
We are mainly interested in the following question:
For a fixed memory budget and fixed number of gradient descent iterations, how quickly does our proposed method optimize the training loss compared to standard SGD with smaller mini-batch?
Reducing the mini-batch size will also reduce computational costs, while RAD will only reduce memory costs. Theoretically method could reduce computational costs slightly, but this is focus. We only consider the memory/gradient variance tradeoff while avoiding adding significant overhead on top of vanilla reverse-mode (as is the case for checkpointing).
Results are shown in Figure 6. Our feedforward network full-memory baseline is trained with mini- batch size of 150. For RAD keep mini-batch size of 150, and try 2 different configurations. For "same sample", we sample with replacement a 0.1 fraction of activations, and the same activations are sampled for each mini-batch element. For "different sample", we sample a 0.1 fraction of activations, independently for each mini-batch element. Our "reduced batch" experiment is trained without RAD with mini-batch size of 20 for CIFAR-11 10 and for MNIST. This achieves similar memory budget as RAD with mini-batch size 150. Details of this calculation and of hyperparameters are in the appendix.

---

Published as a conference paper at ICLR 2021
Tartry Low Itruiin StalCorsfer CIFMA-33
Transn va hastto for Sma IFCNe MH6T
Inaesg Los en hurztorn for EdM on Sequanna WMMS
Reds ond batch
12
Same Sample
Dillweet Sampll
127
07
12-4
12.4
3.00
13.00
(2:00
3000
200
7000
1080
1900
100000
(a) CIFAR- 10 Curves
(b) MNIST Curves
(c) Srauuential-NNISS Curves
Figure 6: Training curves for neural networks. The legend in (c) applies to all plots. For the convolutional and fully connected neural networks, the loss decreases faster using activation sampling, compared to reducing the mini-batch size further to match the memory usage. For the fully connected NN on MNIST, it is important to sample different activations for cach mini-batch element, since otherwise only part of the weight vector will get updated with each iteration. For the convolutional NN on CIFAR-I0. this is not an issue due to weight tying. As expected, the full memory baseline converges quicker than the low memory versions. For the RNN on siquential-MNISS sampling different activations at each time-step matches the performance obtained by reducing the mini-bated size.
Fraction of activations
Baseline (1.0)
0.8
0.5
0.3
0.1
0.05
ConvNet Mem
23.08
19.19
12.37
7.82
3.28
2.14
Fully Connected Mem
2.69
2.51
2.21
2.00
1.80
1.75
RNN Mem
47.93
39.98
25.85
16.43
7.01
4.66
Table 1: Peak memory (including weights) in MBytes used during training with 150 mini-batch elements. Reported values are hand calculated and represent the expected memory usage of RAD under an efficient implementation.
For the feedforward networks we tune the learning rate and lg regularization parameter separately for each gradient estimator on a randomly held out validation set. We train bootstrapped
MSE el Stecheste Gradents from Full Gorlee (Rectaged ever 1000 becknl
Activtial Fussios (4 N2)
Back ID) Higt Aciivaae F'adios 8A (2 M31
kuaik 2t Beat we 1aar U9.
with the best performing hyperparameters versions of the full training set to measure variability in training. Details are in the appendix, including plots for train/test accuracy/loss, and wider range of fraction of activations sampled. All feedforward models are trained with Adam.
In the RNN case, we also run baseline, "same sample", "different sample" and "reduced batch" experiments. The "reduced batch" experiment used a mini-batch size of 21, while the others used a mini-batch size of 150. The learn- ing rate was fixed at 10 for all gradient estimators, found via a coarse grid search for the largest learning rate for which optimization did not diverge. Although we did not tune the learning rate separately for each estimator, still expect that with a fixed learning rate, the lower variance estimators should perform better. When sampling, we sam- ple different activations at each time-step. All recurrent models are trained with SGD without momentum.
Figure 5: We visualize the gradient noise for each stochastic gradient method by com- puting the full gradient (over all mini-batches the training set) and computing the mean squared error deviation for the gradient es- timated by each method for each layer in the convolutional net. RAD has significantly less variance vs memory than reducing mini- batch size. Furthermore, combining RAD with an increased bani-batch size achieves similar variance to baseline 1 150 mini-hatch elements while saving memory.
5 CASE STUDY: REACTION-
PDE- -OE-CONSTRAINED OPTIMIZATION
Our second application is motivated by the observation that many scientific computing problems involve a repeated or iterative computation resulting in a layered computational graph. We may apply RAD to get stochastic estimate of the gradient by subsampling paths through the compu- tational graph. For certain problems, we can leverage problem structure to develop a low- now-memory consider
stochastic gradient estimator without exploding variance. To illustrate this possibility the optimization of a linear reaction-diffe diffusioo PDE on a square domain with Dirichlet boundary conditions, representing the production and diffusion of neutrons in fission reactor (McClarren,
7

---

Published as a conference paper at ICLR 2021
2018). Simulating this process involves solving for potential 6(z, 3.2. varying in two spatial coordinates and in time. The solution obeys the partial differential equation:
06(x, v, DV (x,,1, Cl ys (,0)),,,,,
dL
We discretize the PDE in time and space and solve rule Φr+1 = M Φe + ALC; @ Φ., where M summarizes the discretization of the PDE in space. The exact form is available in the appendix. The initial condition Φo - - sin (#r)) sin (zy), with & = 0 on the boundary of the domain. The loss function is the time-averaged squared error between Φ and a time-dependent target, I, = -/1// E, (0) &f 12- The target is 6t target Φu + s/4sin (тt) sin (2T..) sin (Ty). The source C' is given by seven-term Fourier series and with coefficients given by R., where @ is the control parameter to be optimized. Full simulation details are provided in the appendíx.
spatial grid using explicit update
The gradient is aL Ei Beb,
ØL Ei (II)=) 0Φ,+1) Soh 3Ø; 8C,
SC.-1
As the reaction-diffusion PDE is
80
linear and explicit, 0 TRN2XN known and independent of Φb. We avoid storing C at each timestep by recomputing C from 0 and t. This permits low-memory stochastic gradient estimate without exploding variance by sampling from 0L/a0, E N and the diagonal matrix 0p./0C.... dh with the unbiased estimator replacing 00
E
Sper
aL DΦ.
Σ
III
pept dΦ;
Sp
Op:
8C
(8)
[C;-1]
00
j=i
This estimator can reduce memory by as much as 99% without harming optimization; see Figure 7b.
6 RELATED WORK
Approximating gradients and matrix opera- tions Much thought has been given to the approximation of general gradients and Jacobians. We draw inspiration from this literature, although our main objective is designing an unbiased gradient estimator, rather than an approximation with bounded accuracy. Abdel- -hhalik al. (2008) accelerate Jacobian accumulation via random projections, in similar manner to randomized methods for SVD and matrix multiplication. Choromanski & Sindhwani (2017) recover Jacobians in cases where AD is not available by performing a small number of function evaluations with random input perturbations and leveraging known structure of the Jacobian (such as sparsity and symmetry) via compressed sensing.
(a) Visualization of sampling
Training Loss Ittrations for Reaction Diffusion Equation
Menory
0.001
0.002
Other work aims to accelerate neural network training by approximating operations from the forward and/or backward pass. Sun et al. (2017) and Wei et al. (2017) backpropagate sparse gradients, keeping only the top k elements of the adjoint vector. Adelman & Silberstein (2018) approximate matrix multiplications and convolu- tions in the forward pass of neural nets nets using column- row sampling scheme similar to our subsampling scheme. Their method also reduces the computational cost of the backwards pass but changes the objective landscape.
0. 0.006
0.01
0.1
10
1.0
200
100
020
500
(b) RAD curves
Figure 7: Reaction-diffusion PDE expt. (b)
RAD saves up 99% memory without significant slowdown convergence.
Related are invertible and reversible transformations, which remove the need to save intermediate variables on the forward pass, these can be recomputed on the backward pass. Maclaurin et al. (2015) use this idea for hyperparameter optimization, reversing the dynamics of SGD with momentum to avoid the expense of saving model parameters each training iteration. Gomez al. (2017) introduce reversible ResNet (He et al., 2016) to avoid storing activations. Chen et al. (2018) introduce Neural ODEs, which also have constant memory cost function of depth.
8

---

Published as a conference paper at ICLR 2021
Limited-memory learning and optimization Memory is major bottleneck for reverse- AD, and much work aims to reduce its footprint. Gradient checkpointing is perhaps the most well known, and has been used for both reverse-mode AD (Griewank & Walther, 2000) with general layerwise computation graphs, and for neural networks (Chen et al., 2016). In gradient checkpointing, some subset of intermediate variables are saved during function evaluation, and these are used to re-compute downstream variables when required. Gradient checkpointing achieves sublinear memory cost with the number of layers in the computation graph, at the cost of constant-factor increase in runtime.
Stochastic Computation Graphs Our work is connected to the literature on stochastic estimation of gradients of expected values. of the expected outcome of stochastic computation graph. The distinguishing feature of this literature (vs. the proposed RAD approach) is that it uses stochastic estimators of an objective value to derive stochastic gradient estimator, i.e., the forward pass randomized. Methods such as REINFORCE (Williams, 1992) optimize an expected return while avoiding enumerating the intractably large space of possible outcomes by providing an unbiased stochastic gradient estimator, i..., by trading computation for variance. This is also true of mini-batch SGD, and methods for training generative models such as contrastive divergence (Hinton, 2002), and stochastic optimization of evidence lower bounds (Kingma & Welling. 2013). Recent approaches have taken intractable deterministic computation graphs with special structure, i.e. involving loops or the limits of a series of terms, and developed tractable, unbiased, randomized telescoping ceries-based estimators for the graph's 's output, which naturally permit tractable unbiased gradient estimation (Tallec & Ollivier, 2017; Beatson Adams, 2019; Chen et al., 2019; Luo et al., 2020).
7 CONCLUSION
We present framework for randomized automatic differentiation. Using this framework, construct reduced- reduced-memory unbiased estimators for optimization of neural networks and linear PDE. Future work could develop RAD formulas for new computation graphs, c.g., using randomized rounding to handle arbitrary activation functions and nonlinear transformations, integrating RAD with the adjoint method for PDEs, or exploiting problem-specific sparsity in the Jacobians of physical simulators. The randomized view on AD we introduce may be useful beyond memory savings: we hope it could be useful tool developing reduced-come computtaiooittiio stochastic gradient methods or achieving tractable optimization of intractable computation graphs.
ACKNOWLEDGEMENTS
The authors would like to thank Haochen Li for early work on this project. We would also like to thank Greg Gundersen, Ari Seff, Daniel Greenidge, and Alan Chung for helpful comments the manuscript. This work is partially supported by NSF IIS-2007278.
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: A system for large-scale machine learning. In 12th USENIX Symposium on Operating Systems Design and Implementation (OSDI 16), pp. 265-283, 2016.
Hany KhallkK Paul D Hovland, Andrew Lyons, Tracy E Stover, and Jean Utke. low rank approach to automatic differentiation. In Advances in Automatic Differentiation, pp. 55-65. Springer, 2008.
Menachem Adelman and Mark Silberstein. Faster neural network training with approximate tensor operations. arXiv preprint 18Xiv:8805.0800 2018.
Friedrich L Bauer. Computational graphs and rounding error. SIAM Journal on Numerical Analysis, 11(1):87-96, 1974.
Atilim Gunes Baydin, Barak A Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: a survey. Journal of Machine Learning Research, 18(153), 2018.

---

Published as a conference paper at ICLR 2021
Alex Beatson and Ryan P Adams. Efficient optimization of loops and limits with randomized telescoping sums. In International Conference on Machine Learning, 2019.
James Bergstra, Olivier Breuleux, Frédéric Bastien, Pascal Lamblin, Razvan Pascanu, Guillaume Desjardins, Joseph Turian, David Warde-Farley, and Yoshua Bengio. Theano: CPU and GPU math expression compiler. In Proceedings of the Python for Scientific Computing Conference (SciPy), volume 4, 2010.
Christian Bischof, Alan Carle, George Corliss, Andreas Griewank, and Paul Hovland. ADIFOR- generating derivative codes from Fortran programs. Scientific Progrumming, 1(1):11-29) 1992.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 2018.
Tian Qi Chen, Jens Behrmann, David Duvenaud, and Jörn-Henris Jacobsen. Residual flows for invertible generative modeling. In Advances in Neural Information Processing Systems, pp. 9913-9923, 2019.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174 2016.
Krzysztof M Choromanski and Vikas Sindhwani. On blackbox backpropagation and Jacobian sensing. In Advances Neural Information Processing Systems, pp. 6521-6529, 2017.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(ull):121-2155) 2011.
Conal Elliott. The simple essence automatic differentiation. Proceedings of the ACM on Program-
ming Languages, 2(ICFP):70, 2018
Aidan N Gomez, Mengye Ren, Raquel Urtasun, and Roger B Grosse. The reversible residual network: Backpropagation without storing activations. In Advances in Neural Information Processing Systems, pp. 2214 2224,, 2017.
A Griewank and U Naumann. Accumulating Jacobians by vertex, edge, or face elimination. cari 2002. In Proceedings of the 6th African Conference on Research in Computer Science, INRIA, France, pp. 315-373333333 383, 2002.
Andreas Griewank and Andrea Walther. Algorithm 799: revolve: an implementation of check- pointing for the reverse or adjoint mode of computational differentiation. ACM Transactions on Mathematical Software (TOMS), 26(1):19-45, 2000.
Andreas Griewank and Andrea Walther. Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, volume 105. SIAM, 2008
Laurent Hascoet and Valérie Pascual. The Tapenade automatic differentiation tool: Principles, model, and specification. ACM Transactions on Mathematical Software (TOMS), 39(3):20, 2013.
Kaiming Hc. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 770-778, 2016.
Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 1771)1177 2002.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1442, 1997.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference Learning Representations, 2015.
10

---

Published as a conference paper at ICLR 2021
Diederik P Kingma and Max Welling. arXiv:13/2.6114, 2013.
Auto-encoding variational Bayes. arXiv preprint
Quoc V Le, Navdeep Jaitly, and Geoffrey E Hinton. A simple way to initialize recurrent networks of rectified linear units. arXiv preprint arXiv: 1eXiv::504.00990 00941, 2015.
Yucen Luo, Alex Beatson, Mohammad Norouzi, Jun Zhu, David Duvenaud, Ryan P Adams, and Ricky TQ Chen. Sumo: Unbiased estimation of log marginal probability for latent variable models. In International Conference on Learning Representations, 2020.
Dougal Maclaurin, David Duvenaud, and Ryan P Adams. Autograd: Effortless gradients in numpy. URL peetps://githe cww/HIPS HwW.HIISSSSaaaaaaaa autograd.
Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient-based hyperparameter optimization through reversible learning. In International Conference on Machine Learning, pp. 2112-2122 2015.
Ryan McClarren. Computational Nuclear Engineering and Radiological Science Using Python: Chapter 18 One-Group Diffusion Equation. Academic Press, 2018
Uwe Naumann. Optimal accumulation of Jacobian matrices by elimination methods on the dual computational graph. Mathematical Programming. 99(3):399-421) 2004.
Uwe Naumann. Optimal Jacobian accumulation is NP-complete. Mathematical Programming, 112 (2):427-441, 2008
Samuel I. Smith and Quoc Le. Bayesian perspective on generalization and stochastic gradient descent. In International Conference on Learning Representations, 2018.
Xu Sun, Xuancheng Ren, Shuming Ma, and Houfeng Wang. meprop: Sparsified back propagation for accelerated deep learning with reduced overfitting. In Proceedings of the 34th International Conference Machine Learning, pp. 2299-3308. JMLR. org, 2017
Corentin Tallec and Yann Ollivier. Unbiasing truncated backpropagation through time. arXiv preprint arXiv:1705.08200 2017.
Bart van Merrienboer, Dan Moldovan, and Alexander Wiltschko. Tangent: Automatic differentiation using source-code transformation for dynamically typed array programming. Advances in Neural Information Processing Systems, pp. 6256-6265, 2018.
Andrea Walther and Andreas Griewank. Getting started with ADOL-C. Combinatorial Scientific Computing, (09061)::8) 188-2202, 2009.
Bingzhen Wei, Xu Sun, Xuancheng Ren, and Jingjing Xu. Minimal effort back propagation for convolutional neural networks. arXiv preprint arXiv:1709 05804, 2017.
Ronald J Williams. Simple statistical follong algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):229-22) 256, 1992.

---

Published as a conference paper at ICLR 2021
APPENDIX A: NEURAL NETWORK EXPERIMENTS
RANDOM PROJECTIONS FOR RAD
As mentioned in Section 3.2 (around Equation 5) of the main paper, we could also use different matrices P that have the properties
EP la
and
P P= RR BR tecck,k< <dd.
In the appendix we report experiments of letting be a matrix of iid Rademacher random variables, scaled by vk. P RR defined in this way satisfies the properties above. Note that this would lead additional computation: The Jacobian or imput vector would have be fully computed, and then multiplied by R and stored. In the backward pass, it would have to be multiplied by R We report results as the "project" experiment the full training/test curves the following sections. We see that it performs competitively with reducing the mini-batch size.
ARCHITECTURES USED
use three different neural network architectures for our experiments: one fully connected feedfor- ward, convolutional feedforward, and one recurrent.
Our fully-connected architecture consists of:
1. Input: 784-dimensional flattened MNIST Image
2. Linear layer with 300 neurons (+ bias) (+ ReLU)
3. Linear layer with 300 neurons (+ bias) (+ ReLU)
4. Linear layer with 300 neurons (+ bias) (+ ReLU)
5. Linear layer with 10 neurons (+ bias) (+ softmax)
Our convolutional architecture consists of:
1. Input: 3 32 32-dimensional CIFAR-10 Image
2. 5 x 5 convolutional layer with 16 feature maps (+ 2 zero-padding) (+ bias) (+ ReLU)
3. 5 x convolutional layer with 32 feature maps (+ 2 zero-padding) (+ bias) (+ ReLU)
4. average pool 2-d
5. 5 5 convolutional layer with 32 feature maps (+ 2 zero-padding) (+ bias) (+ ReLU)
6. 5 x 5 convolutional layer with 32 feature maps (+ 2 zero-padding) (+ bias) (+ ReLU)
7. 2 x average pool 2-d (+ flatten)
8. Linear layer with 10 neurons bias) (+ softmax)
Our recurrent architecture was taken from Le et al. (2015) and consists of:
1. Input: sequence of length 784 of 1-dimensional pixels values of flattened MNIST image.
2. A single RNN cell of the form
he = DeLUWW + ban W Wanhe-1 + bua)
where the hidden state (/t) dimension is 100 and x, is the 1-dimensional input.
3. An output linear layer with 10 neurons (+ bias) (+ softmax) that has as input the last hidden state.
CALCULATION OF MEMORY SAVED FROM RAD
For the baseline models, we assume inputs to the linear layers and convolutional layers are stored in 32-bits per dimensions. The ReLU derivatives then recalculated the backward pass.
For the RAD models, we assume inputs are sampled or projected to 0.1 of their size (rounded up) and stored in 32-bits per dimension. Since ReL.U derivatives not exactly be calculated now,
12

---

Published as a conference paper at ICLR 2021
assume they take 1-bit per dimension (non-reduced dimension) store. The input to the softmax layer is not sampled or projected.
In both cases, the average pool and bias gradients does not require saving since the gradient is constant.
For MNIST fully connected, this gives (per mini- element memory):
Baseline: (784 + 300 + 300 + 300 + 10) 32 bits = 6.776 kBytes
RAD 0.1: (79 30 + 30 30 + 10) 32 bits + (300 + 300 + 300) 1 bits 828.5 bytes
which leads to approximately 8x savings per mini-batch element.
For CIFAR-10 convolutional, this gives (per mini-batch element memory):
Baseline: (3-32-32-4 16 32-32-2--222 4-3 16 16 4-- 16 16 32. 8- 10) 322 bits 151.59 kBytes RAD 0.1: (308 1639 + 820 + 820 205 + 10)-33 bits + (16384 + 8192 + 8192 + 2048) 1 bits 19.56 kBytes
which leads to approximately 7.5x savings per mini batch element.
For Sequential-MNIIT RNN, this gives (per mini-batch element memory):
Baseline: (784 (1 + 100) + 100 + 10) -32 bits = 317.176 kBytes
RAD 0.1: (784 (1 + 10) 10+ + 10) -3 bits + (784- 100)- bits 44.376 kBytes
which leads to approximately 7.15x savings per mini-batch element.
FEEDFORWARD NETWORK TRAINING DETAILS
We trained the CIFAR-10 models for 100, 000 gradient descent iterations with a fixed mini-batch size, sampled with replacement from the training set. We lower the learning rate by 0.6 every 10. 000 iterations. We train with the Adam optimizer. We center the images but do not use data augmentation. The MNIST models were trained similarly, but for 20, 000 iterations, with the learning rate lowered by 0.6 every 2, 000 iterations. We fixed these hyperparameters in the beginning and did not modify them.
We tune the initial learning rate and & weight decay parameters for each experiment reported in the main text for the feedforward networks. For each experiment (project, same sample, different sample, baseline, reduced batch), for both architectures, we generate 20 (weight decay, learning rate) pairs, where each weight decay is from the loguniform distribution over 0.0000001 0.001 and learning rate from loguniform distribution over 0.00001 - 0.01.
We then randomly hold out validation dataset of size 5000 from the CIFAR-10 and MNIST training sets and train each pair on the reduced training dataset and evaluate on the validation set. For each experiment, we select the hyperparameters that give the highest test accuracy.
For each experiment, we train each experiment with the best hyperparameters 5 times on separate bootstrapped resamplings of the full training dataset (50, 000 for CIFAR-10 and 60, 000 for MNIST), and evaluate on the test dataset (10, 000 for both). This is to make sure the differences we observe across experiments are not due to variability in training. the main text we show 3 randomly selected training curves for each experiment. Below we show all 5.
All experiments were run on a single NVIDIA K80 or V100 GPU. Training times were reported on V100.
RNN TRAINING DETAILS
All RNN experiments were trained for 200,000 iterations (mini-batel updates) with a fixed mini-bate size, sampled with replacement from the training set. We used the full MNIST training set of 60,000 images whereby the images were centered. Three repetitions of the same experiment were performed with different seeds. Hyperparameter tuning was not performed due to time constraints.
13

---

Published as a conference paper at ICLR 2021
The hidden-to-hidden matrix (Wha) is initialised with the identity matrix, the input-to-hice so-hidden matrix (Wa)) and hidden-to-output (last hidden layer to softmax input) are initialised with a random matrix where each element is drawn independently from (O, 0.001) distribution and the biases (6i.. byr)) are initialised with zero.
The model evaluated the test set of 10,000 images every 400 iterations and the entire training set every 4000 iterations.
For the "sample", "different sample", "project" and "different project" experiments different activa- tions/random matrices were sampled at every time-step of the unrolled RNN.
All experiments were
single NVIDIA K80 V100 GPU.
The average running times for each experiment are given in Table 2. Note that we did not optimise implementation for speed and these running times be reduced significantly.
Table 2: Average running times for RNN experiments Sequential-MNIST.
Experiment
Running Time (hrs)
GPU
Baseline
34.0
V100
Small batch
16.0
V100
Sample
96.0
K80
Different Sample
110.0
K80
Project
82.0
K80
Different Project 89.0
K80
IMPLEMENTATION
The code is provided on GitHubl, Note that we did not optimize the code for computational efficiency; only implemented our method to demonstrate the effect it has the number of gradient steps to train. Similarly, we did not implement all of the memory optimizations that we account for in our memory calculations: in particular in our implementation
did not take advantage of storing ReLU derivatives with 1-bit or the fact that average pooling has a constant derivative. Although these would have be implemented in a practical use-case, they are not necessary in this proof of concept.
httt::///iittbl..oo////ottttttia ttt cDi fferent i
zedkutooatt
14

---

Published as a conference paper at ICLR 2021
FULL TRAINING/TEST CURVES FOR MNIST AND CIFAR-10
Trairing los Ineatioon for SmalCencllt CIFAR.10 10
MMST
109
10
10-2
10"1
10-2
Projct
Freact
Rabood back
Baokie
Baudies
Sure Saople
Sie Sample
10'
Offoren Sample
Diferst Sampl
2000
4003
0000
83030
(00033
1550
2000
Traiiing Nocuracy boratons SmalCConNet
Taining Accuracy vs Ieeatiore for SmallF on MNIST
100.00
66
20:00
99.95
Pegjer
95.75
Projec
Reducdd batch
Rediced batch
Blue line
99.00
Save Scrple
Sampl
98:25
2000
40003
3000
1790)
Test Loss ltaations for Sma Conwheet CIFAR11 10
Text Los Iteration for SmalF MNIST
3.8.
8
Poçect
Rakood batch
3.5 10
Bandda
Sine
2x
10
Sane Surple
BiFlowe Sample
70000
1060
00001
4003C
1930
5000
(0000)
3000
SmallConoNee CIFAR.10
Accuracy SMalFCNee
Propet
66-
Heducad batch
Roolive
Trairing tine vs lteations for SmallCeneNet on CIFNR 10
Traning time vs Iterztoons for Sma IFCNet MNIST
12.5
20
15.0
15
12.
10
5.0
2.5
25000
00000
50COC
100000
5000
10000
35300
20000
Figure 8: Full train/test curves and runtime per iteration. Note that training time is heavily dependent on implementation, which we did not optimize. In terms FLOPs, "project" should significantly higher than all the others and "reduced batch" should be significantly smaller. The "baseline". "same sample", and "different sample" should theoretically have approximately the same number of FLOPs.
15

---

Published as a conference paper at ICLR 2021
FULL TRAINING/TEST CURVES FOR RNN ON SEQUENTIAL-MNEE
Traiiig
Sequ
2x
Base ine
Same Sample
Sample
Project
Project
00000
10000
390000
2000
Taining Arcurary heestione Sequ Sequemial-AMN
70
50
Bosekine
Same Sample
Sample
Project
30
Project
50000
190000
10000
Test Loss Ieeatioos IRNN Sequ entiaal MINIST
Baseire
Reduced batch
Same Sample
Different Sample
Project
Diffcren Project
103000
150000
Teat. Accurscy va lterations for IRNN on Sequertial MNITT
50
Baselre
Reduced batch
Same Sample
30
Different Sample
Project
Different Project
Toiniing time vs berstons for IRNN mn Seguentil-MMS MNIST
750
Easeine
Redioeel
530
Same Simple
Different Sample
Project
Different Project
230
Figure 9: Full train/test curves and runtime per 400 iterations. We also include results for random projections with shared and different random matrices for each mini-batch element.
16

---

Published as a conference paper at ICLR 2021
S SWEEP OF FRACTIONS SAMPLED FOR MNIST AND CIFAR-10
Traning Loss as hentions for SmallComhiet on CFFR41
Training
1f
10
10
2.05
11-
50000
20000
5009C
50000
100000
190.0
Training Accuracy Iterations for SmallCom/Ne CIFAR:10
100.00
Training Accuracy lterations for SmalF on MNS
363
0.25
35
50.00
55
00.00
51
06.50
52
Uand
92
20000
00300
50000
00.00
10000
06.00
5000
7000
10000
12500
19030
17500
20030
Iaertioon SmallConolis
1x 109
00
3x
Buscire
0.55
Test Accuracy vs herccios for ConeNet or CIFAR
Test Accuracy Ieeratinss for SmellFCNet on MNIST
100
2.06
33
3.5
56
20003
40030
COCC
8000
10000
50
3002
30003
19000
20030
Trainng
Triining SmalFCNis
11
10
22
Figure 10: Full train/test curves and runtime per iteration for various fractions for the "different sample" experiment. Note that the reason 0.8 does not quite converge to bascline in the training curve is because we sample with replacement. This is an implementation detail; our method could be modified to sample without replacement, and at fraction 1.0 would be equivalent baseline. The weight decay and initial learning rate for the RAD experiments above are all the same as the ones tuned for 0.1 fraction "different sample" experiment. The baseline experiments tuned for baseline.
17

---

Published as a conference paper at ICLR 2021
APPENDIX B: REACTION-DIFFUSION PDE
The reaction-diffusion equation a linear parabolic partial differential equation. In fission reactor analysis, is called the one-group diffusion equation one-speed diffusion equation, shown below.
0Φ Ōt - DV + Co+5
Here & represents the neutron flux, D is a diffusion coefficient, and C. and S are source terms removal of neutron flux. In this paper, solve the one-speed
related the local production diffusion equation in two spatial dimensions on the unit square with the condition that á = 0 on the boundary. We assume that D is constant equal to 1/4, C(z,, M, i, 0) is function of control parameters described below, and S is zero. We discretize Φ on regular grid in space and time, which motivates the notation & > Φt. The grid spacing is Ar 1/32 and the timestep is At 1/4096. We simulate from t = 0 to t = 10. We use the explicit forward- homeard-time, centered-space (FTCS) method to timestep Φb The timestep chosen to satisfy the stability criterion, Dát/(an)² S x- In matrix notation, the FTCS update rule can be written Φ++1 MΦ. + AIC, @ Φr, in index notation it can be written as follows:
0 + (Ar)²
DAt
His + -
si5t1 0,5-1 -50-1-40))) + AICH ON
The term Co in the ope-speed diffusion equation relates to the local production or removal neutrons due nuclear interactions. In real fission reactor, C' is a complicated function of the material properties of the reactor and the heights of the control rods. We make the simplifying assumption that C' can be described by 7-term Fourier series and written below. Physically, this equivalent to the assumption that the material properties of the reactor are constant in space and time, and the heights of the control rods are sinusoidally varied in æ and &o initialized so that the reactor begins in stable state, the other parameters are initialized from uniform distribution.
C(a,y, y, 0, + 01 sin(xt) + 02 cos(mt) sin (it( + 04 (inT) cos(m() + O3 cos(2 sin(T() + 0g cos(2 cos(TT)
The details of the stochastic gradient estimate and optimization are described the main text. The Adam optimizer is used. Each experiment of 800 optimization iterations runs in about 4 hours GPU.
18

---

Published as a conference paper at ICLR 2021
APPENDIX C: PATH SAMPLING ALGORITHM AND ANALYSIS
Here we present an algorithm for path sampling and provide a proof that leads to an unbiased estimate for the gradient. The main idea is sample edges from the of outgoing edges for each vertex in topological order, and scale appropriately. Vertices that have no incoming edges sampled can be skipped.
Algorithm 1 RMAD with path sampling
1: Inputs:
2:
G (V, E)- Computational Graph. d, denotes outdegree, v.suce successor set of vertex v.
3: y Output vertex
4:
è = (01. .u,...,.... 6m)) CV- Input vertices
5: :: > 0 Number of samples per vertex
6: Initialization:
7:
Q(e) 0, Ve
8: for topological order; synchronous with forward computation do
9: if' No incoming edge of has been sampled then
10:
Continue for times do
12: Sample from [d.] uniformly.
Q(v, v.succlil) + Qrv, v.suceli]) + @v.suceli]
80
14:
Run backpropagation from y to θ using Q as intermediate partials.
15:
Output: Vey
The main storage savings from Algorithm will come from Line 9, where we only consider a vertex if it has an incoming edge that has been sampled. In computational graphs with a large number of independent paths, this will significantly reduce memory required, whether we record intermediate variables and recompute the LCG, or store entries of the LCG directly.
To see that path sampling gives an unbiased estimate, we use induction on the vertices in reverse topological order. For every vertex 2, we denote az
and our approximation for For our
base case, we let dy 1, so Eý = y. For all other vertices we define
Ov
8 d2
Σ D= Oz (2,0)E
(9)
where d, is the out-degree of U/ is sampled uniformly from the set of successors of 2, and Iu=e, is an indicator random variable denoting 0 = V. We then have
Ou ED 02
Ov
E: Σ
d.E
Σ
Ed
(10)
(2
(x,0)EE
assuming that the randomness over the sampling of outgoing edges is independent of é, which must be true because our induction reverse topological order. Since by induction we assumed Ei we have
E8
Σ
(11)
Oz
(
which completes the proof.
19


Reviews:
[{"review_id": "wMCL09Fa7FE", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In the context of deep learning, back-propagation is stochastic in the sample level to attain bette efficiency than full-dataset gradient descent. The authors asked that, can we further randomize the gradient compute within each single minibatch / sample with the goal to achieve strong model accuracy. In modern deep learning, training memory consumption is high due to activation caching. Thus this randomized approach can help attain strong model accuracy under memory constraints.\n\nThe authors proposed a general framework for randomized auto differentiation to achieve unbiased gradient estimators. This general framework allows randomization at different granularity such as layer level and individual neuron level. It also includes conventional minibatch gradient estimators as a special case at the sample/minibatch level for randomization. The memory saving here is achieved by trading off gradient variance for activation memory saving. \n\nEmpirically, the authors show that for 1) convolution nets on MNIST and CIFAR and 2) RNN on sequential-MNIST, under the same memory budget, neuron-level randomized gradient estimator can achieve higher model accuracy than conventional SGD with smaller minibatch size. \n\nStrong point: This paper is well written with novel thoughts on the fundamental aspects of auto-diff when applied to deep learning (and also to PDE as demoed in section 5.). It can also provide new options for practitioners to train models with high accuracy under memory constraints. Thus I recommend to accept this paper.\n\nI have the following comments / questions on the technical aspects. I only raise these questions up for improvement on the paper; they are not concerns on the quality of the current version. Nonetheless,  I am happy to raise the score if the authors can demonstrate results on these aspects.\n\n1. The authors mentioned about leveraging model specific structures to control/reduce gradient variance for fine-grained randomization (such as at the neuron level). Specifically, they considered using the randomized activation sparsification only for activation memory-consuming layers. I was wondering if other model structures can also help here. E.g. can we sample at the full-channel level for convolutional layers or column / row level for linear layers. Would this has a significant impact of the attained model accuracy?\n\n2. The main focus on practical implications in this paper is about high accuracy training under memory constraints. However, I was wondering how the authors think about the implication on compute, especially related to the sparse training trend. E.g. can we also make the forward compute itself also sparse and still minimally influence model accuracy?\n\n\nNITs: \n\n1. First letter capitalization for figure 3 at the beginning of section 4.\n\n2. Would it be possible to provide some preliminary in appendix on the solution to the section 5 PDE example.  \n\n3. In section 3.3, the discussion on two extremal cases needs a more precision in text. Assuming the sampling on each connection (i.e. segment on paths) are independent, I think both case will have variance exponential in terms of depth. Currently it reads like only the fully connected case have exponentially large variance.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "ag4Yp4_sjfV", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes a novel approach to reduce memory in backpropagation by sampling the paths in DAG. The paper developed and proved the proposed method in a general framework (which is nice!). Still, I feel that the explanation for applying this method to neural networks is somewhat lacking. I summarize the missing parts in the following.\n\n(1) For fully-connected layers (matrix multiplication), the explanations in Section 3.2 and Section 4.2 are not consistent --- Section 3.2 suggests sampling from the paths (connections), while Section 4.2 suggests sampling from the vertices (activations). I believe in practice sampling from activations is used (correct me if I am mistaken), which indeed implies that random matrices {P_1, ..., P_n} are not independent (I guess it forms a Markov chain?). Therefore, the analysis in Section 3.2 is not applicable. I think an analysis using dependent sampling (ancestral sampling) is needed in this case.\n\n(2) The explanation for convolutional layers is not sufficient. The paper mentions that the compact structure could be exploited multiple times, but exactly how it is somewhat missing. I am particularly wondering whether the activations (pixels) or the feature maps (channels) are sampled. If the activations are sampled, I think it is not friendly to the CUDA kernels as the activations within the same channel are computed together; And if the feature maps are sampled, it again violates the analysis in Section 3.2 (since some paths must be chosen together as a group).\n\n(3) It is not clear how the proposed method is compatible with normalization layers. In normalization layers (e.g., batch normalization), the statistics of the activations are computed. However, since the activations now are sampled, it is not explained how the normalization layers should be modified accordingly to preserve unbiasedness.\n\n(4) It is not clear how the proposed method is compatible with randomized operations in neural networks (e.g., dropout, reparameterization, etc.) I am wondering if the combination is still unbiased.\n\nGiven that the technical explanation is lacking for the moment, it is hard for me to judge this paper's correctness. I will temporarily give a reject. I am very willing to increase my score if the authors address the confusion above.", "extra_scores": {}, "extra_text": {}}, {"review_id": "9F1OFkeE1IC", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: \n\nThe paper proposes to subsample the computational graph to obtain an unbiased gradient estimator with less memory requirement, in the same spirit as minibatching reducing the memory of the full-batch gradient descent. The authors propose to do so by modifying the forward pass of a linear layer to only save a subset of the hidden units to estimate the gradient. \n\n\nContributions: \n\nJust like minibatching is introduced to reduce the memory requirement of full-batch gradient descent, this paper introduces another dimension to trade off between noisy gradient estimate and memory by subsampling the computational graph for computing a noisy gradient. The experiments, albeit small-scale, clearly demonstrate the effect of using SGD with randomized AD to train standard neural networks. The application on stochastic optimization of PDE is also very interesting and relevant. The paper is very well written and explains the idea very clearly. I look forward to seeing future work using this principle to train larger scale models as well as how this work will enable the development of many more research ideas the same way minibatch SGD acts as a workhorse of deep learning / machine learning. \n\n\nAdditional details:\n\nIn the exposition of the “alternative SGD scheme with randomized AD“, it might help to understand the sources of the variance of the gradient estimator by decomposing it as Var(x_ij) = Var(x_i) + E[Var(x_ij | x_i)], where i is the index for minibatching and j is for subgraph sampling. It will also help to understand the different curves of Fig 5.  \n\nIs there any speculation of why the drop in performance is more drastic in MLP and CNN than RNN in Fig 5? \n\nI’d like to see some negative results where the variance is too high and harms training, to get a sense of how important it is to tailor the sampling strategies for different computational graphs.  \n\nPerhaps explain what M in section 5 is. I had to go to the appendix to figure out that it is a matrix that summarizes the discretization of the space and time. \n\nRelated work: Perhaps in the discussion of constant-memory via invertible transformation, Neural ODE can be included since it’s constant memory the same way as Gomez et al 17’. \n\n--- post rebuttal --\n\nThanks for the response. I think the contribution of this work is solid and I vote for clear acceptance. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "2EI6ztnbJkq", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors introduce the novel idea of producing unbiased gradient estimates by Monte Carlo sampling the paths in the autodiff linearized computational graph.\nBased on sampling paths it is possible to save memory due to not having to store the complete linearized computational graph (or the intermediate variables necessary to reconstruct it). \nMemory is the main bottleneck for reverse mode autodiff for functions with lots of floating point operations (such as a numerical integrator that performs many time steps, or a very deep neural network).\nThe authors' idea can therefore potentially enable the gradient based optimization of objective functions with many floating point operations without check pointing and recomputing to reverse through the linearized computational graph.\nThe tradeoff made is the introduction of (additional) variance in the gradient estimates.\n\nThe basic idea is simple and elegant:\nThe linearized computational graph of a numerical algorithm is obtained by\na) having the intermediate variables of the program as vertices\nb) drawing directed edges from the right-hand side variables of an assignment to its left-hand side variable\nc) labeling the edges by the (local) partial derivatives of assignments' left-hand side with respect to their right-hand side.\n\nThe derivative of an \"output\" y with respect to an \"input\" x of the function is the sum over all paths from x to y through the linearized computational graph taking the product of all the edges in the path.\nThe sum over all paths corresponds to the expectation of a uniform distribution over the paths times the number of paths.\nThat expectation can be Monte Carlo sampled.\n\nThe authors suggest a way of producing the path sampling based on taking a chained matrix view of the computation graph (see e.g https://arxiv.org/abs/2003.05755) and injecting low rank random matrices.\nDue to the fact that the expectation of the product of independent random variables is equal to the product of the expectations this injection is unbiased as well if the injected matrices have the identity matrix as expectation.\n\nTo take advantage of the simple idea it is in practice necessary to consider the concrete shape of the computational graph at hand in order to decide where to best randomize and save memory without letting variance increase too much.\n\nThe authors present a neural network case study where they show that for some architectures the suggested approach has a better memory to variance trade off than simply choosing a smaller mini-batch size.\nFurthermore, they present a 2D PDE solver case study where their approach can save a lot of memory and still optimize well.\n\nI recommend to accept the paper.\n\nRemarks:\n\nI would love to see a more in depth analysis of the variance for example for simple but insightful toy examples.\n\nFor exampl simple sketching with random variates v with E[vv^T] = I can be used to obtain an unbiased gradient estimate via E[gvv^T] = g, i.e. by evaluating a single forward-mode AD pass (or just a finite difference perturbation).\nBut of course the gradient estimate has such a high variance so as to not give any advantage over finite difference methods (since with each sample / evaluation we are only capturing one direction in the input space).\nWe are not gaining the usual benefit of reverse mode autodiff of getting information about the change of the function in all directions.\n\nIn order for paths to be efficiently Monte Carlo-friendly it is probably necessary that they are correlated with other paths.\nIn practice this will perhaps have something to do with e.g. the regularity of the PDE solution (the gradient with respect to the solution is similar to that of its neighborhood).\n\nA simple example (ODE integrator):\n\np1 = x\np2 = x\nfor i in range(n):\n\tp1 = p1 + h * sin(p1)\n\tp2 = p2 + h * sin(p2)\n\ny = 0.5 * (p1 + p2)\n\nThe two paths in the program compute exactly the same values so leaving one path out randomly does not make any difference at all (if we correctly re-weight the estimate).\n\nMini-batches are often like that: Independent samples from the same class give correlated computations, hence the variance is related to the variance in the data.\n\nBut if two paths involve completely independent and uncorrelated computations then the variance is such that we do not gain anything.\nWe need at least two gradient steps to incorporate the information from both paths.\nSince we do not systematically cycle through them but sample randomly, we are actually going to be less efficient.\n\nIn terms of arguing about memory savings for machine learning applications it would be interesting to see a case study with a large scale architecture that does not fit into memory.\n\nThe random matrix injection section could be clarified by moving the sentence \"the expectation of a product of independent random variables is the product of their expectation\" further to the front and state clearly the idea that:\nE[A PP^T B QQ^T C] = A E[PP^T] B E[QQ^T] C = A I B I C = A B C\n\nIn the PDE example you could clarify the notation used to properly distinguish between the continuous and the discretized solution.\n\nAlso the PDE constrained optimization problem is not inherently stochastic (as can be argued for the empirical risk minimization setting in machine learning).\nTherefore, it is possible to use non-SGD methods with linear or even super-linear convergence rates (quasi-Newton methods).\nSGD with unbiased gradients has a sublinear rate of convergence.\nBut the ideas of the paper are of course still useful even when trying to find the optimum up to machine precision in finite time.\nWe can first use the RAD SGD approach in the early optimization and then go to the deterministic setting later in the optimization.\n\n- Page 3: eqn equation 1 -> Equation 1\n- Page 6: figure 5 -> Figure 5\n- Throughout: Perhaps clarify the meaning of batch vs mini-batch (in other papers batch can refer to full-batch)\n- Figure 5 (a) has blue curves but blue is not in the legend of Figure 5 (c)\n- Page 8: backpropogate -> backpropagate \n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers agree that this is an interesting and original paper that will be of interest to the ICLR community, and is likely to lead to follow up work.
Decision: Accept (Oral)