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: A main problem with RNN is to update all hidden dimensions in each time step. The authors proposed selective-activation RNN (SA-RNN), which modifies each state of RNN by adding an update coordinator which is modeled as a lightweight neural network. The coordinator, based on the incoming data, makes a discrete decision to update or not update each individual hidden dimension. A multi-objective optimization problem is defined to both solving a sequential learning task and minimizing the number of updates in each time step. The authors evaluated their networks on three public benchmark datasets and achieved good results compared to the state-of-the-art ones.