Learning by Directional Gradient Descent

1 INTRODUCTION . A key requirement of intelligence is the ability to integrate information over time , by constructing an internal representation of state from sequences of observations . This is critical for an agent to predict or control its future experience , whenever its environment is not fully observable ( i.e . whenever individual observation are not sufficient for making the required predictions , or decisions ) . An efficient representation must construct state incrementally through an update function st+1 = update ( st , ot ) with the previous state st and the last observation ot as inputs . Since observations may include lots of redundant or irrelevant information , state representations must be tailored to a given task , encoding all and only the aspects of past experience required to perform such task . To learn state representations from data , it is useful to be able to assess how changing the update function affects the long term performance of the agent . If the update function is parameterized by a differentiable function such as a recurrent neural network , we can measure the effect of changing each parameter in terms of a suitable loss function , and then update the parameter in the direction of steepest descent of the loss function , by using a suitable variant of ( stochastic ) gradient descent . Unfortunately , no existing method for computing the steepest descent direction of a recurrent function satisfies all the basic desiderata for a general approach to the problem . Prior methods either use too much memory ( BPTT , Rumelhart et al . 1986 ; Werbos 1988 ) , too much computation ( RTRL , Williams & Zipser 1989 ) , have too much variance ( UORO , Tallec & Ollivier 2017 ; SPSA , Spall et al . 1992 ) , or too much bias ( DPG , Silver et al . 2014 ; synthetic gradients , Jaderberg et al . 2017b ) . In this paper we propose an algorithm to efficiently perform near-steepest descent of recurrent functions , using the same computational and memory complexity as the forward computation of the recurrent function . The proposed approach is based upon the observation that the directional derivative of a recurrent function along any arbitrary direction u can be computed efficiently and then can be used to construct a descent direction . When applied to recurrent networks , this observation allows the directional derivative of a recurrent function to be computed efficiently in a fully online manner , with computational cost the same as the forward computation of the function . What direction u is most effective ? In the special case where u matches the direction of the true gradient , we recover the steepest descent update . This suggests the following algorithm : first approximate the true gradient , and then follow a descent direction along the direction of the estimated gradient , in order to reduce the loss . We analyse and evaluate empirically random choices of direction , and also investigate directions based upon several methods of gradient approximation . 01 DeepMind , London , UK , 2 University College London , 3 Mila , University of Montreal . Corresponding author : anirudhgoyal9119 @ gmail.com 2 METHODS FOR APPROXIMATING THE GRADIENT . In this section , we first give a brief overview of common methods for approximating true gradient . Notation . We consider a recurrent network represented by a differentiable function ( lt , xt ) = f ( xt−1 , ot , w ) . The recurrent network maintains at each time-step t an internal `` hidden '' state xt of size n , and receives an observation ot . The recurrent function is parameterized by a vector w of size p. At every time-step there is a loss lt. For our purposes , any output of the network ( such as its predictions or actions ) may be considered to be part of the internal state xt that leads to a loss lt . The objective is to minimise the total loss L = ∑T t=1 lt over T time-steps . A key requirement for efficient optimization is the computation of the gradient of the total loss with respect to the parameters , ∂L ∂w = ∑T t=1 ∂lt ∂w . Backpropagation-through-time ( BPTT ) . Backpropagation-through-time ( BPTT ) uses backward differentiation to compute the gradient of the cumulative loss with respect to the parameters of a recurrent neural network . Specifically , the gradient can be computed by exchanging indices to consider the effect of state xt on all future losses lt , ... , lT : ∂L ∂w = T∑ t=1 ∂lt ∂w = T∑ t=1 t∑ j=1 ∂lt ∂xj ∂xj ∂w = T∑ t=1 T∑ j=t ∂lj ∂xt ∂xt ∂w . ( 1 ) The inner term may be accumulated by backpropagation of error , T∑ j=t ∂lj ∂xt = ∂lt ∂xt + T∑ j=t+1 ∂lj ∂xt+1 ∂xt+1 ∂xt . This requires just O ( p ) computation per time-step ; however , the entire sequence must first be processed forward to compute the output . In practice , the memory-requirements of BPTT limit the computation of gradients to short windows , and limit the choice of neural network architecture to those with modest memory requirements . If the gradient of large neural networks must be computed over long horizons then BPTT may not be an appropriate choice of algorithm . Synthetic Gradients and Direct Policy Gradients . We now consider how one may use ideas drawn from reinforcement learning to estimate the future loss and hence allow for efficient online computation of gradients . One may estimate the loss using a value function ṽ ( xt ) ≈ ∑T j=t lj and use the estimated loss in place of the true loss , ∂L ∂w = T∑ t=1 T∑ j=t ∂lj ∂xt ∂xt ∂w ≈ T∑ t=1 ∂ṽ ( xt ) ∂xt ∂xt ∂w , ( 2 ) where the loss may be estimated by bootstrapping from the next time-step , ṽ ( xt ) ≈ lt + ∑T j=t+1 lj ≈ lt + ṽ ( xt+1 ) . Alternatively , one may estimate the backpropagated gradient ṽ′ ( xt ) ≈ ∑T j=t ∂lt+j ∂xt and use the estimated gradient in place of the true gradient , ∂L ∂w ≈ T∑ t=1 ṽ′ ( xt ) ∂xt ∂w , where the backpropagated gradient may be estimated by bootstrapping from the next time-step , ṽ′ ( xt ) ≈ ∂lt∂xt + ∑T j=t+1 ∂lj ∂xt+1 ∂xt+1 ∂xt ≈ ∂lt∂xt + ṽ ′ ( xt+1 ) ∂xt+1 ∂xt . It is worth noting that the former approach is equivalent to the Deterministic Policy Gradient ( DPG ) ( Silver et al. , 2014 ) in a reinforcement learning problem where the state is considered to be the input state st = xt−1 , the action is the next state at = xt , the policy is the function mapping from state xt−1 to action xt parameterized by weights w , and the reward is the negative loss rt = −lt ( Jaderberg et al. , 2017b ) . In practice , both deterministic gradients and synthetic gradients introduce significant bias into the estimate of the gradient . This has the downside that stochastic gradient descent , applied to the estimated gradient , may not actually descend the loss and may fail to converge . Real-time recurrent learning ( RTRL ) . We now consider methods that compute the gradient of a recurrent function by forward differentiation . Given a function z = hT ( ... ( h1 ( y ) ) ... ) with derivative dz dy = JT ... J1 where Jt is the Jacobian of function ht , one may instead compute the derivative dz dy = ( JT ... ( J2 ( J1 ) ) ) by processing the chain rule in a forward direction and caching the intermediate quantity dhtdy = Jt ... J1 . This is only computationally efficient when y is smaller than z ; this is the opposite case to usual in deep learning where we have many parameters for y and a scalar loss for z. Real-time recurrent learning ( RTRL , Williams & Zipser , 1989 ) applies forward differentiation to compute the gradient of the cumulative loss with respect to the parameters of recurrent functions . This is achieved by caching the sensitivity of state to weights dxtdw , i.e . the total derivative of state to weights via all previous states , ∂L ∂w = T∑ t=1 ∂lt ∂w = T∑ t=1 ∂lt ∂xt dxt dw . The sensitivity may be accumulated via forward accumulation , dxt+1 dw = ∂xt+1 ∂w + ∂xt+1 ∂xt dxt dw . The size of the matrix dxtdw is n× p and the computational cost of multiplying the Jacobian matrices in the second term is O ( pn2 ) , leading to a costly overall algorithm that is impractical for large-scale deep learning . 3 DEEP ONLINE DIRECTIONAL GRADIENT ESTIMATE ( DODGE ) . After explaining a directional derivative , we will introduce our online directional gradient estimate . 3.1 DIRECTIONAL DERIVATIVE . A basic property of gradients is that the directional derivative can be computed efficiently . One may efficiently compute the Jacobian-vector product dzdy · u = ( Jn ... ( J2 ( J1 · u ) ) ) for any tangent vector u of size p. A geometric interpretation of this product is the directional derivative of the gradient , i.e . dz dy · u describes how much the gradient descends in the direction of the vector u . 3.2 GRADIENT ESTIMATE . When applied to recurrent networks , the directional derivative of a recurrent function can be computed efficiently with computational cost of justO ( p ) per time-step , i.e . the same as the forward computation of the function . ∂L ∂w · u = T∑ t=1 ∂lt ∂xt ( dxt dw · u ) dxt+1 dw · u = ∂xt+1 ∂w · u+ ∂xt+1 ∂xt ( dxt dw · u ) For any unitary u , calculating the directional derivative along that direction provides a valid descent direction for updating weights , ∆w ∝ ( ∂L ∂w · u ) u . The direction of steepest descent may be computed from many such tangent vectors , for example ∂L ∂w = ∑p i=1 ( ∂L ∂w · ei ) ei for the unit bases e1 , ... , ep . This requires p separate directional derivatives . Alternatively , random directions such as ui ∼ { −1 , +1 } p may be sampled , leading to an accumulated gradient estimate ∂L∂w ≈ 1 n ∑n i=1 ( ∂L ∂w · ui ) ui , but at the cost of introducing significant variance . However , if the tangent vector is chosen to be the gradient direction itself , u = g‖g‖ where g = ∂L ∂w , then only one such directional derivative is required , ( g · u ) u = ( g · g ‖g‖ ) g ‖g‖ = g. We are often interested in the expectation of the gradient , E [ g ] rather than the sample gradient . For example , stochastic gradient descent ( SGD ) averages over sample gradients to descend in the direction of the overall gradient . When using directional gradient descent , one may either compute the directional derivative of an individual sample gradient , and then average , or compute the directional derivative in the direction of the expected gradient u = ḡ‖ḡ‖ where ḡ = E [ g ] , and then average : E [ ( g · u ) u ] = E [ ( g · ḡ ‖ḡ‖ ) ḡ ‖ḡ‖ ] = ( ḡ · ḡ ‖ḡ‖ ) ḡ ‖ḡ‖ = ḡ . This suggests a novel strategy of estimating the expected gradient ḡ and using the estimated gradient as a descent direction u in ( ∂L∂w · u ) u . Note that , regardless of the quality of the gradient estimate , the algorithm will always follow a valid descent direction.A valid descent direction has a non-negative dot-product with the true gradient . The quality of the gradient solely determines the steepness of descent in expectation . In the next section , we consider methods that provide such an estimate of the gradient . We refer to this family of algorithms as Deep Online Directional Gradient Estimate ( DODGE ) and they can be understood as near-steepest descent algorithms . def dodge ( loss_function , params , direction ) : # Compute the Jacobian-vector product . loss , loss_jvp = jax.jvp ( loss_function , [ params ] , [ direction ] ) grad_estimate = loss_jvp * direction return loss , grad_estimate Listing 1 : DODGE implemented in JAX .

Summary: This paper proposes to adapt the RTRL technique when training RNNs to make it usable in practice. That is, instead of computing the full gradient of the loss $L$ according to all parameters, the proposed method "DODGE" consists in computing the gradient of $L$ in *one* or a small number of directions in the parameter space. These "directional gradients" of $L$ are much easier to compute than the full gradient when using RTRL. More precisely, the user can choose a subspace $\mathcal{P}'$ (of dim. $p'$) of the space of parameters $\mathcal{P}$ (of dim. $p$), and compute the projection of the true gradient $g^*$ of $L$ on the space $\mathcal{P}'$, *with a computational cost proportional to $p'$*. This is advantageous when $p' \ll p$.  Notably, this method can be used in two ways:  * estimating the true gradient by computing its projection on a random low-dimensional subspace of the parameters;  * improve an existing estimation $\hat{g}$ of the true gradient $g^*$ by computing the projection of $g^*$ on $\hat{g}$: this projection has better properties than $\hat{g}$.