Cutting Down Training Memory by Re-fowarding

1 INTRODUCTION . The standard DNN training process consists of two alternated stages : forward and backward . Fig . 1 ( a ) illustrates an example of feed-forward neural networks . In the forward stage , the network takes an input tensor , [ BatchSize× Channel ×Width×Height ] , and computes the tensors at each layer until producing the output . In the backward stage , difference between the output and ground truth is passed back along the network to compute the gradients at each layer . The regular training approach saves tensors at all layers during forward , because they are all needed to compute gradients during backward . The total memory cost is the sum of cost over all layers . In popular backbone DNNs for feature extraction of images , such as AlexNet ( Krizhevsky et al . ( 2012 ) ) , VGG ( Simonyan & Zisserman ( 2014 ) ) and ResNet ( He et al . ( 2016 ) ) , the memory cost increases quadratically with the input image resolution and network depth . For example , given an median size input tensor of ( 32 , 3 , 224 , 224 ) , ResNet101 requires around 5000 MB . In more challenging tasks , DNNs that detect small objects and large number of object categories require input image resolution of more than 600 × 600 ( Ren et al . ( 2015 ) ; Singh et al . ( 2017 ) ; Redmon & Farhadi ( 2018 ) ) . The memory issue is worse for video-based DNNs , such as CDC ( Shou et al . ( 2017 ) ) , C3D ( Ji et al . ( 2013 ) ) and 3D-ResNet ( Hara et al . ( 2017 ) ) . To model complex activities in video , the input tensor may contain 64 frames . Moreover , DNN training takes much more memory than testing . In order to train DNNs with large databases and big learning rate , the batch size can be up to 64 . In training DNN compositions , such as Generative adversarial networks ( GANs ) , multiple generator and discriminator networks are simultaneously stored in GPU memory . Existing efforts to address memory issues presented three main approaches : ( 1 ) Better single GPUs . Recent GPUs provide larger memory at the expense of exponentially growing price and power consumption . For instance , from TitanXp , Quadro P6000 to Tesla V100 , for 1-2.7 times increase in memory , the prices increase 2.8-8.5 times . ( 2 ) Parallelization among multiple GPUs ( Dean et al . ( 2012 ) ; Shi et al . ( 2009 ) ; Langford et al . ( 2009 ) ; Mcdonald et al . ( 2009 ) ; McDonald et al . ( 2010 ) ; Zinkevich et al . ( 2010 ) ; Agarwal et al . ( 2014 ) ; Agarwal & Duchi ( 2011 ) ) , which requires expensive clusters , introduces substantial I/O cost , and does not reduce the total memory cost . ( 3 ) Low-level heuristic techniques . Optimization of computation graphs ( Aho et al . ( 1986 ) ) , which merges inplace operations into non-inplace operations to cut down memory . Liveness analysis ( Aho et al . ( 1986 ) ) , which dynamically recycles garbage tensors in training epochs . These approaches are specific to certain DNN structures , data and tasks . To address above issues , we propose a fundamental approach that explores trade-off between memory and computation power of GPUs . Note that recent affordable GPUs , although limited in memory ( 12GB ) , provide exceptional improvement in GPU cores and FLOPS . Trading computational time for memory is a very attractive solution that make it possible to train very heavy DNNs with finite GPU memory . Our approach only saves tensors at a subset of layers during the first forward , and conduct only extra local forwards to compute the missing tensors needed during backward . We call the extra forward process as Re-forwarding . The total memory cost is the sum of ( 1 ) the cost at the subset of layers and ( 2 ) the maximum memory cost among local re-forwards . Training with Reforwarding , see Fig . 1 ( b ) , leads to substantial memory reduction . We propose sophisticate theories and efficient algorithms that achieve the optimal memory solution of arbitrary computation graphs . 2 RELATED WORK . To alleviate the memory pressure from a single GPU processor , many researchers utilized the wellestablished techniques for distributed computation ( Dean et al . ( 2012 ) ; Shi et al . ( 2009 ) ; Langford et al . ( 2009 ) ; Mcdonald et al . ( 2009 ) ; McDonald et al . ( 2010 ) ; Zinkevich et al . ( 2010 ) ; Agarwal et al . ( 2014 ) ; Agarwal & Duchi ( 2011 ) ) . These techniques distribute memory pressure to possibly infinite GPUs or server clusters , but do not reduce the total memory cost of DNNs . Other researchers reduced the memory on finite hardware by optimizing computation graph of DNN and performing liveness analysis . The computation graph of DNNs describes the dependencies of tensors among layers . Liveness analysis recycles garbage to manage memory . These ideas were originated from compiler optimization ( Aho et al . ( 1986 ) ) and has been widely adopted by deep learning frameworks : Theano ( Bastien et al . ( 2012 ) ; Bergstra et al . ( 2010 ) ) , MXNet ( Chen et al . ( 2015 ) ) , Tensorflow ( Abadi et al . ( 2016 ) ) and CNTK ( Yu et al . ( 2014 ) ) . Some other techniques efficiently swap data between CPU and GPU ( Wang et al . ( 2018 ) ; Rhu et al . ( 2016 ) ) . These techniques usually cost extra I/O time and still do not actually reduce the total memory cost . The closest work to our approach , Chen et al . ( Chen et al . ( 2016 ) ) , uses the gradient checkpoints ( similar to the subset of layers in Re-forwarding ) . However , ( Chen et al . ( 2016 ) ) only worked on linear computation graph via a heuristic algorithm . Our approach generates optimal solutions for both linear and arbitrary computation graphs . Our algorithm reduces training memory by manipulating high-level tensors , therefore is generalizable to any DNNs and their compositions . All previous techniques are compatible to our approach and can further improve the memory efficiency of DNN training . 3 LINEAR COMPUTATION GRAPH ( LCG ) . Denote a computation graph as G = ( E , V ) . E = { ei } and V = { vi } are the edges and vertices in the computation graph , respectively . In deep neural networks , the vertices represent the tensors and the edges represent operations . Denote function l ( · ) as a measure of memory cost . VR is the subset of vertices saved during the first forward . l ( vi ) is defined as the memory cost of storing vertex vi . For two adjacent vertices vi and vj in set VR , the memory cost during re-forwarding from vi to vj is defined as l ( vi , vj ) = ∑j−1 t=i+1 l ( vt ) , which is the sum of cost over all the vertices between vi and vj . Using these notations , the memory cost of training with re-forwarding is formulated as min VR ∑ i l ( vi ) + max j l ( vj , vj+1 ) , ( 1 ) where the first term is the sum of the memory cost of all the stored tensors , and the second term is the maximal cost among the re-forwards . For easy illustration , we start by formulating Re-forwarding on Linear Computation Graphs ( LCG ) ( Fig . 2 ( a ) ) . For LCGs , Eqn . 1 can be solved in two cases . Case ( 1 ) LCG with Identical Vertex Cost : Suppose a LCG has N vertices , each of which has the same cost l ( vi ) = 1N and the total cost is 1 . Obviously , the optimal solution is reached when vertices in VR are distributed evenly in the LCG . Suppose the number of vertices in VR is k. The total cost is then kN + 1 k . The optimal solution of Eqn . 1 is k = √ N , and the optimal total cost is 2√ N . Case ( 2 ) LCG with Non-identical Vertex Cost : When the assumption of identical cost does not hold , the solution to Eqn . 1 does not have an analytic form . Denote the maximal Re-forward cost max j l ( vj , vj+1 ) as a constant C , and the solution to Eqn . 1 is reduced to solving for min VR ∑ i l ( vi ) . Algorithm 1 Linear Computation Graph ( LCG ) Solver 1 : for each vertex pair ( vi , vj ) in G do 2 : Set the maximal term as l ( vi , vj ) 3 : Construct Accessibility Graph 4 : Find the shortest path in the Accessibility Graph as the solution 5 : Compute the actual total cost of the solution 6 : Save the solution if it ’ s better . 7 : Suppose the actual max term of this solution is B , and l ( vi , vj ) = C , skip the loops where B ≤ l ( vi , vj ) < C All the Re-forward costs in an optimal solution satisfy the constraint l ( vj , vj+1 ) ≤ C. We solve Eqn . 1 by constructing a new graph , called Accessibility Graph GA = ( EA , V ) . The edges of GA , called Accessibility Edge eAij , exists between vertex vi and vj if and only if l ( vi , vj ) ≤ C. Now the problem of solving min VR ∑ i l ( vi ) is equivalent to finding the shortest path from the source vertex and the target vertex in the Accessibility Graph . Notice that in the optimal solution , the max term equal the one maximal term among all l ( vi , vi+1 ) terms . To traverse all possible max terms , we can simply compute the loss of every vertex pair and use it as a possible max term . Given a max term C , suppose the actual max term of the solution under C is B and B < C. It ’ s obvious that for all the max terms B ≤ max < C , the solution would be the same solution . Therefore , these max terms can be skipped . Algorithm 1 summarizes the process for searching an optimal solution for LCG . Suppose there are N vertices in the computation graph , the time complexity of Algorithm 1 is O ( N4 ) 1 . 4 ARBITRARY COMPUTATION GRAPH ( ACG ) . As generalization of DNNs with LCG , we present theory2 and algorithms for DNNs with Arbitrary Computation Graphs ( ACG ) , in particular the acyclic directed graphs ( Fig . 2 ( b ) ) . 4.1 ASSUMPTION FOR OPTIMALITY . The optimal solution of Re-forwarding corresponds to an optimal division of ACG , such that memory cost ( Eqn . 1 ) is minimum . We denote that an ACG is divided into end-to-end segments by a set of vertices . These end-to-end segments can have multiple endpoint vertices , for example , multiple source vertices and multiple target vertices . In this paper , as an assumption and also for simplification , these end-to-end segments are narrowed down to those with only one source vertex and one target vertex . Another assumption in the case of ACG is imposed on the operation that has multiple inputs : one can compute the gradients of output with respect to the gradients of inputs without using the current value of inputs . Examples of operations that meet this assumption are : concatenation ( the gradient of output is also the concatenation of the gradient of input ) , add ( the gradient of output equals the gradient of input ) , etc . An example that breaks this assumption is multiplication ( the gradient of input depends on the input ) . Fortunately , most of the popular networks meet this assumption . A simple way to remove this assumption is to store all the input tensors of this multi-input operation . However , this is not modeled by our loss function and may lead to sub-optimal solution . In summary , there are only two assumptions in our approach : ( 1 ) the segment in a solution only has two endpoints ( source and target ) . ( 2 ) the multi-input operation can compute the gradients of output without using the current value of input . Under these two assumptions , our approach is optimal for ACGs .

Summary: The paper suggests a method to reducing the space consumption of training neural nets, in exchange for additional training time. The method stores in memory only a subset of the intermediate tensors computed in the forward step, and then in the backward step it re-computes the missing tensors as they are needed by interpolating forward (again) between the stored tensors. The paper also gives a combinatorial algorithm for choosing which tensors to store on a given computation DAG annotated with vertex costs.