Practical Real Time Recurrent Learning with a Sparse Approximation

1 Introduction . Recurrent neural networks ( RNNs ) have been successfully applied to a wide range of sequence learning tasks , including text-to-speech ( Kalchbrenner et al. , 2018 ) , language modeling ( Dai et al. , 2019 ) , automatic speech recognition ( Amodei et al. , 2016 ) , translation ( Chen et al. , 2018 ) and reinforcement learning ( Espeholt et al. , 2018 ) . RNNs have greatly benefited from advances in computational hardware , dataset sizes , and model architectures . However , the algorithm used to compute their gradients in almost all practical applications has not changed since the introduction of Back-Propagation Through Time ( BPTT ) . The key limitation of BPTT is that the entire state history must be stored , meaning that the memory cost grows linearly with the sequence length . For sequences too long to fit in memory , as often occurs in domains such as language modelling or long reinforcement learning episodes , truncated BPTT ( TBPTT ) ( Williams & Peng , 1990 ) can be used . Unfortunately the truncation length used by TBPTT also limits the duration over which temporal structure can be realiably learned . Forward-mode differentiation , or Real-Time Recurrent Learning ( RTRL ) as it is called when applied to RNNs ( Williams & Zipser , 1989 ) , solves some of these problems . It doesn ’ t require storage of any past network states , can theoretically learn dependencies of any length and can be used to update parameters at any desired frequency , including every step ( i.e . fully online ) . However , its fixed storage requirements are O ( k · |θ| ) , where k is the state size and |θ| is the number of parameters θ in the core . Perhaps even more daunting , the computation it requires is O ( k2 · |θ| ) . This makes it impractical for even modestly sized networks . The advantages of RTRL have led to a search for more efficient approximations that retain its desirable properties , whilst reducing its computational and memory costs . One recent line of work introduces unbiased , but noisy approximations to the influence update . Unbiased Online Recurrent Optimization ( UORO ) ( Tallec & Ollivier , 2018 ) is an approximation with the same cost as TBPTT – O ( |θ| ) – however its gradient estimate is severely noisy ( Cooijmans & Martens , 2019 ) and its performance has in practice proved worse than TBPTT ( Mujika et al. , 2018 ) . Less noisy approximations with better accuracy on a variety of problems include both Kronecker Factored RTRL ( KF-RTRL ) ( Mujika et al. , 2018 ) and Optimal Kronecker-Sum Approximation ( OK ) ( Benzing et al. , 2019 ) . However , both increase the computational costs to O ( k3 ) . The last few years have also seen a resurgence of interest in sparse neural networks – both their properties ( Frankle & Carbin , 2019 ) and new methods for training them ( Evci et al. , 2019 ) . A number of works have noted their theoretical and practical efficiency gains over dense networks ( Zhu & Gupta , 2018 ; Narang et al. , 2017 ; Elsen et al. , 2019 ) . Of particular interest is the finding that scaling the state size of an RNN while keeping the number of parameters constant leads to increased performance ( Kalchbrenner et al. , 2018 ) . In this work we introduce a new sparse approximation to the RTRL influence matrix . The approximation is biased but not stochastic . Rather than tracking the full influence matrix , we propose to track only the influence of a parameter on neurons that are affected by it within n steps of the RNN . The algorithm is strictly less biased but more expensive as n increases . The cost of the algorithm is controlled by n and the amount of sparsity in the Jacobian of the recurrent cell . We study the nature of this bias in Appendix C. Larger n can be coupled with concomitantly higher sparsity to keep the cost fixed . This enables us to achieve the benefits of RTRL with a computational cost per step comparable in theory to BPTT . The approximation approaches full RTRL as n increases . Our contributions are as follows : • We propose SnAp – a practical approximation to RTRL , which is is applicable to both dense and sparse RNNs , and is based on the sparsification of the influence matrix . • We show that parameter sparsity in RNNs reduces the costs of RTRL in general and SnAp in particular . • We carry out experiments on both real-world and synthetic tasks , and demonstrate that the SnAp approximation : ( 1 ) works well for language modeling compared to the exact unapproximated gradient ; ( 2 ) admits learning temporal dependencies on a synthetic copy task and ( 3 ) can learn faster than BPTT when run fully online . 2 Background . We consider recurrent networks whose dynamics are governed by ht = fθ ( ht−1 , xt ) where ht ∈ Rk is the state , xt ∈ Ra is an input , and θ ∈ Rp are the network parameters . It is assumed that at each step t ∈ { 1 , ... , T } , the state is mapped to an output yt = gφ ( ht ) , and the network receives a loss Lt ( yt , y∗t ) . The system optimizes the total loss L = ∑ t Lt with respect to parameters θ by following the gradient ∇θL . The standard way to compute this gradient is BPTT – running backpropagation on the computation graph “ unrolled in time ” over a number of steps T : ∇θL = T∑ t=1 ∂L ∂ht ∂ht ∂θt = T∑ t=1 ( ∂L ∂ht+1 ∂ht+1 ∂ht + ∂Lt ∂ht ) ∂ht ∂θt ( 1 ) The recursive expansion ∂L∂ht = ∂L ∂ht+1 ∂ht+1 ∂ht + ∂Lt∂ht is the backpropagation rule . The slightly nonstandard notation θt refers to the copy of the parameters used at time t , but the weights are shared for all timesteps and the gradient adds over all copies . 2.1 Real Time Recurrent Learning ( RTRL ) . Real Time Recurrent Learning ( Williams & Zipser , 1989 ) computes the gradient as : ∇θL = T∑ t=1 ∂Lt ∂ht ∂ht ∂θ = T∑ t=1 ∂Lt ∂ht ( ∂ht ∂θt + ∂ht ∂ht−1 ∂ht−1 ∂θ ) ( 2 ) This can be viewed as an iterative algorithm , updating ∂ht∂θ from the intermediate quantity ∂ht−1 ∂θ . To simplify equation 2we introduce the following notation : have Jt : = ∂ht∂θ , It : = ∂ht ∂θt andDt : = ∂ht∂ht−1 . J stands for “ Jacobian ” , I for “ immediate Jacobian ” , and D for “ dynamics ” . We sometimes refer to J as the “ influence matrix ” . The recursion can be rewritten Jt = It +DtJt−1 . Cost analysis Jt is a matrix in Rk×|θ| , which can be on the order of gigabytes for even modestly sized RNNs . Furthermore , performing the operationDtJt−1 involves multiplying a k× k matrix by a k × |θ| matrix each timestep . That requires |θ| times more computation than the forward pass of the RNN core . To make explicit just how expensive RTRL is – this is a factor of roughly one million for a vanilla RNN with 1000 hidden units . 2.2 Truncated RTRL and stale Jacobians . In analogy to Truncated BPTT , one can consider performing a gradient update partway through a training sequence ( at time t ) but still passing forward a stale state and a stale influence Jacobian Jt rather than resetting both to zero after the update . This enables more frequent weight updating at the cost of a staleness bias . The Jacobian Jt becomes “ stale ” because it tracks the sensitivity of the state to old parameters . Experiments ( section 5.2 ) show that this tradeoff can be favourable toward more frequent updates in terms of data efficiency . In fact , much of the RTRL literature assumes that the parameters are updated at every step t ( “ fully online ” ) and that the influence Jacobian is never reset , at least until the start of a new sequence . All truncated BPTT experiments in our paper pass forward a stale state if an update is done before the end of the sequence . 2.3 Sparsity in RNNs . One of the early explorations of sparsity in the parameters of RNNs ( i.e . many entries of θ are exactly zero ) was Ström ( 1997 ) , where one-shot pruning based on weight magnitude with subsequent retraining was employed in a speech recognition task . The current standard approach to inducing sparsity in RNNs ( Zhu & Gupta , 2018 ) remains similar , except that magnitude based pruning happens slowly over the course of training so that no retraining is required . Kalchbrenner et al . ( 2018 ) discovered a powerful property of sparse RNNs in the course of investigating them for text-to-speech – for a constant parameter and flop budget sparser RNNs have more capacity per parameter than dense ones . This property has so far only been shown to hold when the sparsity pattern is adapted during training ( in this case , with pruning ) . Note that parameter parity is achieved by simultaneously increasing the RNN state size and the degree of sparsity . This suggests that training large sparse RNNs could yield powerful sequence models , but the memory required to store the history of ( now much larger ) states required for BPTT becomes prohibitive for long sequences . In this paper , we use a fixed sparsity pattern rather than pruning ( see Appendix B ) , for simplicity . In particular , we pick uniformly at random which indices of weight matrices to force to zero and hold this sparsity pattern constant over the course of training . 3 The Sparse n-Step Approximation ( SnAp ) . Ourmain contribution in this work is the development of an approximation to RTRL called the Sparse n-Step Approximation ( SnAp ) which reduces RTRL ’ s computational requirements substantially . SnAp imposes sparsity on J even though it is in general dense . We choose the sparsity pattern to be the locations that are non-zero after n steps of the RNN ( Figure 1 ) . We also choose to use the same pattern for all steps , though this is not a requirement . This means that the sparsity pattern of Jt is known and can be used to reduce the amount of computation in the product DtJt−1 . See Figure 2 for a visualization of the process . The costs of the resulting methods are compared in Table 1 . We note an alternative strategy would be to perform the full multiplication ofDtJt−1 and then only keep the top-k values . This would reduce the bias of the approximation but increase its cost . More formally , we adopt the following approximation for all t : ( Jt ) ij ≈ { ( Jt ) ij if ( θt ) j influences hidden unit ( ht+n ) i 0 otherwise 3.1 Sparse One-Step Approximation ( SnAp-1 ) . Even for a fully dense RNN , each parameter will in the usual case only immediately influence the single hidden unit it is directly connected to . This means that the immediate Jacobian It tends to be extremely sparse . For example , a Vanilla RNN will have only one nonzero element per column , which is a sparsity level of k−1k . Storing only the nonzero elements of It already saves a significant amount of memory without making any approximations ; It is the same shape as the daunting Jt matrix whereas the nonzero values are the same size as θ . It can become more dense in architectures ( such as GRU and LSTM ) which involve the composition of parameterised layers within a single core step ( see Appendix A for an in-depth discussion of the effect of gating architectures on Jacobian sparsity ) . In the Sparse One-Step Approximation , we only keep entries in Jt if they are nonzero in It . After just two RNN steps , a given parameter has influenced every unit of the state through its intermediate influence on other units . Thus only SnAp with n = 1 is efficient for dense RNNs because n > 1 does not result in any sparsity in J ; for dense networks SnAp-2 already reduces to full RTRL . ( N.b . : SnAp-1 is also applicable to sparse networks . ) Figure 1 depicts the sparse structure of the influence of a parameter for both sparse and fully dense cases . SnAp-1 is effectively diagonal , in the sense that the effect of parameter j on hidden unit i ismaintained throughout time , but ignoring the indirect effect of parameter j on unit i via paths through other units i′ . More formally , it is useful to define u ( j ) as the one component in the state ht connected directly to the parameter j ( which has at the other end of the connection some other entry i′ within ht−1 or xt ) . Let i = u ( j ) . The imposition of the one-step sparsity pattern means only the entry in row i will be kept for column j in Jt . Inspecting the update for this particular entry , we have ( Jt ) ij = ( It ) ij + n∑ m=1 ( Dt ) im ( Jt−1 ) mj = ( It ) ij + ( Dt ) ii ( Jt−1 ) ij ( 3 ) The equality follows from the assumption that ( Jt−1 ) mj = 0 if m 6= i. Diagonal entries in Dt are thus crucial for this approximation to be expressive , such as those arising from skip connections .

Summary: This work presents a method, named SnAp, that takes Real-time Recurrent Learning derivations and proposes to approximate its computations with sparse approximations to make them more computational tractable. The method is an alternative to overcome the truncation in backpropagation through time (BPTT) over long term temporal structure. The method assumes a sparsity pattern on the parameters of the network that leads to the relationship on how the gradients could be updated. Finally, the method is evaluated against BPTT, UORO and RTRL on character-level language modeling and a copy task.