An Automatic Operation Batching Strategy for the Backward Propagation of Neural Networks Having Dynamic Computation Graphs

1 INTRODUCTION . Recent years , neural networks ( NN ) have been applied and shown their great effects on a lot of Natural Language Processing ( NLP ) topics like Sentiment Classification ( Socher et al. , 2013 ; Tai et al. , 2015 ) , Named Entity Recognition ( NER ) ( Lample et al. , 2016 ) , Machine Translation ( MT ) ( Sutskever et al. , 2014 ; Wu et al. , 2016 ; Eriguchi et al. , 2016 ) , Question Answering ( QA ) ( Kumar et al. , 2016 ) and so on . It has also become crucial to improving the speed of training neural networks for the research and development of deep learning models and applications since the neural networks as well as the training sample datasets become larger and larger , which leads to days of training time . Taking Machine Translation as an example topic , Table . 1 summarized both the accuracy performance and computing performance of each state-of-the-art work on WMT 2014 English-to-French Machine Translation task since 2014 . We can see that the training phase of machine translation is time-consuming and it takes at least days to finish even accelerated by multiple GPUs . Therefore , it ’ s becoming more and more important to take the computing performance of training neural networks into consideration . In order to maximize the use of computing resources provided by modern parallel processors such as multi-core CPUs and many-core GPUs , batching is widely used in the existing implementations of neural networks for deep learning . This is because the operations batched together can be evaluated simultaneously on modern CPUs and GPUs and the evaluation is much faster than evaluating the operations one by one . For example , programmers can execute a single matrix-matrix multiplication using an optimized BLAS call instead of multiplying tens or hundreds of different vectors by the same matrix one by one and the former implementation is much faster on both CPUs and GPUs . Usually , it ’ s the programmers who implement the neural networks in charge of crafting efficient batching by hand . This is easy in some cases where the computation graphs of neural networks are stable and the inputs and outputs are represented as fixed sized tensors ( e.g. , training LeNet-5 ( LeCun et al. , 1998 ) with MNIST dataset in which all the images have the same fixed size ) . In these cases , the programmers write code for the computation of processing an individual training sample and add another leading dimension to the tensors representing the inputs and outputs so that multiple tensors can be packed into and represented as only one . Then the program will process the new tensor which actually contains the data from multiple inputs as an individual one using the high performance computing technology like SIMD instructions or BLAS library to execute it efficiently in parallel ( the correctness is guaranteed by the batch learning which allows multiple inputs to be processed using the same parameters ) . However , the batching tasks tend to become difficult for programmers to finish it manually when the inputs have different structures or the computation graphs of the neural networks is unstable during the training phase . Neural networks having dynamic computation graphs are common in models or application for NLP . For example , the inputs of neural networks based NLP model are always sentences represented as word sequences with different lengths or parse trees with different structure ( Socher et al. , 2011 ) , and the corresponding unfolded computation graphs of recurrent or recursive neural networks ( Elman , 1990 ) vary when processing inputs with different structures . In some cases where the inputs are sequences with different lengths , programmers can pad the shorter ones with the same special symbols so that all the sequences have the same lengths . This approach , which is called padding , is widely used in existing neural network based NLP models and helps with utilizing batching in the execution . However , the strategy padding strategy still has limitations : the computation on the padded symbol is considerable but totally a waste ( Qiao et al. , 2018 ) and it is also difficult to apply this strategy on neural networks with more complex architectures , like Tree-LSTM ( Tai et al. , 2015 ) . Indeed , many implementations of neural networks having dynamic computation graphs choose to avoid batching data and operations entirely or just use single instance training . In order to apply more efficient batching without programmer ’ s manual work , researchers tried to develop automatic batching strategies that can instruct the deep learning frameworks or libraries to apply the operation batching by itself . The programmer should only write the code that defines the computation of processing one individual instance and leave the details of operation batching to the frameworks or libraries . Google proposed and implemented their automatic operation batching strategy in TensorFlow Fold ( Looks et al. , 2017 ) . Another strategy called Agenda-based batching was also proposed in Neubig et al . ( 2017b ) and implemented using DyNet toolkit ( Neubig et al. , 2017a ) . Both of these operation auto-batching strategies bring huge gains in computing performance of neural networks having dynamic computation graphs . Unfortunately , there are still limitations in both of these strategies . In the training phase , the autobatching strategy generates an execution order of the operations in the computation graphs and this execution order determines which operations can be batched and executed on the same step . Usually , the execution order is generated according to the dependency between the operations in the forward propagation of computation graph . In the backward propagation , the framework just executes the operations ( calculating the derivative of each node ) using a reverse of that generated execution order . This scheme leads to missing some batching opportunities in the backward propagation . For example , there are some operations with the same type share the same parameters but can ’ t be batched in the forward propagation . But the calculations of the derivatives from the results of those operations to the shared parameters can be executed together . If we are able to utilize this kind of operation batching opportunity in the backward propagation , we can achieve further improvement of the training speed of neural networks having dynamic computation graphs . In this paper , we propose our automatic operation batching strategy for the backward propagation of neural networks having dynamic computation graphs . This strategy explores and utilizes the operation batching opportunities described as above in the backward propagation . It also helps to execute those new batched operations with efficient memory accesses . The combination of our strategy and the existing ones can provide further improvement in the speed of the training phase . 2 BACKGROUND . 2.1 OPERATION BATCHING . Operation batching is the most common way to enable parallelism in deep learning . It groups the same operations from multiple training samples together and executes them simultaneously . Operations from the same training can also be batched if they have the same type and no dependency on each other . Since modern parallel processors ( CPUs and GPUs ) have powerful SIMD processing instructions , operation batching often brings large gains in the computation efficiency . Figure 1 illustrates the basic concept of operation batching . The function y = tanh ( Wx + b ) is called three times for the independent input vector x1 , x2 , x3 . For efficient execution , we can concatenate the three vectors into a matrix X and broadcast the bias vector b into a matrix B , then execute the affine transformation as Y = tanh ( WX + B ) . The result Y consists of three column vectors y1 , y2 , y3 , which are the corresponding results of the function calls for x1 , x2 , x3 . In this case , the operation batching brings benefits for computation efficiency from three aspects : First , the matrix W is loaded from the lower level memory to the upper-level memory for only once , which results in a better locality to the memory access behavior ; Second , the element-wise activation function tanh is called only once , which asks for fewer times of computation kernel launch ; Third , the matrix-matrix product between W and X can be implemented using an optimized BLAS library , which makes efficient use of hardware resources . 2.2 MANUAL OPERATION BATCHING . It is not easy to apply the operation batching on the neural networks having dynamic computation graphs . Figure 2 illustrates the batch learning using a basic vanilla Recurrent Neural Network ( RNN ) which takes a sequence of word embedding vectors as the input . The final hidden state of this RNN is used to make a prediction and calculate the training loss . In the example shown in the left part , a computation graph of mini-batch learning is shown . The batch size is 3 and each training instance has a different length . An ideal set of batched operations for this given example is ( Op11 , Op 2 1 , Op31 ) , ( Op 1 2 , Op 2 2 , Op 3 2 ) , ( Op 1 3 , Op 3 3 ) , ( Op 1 4 ) , ( L 1 , L2 , L3 ) , where Opit represents the operations of RNN cell at the t th step in the i th instance . The operations in the same batch share the same operation type and the same size of inputs and outputs , meanwhile , they don ’ t have dependencies on each other . After batching operations like that , we can execute them in an efficient way as described above . Nevertheless , it ’ s really difficult to implement an ideal batching for every mini-batch since the length distributions of sequences vary in each mini-batch since we can ’ t count on the programmers to manually implement that with hard code for each one . Besides the operation batching , the multiprocessing , which means distributing training instances from a mini-batch to multiple threads or processes then synchronizing after processing each mini-batch , is another way to enable parallel processing . The multi-processing is easy to implement but unfortunately , its execution is generally inefficient since the similar operations across the training instances are performed separately and sequentially . The benefit of operation batching , which helps to process training instances in parallel and taking advantage of data-parallelism in the training phase , is still attractive . One most common way to make use of the operation batching for training RNN-based models is to pad the input sequences with dummy elements so that all the training instances from a mini-batch have the same length . The right part of Figure 2 illustrates the corresponding computation graph after padding the training instances . In the practical usage of deep learning libraries or frameworks , it is always required to pad all the instances of training dataset so that they have the same length . After the padding , the computation graph of each mini-batch has the same structure and can be treated as a static one . The programmers can write the code for manual operation batching once and repeatedly use it to apply operation batching on every mini-batch . In order to keep the correctness of execution , the intermediate results and resulting loss from the computation at the dummy elements are masked out using a mask matrix . Even though the padding strategy makes it possible to apply operation batching for the input sequences with different lengths from a mini-batch , it still has three disadvantages : First , the programmers are still in charge of implementing padding and masking , which increases the possibility of inefficient implementation and presenting bugs ; Second , the computation on the dummy elements is totally unnecessary leads to the waste of computing resource ; Third , for more complex neural networks having dynamic computation graphs like Tree-LSTM , it ’ s almost impossible to apply the padding strategy on each training instance .

Summary: This work describes a shortcoming in existing dynamic batching strategies, namely that they operate only on the forward pass while some operations can be batched only in the backward pass. For example, the gradient of the transition matrix in a RNN consists of the sum of partial derivatives over each time step; the terms of this sum and the summation can be batched into a single matrix multiplication. The authors implement this batching strategy in DyNet and show empirically that it can lead to decent (0-25%) speedups.