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: This work is mainly concerned with the time cost of LSTM computation. Inspired by the sparsity of neural networks, the authors propose to structure dropout patterns by discarding the same collections of neural units within a training batch. They have also carefully studied how dropout-induced sparsity propagated through different training phrases and how it was utilized in each of them. The experiments are very solid, with the investigations on LM, NMT, and NER.