Structured in Space, Randomized in Time: Leveraging Dropout in RNNs for Efficient Training

1 Introduction . Recurrent Neural Networks ( RNNs ) are an important class of machine learning approaches for analyzing sequential data including language , speech , time-series data , etc . Given an input training sequence , RNNs cells pass information from the previous time-step to the current time-step , in the form of a hidden-state , which helps the network to learn inherent temporal constructs . To further enhance the ability to learn from past information , Long Short-Term Memory ( LSTM ) [ Hochreiter and Schmidhuber , 1997 ] cells have been proposed , which consist of multiple internal gates and a cell-state in addition to the hidden-state . Such LSTM cell-based RNNs have been successfully deployed in a variety of application domains , e.g. , language modelling [ Mikolov and Zweig , 2012 ] , automatic speech recognition [ Hannun et al. , 2014 , Amodei et al. , 2016 , Graves et al. , 2013 ] , sentiment analysis [ Ma et al. , 2018 ] , and machine translation [ Luong et al. , 2015 ] . However , training such RNNs is a time- and compute-intensive task . This is because , for each input sequence , a 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . given RNN is logically unrolled into a sequence equal to the sequence length , and computation is performed following the dependencies . Once the training loss is calculated , it is back-propagated over the sequence in the reverse direction to compute the weight gradients , which is known as back-propagation through time ( BPTT ) [ Werbos , 1990 ] . One of the promising ways of reducing the effective amount of computation is to exploit sparsity present in data structures . Sparsity refers to the fraction of zero-valued operands to the total size of the operands involved in a computation step . Thus , efficiently skipping zero-valued operands from computation leads to faster training time and energy efficiency . Convolutional Neural Networks ( CNNs ) , owing to their prominent use of ReLU [ Nair and Hinton , 2010 ] as the activation function , generate significant dynamic sparsity associated with the activation and neuron gradient values , which can be leveraged during training . However , due to the non-ReLU activation functions , e.g. , tanh and sigmoid , dynamic sparsity in LSTM based networks is not inherent in nature . Such activation functions do not directly induce zero values , and hence the opportunity of exploiting sparsity is only limited to thresholding and approximation-based approaches [ Gupta et al. , 2019 ] . Although such techniques have been shown to work relatively well during the inference phase ( after the network is fully trained ) , training still relies on high precision arithmetic operations . Also , the derivatives of the tanh and sigmoid functions do not lend themselves to sparsity as in ReLU . For the same reasons , the back-propagated gradient values are not sparse either , thus ruling out the opportunity for sparsity exploitation during the backward pass . Motivated by these observations , in this work , we explore the opportunity of using the dropoutinduced zero values as a source of sparsity in LSTMs . Dropout , a widely used regularizing technique , works by randomly dropping activation values in the forward pass . Since the dropped activation map becomes sparse in nature , we can take advantage of this sparsity towards achieving computation speedup . However , such dropout masks are sampled randomly , which creates bottlenecks in accelerating the resultant sparse computation in a general-purpose SIMD hardware owing to the storage and memory access overheads . Therefore , we propose to structure such dropout patterns which makes it amenable for hardware acceleration . We distinguish between sparsity types , namely , input and output sparsity , depending on whether the input operands are zero-valued or an output operand is going to be zero . Since the dropout mask can be sampled ahead of time , exploiting output sparsity may seem straightforward . However , applying it naively to the hidden state can result in unintended consequences . This is because the sparsity also propagates to the cell states due to the inherent computational dependencies , resulting in possible loss of functional performance ( e.g . accuracy ) . Further , we carefully analyze the computational flow in the backward pass of LSTM to identify sparse operands . We also extend the sparsity technique to SIMD GPU-based matrix-multiplication using shared memory , and show that that our proposed approach is exploitable in existing software based techniques . Our proposed sparsity pattern is also well-suited to be leveraged in systolic array-based computations , resulting in energy-efficiency benefits of dense matrix operations combined with high-performance sparse operations . In summary , this paper makes the following key contributions : • We outline a generic framework for applying dropout in the context of LSTMs , considering uniformity of dropout patterns within a batch and across time steps . In addition to commonly used non-recurrent ( NR ) direction , we also extend dropout structural pattern to recurrent ( NR+RH ) directions to maximize the scope for potential training speedup . • We provide detailed opportunities for exploiting the uniform sparsity in forward as well as backward passes of LSTMs based on sparsity propagation during training . We also identify the sparsity types based on input/output sparsity , and leverage it towards faster computation operation . • We conduct extensive experiments with three representative applications including Language Modelling , Machine Translation , and Named Entity Recognition to evaluate the general applicability and performance benefits of our proposed approach . Our results show that uniform dropout pattern within a batch and random in time steps achieves significant training speedup ranging from 1.23x-1.64x , while not compromising on the target accuracy metric . 2 Related Work . Dropout and DropConnect . Several prior works study the scope of Dropout towards reducing model over-fitting and achieve regularization . Hinton et al . [ 2012 ] first proposed the idea of dropout , which randomly removes the neurons in a given layer of a feed-forward neural network for a given probability p. Wan et al . [ 2013 ] on the other hand proposed DropConnect , which , instead of neurons , randomly drops the weight connections between two layers . Pham et al . [ 2014 ] and Zaremba et al . [ 2014 ] applied Dropout towards LSTM RNN regularization , where the dropout operator is specifically applied on the non-recurrent direction in LSTM . Kingma et al . [ 2015 ] and Gal and Ghahramani [ 2016 ] explored dropout strategies based on Bayesian Approximation . Gal and Ghahramani [ 2016 ] in particular also extended Dropout to the recurrent direction of LSTMs . Moon et al . [ 2015 ] have proposed using dropout on the LSTM cell-state instead of the hidden-states ; however , their approach is application-specific towards automatic speech recognition ( ASR ) tasks and doesn ’ t discuss generality . Semeniuta et al . [ 2016 ] investigated a variant of recurrent dropout that does not cause a loss of long term memory by letting the dropout mask operate directly on one of the internal state vectors . Krueger et al . [ 2016 ] proposed to regularize RNNs by randomly preserving hidden units from the previous time-step , instead of dropping them out . Merity et al . [ 2017 ] used DropConnect along hidden-to-hidden weight matrices as a form of recurrent regularization for LSTMs , while using Dropout on the non-recurrent dimension . However , all of these prior works focus on achieving better regularization to avoid over-fitting while our approach seeks to ( a ) act as a simple , plug-in replacement in existing LSTM applications without requiring to alter other hyper-parameter settings , and ( b ) reduce the overall training time by structuring the dropout pattern itself , while achieving the same degree of regularization effect . Weight Sparsity in Deep Neural Networks . Sparsity in Deep Neural Nets has also been extensively explored , and the work in this area can be categorized as unstructured or structured types . In particular , unstructured pruning approaches [ Han et al. , 2015a , b , Dai et al. , 2019 , Mao et al. , 2017 , Narang et al. , 2017a ] result in random sparsity in weight matrices , which is difficult to accelerate on general-purpose hardware due to storage and memory access overheads . This has motivated structured sparsity-based approaches for both CNNs [ Yu et al. , 2017 , Wen et al. , 2016 , Narang et al. , 2017a ] and RNNs [ Lu et al. , 2016 , Wen et al. , 2017 , Narang et al. , 2017b ] . However , these works mainly aim to improve model size and speedup during model inference ; in contrast , our work is aimed at achieving speedup during the training phase . We also believe our approach can be adapted in structured pruning approaches as well , towards faster training on general-purpose hardware . 3 A General Framework of Using Dropout for Efficient LSTM Training . In this section , we introduce a generic dropout framework to distinguish different possibilities of dropout application followed by the scope of the dropout-induced sparsity throughout different phases of training . We first highlight the computational dynamics of an LSTM cell [ Hochreiter and Schmidhuber , 1997 ] , which is a variant of RNN , consisting of four different types of gates : input gate : it = sigmoid ( i∗t ) , i ∗ t = h l−1 t Wi + h l t−1Ui + bi ( 1 ) forget gate : ft = sigmoid ( f∗t ) , f ∗ t = h l−1 t Wf + h l t−1Uf + bf ( 2 ) output gate : ot = sigmoid ( o∗t ) , o ∗ t = h l−1 t Wo + h l t−1Uo + bo ( 3 ) input modulation gate : gt = tanh ( g∗t ) , g ∗ t = h l−1 t Wc + h l t−1Uc + bc ( 4 ) where hlt notation refers to hidden-state from layer l and time step t. Based on the above four internal gate states , the following two output states are computed : cell state : clt = ft ct−1 + it gt ( 5 ) hidden state : hlt = ot tanh ( clt ) ( 6 ) Equations ( 1 ) through ( 4 ) form the core of an LSTM cell and are also the computationally intensive part of the operation , during both the forward and backward passes of training . Parameters { W∗ , U∗ , b∗ } constitute the weight of the network , where W∗ and U∗ refer to the input-to-hidden and hidden-to-hidden weight projection matrices , respectively , and b∗ is the bias . Each weight matrix is given by size H ×H , where H is the LSTM hidden size . 3.1 Reframing the Scope of Dropout Structuring in LSTMs . Considering an LSTM cell with hidden state dimension H and an input batch size of B , at each time step t , the cell generates a hidden-state matrix ( ht ) and cell-state matrix ( ct ) , each of size B ×H . Each row of this matrix corresponds to the feature vector of the corresponding input . The hidden state ( ht ) is used as recurrent input to the next time-step t+ 1 , and also as input to the subsequent LSTM cell/layer as a non-recurrent input . The idea of regularizing recurrent neural networks by using dropouts has initially been introduced by Zaremba et al . [ 2014 ] . The two key aspects of the said work were – ( a ) dropouts were only used along the non-recurrent direction , and ( b ) at different time steps , it was sampled using different dropout masks making it time-variant . Subsequent work from Gal and Ghahramani [ 2016 ] has shown that dropout can be extended to the recurrent direction as well , thereby increasing network regularization scope . Similar to the work of Zaremba et al . [ 2014 ] , the dropout mask itself was random . However , unlike Zaremba et al . [ 2014 ] , the sampled dropout masks remained the same across time steps . Given this possible choice of application of dropout , we propose a unifying framework considering the two dimensions of dropout pattern – dropout within a batch and dropout across time steps . Each dimension in turn offers two possible choices , resulting in four possible combinations , as summarized in Fig . 1 ( a ) and illustrated in Fig . 1 ( b ) . Case-I refers to the scenario where dropout is random within a batch and also different across time-steps . This refers to the most common application of dropout including Zaremba et al . [ 2014 ] ’ s method . Case-II on the other hand refers to the scenario in which dropout pattern is random within a batch , but the same pattern is repeated across time steps . This approach is similar to Gal and Ghahramani [ 2016 ] ’ s work and has also been adopted later by Merity et al . [ 2017 ] in AWD-LSTM . Case-III refers to a dropout pattern that is structured within a batch , but varies across time steps . Finally , Case-IV is the most restricted form of dropout , where the pattern is not only structured within a batch , but also remains the same across time-steps . These different dropout patterns exhibit different trade-offs . For example , while a mask , which is structured in a batch ( B ) and same across time-steps ( T ) , is most amenable for hardware acceleration due to least metadata overhead , it also offers the least variability to the training model underneath . In this work , we stick to Case-III , where a structure in dropout mask within batch allows us to take advantage of computation speed-up , while variation across time-steps injects sufficient randomness into the training process , helping to retain the functional accuracy . Also , at a given time step , it is possible to view dropout being applied to only non-recurrent dimension ( NR ) or both nonrecurrent as well as recurrent-hidden ( NR+RH ) dimensions . Therefore , we can further distinguish those cases as NR+ST or NR+RH+ST , and so on . In this study , we explore both NR+ST and NR+RH+ST configurations under Case-III towards maximizing potential training speedup , while retaining functional accuracy .

Summary: The authors present a generic framework to applying dropout as a sparsity technique to LSTM to reduce training time and as a plug-in technique w/o hyperparameters adjustment when comparing to previous similar techniques that focus on regularization. The dropout technique is randomized in timesteps but structured in the batch as a dropout mask that saves computation time on the dropout mask (does not compute), as well as inducing the regularization effect. This technique is applied to both the non-recurrent part of LSTM and the recurrent-hidden layer.