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: The paper proposes a new gradient-based learning algorithm for recurrent neural networks, making use of the directional derivative along a candidate direction. The directional derivative serves the purpose of improving the usefulness of a given candidate direction for gradient-based parameter updates by computing the projection of the gradient along the given direction. The candidate direction can come from various sources such as randomly sampled directions, truncated backpropagation through time (BPTT), the “Reptile” meta-learning approach by Nichol et al. (2018), or synthetic gradients. The authors call this technique “deep online directional gradient estimate” (DODGE) and demonstrate it in several experiments including a copy task (Graves et al. 2014) and a NeRF task (Mildenhall et al. 2020).