Title: CUTTING DOWN TRAINING MEMORY BY RE-FOWARDING

Abstract: Deep Neutral Networks(DNNs) require huge GPU memory when training on modern image/video databases. Unfortunately, the GPU memory as a hardware resource is always finite, which limits the image resolution, batch size, and learning rate that could be used for better DNN performance. In this paper, we propose a novel training approach, called Re-forwarding, that substantially reduces memory usage in training. Our approach automatically finds a subset of vertices in a DNN computation graph, and stores tensors only at these vertices during the first forward. During backward, extra local forwards (called the Re-forwarding process) are conducted to compute the missing tensors between the subset of vertices. The total memory cost becomes the sum of (1) the memory cost at the subset of vertices and (2) the maximum memory cost among local re-forwards. Re-forwarding trades training time overheads for memory and does not compromise any performance in testing. We propose theories and algorithms that achieve the optimal memory solutions for DNNs with either linear or arbitrary computation graphs. Experiments show that Re-forwarding cuts down up-to 80% of training memory on popular DNNs such as Alexnet, VGG, ResNet, Densenet and Inception net.

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 × W idth × 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 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 = {e i } and V = {v i } 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. V R is the subset of vertices saved during the first forward. l(v i ) is defined as the memory cost of storing vertex v i . For two adjacent vertices v i and v j in set V R , the memory cost during re-forwarding from v i to v j is defined as l(v i , v j ) = j-1 t=i+1 l(v t ), which is the sum of cost over all the vertices between v i and v j . Using these notations, the memory cost of training with re-forwarding is formulated as

EQUATION

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.

Figure 2 : (a) Linear Computation Graph (LCG). "s" denotes the start vertex,"t" denotes the end vertex. (b) Arbitrary Computation Graph (ACG). The structure between "s" and "t" vertices may contain arbitrary branches and connections.

Case(1) LCG with Identical Vertex Cost: Suppose a LCG has N vertices, each of which has the same cost l(v i ) = 1 N and the total cost is 1. Obviously, the optimal solution is reached when vertices in V R are distributed evenly in the LCG. Suppose the number of vertices in V R is k. The total cost is then k N + 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(v j , v j+1 ) as a constant C, and the solution to Eqn. 1 is reduced to solving for min

V R i l(v i ).

Algorithm 1 Linear Computation Graph (LCG) Solver 1: for each vertex pair

(v i , v j ) in G do 2:

Set the maximal term as l(v i , v j ) Save the solution if it's better.

7::

Suppose the actual max term of this solution is B, and l(v i , v j ) = C, skip the loops where

B ≤ l(v i , v j ) < C

All the Re-forward costs in an optimal solution satisfy the constraint l(v j , v j+1 ) ≤ C. We solve Eqn. 1 by constructing a new graph, called Accessibility Graph G A = E A , V . The edges of G A , called Accessibility Edge e A ij , exists between vertex v i and v j if and only if l(v i , v j ) ≤ C. Now the problem of solving min V R i l(v i ) 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(v i , v i+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(N 4 )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. Definition 1. Closed Set: A set s containing vertices and edges is a closed set if and only if it satisfies the following three properties: 1. All the vertices of s have a common ancestor v i and a common descendent v j ; 2. Denote the vertex subset of s as V , edge subset as E, and the set of edges between two arbitrary vertices of V ∪ {v i , v j } is E , the edge from v i to v j (if exists) as e ij . E must either be E or E -{e ij }; 3. An arbitrary v 1 ∈ V doesn't have edge with another arbitrary v 2 / ∈ V ∪ {v i , v j }. For multiple valid closed sets between v i and v j , we denote the largest one as

4.2. DEFINITION AND THEOREM:

s ij Definition 2. [s ij ] = s ij ∪ {v i , v j }. [s ij ) = s ij ∪ {v i }. (s ij ] = s ij ∪ {v j }

In the definition of Closed Set, property 1 corresponds to the two endpoint assumption in section 4.1 where the two endpoints become v i and v j in the definition. Property 2 confines the edge subsets of s to be one of two cases: E or E -{e ij }. Both cases are valid although they have different edges. Property 3 guarantees the independence of such a set s, meaning that the vertices within s have no connections with other vertices outside s ∪ {v i , v j }. As there might be multiple valid closed sets between v i and v j , which corresponds to the Branched Closed Set in Definition 5, we denote the largest closed set between v i and v j as s ij and denote smaller closed set with an extra superscript, such as s 1 ij . Definition 3. Splitting Vertex: A vertex v t ∈ s ij is a splitting vertex of s ij if and only if s it exists, s tj exists and s ij = s it ∪ s tj ∪ {v t } and s it ∩ s tj = ∅ Definition 4. Splittable Closed Set (Type 1): closed set with at least 1 splitting vertex.

The definition of Splitting Vertex is to describe whether a closed set can be divided into two linearly arranged closed set. A closed set is splittable if it has at least 1 splitting vertex and is defined as Closed Set Type 1. Definition 5. Branched Closed Set (Type 2): A closed set is branched if it has 0 splitting vertex and can be divided into branches: 

s ij = s 1 ij ∪ s 2 ij and s 1 ij ∩ s 2 ij = ∅

[s pq ]} is a maximal split of non-branched s ij if [s ij ] = ∪{[s pq ]} and ∀s ab , s cd ∈ {[s pq ]}, s ab ∩ s cd = ∅ and ∃{[s pq ]} {[s pq ]} such that ∪{[s pq ]} = [s kt ] [s ij ]

Definition 8. Division of Closed Set: For type 1, its division is the linear segments separated by all its splitting vertices; for type 2, its division is all its branches, any of which cannot be divided into more branches; for type 3, its division is its maximal split.

For closed set type 1, it can be divided into linearly arranged segments. For closed set type 2, it can be divided into branches. So here we investigate the division of closed set type 3. As we don't want trivial division, for example, division that is formed by every edge in the closed set, we define Maximal Split to describe the split such that each member of the split is as large as possible. An example of maximal split is shown in Fig. 4 (c ). In the definition of maximal split, the term maximal is implied by saying that any subset of this split cannot be combined into a single closed set. If it can, then the maximal split will be formed by this larger closed set and all the rest of the previous split. For closed set type 3, we use its maximal split as its division. Definition 9. Division Tree: Division tree is a representation of a computation graph, where the root node is the whole computation graph, the leaf nodes are all the single tensors in the computation graph, and for a non-leaf node, its children is the members of its division.

With the division of 3 types of closed sets, the computation graph can be reorganized into a division tree (Figure 5 ) where a non-leaf node would be a closed set and its children would be its corresponding division. The root node is the whole computation graph, the largest closed set, and the leaf nodes would be single tensors in the computation graph. With division tree, we can apply divide-and-conquer to search for optimal solution.

Figure 5 : In this tree, the root node is the whole computation graph. All the leaf nodes are single tensors. Every other node except root and leaves is a member of the division of its parent.

Theorem 1. The division tree of a computation graph is unique and complete.

The uniqueness of the division tree indicates that the optimal solution of the division tree would also be the optimal solution of the whole computation graph. The completeness indicates that the division tree has included all the possible members of solution and represents the whole search space for the optimal solution. Theorem 1 is proved in the appendix.

4.3. ALGORITHM:

We search optimal solutions for ACGs by solving several sub-problems using Algorithm 2-4 respectively. Based on these components, we present our final solver as Algorithm 5.

Algorithm 2 judges whether a vertex is a splitting vertex of a closed set. This algorithm mainly follows the Definition 3 and uses vertex set to check the property of a splitting vertex. With this algorithm, we can judge whether a closed set is type 1 and get its division if it is. Suppose there are N vertices in s ij , the time complexity of Algorithm 2 is O(N 2 ). Algorithm 2 Judge whether a vertex v t is a splitting vertex of closed set s ij 1: Let {v in } be the vertices of all the vertices within [s ij ] that have paths to v t . Let {v out } be the vertices of all the vertices within [s ij ] that have paths from v t .

2: if {v in }∪{V out }∪{v t } = {v|v ∈ [s ij ]} and {v in }∩{V out } = ∅ and ∃v 1 ∈ {v in }, v 2 ∈ {v out }, v 1 , v 2 have connections then 3:

Return true 4: else 5:

Return False:

Algorithm 3 examines whether a closed set is branched. It uses a growing algorithm to check whether an independent subpart of this closed set can form a closed set. If a non-trivial closed set s ij has an edge from v i to v j , then it's branched because this edge itself can be treated as a closed set. Combined with Algorithm 2, we can know the type of a closed set and get its division if it's type 2. Suppose there are N vertices in s ij , the time complexity of Algorithm 3 is O(N 2 ).

Algorithm 4 addresses the problem of finding the maximal split, the division of a closed set type 3 s ij . First get all the possible closed sets within s ij and use a property of maximal split to judge whether this closed set is a member of the maximal split. The property is: there cannot exist another closed set s ab s ij but contains any member of this maximal split. This property is proved in Lemma 6 of the appendix. Suppose there are N vertices in s ij , the time complexity of Algorithm 4 is O(N 4 ).

Algorithm 5 is the solver for ACGs. First, the division tree of the computation graph is built. Similar to the linear solver, a max term list is formed by the cost of all the possible closed sets for traverse. Given a max term, we propose a greedy idea: for a closed set, never expand it unless the its cost exceed the max term. In other word, if the max term doesn't allow a leap over this closed set, we expand it, otherwise, do not expand it. Because once expanded, some cost of other vertices inside this closed set might be introduced, and the cost will never be smaller than unexpanded. If some children of the closed set type 1 are expanded, the rest reforms a few linear segments and still can be solved by the linear solver. If some children of the closed set type 2 or 3 are expanded, the other Algorithm 3 Judge whether s ij is branched 1: if s ij has at least 1 vertex then 2:

if s ij includes an edge from v i to v j then For any v t ∈ s ij , v t ∈ s that has connection to any v k ∈ s, add v t to s.

8::

if No more vertex can be added to s then 

(v k , v t ) except (v i , v j ) in [s ij ] do 2:

For all the vertices {v} that have paths from v k and have paths to v t .

3:

if ∃v 2 ∈ {v} and v 2 = v k , v t , v 2 has connection to a v 1 ∈ {v} then 4:

Form a closed set s kt with all these vertices. 5: for each formed closed set s kt do 6:

If there doesn't exist a s ab such that s kt s ab s ij , put s kt into the maximal split. members remain unexpanded and need no changes. Suppose there are N vertices in computation graph, the time complexity of Algorithm 5 is O(N 4 ). Algorithm 5 Arbitrary Computation Graph (ACG) Solver 1: Get all possible closed set and their costs. Use their costs to form the max term list. 2: Reorganize the computation graph into a division tree: from the root node (the computation graph), build its children from its division, until all the leaf nodes are single tensors. 3: for each possible max term m in max term list {m} do 4:

if current closed set is type 1 then

5::

For all the children that have cost larger than current max term. Expand them and solve the next level.

6::

All the expanded children have separated the current closed set to linear segments. Solve all the linear segments with current max term. For all the children that have cost larger than current max term. Expand them and solve the next level. 9:

All the other members remain unexpanded. 10: Summarize the total loss, save the current solution if it's better.

5. EXPERIMENT:

We evaluated Re-forwarding on two main groups of neural networks (1) networks with linear structures, such as Alexnet (Krizhevsky et al. (2012) ) and vgg series (Simonyan & Zisserman (2014) ). ( 2) networks with non-linear structures, such as Resnet series (He et al. (2016) ), Densenet series (Huang et al. (2017) ) and Inception net (Szegedy et al. (2016) ). For each network in Table 1 , an computation graph is built such that every vertex is a Float32 tensor, every edge is an operation, and the memory cost of a vertex is its tensor size (measured in MB). We compared Re-forwarding with Chen (Chen et al. ( 2016)) and the regular training approach. Note that Chen et al. ( 2016) only worked on linear computation graphs. To compare with (Chen et al. (2016) ) on non-linear networks, we manually re-organized all the non-linear computation graphs into linear computation graphs with their splitting vertices, and fed them to Chen et al. (2016) (MB)"). Our Re-forwarding approach directly works on arbitrary computation graphs. We have also included a customized network ("CustomNet"), on which even the manual version of Chen's approach is not applicable. Our approach directly works on all networks. The computation graph of this network is visualized in the appendix.

All experiments were conducted in Pytorch. To remove irrelevant GPU memory cost, such as model and Pytorch CUDA interface cost, all training memory costs were measured with two different input sizes and compute the difference between two measurements. For example, to measure the memory cost of Alexnet with input size [BatchSize, Channel, W idth, Height]= [16, 3, 224, 224] , we first record the training memory of input [16, 3, 224, 224] as r 1 , and input [32, 3, 224, 224] as r 2 . The actual memory cost given [16, 3, 224, 224] input is measured as r 2 -r 1 . To use existing DNN implementations, the input of Inception net is [BatchSize, 3, 300, 300] , and the input of all other networks is [BatchSize, 3, 224, 224] . We also measure the training time (time of 1 training iteration) for the regular approach and our approach. Each time is measured as the average of 20 iterations. Our approach has the same training time as Chen's approach and its manual version, see "Space Efficient Training Time" in Table 1 .

Table . 1 shows that Re-forwarding cuts down huge amount of memory from the regular approach at reasonable time overheads: 26% space off and 40% time overhead for Alexnet, around 40% space off and 40% time overhead for Vgg series. For Resnet series, the deeper network, the more memory was cut down. On the deepest Resnet152, 80% space off was achieved with only 39% time overhead. For Densenet series, more than 80% space off was achieved with around 40% time overhead. Notice that, Chen et al. ( 2016) only works on linear networks. Its results on non-linear networks were manually synthesized. Re-forwarding directly works on non-linear networks and constantly outperformed Chen et al. (2016) and its "manual" version. This supports our claim that Re-forwarding is optimal.

6. CONCLUSION:

Re-forwarding is a fundamental approach that explores trade-off between memory and computation power of GPUs. By saving tensors at a subset of layers during forward, and conducting extra local forwards for backward, Re-forwarding makes it possible to train very heavy DNNs with finite GPU memory. To our knowledge, our theoretical and algorithmic results are the first top-down work that achieve an optimal memory solution for arbitrary computation graphs in DNNs. Re-forwarding can be further embedded and optimized with any low-level techniques such as distributed computing, GPU/CPU swapping, computation graph optimization and liveness analysis.

A PROOF:

A.1 LEMMAS

Lemma 1. If s ij ∩ s kt = ∅ and s ij ⊂ s kt and s kt ⊂ s ij , then s ij ∩ s kt = s kj or s ij ∩ s kt = s it Proof. Let [s ij ] ∩ [s kt ] = s = {v, e}.

Let v p be the source vertex of s, v q be the target vertex of s. If v p = v i and v p = v k and v q = v j and v q = v t , then v i , v k has path to v p and v j , v t has path from v q . Therefore, v p has at least 2 immediate parents

v a , v b with v a ∈ [s ij ], v a ∈ [s kt ], v b ∈ [s kt ], v b ∈ [s ij ].

If so, the independence of s ij and s kt is violated. Therefore, v p must be v i or v k .

Same on v q , v q must be v j or v t .

If v p = v i , v q = v j , then s ij ⊂ s kt . If v p = v k , v q = v t , then s kt ⊂ s ij . Therefore, v p = v i , v q = v t or v p = v k , v q = v j . Suppose v p = v k , v q = v j , let's prove s is a closed set. With s ⊂ s kt , ∀v 1 ∈ s, v 1 has no edge with v 2 ∈ [s kt ]. With s ⊂ s ij , ∀v 1 ∈ s, v 1 has no edge with v 2 ∈ [s ij ]. Therefore, ∀v 1 ∈ s, v 1 has no edge with v 2 ∈ [s]

. The independence of s is guaranteed.

In the discussion before, we can see the source vertex v p of s must be either v i or v k . If v i and v k are both the source vertices of s,

then v i ∈ [s kt ] and v k ∈ [s ij ]

, v i has path to v k and v k has path to v i , which will force v i = v k because the s ij , s kt is acyclic. Same on v q , s can only have 1 source vertex and 1 target vertex. Therefore, s is closed set.

Therefore,

s ij ∩ s kt = s kj or s ij ∩ s kt = s it .

Lemma 2. The intersection of two closets s = s i ∩ s j = ∅ is also a closet Proof. Given the independence of s i and s j , the independence of s is obvious. The remaining thing is whether s only has 1 source vertex and 1 target vertex. In the proof of Lemma 1, we can see any source or target vertex of s will eventually become source or target vertex of s i and s j . With simple discussion, we can have this lemma.

Lemma 3. If s ij ∩ s kt = s kj = ∅, then v k is the splitting vertex of s ij and v j is the splitting vertex of s kt

Proof. Let's first prove that v k is the splitting vertex of s ij .

Let s = s ij -[s kj ). Obviously, s ij = s ∪ s kj ∪ {v k } and s ∩ s kj = ∅. We only need to prove that s is closed set. For convenience, let's denote

[s] = s ∪ {v i , v k } v i is obviously the only source vertex of [s] because v i is source vertex of [s ij ]. We discuss the target vertex here. If v k is not the target vertex of [s], as v k ∈ [s], v k must have path to the target vertex v of [s] and v also has path to v j as v ∈ s ij . Because v / ∈ [s kj ],

in the path from v to v j , there exists an edge that connects a vertex v 1 ∈ s with a vertex v 2 ∈ s kt which violates the independence of s kt . Therefore, the target vertex of [s] can only be v k .

As s ⊂ [s ij ), ∀v 1 ∈ s, v 1 has no edge with v 2 ∈ [s ij ). As s kj is close, ∀v 1 ∈ s, v 1 has no edge with v 2 ∈ s kj . ∀v 1 ∈ s, v 1 can only have edge with v 2 ∈ [s]. Thus the independence of s is guaranteed. Therefore, s is closed set, v k is the splitting vertex of s ij . Same on v j , v j is the splitting vertex of s kt Lemma 4. If s ij has n splitting vertices {v 1 , v 2 , ..., v n }, then s ij = s i1 ∪ s 12 ∪ ... ∪ s nj ∪ {v 1 , v 2 , ..., v n } Proof. If n = 2, the splitting vertices are v 1 , v 2 , s ij = s i1 ∪ s 1j ∪ {v 1 } = s i2 ∪ s 2j ∪ {v 2 }. Let v 1 ∈ s i2 , v 1 = v 2 , then s 1j ∩ s i2 = s 12 = ∅. According to Lemma 3, v 1 is splitting vertex of s i2 and v 2 is splitting vertex of s 1j . Therefore, s ij = s i1 ∪ s 12 ∪ s 2j ∪ {v 1 , v 2 }.

For n > 2, the lemma can be proved by repetitively using the conclusion in n = 2. Lemma 5. If the non-branched s ij has a maximal split {[s pq ]}, and |{[s pq ]}| > 2, denote {v} as all the endpoint vertices of [s] ∈ {[s pq ]}. Then ∀v ∈ {v}, v = v i , v j , v is the endpoint vertex of at least 3 members of the maximal split.

Proof. If v b is the endpoint vertex of only 2 members of the maximal split, suppose the 2 members are s ab and s bc . If so, s ab and s bc can be merged into s ac . If s ac = s ij , this violates the definition of maximal split. Otherwise, it violates the condition that s ij is non-branched and |{[s pq ]}| > 2. It is impossible that the 2 members are s ab and s cb because in this way v b has no path to v j and violates the definition of closed set. If v b is the endpoint vertex of only 1 member of the maximal split, then v b must be either v i or v j . Therefore, this lemma is proved. Lemma 6. Any member of a maximal split can not be the subset of another closed set s s ij .

Proof. Suppose the source vertex of s is v 1 and target vertex is v 2 , a member s xy of the maximal split is inside s. Suppose a member s ab of the maximal split has its source vertex v a inside s and target vertex v b outside s. Then the boundary vertex (the vertex that has edges to the non-overlapping parts of both sets) must be v 2 , otherwise the independence of s will be violated. Notice that v 2 is inside s ab and the independence of s ab needs to be guaranteed, for ∀v p ∈ s, v p / ∈ s ∩ s ab , v q ∈ s ∩ s ab , v p has no edge with v q . Therefore, v a is a splitting vertex of s.

Similarly, if s ba has its target vertex v a inside s and source vertex v b outside s, the boundary vertex must be v 1 and v a is a splitting vertex of s.

For the closed set s, from the discussion above, we know that there are at most 2 members of the maximal split that can overlap with s. Other members must be either completely inside s or completely outside s. Let's discuss the number of members that overlaps with s.

If there are 0 member that overlaps with s, s is the union of a subset of members of the maximal split, which violates the definition of maximal split.

If there is 1 member that overlaps with s, suppose the corresponding splitting vertex is v b , and the boundary vertex is actually v 2 . Then s 1b is a closed set containing s xy and corresponds to the situation of 0 member overlapping. s 1b is the union of a subset of members of the maximal split, and violates the definition of maximal split.

If there are 2 members that overlaps with s, suppose they generate two different splitting vertex v a and v b . Then s ab is a closed set containing s xy and corresponds to the situation of 0 member overlapping. s ab is the union of a subset of members of the maximal split, and violates the definition of maximal split.

If they generate the same splitting vertex v b , from lemma 5, v b is also the endpoint vertex of at least 1 other member s ab which has to be inside s. Suppose the two overlapping members are s cb that contains v 1 , and s bd that contains v 2 . As the source vertex of s, v 1 has path to v b and v 1 has path to v a , which implies v b has path to v a . As the target vertex of s, v 2 has path from v b and v 2 has path from v a , which implies v b has path from v a . This conflicts with the fact that s is acyclic. Therefore, this case is not possible.

Therefore, this lemma is proved.

Lemma 7. If non-branched s ij has at least 1 vertex but has 0 splitting vertex, then its maximal split has length > 2 Proof. As s ij is not branched, the members of its maximal split cannot have the starting vertex as v i and the ending vertex as v j at the same time. If s ij has at least 1 vertex, and its maximal split has length 2, then its maximal split must be {[s ik ], [s kj ]}, and v k will be the splitting vertex of s ij , which violates that s ij has no splitting vertex.

B.3 ALGORITHM 3:

Suppose there are N vertices in the closed set s ij . The most time consuming part will be from step 5 to step 13. Other steps are O(1). In step 5 to step 13, every edge between two vertices in s ij is at most visited once and there are O(N 2 ) edges. Therefore, the overall time complexity of Algorithm 3 would be O(N 2 ).

In our implementation, an array of length N is used to represent the vertex set s = {v k }. Therefore, the space complexity would be O(N ).

B.4 ALGORITHM 4:

Suppose there are N vertices in the closed set s ij and there are O(N 2 ) vertex pairs. For each vertex pair, the connection check in step 2-4 will cost O(N 2 ), similar to the connection check in Algorithm 2. Thus step 1-4 will cost O(N 4 ). In our implementation, for each vertex in the closed set s ij , we select the largest formed closed set s kt that contains this vertex. The closed set number is then reduced to O(N ) and step 5-6 can be done in O(N 3 ). Therefore, the overall time complexity of Algorithm 4 would be O(N 4 )

As O(N 2 ) closed sets can be formed in step 1-4 and each closed set is a smaller DAG with O(N ) vertices and cost O(N 2 ) space, the space complexity would be O(N 4 ) for all these closed sets.

B.5 ALGORITHM 5:

Step 1 is similar to step 1-4 in Algorithm 4 with s ij being the whole computation graph. Therefore, the overall time complexity for step 1 is O(N 4 ).

In step 2, the complexity of building division tree is related to the complexity of getting the division of a closed set. For closed set type 1, Algorithm 2 is called for each vertex to get all splitting vertices. Thus getting the division of closed set type 1 cost O(N 3 ) time. For closed set type 2, Algorithm 3 is used to solve for its division and costs O(N 2 ) time. For type 3, Algorithm 4 is called to solve for its division. Notice that we have already stored all possible closed sets in step 1, step 1-4 in Algorithm 4 can be skipped and thus the time complexity of getting the division of closed set type 3 is reduced to O(N 3 ). Therefore, getting the division of an arbitrary closed set costs O(N 3 ) time. In depth i of the division tree, suppose there are k closed sets, and the number of vertices of jth closed sets is a j . To build depth i + 1 of the division tree, we need to get the division of all these closed sets, which will cost j O(a 3 j ). As j a j ≤ N , we have j O(a 3 j ) ≤ O(N 3 ). As the depth of division tree is at most N , the overall time complexity of step 2 would be O(N 4 ).

For step 3-10, if the computation graph is linear, the ACG solver will reduce to LCG solver and has complexity O(N 4 ). If the computation graph is non-linear, the length of {m} would be O(N 2 ) for there are O(N 2 ) vertex pair. For a max term m, from step 4-10, the actual time costing part will be step 6 which calls the LCG solver, and other steps would be O(1). Suppose the LCG solver is called k times, solving problems of a 1 , a 2 , ..., a k vertices. The total complexity of this would be O(a 4 1 ) + O(a 4 2 ) + ... + O(a 4 k ). Notice that a 1 + a 2 + ... + a k ≤ N for the fact that any vertex in the computation graph would not be solved twice by LCG solver, we have a 4 1 + a 4 2 + ... + a 4 k ≤ N 4 . Therefore the time complexity of step 3-10 is O(N 4 ).

Step 1 would cost O(N 4 ) space to store all the possible closed sets. Step 2 would cost O(N 2 ) space for the division tree. Step 3-10 would cost O(N 2 ) space for calling LCG solver.

In conclusion, the overall time complexity of Algorithm 5 is O(N 4 ) and the overall space complexity of Algorithm 5 is O(N 4 ). C RUNTIME AND THEORETICAL ANALYSIS

The number of vertices in the computation graph and the runtime of ACG Solver (Algorithm 5) for each network are listed in Table 2 . All the runtimes were measured on a single core of CPU i7-8700.

Notice that the runtime is measured on only 1 cpu core, it can be massively reduced by parallelization on multiple cpu cores. The runtime can also be further reduced through a better implementation as our implementation is a prototype.

Although it might be concerning that the runtime is too much for some deep networks, it is still relatively small compared to training processes which might cost days or even weeks. More importantly, solving the optimal solution for a network is an one-time effort. The optimal solutions for all popular networks will be released online for people to use without taking the time to run ACG solver.

To see how well the reality matches with the theory, we also compare the measured memory cut off and theoretical memory cut off (given by Algorithm 5) in Table 2 . Observe that all the measured memory cut off are slightly lower than theoretical memory cut off. This is because, in implementation, we assume that the whole input tensors of each operation are always stored for backward.

In reality, some operations only need to store small tensors for backward. For example, batchnormalization only needs a few statistics for backward and doesn't need the whole input tensor.

D VISUALIZATION:

We visualize the computation graph of Alexnet, vgg11, vgg13, vgg16 , vgg19 and CustomNet and the solution of our approach (in green) and the solution of Chen's approach (in red). In the computation graphs, the cost of each vertex and the actual operation of each edge are also marked. The cost of each vertex is the size of this tensor during forward given the input as [1, 3, 224, 224] ([1, 3, 300, 300] for inception v3). For example, in Alexnet, the input is [1, 3, 224, 224] and thus the source vertex has the cost 150528 = 1 × 3 × 224 × 224. After 2D convolution and relu, the tensor becomes [1, 64, 55, 55] and thus the second vertex has the cost 193600 = 1×64×55×55. Figure 10 : Endpoint vertices found on CustomNet, Selected vertices(green) of our approach, Chen's approach not applicable

:

More detailed complexity analysis is in the appendix due to space limitation

All proofs are in the appendix due to space limitation.

Figures and Tables:

Figure 1: Regular Training Approach vs. Re-forwarding (our). (a) The regular approach saves all tensors during forward, and uses these tensors to compute gradients during backward. (b) Reforwarding (our) saves a subset of tensors during the first forward, and conducts "Re-forward" to compute tensors for gradients during backward.

Figure 3: Closed Set Examples: (a) Closed set in a graph. there cannot exist a closed set between v 2 and v 4 because v 3 depends on v 1 . There can exist a closed set between v 1 and v 3 because v 2 doesn't depend on any other vertex. (b) Splittable Closed Set (Type 1). v 2 is the splitting vertex of s 13 . (c) Branched Closed Set (Type 2). (d) Non-branched Closed Set (Type 3).Definition 1. Closed Set: A set s containing vertices and edges is a closed set if and only if it satisfies the following three properties: 1. All the vertices of s have a common ancestor v i and a common descendent v j ; 2. Denote the vertex subset of s as V , edge subset as E, and the set of edges between two arbitrary vertices of V ∪ {v i , v j } is E , the edge from v i to v j (if exists) as e ij . E must either be E or E -{e ij }; 3. An arbitrary v 1 ∈ V doesn't have edge with another arbitrary v 2 / ∈ V ∪ {v i , v j }. For multiple valid closed sets between v i and v j , we denote the largest one ass ij Definition 2. [s ij ] = s ij ∪ {v i , v j }. [s ij ) = s ij ∪ {v i }. (s ij ] = s ij ∪ {v j }

Definition 6. Non-branched Closed Set (Type 3): A closed set s ij is non-branched if it has 0 splitting vertex and no branch: ∃s 1 ij s ij Among closed set with no splitting vertex, we categorize closed set with branches as Closed Set Type 2, and closed set without branches as Closed Set Type 3. The examples of different types of closed set are shown in Fig. 3.

Figure 4: Example divisions of three types of closed sets. Members of a division are colored differently. (a) Division of closed set type 1. The division is {[s 12 ], [s 23 ]} (b) Division of closed set type 2. The division is {[s 1 12 ], [s 2 12 ], [s 3 12 ]} (c) Division of closed set type 3. The division is {[s 12 ], [s 13 ], [s 23 ], [s 24 ], [s 34 ]} Definition 7. Maximal Split: {[s pq ]} is a maximal split of non-branched s ij if [s ij ] = ∪{[s pq ]} and ∀s ab , s cd ∈ {[s pq ]}, s ab ∩ s cd = ∅ and ∃{[s pq ]} {[s pq ]} such that ∪{[s pq ]} = [s kt ] [s ij ] Definition 8. Division of Closed Set:For type 1, its division is the linear segments separated by all its splitting vertices; for type 2, its division is all its branches, any of which cannot be divided into more branches; for type 3, its division is its maximal split.

set s = {v k }. v k ∈ s ij is a randomly chosen vertex.

Find the maximal split of a non-branched s ij with 0 splitting vertex 1: for each vertex pair

Figure 6: Endpoint vertices found on Alexnet

(see Table 1 "Chen et al. (2016) manual Training memory usage and time overhead of the regular, Chen et al. (2016), Chen et al. (2016) manual and Re-forwarding (ours) approach on linear and non-linear computation graph.

Table content: <table><tr><td>Linear network</td><td>Regular (MB)</td><td>Chen et al. (2016) (MB)</td><td>Re-forwarding (ours) (MB)</td><td>Memory Cut off (ours)</td><td>Regular Training Time (s)</td><td>Space Efficient Training Time (s)</td><td>Time Overhead</td></tr><tr><td colspan="2">Alexnet batch 1024 3550</td><td>3108</td><td>2620</td><td>26%</td><td>1.295</td><td>1.816</td><td>40%</td></tr><tr><td>Vgg11 batch 64</td><td>2976</td><td>2292</td><td>1802</td><td>39%</td><td>0.606</td><td>0.819</td><td>35%</td></tr><tr><td>Vgg13 batch 64</td><td>4152</td><td>2586</td><td>2586</td><td>38%</td><td>1.020</td><td>1.333</td><td>31%</td></tr><tr><td>Vgg16 batch 64</td><td>4470</td><td>2894</td><td>2586</td><td>42%</td><td>1.307</td><td>1.696</td><td>30%</td></tr><tr><td>Vgg19 batch 64</td><td>4788</td><td>2894</td><td>2502</td><td>48%</td><td>1.593</td><td>2.060</td><td>29%</td></tr><tr><td>Non-linear network</td><td>Regular (MB)</td><td>Chen et al. (2016) manual (MB)</td><td>Re-forwarding (ours) (MB)</td><td>Memory Cut off (ours)</td><td>Regular Training Time (s)</td><td>Space Efficient Training Time (s)</td><td>Time Overhead</td></tr><tr><td colspan="2">Resnet18 batch 256 5402</td><td>2898</td><td>2898</td><td>46%</td><td>1.144</td><td>1.599</td><td>40%</td></tr><tr><td colspan="2">Resnet34 batch 128 3900</td><td>1936</td><td>1544</td><td>60%</td><td>1.041</td><td>1.419</td><td>36%</td></tr><tr><td>Resnet50 batch 64</td><td>5206</td><td>2332</td><td>1798</td><td>65%</td><td>0.740</td><td>1.027</td><td>40%</td></tr><tr><td colspan="2">Resnet101 batch 32 3812</td><td>1216</td><td>970</td><td>75%</td><td>0.624</td><td>0.853</td><td>37%</td></tr><tr><td colspan="2">Resnet152 batch 16 2810</td><td>636</td><td>564</td><td>80%</td><td>0.450</td><td>0.628</td><td>39%</td></tr><tr><td colspan="2">Densenet121 batch 32 3984</td><td>1012</td><td>776</td><td>81%</td><td>0.558</td><td>0.789</td><td>42%</td></tr><tr><td colspan="2">Densenet161 batch 16 3658</td><td>744</td><td>616</td><td>83%</td><td>0.511</td><td>0.708</td><td>39%</td></tr><tr><td colspan="2">Densenet169 batch 32 4826</td><td>998</td><td>848</td><td>82%</td><td>0.714</td><td>1.022</td><td>43%</td></tr><tr><td colspan="2">Densenet201 batch 16 3164</td><td>600</td><td>582</td><td>82%</td><td>0.449</td><td>0.651</td><td>45%</td></tr><tr><td colspan="2">Inceptionv3 batch 32 2976</td><td>1026</td><td>910</td><td>69%</td><td>0.563</td><td>0.763</td><td>35%</td></tr><tr><td colspan="2">CustomNet batch 64 3233</td><td>Not Applicable</td><td>1353</td><td>58%</td><td>1.226</td><td>1.648</td><td>34%</td></tr></table>

Actual Runtime of ACG Solver and Theoretical Analysis Linear network Number of vertices Runtime (s) Measured Memory Cut off Theoretical Memory Cut off

Table content: <table><tr><td>Alexnet</td><td>12</td><td>0.03</td><td>26%</td><td>42%</td></tr><tr><td>Vgg11</td><td>17</td><td>0.09</td><td>39%</td><td>50%</td></tr><tr><td>Vgg13</td><td>19</td><td>0.15</td><td>38%</td><td>47%</td></tr><tr><td>Vgg16</td><td>22</td><td>0.26</td><td>42%</td><td>51%</td></tr><tr><td>Vgg19</td><td>25</td><td>0.44</td><td>48%</td><td>53%</td></tr><tr><td colspan="5">Non-linear network Number of vertices Runtime (s) Measured Memory Cut off Theoretical Memory Cut off</td></tr><tr><td>Resnet18</td><td>51</td><td>0.09</td><td>46%</td><td>63%</td></tr><tr><td>Resnet34</td><td>91</td><td>0.53</td><td>60%</td><td>73%</td></tr><tr><td>Resnet50</td><td>125</td><td>1.27</td><td>65%</td><td>75%</td></tr><tr><td>Resnet101</td><td>244</td><td>12.40</td><td>75%</td><td>81%</td></tr><tr><td>Resnet152</td><td>363</td><td>59.34</td><td>80%</td><td>84%</td></tr><tr><td>Densenet121</td><td>306</td><td>293.75</td><td>81%</td><td>81%</td></tr><tr><td>Densenet161</td><td>406</td><td>1537.82</td><td>83%</td><td>84%</td></tr><tr><td>Densenet169</td><td>426</td><td>2356.47</td><td>82%</td><td>84%</td></tr><tr><td>Densenet201</td><td>506</td><td>7335.35</td><td>82%</td><td>86%</td></tr><tr><td>Inceptionv3</td><td>219</td><td>7.68</td><td>69%</td><td>76%</td></tr><tr><td>Custom</td><td>35</td><td>0.05</td><td>58%</td><td>68%</td></tr></table>

Peer review (iclr2019):
 I took a look at the revision. I am glad to see that the authors clarified the meaning of "optimality" and added time complexity for each algorithm. The complexities of the algorithms do not seem great (3rd or 4th order polynomial of N) as they appear to be checking things exhaustively, but perhaps they are not a big issue since the decomposition algorithm is run only once and usually N is not huge. I wonder if more efficient algorithms exist. It is also nice to see that the time overhead for one training step is not huge (last column of Table 1). While I still think it is better to see more complete training curves, the provided time overhead is a proxy. I hope the authors can further improve the algorithm description, for example, the pseudo-code of Algorithm 1 is very verbal/ambiguous, and it would be better to have more implementation-friendly pseudo-code. Despite the above-mentioned flaws, I think this work is still valuable in handling the memory consumption of arbitrary computation graph in a principled manner. ============================================================================= This paper presents a method for reducing the memory cost of training DNNs. The main idea is to divide the computational graph into smaller components (close sets), such that the dependency between components is low, and so one can store only tensors at key vertices in each component during the forward pass. In the backward pass, one needs to re-forward the data within each close set for computing the gradient. The idea is quite intuitive and the example in linear computational graph in Section 3 clearly demonstrates the basic idea. The main technical development is in the case of arbitrary computation graph (Section 4), where the authors explain how to divide the computational graph into different types of close sets. I have not read the proofs in appendix, but it is clear that the technical development is mainly in discrete math rather than machine learning. For this part, my suggestions are: 1. give better algorithm description: what are the inputs and outputs of each algorithm, for the examples in Figure 2-4 (and perhaps with costs for each vertex), which vertices do the final algorithm decide to store 2. analyze the complexity of each algorithm 3. provide clear definition of "optimality" and its proof: the authors mentioned in a few places about the method being "optimal", but the paper needs to be crystal clear about "for what problem is the method optimal", "optimal in what sense (the objective)". The more concerning part is the experimental evaluation. While I believe that the proposed method can save memory cost, there is no result on the additional time cost for re-forwarding. Therefore, it is not clear if it is worth the extra complication of using re-forwarding in an end-to-end sense: the authors motivated in Section 1 that saving memory cost allows one to use larger mini-batches, and an important benefit of large mini-batch is to accelerate training, see, e.g., Goyal et al. Accurate, large minibatch SGD: Training imagenet in 1 hour. Hoffer et al. Train longer, generalize better: Closing the generalization gap in large batch training of neural networks. It is less desirable to use re-forwarding it if it causes significant latency in the backward pass and training. An good demonstration of the memory vs. time tradeoff would be the training curves of loss vs. running time. The other detail I would appreciate is how the authors implemented the algorithm, and whether the proposed method is likely to be deployed on popular ML frameworks, which would improve the significance of the work.