Randomized Automatic Differentiation

1 INTRODUCTION . Deep neural networks have taken center stage as a powerful way to construct and train massivelyparametric 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 overparameterization , 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 a 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 PyTorch and TensorFlow have received huge dividends from a 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 are used for stochastic optimization . We ask : Why spend resources on exact gradients when we ’ re 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 . In fact , empirical evidence supports the hypothesis that the regularizing effect of gradient noise assists model generalization ( 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 . 1Department of Computer Science 2Department of Astrophysical Sciences from math import sin , exp def f ( x1 , x2 ) : a = exp ( x1 ) b = sin ( x2 ) c = b ∗ x2 d = a ∗ c return a ∗ d a exp ( x1 ) b sin ( x2 ) c b * x2 d a * c f a * d a b c d f exp ( x1 ) cos ( x2 ) x2 b c a d a a b c d f exp ( x1 ) cos ( x2 ) x2 b c a d a 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 a 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 equations . But in stochastic optimization we can make do with inaccurate gradients , as long as our estimator is unbiased and has reasonable variance . We ask the same question that motivates mini-batch SGD : why compute an 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 to 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 . 2 AUTOMATIC DIFFERENTIATION . Automatic ( or algorithmic ) differentiation is a family of techniques for taking a program that computes a differentiable function f : Rn → Rm , and producing another program that computes the associated derivatives ; most often the Jacobian : J [ f ] = f ′ : Rn → Rm×n . ( For a 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 a loss function that produces a 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 a 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 ( e.g. , Bischof et al . ( 1992 ) ; Hascoet & Pascual ( 2013 ) ; van Merrienboer et al . ( 2018 ) ) , execution tracing ( e.g. , Walther & Griewank ( 2009 ) ; Maclaurin et al . ) , manipulation of explicit computational graphs ( e.g. , Abadi et al . ( 2016 ) ; Bergstra et al . ( 2010 ) ) , and category-theoretic transformations ( Elliott , 2018 ) . AD implementations 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 , run-time , or via an embedded domain-specific language , all AD approaches can be understood as manipulating the linearized computational graph ( LCG ) to collapse out intermediate variables . Figure 1 shows the LCG for a simple example . These computational graphs are always directed acyclic graphs ( DAGs ) with vertices as variables . Let the outputs of f be yj , the inputs θi , and the intermediates zl . AD can be framed as the computation of a partial derivative as a sum over all paths through the LCG DAG ( Bauer , 1974 ) : ∂yj ∂θi = Jθ [ f ] j , i = ∑ [ i→j ] ∏ ( k , l ) ∈ [ i→j ] ∂zl ∂zk ( 1 ) where [ i→ j ] indexes paths from vertex i to vertex j and ( k , l ) ∈ [ i→ j ] denotes the set of edges in that path . See Figure 1d 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 a 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 as 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 to structures that are difficult to reason about . Often we prefer to manipulate vectors and matrices , and we can instead let each intermediate zl represent a dl dimensional vector . In this case , ∂zl/∂zk ∈ Rdl×dk represents the intermediate Jacobian of the operation zk → zl . Note that Equation 1 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 , Jθ [ f ] . Reverse-mode AD efficiently calculates Jθ [ f ] , but requires the full linearized computational graph to either be stored during the forward pass , or to be recomputed during the backward pass using intermediate variables recorded during the forward pass . For large computational graphs this could provide a large memory burden . The most common technique for reducing the memory requirements of AD is gradient checkpointing ( Griewank & Walther , 2000 ; Chen et al. , 2016 ) , which saves memory by adding extra forward pass computations . Checkpointing is effective when the number of `` layers '' in a computation graph is much larger than the memory required at each layer . We take a different approach ; we instead aim to save memory by increasing gradient variance , without extra forward computation . Our main idea is to consider an unbiased estimator Ĵθ [ f ] such that EĴθ [ f ] = Jθ [ f ] which allows us to save memory required for reverse-mode AD . Our approach is to determine a 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 formula each Jacobian entry is expressed as a sum over paths in the LCG . A simple strategy is to sample paths uniformly at random from the computation graph , and form a Monte Carlo estimate of Equation 1 . Naïvely 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 as reverse-mode automatic differentiation can then be applied to the sparsified LCG . 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 . A single path in the vectorized computation graph represents many paths in the underlying scalar computation graph . As an example , Figure 2c is a vector representation for Figure 2b . For this example , ∂y ∂θ = ∂y ∂C ∂C ∂B ∂B ∂A ∂A ∂θ ( 2 ) where A , B , C are vectors with entries ai , bi , ci , ∂C/∂B , ∂B/∂A are 3× 3 Jacobian matrices for the intermediate operations , ∂y/∂C is 1× 3 , and ∂A/∂θ is 3× 1 . We now note that the contribution of the path p = θ → a1 → b2 → c2 → y to the gradient is , ∂y ∂C P2 ∂C ∂B P2 ∂B ∂A P1 ∂A ∂θ ( 3 ) where Pi = eieTi ( outer product of standard basis vectors ) . Sampling from { P1 , P2 , P3 } and right multiplying a Jacobian is equivalent to sampling the paths passing through a vertex in the scalar graph . In general , if we have transition B → C in a vectorized computational graph , where B ∈ Rd , C ∈ Rm , we can insert a random matrix P = d/k ∑k s=1 Ps where each Ps is sampled uniformly from { P1 , P2 , . . . , Pd } . With this construction , EP = Id , so E [ ∂C ∂B P ] = ∂C ∂B . ( 4 ) If we have a matrix chain product , we can use the fact that the expectation of a product of independent random variables is equal to the product of their expectations , so drawing independent random matrices PB , PC would give E [ ∂y ∂C PC ∂C ∂B PB ] = ∂y ∂C E [ PC ] ∂C ∂B E [ PB ] = ∂y ∂C ∂C ∂B ( 5 ) Right multiplication by P may be achieved by sampling the intermediate Jacobian : one does not need to actually assemble and multiply the two matrices . For clarity we adopt the notation SP [ ∂C/∂B ] = ∂C/∂BP . This is sampling ( with replacement ) k out of the d vertices represented by B , and only considering paths that pass from those vertices . The important properties of P that enable memory savings with an unbiased approximation are EP = Id and P = RRT , R ∈ Rd×k , k < d . ( 6 ) We could therefore consider other matrices with the same properties . In our additional experiments in the appendix , we also let R 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 a two-level sampling scheme . We can both sample paths from the computational graph where each vertex on the path corresponds to a 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 B → C where B , C ∈ Rd . The Jacobian is d × d. Thankfully many common operations are element-wise , leading to a diagonal Jacobian that can be stored as a d-vector . Another common operation is matrix-vector products . Consider Ab = c , ∂c/∂b = A . Although A has many more entries than c or b , in many applications A is either a 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 or 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 a compactly-representable Jacobian and for which this is not possible .

Summary: 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.