Reducing Computation in Recurrent Networks by Selectively Updating State Neurons

1 INTRODUCTION . Recurrent Neural Networks ( RNN ) are the state-of-the-art approach to many sequential learning problems including speech recognition ( Graves et al. , 2013 ) , machine translation ( Bahdanau et al. , 2015 ) , and sequence generation ( Graves , 2013 ; Xu et al. , 2015 ) . However , RNNs typically rely on computationally-taxing updates to their entire hidden state at each timestep , a cost that grows with hidden state size . As demonstrated by the success of gating mechanisms such as the GRU ( Cho et al. , 2014 ) and LSTM ( Hochreiter & Schmidhuber , 1997 ) , all dimensions rarely need to be re-computed from scratch at each timestep . By discretely selecting which dimensions to update at each timestep via a learned update pattern , RNNs with a large hidden state can be trained with lower computational requirements ( Bengio et al. , 2013 ) , inference in long RNNs can be expedited ( Campos et al. , 2018 ) , and hidden representations may be made more robust to misleading inputs such as outliers or noise . Selective neuron activation in RNNs has recently gained attention in the literature ( Koutnik et al. , 2014 ; Neil et al. , 2016 ; Shen et al. , 2019 ; Jernite et al. , 2017 ; Campos et al. , 2018 ) . The most popular methods hand-craft specific update patterns , dictating which dimensions of the hidden state will update at which timesteps according to prior knowledge of a task ( Koutnik et al. , 2014 ; Neil et al. , 2016 ) . This imposes undue challenges in implementation , limits extensibility , and ignores the data-driven curation of information-flow through the RNN , a signature property of recurrent memory cells ( Hochreiter & Schmidhuber , 1997 ; Cho et al. , 2014 ) . More recent methods learn to react to input data but impose strict relationships between the update patterns across both hidden dimensions and time ( Shen et al. , 2019 ; Jernite et al. , 2017 ; Campos et al. , 2018 ) . While applicable to tasks with clear hierarchical components , such as modeling character-level text ( Chung et al. , 2017 ) , these assumptions limit the expressiveness of the learned update patterns . Specifically , we study the problem of generating a binary update-pattern for the hidden states learned by an RNN . The learned update-pattern defines which dimensions of the hidden state to update at each timestep , similar to the motivation for Residual Networks ( He et al. , 2016 ; Wang & Tian , 2016 ) and Highway Networks ( Srivastava et al. , 2015 ; Zilly et al. , 2017 ) . Ideally , only a small subset of the hidden state ’ s dimensions needs to be updated at each timestep , especially with high-dimensional hidden states . In this way , representations can be learned while both solving a sequential learning task and minimizing the number of updates . This results in a reduction of the compute time . A solution to this multi-objective optimization problem should have a comparable accuracy to a traditional constantly-updating RNN but save the majority of computation steps along the way , ultimately accelerating inference and training ( Neil et al. , 2016 ) . Despite the potential for reducing computation required by RNNs , learning said update patterns is a challenging problem . First , binary-output neurons making discrete decisions ( whether or not to update a hidden state dimension , for example ) in the interior of a neural network is a classic challenge to gradient-based learning . This is because such decisions are non-differentiable by nature and therefore backpropagation can not be directly used to update the weights . Second , the quality of a learned update pattern is unsupervised and thus the only feedback is task-specific . This discourages a priori assumptions of the update patterns . To address the aforementioned challenges , we propose the selective activation RNN , or SA-RNN , which parameterizes a distribution of update-likelihoods , one per hidden state dimension , from which update-decisions can be made at each timestep . We augment an RNN with an update coordinator that adaptively controls which coordinate directions to update in the hidden state on the fly . The coordinator is modeled as a lightweight neural network that observes incoming data at each timestep and makes a discrete decision of whether or not enough information is stored in each individual hidden dimension to warrant an update . Subsequently , each hidden dimension is either computed by the RNN or copied from the previous timestep . The coordinator ’ s architecture is kept as simple as possible so the complexity of the RNN can scale without simply outsourcing computation to another network , similar to the controller in Ha & Schmidhuber ( 2018 ) . Most notably , in contrast to other recent approaches ( Koutnik et al. , 2014 ; Jernite et al. , 2017 ; Neil et al. , 2016 ; Campos et al. , 2018 ; Shen et al. , 2019 ; Liu et al. , 2018 ) we impose no assumptions of which individual hidden dimensions should update ( or not update ) together . Instead , we show that using an entirely-learned approach still results in complex task-specific update patterns . On three publicly-available datasets , we show that our low-bias approach achieves higher accuracy with far fewer updates than recent state-of-the-art methods ( Koutnik et al. , 2014 ; Jernite et al. , 2017 ; Neil et al. , 2016 ; Campos et al. , 2018 ) . These results indicate that predicting RNN update-patterns solely with respect to a task is not only feasible and low-bias , but is also favorable in a variety of settings . 2 RELATED WORK . Recurrent neuron update patterns have gained much interest in recent literature ( Koutnik et al. , 2014 ; Jernite et al. , 2017 ; Neil et al. , 2016 ; Chung et al. , 2017 ; Shen et al. , 2019 ; Liu et al. , 2018 ) . All of these methods boast fewer updates to the hidden states than standard RNN architectures . However , there are several limitations of these methods , two of which are summarized as follows . First , the most popular methods rely on extensively-handcrafted update patterns consisting of periodic neuron activations ( Koutnik et al. , 2014 ; Neil et al. , 2016 ; Liu et al. , 2018 ) . This requires either prior knowledge of sampling frequencies or seasonal patterns present in the data , reducing the potential extension to many sequential learning problems . Additionally , these input-agnostic periodic updates are fixed prior to learning . The choice of update periods heavily impacts the performance of the model , and sequences with irregular information flow can not be modeled without massive state representations . Second , the most recent works allow for data-reactive update patterns ( Jernite et al. , 2017 ; Shen et al. , 2019 ; Campos et al. , 2018 ) but assume temporal hierarchies in the input sequences and study settings where this effect is exceedingly obvious ( for example , character-level sentence modeling ( Chung et al. , 2017 ) ) . In many real-world settings , temporal hierarchies are often subtle and forcing this assumption into the architectural design may limit its applications . Additionally , our approach is related to conditional computation , which predicts subsets of neural networks to activate depending on input data ( Bengio et al. , 2015 ; Shazeer et al. , 2017 ; Cheng et al. , 2017 ) . In many cases , when a particular concept can be represented using only the sub-network of a large neural network , computation can be preserved by learning the structure of said sub-network ( Schmidhuber , 2012 ) and activating it accordingly . 3 SELECTIVE NEURON ACTIVATION FOR RNNS . We introduce the Selective-Activation RNN , or SA-RNN , a broadly-applicable augmentation to RNNs which minimizes the computation required for RNNs by facilitating unimpeded informationflow across timesteps for individual dimensions of the hidden state . At its core , SA-RNN learns a data-driven strategy for discretely reading and writing information to the latent state space through the learned parameterization of an update-likelihood distribution . Despite leaving hidden dimension update patterns independent from one another , complex strategies still arise naturally depending on the sequential learning task at hand . In this section , we describe the training process of SA-RNN with D-dimensional hidden states on sequences of length T for input data x with V variables . We omit biases from affine transformation equations and use notation for one training instance for ease of readability . An overview of the forward pass through SA-RNN is shown in Figure 1 . 3.1 COMPUTING HIDDEN STATES . RNNs compute a sequence of hidden states one timestep at a time ( Elman , 1990 ) , each computed by a parametric recurrence function R ( · ) : ht = R ( ht−1 , xt|θr ) . The result is a sequence of vector representations H = { h1 , . . . , hT } where each ht ∈ RD represents temporal dynamics of the time series up to timestep t with respect to a task , preserving not only temporal dependencies but also the ordering of the inputs . A popular and powerful augmentation to the RNN , as it was originally proposed , is the Gated Recurrent Unit ( GRU ) ( Cho et al. , 2014 ) , which adds a series of gates between ht−1 and ht to alleviate the vanishing gradient problem ( Bengio et al. , 1994 ) : rt = σ ( Wrht−1 + Urxt ) ( 1 ) zt = σ ( Wzht−1 + Uzxt ) ( 2 ) st = φ ( Wcxt + Uc ( rt ht−1 ) ) ( 3 ) h̃t = ( 1− zt ) ht−1 + zt st ( 4 ) where W s and Us are matrices of learnable parameters of shape D ×D and D × V respectively , xt ∈ RV is the input data at timestep t , represents the element-wise multiplication , σ represents the sigmoid function , and φ represents a non-linearity ( traditionally the hyperbolic tangent function ) . Its design is motivated heavily by the LSTM ( Hochreiter & Schmidhuber , 1997 ) . The GRU performs soft read/write operations , recomputing the entire vector ht at each timestep since gate z ∈ [ 0 , 1 ] D , the space of vectors with values inclusively between 0 and 1 . Instead , we propose that all dimensions do not need to be updated at each timestep , as the position of the hidden state in many dimensions may often encode enough of the modeled input . Note that the output of the recurrence function is referred to as h̃t . In the next section , we describe how to compute ht , the final hidden state for timestep t which is subsequently used for computing ht+1 or the task . 3.2 SELECTIVE NEURON ACTIVATION . To reduce computation required to generate state representations , we assume updating representations to be a sequence of binary decisions – at each timestep a neuron will either be updated or not . Thus , we propose a learned update coordinator , which generates a binary mask for each hidden dimension , forecasting which dimensions need to be updated at the next timestep . First , an update-likelihood ũt is computed for each neuron , informed by both the data observed at the current timestep and the previous update-likelihoods : ũt = σ ( Wuht−1 +Wixt ) where Wu ∈ RD×D is a diagonal matrix of trainable parameters which dictate the linear relationship between ht−1 and ũt . Wu is kept diagonal to maintain relationships between update-decisions of the same dimension while avoiding the extensive computation of a fully-connected layer , similar to the hidden state decay in Che et al . ( 2018 ) . Wi ∈ RD×V encodes the influence of the input data on the current update-likelihood and σ ( · ) represents the hard sigmoid function , bounding ũ according to a slope α . Thus ũt ∈ [ 0 , 1 ] D , with one update-likelihood per dimension of the hidden state . To discretize ũt , allowing information to flow unimpeded , element-wise binarization is applied : ut = binarize ( ũt ) , where ( 5 ) binarize ( a ) = { 1 if a > 0.5 , 0 otherwise . ( 6 ) We apply this final discrete update decision as a binary gating mechanism since ut ∈ { 0 , 1 } D : ht = ut h̃t + ( 1− ut ) ht−1 ( 7 ) As written , this equation requires the pre-computation of h̃t . However , through masking the computation can be directed at only the needed updates upon calculation of ut . Thus when ũnt , the update decision for the n-th dimension in h , is 1 , hnt is updated according to the new information present in h̃nt . We note that this update-decision strategy does not impose the inter-neuron assumptions of Jernite et al . ( 2017 ) ; Shen et al . ( 2019 ) ; Koutnik et al . ( 2014 ) while still allowing such strategies to be learned if they are found to be optimal by the model since decisions are made with respect to previous decisions , similar to Campos et al . ( 2018 ) . We hypothesize that updating neurons together may generally be beneficial since complex temporal dependencies often require representations evolving in blocks of multiple neurons , as discussed in Koutnik et al . ( 2014 ) . Since binary-output neurons are inherently non-differentiable , barring the direct use of backpropagation , we approximate the gradient of the binarization function using the straight-through gradient estimator ( Bengio et al. , 2013 ) trained with slope-annealing ( Chung et al. , 2017 ) : ∂binarize ( x ) ∂x = 1 . ( 8 ) By estimating the gradient in this way we avoid additional loss terms and end up with empiricallyreasonable approximations in comparison to other high-variance methods , such as REINFORCE ( Williams , 1992 ; Chung et al. , 2017 ; Campos et al. , 2018 ) . After computing the sequence of state representations H , they are projected into the output space depending on the task at hand .

Summary: This paper proposes selective activation RNN (SA-RNN), by using an update coordinator to determine which subset of the RNN’s hidden state dimensions should be updated at a given timestep. The proposed loss term is then a sum of the original objective (e.g. classification) and a weighted sum of the probability that each dimension will be updated for each timestep. The method is evaluated on 3 time series datasets: Seizures, TwitterBuzz, Yahoo.