Correcting Momentum in Temporal Difference Learning
Authors: Emmanuel Bengio
Joelle Pineau
Doina Precup
Venue:  (2021)
Keywords: Momentum
Reinforcement Learning
Temporal Difference
Deep Reinforcement Learning

Abstract: A common optimization tool used in deep reinforcement learning is momentum, which consists in accumulating and discounting past gradients, reapplying them at each iteration. We argue that, unlike in supervised learning, momentum in Temporal Difference (TD) learning accumulates gradients that become doubly stale: not only does the gradient of the loss change due to parameter updates, the loss itself changes due to bootstrapping. We first show that this phenomenon exists, and then propose a first-order correction term to momentum. We show that this correction term improves sample efficiency in policy evaluation by correcting target value drift. An important insight of this work is that deep RL methods are not always best served by directly importing techniques from the supervised setting.

Under review as a conference paper at ICLR 2021
CORRECTING MOMENTUM IN TEMPORAL DIFFERENCE LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
A common optimization tool used in deep reinforcement learning is momentum, which consists in accumulating and discounting past gradients, reapplying them each iteration. We argue that, unlike in supervised learning, momentum in Tem- poral Difference (TD) learning accumulates gradients that become doubly stale: not only does the gradient of the loss change due to parameter updates, the loss itself changes due to bootstrapping. We first show that this phenomenon exists, and then propose a first-order correction term to momentum. We show that this correction term improves sample efficiency in policy evaluation by correcting tar- get value drift. An important insight of this work is that deep RL methods are not always best served by directly importing techniques from the supervised setting.
INTRODUCTION
Temporal Difference (TD) learning (Sutton, 1988) is fundamental method of Reinforcement Learn- ing (RL), which works by using estimates of future predictions to learn to make predictions about the present. To scale to large problems, TD has been used in conjunction with deep neural networks (DNNs) to achieve impressive performance (e.g. Mnih et al., 2013 3; Schulman et al., 2017; Hessel al., 2018). Unfortunately, the naive application of TD DNNs has been shown to be brittle (Machado et al., 2018; Farebrother et al., 2018; Packer et al., 2018; Witty et al., 2018), with exten- sions such the n-step TD update (Fedus et al., 2020) or the TD(A) update (Molchanov al., 2016) only marginally improving performance and generalization capabilities when coupled with DNNs.
Part of the success of DNNs, including when applied TD learning, the use of adaptive or accel- erated optimization methods (Hinton et al., 2012; Sutskever et al., 2013; Kingma & Ba, 2015) to find good parameters. In this work we investigate and extend the momentum algorithm (Polyak, 1964) as applied to TD learning in DNNs. While accelerated TD methods have received some attention in the literature, this typically done in the context of linear function approximators (Baxter & Bartlett, 2001; Meyer et al., 2014; Pan et al., 2016; Gupta et al., 2019; Gupta, 2020; Sun et al., 2020), and while some studies have considered the mix of DNNs and TD (Zhang et al., 2019; Romoff et al., 2020), many are limited to a high-level analysis of hyperparameter choices for existing optimization methods (Sarigül& & Avci, 2018; Andrychowicz et al., 2020); or indeed the latter are simply applied as-is train RL agents (Mnih al., 2013; Hessel al., 2018). For extended discussion of related work, we refer the reader to appendix A.
As first step in going beyond the naive use of supervised learning tools in RL, examine mo- mentum. We argue that momentum, especially as it is used in conjunction with TD and DNNs, adds an additional form of bias which can be understood as the staleness of accumulated information. We quantify this bias, and propose a corrected momentum algorithm that reduces this staleness and is capable of improving performance.
1.1 REINFORCHMENT LEARNING AND TEMPORAL DIFFERENCE LEARNING
A Markov Decision Process (MDP) (Bellman, 1957; Sutton & Barto, 2018) M = AS, consists of a state space S. an action space A, a reward function R : S - R and a transition function P(s' s, a)). RL agents usually aim to optimize the expectation of the long-term return, G(S,) = E x where E 10 called the discount factor. Policies (aass) map states to action distributions. Value functions va and Q" map states and states-action pairs to

---

Under review as a conference paper at ICLR 2021
expected returns, and can be written recursively:
V (=((.) = 1E, [G(S:]] - Bn [R(S,, As) + (S ~ x (S.)]
(", = Ex ( We approximate V with Vg. We can train V/8 via regression to observed values of G, but these recursive equations also give rise to the Temporal Difference (TD) update rules for policy evaluation, relying on current estimates of V to bootstrap, which for example in the tabular case is written as:
V(S)) + V(S.) (S.) (R(()) + +V(S++))) where a E 1) is the step size. Alternatively, estimates of this update can be performed by a so-called semi- remi-gradient (Sutton & Barto, 2018) algorithm where the "TD(O) loss" is minimized:
(1)
O 02 V. (Vo, (Sa) (R(S,) + +Vo. (S++1)))
with V meaning we consider V constant for the purpose of gradient computation.
1.2 BIAS AND STALENESS IN MOMENTUM
The usual form of momentum (Polyak, 1964; Sutskever et al., 2013) in stochastic gradient descent (SGD) maintains an exponential moving average with factor 33 of gradients w.r.t. to some objective
J, changing parameters 0, E IR" (t is here SGD time rather than MDP time) with learning rate a:
Pr = 3p4-1 + (1 V
(3)
0: 0, ox/
(4)
We assume here that the objective Jr is dependepn,, as is the case for example in minibatch training or online learning. Note that other similar forms of this update exist, notably Nesterov's accelerated gradient method (Nesterov, 1983), as well as undampened methods that omit (1 - (8 in (3) replace (1 3) with o, found in popular deep learning packages (Paszke al., 2019).
We make the observation that, at time f, the gradients accumulated in j are stale. They were com- puted using past parameters rather than Oc, and in general we'd expect ((.)) + Vo (Ox)). k < t. As such, the update in (4) is a biased update.
In supervised learning where learns a mapping from x to y, this staleness only has one source: changes but the target stays constant. We argue that in TD learning, momentum becomes doubly stale: not only does the value network change, but the target (the equivalent of y) itself changes with every parameter update. Consider the TD objective in (2), when @ changes, not only does V (s) change, but V(s') as well. The objective itself changes, making past gradients stale and less aligned with recent gradients (even more so when there is gradient interference (Liu et al., 2019; Achiam et al., 2019; Bengio et al., 2020), constructive or destructive).
Note that several sources of bias already exist in TD learning, notably the traditional parametric bias (of the bias-variance tradeoff when selecting capacity), as well as the bootstrapping bias (of the error in V(s') when using it as target; using frozen target prevents this bias from compounding). We argue that the staleness in momentum we describe is an additional form of bias, slowing down or preventing convergence. This has been hinted at before, e.g. Gupta (2020) suggests that momentum hinders learning in linear TD(0).
We wish to understand and possibly correct this staleness in momentum. In this work propose answers to the following questions:
Is this bias significant in supervised learning? No, the bias exists but has minimal effect best when comparing to an unbiased oracle.
Is this bias significant in TD learning? Yes, unbiased oracle reveal significant differences.
quantify the bias, and comparisons to an Can we correct μe to remove this bias? Yes, we derive an online update that approximately corrects 14 using only first order derivatives.
Does the correction help in TD learning? Yes, using staleness- correcte momentum improves sample complexity, in policy evaluation, especially in an online setting.
'Interestingly, even in most recent value-based control works (Hessel et al., 2018) a (usually frozen) copy is used for stability, meaning that the target only changes when the copy is updated. This is considered a "trick" which it would he compelling to get rid of, since it slows down learning, and since most recent policy-gradient methods (which still use value function) do not make use of such copics (Schulman et al., 2017).
2

---

Under review as a conference paper at ICLR 2021
2
CORRECTING MOMENTUM
2.1 IDENTIFYING BIAS
Here propose an experimental protocol quantify the bias of momentum. We assume minimizing some objective in the minibatch or online setting. In such a stochastic setting, mo- mentum is usually understood variance reduction method, setting (Botev et al., 2017), but as discussed above, momentum also induces bias in the optimization
approximating the large-batch
We note that u4, the momentum time I, can be rewritten as:
= (1 3)
B Vo,
(5)
==1
and argue that an ideal unbiased momentum fiz would approximate the large batch case by only discounting past minibatches and using current parameters rather than past parameters:
He (1 3)
J.(..)
(6)
Note that the only difference between (5) and (6) is the use of 0g vs 0z. The only way to exactly compute 14 is to recompute the entire sum after every parameter update. We will consider this our unbiased oracle. To compute i empirically, since beyond certain power 3 becomes small, we will use an effective horizon of h = 2/(1 3) steps (i.e. start at i = t h rather than at i = 1).
We call the difference between μ4 and μu the bias, but note that in the overparameterized stochastic gradient case, minor differences in gradients can quickly send parameters in different regions. This makes the direct measure of ||pe /4; uninformative. Instead, we measure the optimization bias by simply comparing the loss of model trained with 14, against that of model trained with 14.
Finally, we note that, in RL, momentum approximating the batch case is related (approximately) replaying an entire buffer at once (instead of sampling transitions). This has been shown to also have interesting forms in the linear case (van Seijen & Sutton, 2015), reminiscent of the correction derived below. We also note that, while the mathematical expression of momentum and eligibility traces (Sutton, 1988) look fairly similar, they estimate a very different quantity (see appendix A).
2.2 CORRECTING BIAS
Here propose way approximate ,c", and so derive approximate correction to the bias in momentum for supervised learning, as well as for Temporal Difference (TD) learning.
We consider here the simple regression and TD(0) cases, for the online (minibatch size 1) case. We show the full derivations, and results for any minibatch size, TD(A) and n-step TD, in appendix B. In least squares regression with loss 82 we can write the gradient 9: as:
ge (6,) = (ye- - (xt)) fe. = 5,,6//2.
(7)
We would like to "correct" this gradient as we move away from 0r. A simple way to do so is to compute the Taylor expansion around 0, of g::
ge (0, + A0) = (0.) + Vo.g. 2B ((1A))1))
(8)
= ge (0,) + (8V, fo. (z.) - (e. & (x)) AA,
(9)
where & is the outer product, 7 the second derivative. We note that the term multiplying Ae commonly known as "the Hessian" of the loss. We also note that Taylor expansions around parameters have a rich history in deep learning (LeCun et al., 1990; Molchanov et al., 2016), and that in spite of its non-linearity, the parameter space of deep network is filled with locally consistent regions (Balduzzi et al., 2017) in which Taylor expansions are accurate.
The same correction computation can be made for TD(0) with TD loss 82, Remember that we write t as the learning time; we denote an MDP transition as (St, az, T', s(), making no assumption on the distribution of transitions:
- -Vo, Vo, = Ve, go (O,) = (Vo,
(10)
ge(0, + A0) (0.) + (Ve, Vo. (s1)) V., (s. + 888 Vo, A0.
(11)
3

---

Under review as a conference paper at ICLR 2021
Here, because we using a semi semi-gradient method, the term multiplying A is not exactly the Hessian: when computing Vod", we hold (s'") constant, but when computing Veg, we need to consider Vs(s") a function of 0 as it affects g. and compute its gradient."
Without loss of generality, let us write equations like (9) and (11) using the matrix Z, E R"
go (O, + A0) 9r (0.) Z 20,
(12)
where the form of Z₁, the "Taylor term", depends the loss (e.g. the Hessian in (9)).
Recall that (6) we define μ² as the discounted sum gradients using 0g for losses Jr, i S t. We call those gradients corrected gradients, 9f go + (0, 0,)). At timestep we update 02 with
ajt, thus we substitute AO =0,- 0: =- > Amm pa in (12) and get:
x 2 (0,) aZ, -Z
1 ~ gi-
(13)
k=
We can approximate the unbiased momentum p using (13), which we denote ja:
A def (1 33)
g 0
(14)
ut - - -
2 p..
(15)
Noting that μ4 can be computed as usual and that the second term of (15) has a recursive form, which we denote Q). (see appendix B), we rewrite á as follows:
ps = Pt
(16)
B1)=1 + ag 0,
(correction term)
(17)
14 (1 B)
gi-b Ji = =(1 -1--3)9 + Bpu-
(normal momentum)
(18)
1=1
G = (1 (1- 8) Σ 3 2, (1- (1) BG---
("momentum" of Taylor terms)
(19)
Algorithmically, this requires maintaining 2 vectors μ and of size the number of parameters, and one matrix Ç of size n x n. In practice, we find that only maintaining the diagonal of Ç can also work and avoid the quadratic growth in
The computation of Z also calls on computing second order derivatives V (e.g. in (9) and (11)), which is impractical for large architectures. In this work, is commonly done due to the usually small magnitude of & (Bishop, 2006, section 5.4), we ignore them and only rely on the outer product of gradients.
Ignoring the second derivatives, computing Z for TD(0) requires 3 backward passes, for V (sV(s'). and (s(s). In the online case g VVV(ss, requiring only backward passes, but in the more general minibatch of size > 2 case, it is more efficient with modern automatic differentiation packages do backward passes than 2m passes (see appendix B.4).
Finally, we note that this method easily extends to forward-view TD(X) and n-step TD methods, all that is needed is to compute appropriately (see appendix B.2).
EXPERIMENTAL RESULTS
We evaluate proposed correction well the oracle (the best our method could perform) common RL benchmarks and supervised problems. We evaluate each method with a variety of
This computation of the gradient Vo(s') may remind the reader of the so-called full gradient or residual gradient (Baird, 1995), but its purpose here is very different: we care about learning using semi graddentssss TD(0) is principled algorithm, but we also care about how this semi-gradient evolves as parameters change, and thus need compute V VoV (s')..
4

---

Under review as a conference paper at ICLR 2021
hyperparameters and multiple seeds for each hyperparameter setting. The full range of hyperparam- eters used as well as architectural details can be found in appendix D.
We will use the following notation for the optimizer used to train models: yu is the usual momen- tum, ie. (3)&(4) and serves as a bascline. 1r' is our oracle, defined in (6). fi is the proposed correction, with updates in (16)-(19), using the outer product approximation of Z. Paing is the proposed correction, but with a diagonal approximation of Z. Throughout our figures, shaded areas bootstrapped 95% confidence intervals over hyperparameter settings (if applicable) and runs.
3.1 SUPERVISED LEARNING
We first test our hypothesis, that there is bias in momentum, in simple regression task and SVHN. For regression, task 4-layer MLP to regress a 1-d function from a I-d input. For illustration we use smooth but non-trivial mixture of sines of increasing frequency in the domain E [-1, 1]:
M
7(1) 0.5 0.5 22 (. + 2)) + 0.82 (9 + 0.4)+
x
0.38 sin (1xx) + 0.32 s38r - 0.1)
(20)
Note that this choice purely illustrative and that our findings extend to similar simple functions. We train the model on fixed sample of 10k uniformly sampled points. We measure the mean squared error (MSE) when training small MLP with momentum SGD versus our oracle momentum ju" and the outer product correction P. As shown in Figure 1, we find that while there is significant difference between the oracle and the baseline (p = 0.001 from Welsh's t-test), the difference is fairly small, and is no longer significant when using the corrected (p RJ 0.1).
0.4
Oracle: μr
Negative Log. Likelihood
Oracle: l
Mean Squared Error
Baseline:
2.0
Baseline:
Corrected: f
Corrected: Adiag
0.3
1.5
0.2
1.0
0.1
0
1000
2000
3000
4000
5000
5000
10000
15000
20000
SGD steps
SGD steps
Figure 1: Regression to (20) with varying mo- mentums (10 seeds per setting). Note that the difference between i and j is not significant, but 4 and yr is.
Figure 2: Classification on SVHN with varying momentums (5 seeds per setting). Dotted lines are test losses. The only significant difference is between the training loss of μ* and j..
We compare training a small convolutional neural network on SVHN (Netzer et al., 2011) with momentum SGD versus our oracle momentum μ and the diagonalized correction momentum Ídiag- The results are shown in Figure 2. We do not find significant differences except between the training loss of μ and μl, and find that the oracle performs worse than the normal and corrected momentum.
From these experiments we conclude that, in supervised learning, there exists quantifiable opti- mization bias to momentum, but that correcting it does not appear to offer any benefit. improves performance only marginally at best, and degrades it at worst. This is consistent with the insight that u* approximates the large batch gradient, and that large batches are often associated with overfitting poor convergence in supervised learning (Wilson & Martinez, 2003; Keskar al., 2016).
3.2 TEMPORAL DIFFERENCE LEARNING
We now test our hypotheses, that there optimisation bias and that we can correct it, on RL prob- lems. First, test policy evaluation of the optimal policy on the Mountain Car problem (Singh
5

---

Under review as a conference paper at ICLR 2021
& Sutton, 1996) with small MLP. We also test Acrobot and Cartpole environments and find very similar results (see appendix C). We then test our method on Atari (Bellemare et al., 2013) with convolutional networks.
Figure shows policy evaluation on Mountain Car using a replay buffer (on-pat con-policy state transitions sampled i.i.d. in minibatches). We compare the loss distributions (across hyperparameters and seeds) at step 5k, and find that all methods are significantly different (p < 0.001) from one another. Figure 4 shows online policy evaluation, i.e. the transitions are generated and learned from once, in-order, and one at a time (minibatch size of 1). There we see that the oracle j²' and full corrected version j are significantly different from the baseline μ (p < 0.001) and diagonalized correction Adiag. as well as from Adiag, while μr" and P are not significantly different (p > 0.1).
This suggests that the ( matrix carries useful off-diagonal temporal information about parameters which co-vary, especially when the data is not used uniformly during learning. We test another possible explanation, which is that performance
degraded in online learning because the batch size is (rather than 16 or 32 as in Figure 3). We find that a batch size of does degrade fdinng performance significantly, but does not fully explain its poor online performance (see appendix Figure 10).
Mean Squared Error, EV Vaya
0.4
Mean Squared Error, EV Vaya
0.4
Oracle: je
Baseline: j
0.3
0.3
Corrected: j
Corrected: Adins
0.2
0.2
Oracle: 1²'
Baseline:
0.1
0.1
Corrected:
Corrected: Pding
0.0
0.0
0
1000
2000
3000
4000
5000
5000
10000
15000
20000
SGD steps
SGD steps
Figure 3: TD(0) policy evaluation on Moun- tain Car with varying momentums on replay buffer. The MSE is measured against a pre- trained 4 (10 seeds per setting). At step 5k, all methods are significantly different.
Figure 4: TD(0) online policy evaluation on Mountain Car, transitions are seen in-order. The MSE measured against a pretrained // (50 seeds per setting). At step 20k, μ²' and are not significantly different.
A central motivation of this work is that staleness in momentum arises from the change in gradients and targets. In theory, this is especially problematic for function approximators which tend to have interference, such DNNs (Fort et al., 2019; Bengio et al., 2020), i.e. where taking gradient step using a single input affects the output for virtually every other input. More interference means that as we change 6, the accumulated gradients computed from past és become stale faster. We test this hypothesis by (1) measuring the drift of the value function in different scenarios, and (2) computing the cosine similarity between the corrected gradients of (13), G and their true value gf (0.).
In Figure 5 we compare the value drift of MLPs with that of linear models with Radial Basis Function (RBF) features. We compute the value drift of the target on recently seen examples, i.e. we compute the average (Va, (s() Vo, ( for the last h 2m (1 (1-8) B) examples, where Nmmb is the minibatch size. We compute the RBF features for s E R² as exp(-ll (lp
uzyll//"* for a regular grid of May E 2 in the input domain. We find that the methods we try (even the oracle) are all virtually identical when using RBFs (see Figure 12 in appendix). We also find, as shown in Figure 5, that RBFs have very little value drift (due to their sparsity) compared to MLPs. This is consistent with our hypothesis that the method we propose is only useful if there is value drift-otherwise, there is no optimization bias incurred by using momentum. We can test this hypothesis further by artificially increasing the width of the RBFs, o² such that they overlap. As predicted, sparsity (increasing interference) increases value drift and increases the gap between our method and the baseline (Figure 6). This drift correlates with performance when changing o² (Figure 7).
find that reducing

---

Under review as a conference paper at ICLR 2021
0.16
Oracle: j
10-1
Baseline: j
MSE, ((VV -ar2
0.14
Corrected: f
Value Drift
MLP
Corrected: Pdiag
10-2
1 RBF
0.12
MLP
RBF
10-3
0.10
á RIFF
Pdiag MLP
0.08
10
Pding RBF
0.06
0
1000
2000
3000
4000
5000
2
3
4
SGD steps
width RBFs, o²
Figure 5: Value drift of (ss) when training with TD(O) on a replay buffer. We see that RBFs being sparse feature representation, the value functions of recently seen data tend not drift (10 seeds per setting). Here o² 1.
Figure 6: MSE as a function of the width, o², of RBF kernels. The larger the kernel, the more value drift our method, it, is able to correct (10 seeds per setting).
1.00
0.16
Oracle: μ
Corrected: P
MSE, (VV 1/122
0.14
Baseline: μ
Cosine Similarity
0.95
Corrected: Adiag
Corrected: j
0.90
0.12
Corrected: Pding
0.10
0.85
0.08
0.80
0.06
0.75
10
10
10
10
1000
2000
3000
4000
5000
Mean value drift
SGD steps
Figure 7: MSE function of mean Value drift of (s'') for RBFs varying kernel size. The lines match the 6² lines of Figure 6, and show the relation between o², drift, and error.
Figure 8: Average cosine similarity of the Tay- lor approximations of with their true value of for recently seen data; Mountain Car, replay buffer policy evaluation (40 seeds per setting).
In Figure 8 we measure the cosine similarity of the Taylor approximations g + Z (0: 0,) with the true gradients of (.,( for the last h 21 - :3) examples. We find that the similarity is relatively high (close to 1) throughout training but that it gets lower as the model converges the true value V This also consistent with our hypothesis that there is change (staleness) in gradients, while also validating the approach of using Taylor approximations to achieve this correction mechanism.
It is also possible to correct momentum for a subset of parameters; we try this on a per-layer basis and find that, perhaps counter-intuine counter-intuitively" than the top layers (close to V), although correcting all layers is the best (see appendix Figure 13).
much better correct the bottom layers (close to a:)
We now apply our method on the Atari MsPacman game. We first do policy evaluation on an expert agent (we use pretrained Rainbow agent (Hessel et al., 201 2018)). Since the architecture required train an Atari agent is too large (n R 4.8M) to maintain n² values, we only use the diagonal version of our method. We also experiment with smaller (n. x 12.8k) models and the full correction,
Although one may expect that changes close to V should produce more drift (it is commonly thought that hootstrapping happens in the last layer top of relatively stable representations). the opposite is consistent with á interacting with interference in the input space, which the bottom layers have learn disentangle.

---

Under review as a conference paper at ICLR 2021
with similar results (see D.8). To (considerably) speed up learning, since the model is large, additionally use per par-marmettrr learning rates as in Adam (Kingma & Ba, 2015), where an estimate of the second moment is used as denominator in the parameter update; denote this combination Adig// /V2 + E. We see in Figure 9 that our method provides significant (p < 0.01) advantage.
Note that our method does not use frozen tar- gets (as is usually necessary for this environ- ment). A simple way to avoid momentum stal- eness and/or target drift in TD(O) is the of frozen targets, i.e. to keep separate 0 to com- pute (s's,, updated (0 0) at large inter- vals. Such a method is central to DQN (Mnih et al., 2013), but its downside that it requires more updates to bootstrap. We find that for pol- icy evaluation, frozen targets are much slower (both in Atari and simple environments) than our baseline (see appendix Figures 16 and 17).
20
Mean Squared Error, E(VV
Corrected: Pdlag/ V +E
18
Baseline (Adam): 4/vV2+ +
16
14
12
10
0
100
200
300
400
500
SGD steps (x (x10))
We finally apply our method to control, first in combination with a policy gradient method, PPO (Schulman et al., 2017), in its policy eval- uation part, and second with a 5-step Sarsa method, but find improvement (or marginal at best) in either setting. As in simpler environ- ments we measure cosine similarity and value drift. We find low similarity (+s 0.1) but a 2x to 3x decrease in drift using our method, suggesting that while our method corrects drift, its effect on policy improvement minor. We suspect that in control, even with a better value function, other factors such as exploration or overestimation come into play which are not addressed by our method.
Figure 9: TD(0) policy evaluation on Atari (MsPacman) with varying momentums (20 seeds) on replay buffer. The MSE is measured against sampled returns GT.
4 DISCUSSION
We propose a method which improves momentum, when applied to DNNs doing TD learning, by correcting gradients for their staleness via approximate Taylor expansion. We show that correct- ing this staleness is particularly useful when learning online using a TD objective, but less so for supervised learning tasks. We show that the proposed method corrects for value drift in the boot- strapping target of TD, and that the proposed approximate Taylor expansion is useful tool that aligns well enough with the true gradients.
Shortcomings of the method In its most principled form, the proposed method requires comput- ing a second-order derivative, which is impractical in most deep learning settings. While we do find experimentally that ignoring its contribution has negligible effect in toy settings, we are unable to verify this for larger neural networks. Compared to the usual momentum update, the proposed method also requires performing two additional backward passes, well as storing 2n + n² addi- tional scalars. This can be improved with a diagonal approximation requiring only 3n scalars, but this approximation is not as precise, which is currently an obstacle for large architectures. While these extra computations improve sample complexity, i.e. we get more out of each sample, they do not necessarily improve convergence speed in computation time (although, our method is suited to GPU parallelism and has reasonable speed even for large n²s,, including in Atari). Sample complex- ity may be particularly important in settings where samples are expensive to get (e.g. real world systems), but less so when they are almost free (e.g. simulations).
Incorporating TD in optimization One meta-result stands out from this work: something is lost when naively applying supervised learning optimization tools to RL algorithms. In particular here, by simply taking into account the non-stationarity of the TD objective, we successfully improve the sample complexity of a standard tool (momentum), and demonstrate the potential value of incorpo- rating elements of RL into supervised learning tools. The difference explored here with momentum is only one of the many differences between RL and supervised learning, and we believe there main plenty of opportunitics to improve deep learning methods by understanding their interaction with the peculiarities of RI..
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Joshua Achiam, Ethan Knight, and Pieter Abbeel. Towards characterizing divergence in deep q- learning. arXiv preprint arXiv:/903.08894. 2019.
Marcin Andrychowicz, Anton Raichuk, Piotr Stañczyk, Manu Orsini, Sertan Girgin, Raphael Marinier, Léonard Hussenot, Matthieu Geist, Olivier Pietquin, Marcin Michalski, Sylvain Gelly, and Olivier Bachem. What matters in on-policy reinforcement learning? a large-scale empirical study, 2020.
Leemon Baird. Residual algorithms: Reinforcement learning with function approximation. In Machine Learning Proceedings /995, pp. 30-37. Elsevier, 1995.
David Balduzzi, Brian McWilliams, and Tony Yeoman. Neural taylor approximations: Con- vergence and exploration rectifier networks. In International Conference Machine Learning, pp. 351-360, 2017.
Jonathan Baxter and Peter L Bartlett. chinite-horizon policy-gradient estimation. Journal of Artifi- cial Intelligence Research, 15:319-350, 2001.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Richard Bellman. A markovian decision process. Journal of mathematics and mechanics, pp. 679- 684, 1957.
Emmanuel Bengio, Joelle Pineau, and Doina Precup. Interference and generalization temporal difference learning. In International Conference on Machine Learning, 2020.
Christopher M Bishop. Puttern recognition and machine learning- springer, 2006.
Alcksandar Botev, Guy Lever, and David Barber. Nesterov's accelerated approximations to regularised update descent. 2017 International Joint Conference on Neural Networks (IJCNN), pp. 1899-1903. IEEE, 2017.
accelrrateee gradient and momentum as
Felix Dangel, Frederik Kunstner, and Philipp Hennig. BackPACK: Packing more into backprop. In International Conference Learning Representations, 2020. URL https::tttssppttt //tppr/ee///peeeee....... net pat//for forum? 1d=3J1=F24
A. P. Engelbrecht. A new pruning heuristic based variance analysis of sensitivity information.
IEEE Transactions on Neural Networks, 12(6):1386-133 399, 2001.
Andries Petrus Engelbrecht. Using the taylor expansion of multilayer feedforward neural networks.
South African Computer Journal, 2000(26):181-188 2000.
Jesse Farebrother, Marlos C Machado, and Michael Bowling. Generalization and regularization in dqn. arXiv preprint arxiv://810.00/23, 2018
William Fedus, Prajit Ramachandran. Rishabh Agarwal, Yoshua Bengio. Hugo Larochelle, Mark Rowland, and Will Dabney. Revisiting fundamentals of experience replay, 2020.
Stanislav Fort. Pawel Krzysztof Nowak. Stanislaw Jastrzebski, and Srini Narayanan. Stiffness: A new perspective on generalization in neural networks, 2019.
Dhawal Gupta. Applicability of momentum in the methods of temporal learning. 2020.
Harsh Gupta, R Srikant, and Lei Ying. Finite-time performance bounds and adaptive learning rate selection for two time-scale reinforcement learning. In Advances in Neural Information Process- ing Systems, pp. 4704 2019.
Babak Hassibi and David G Stork. Second order derivatives for network pruning: Optimal brain surgeon. In Advances in neural information processing systems, pp. 164-171, 1993.

---

Under review as a conference paper at ICLR 2021
Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Thirty-Second AAAI Conference Artificial Intelligence, 2018.
Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent. CSC321, 2012.
Steven Kapturowski, Georg Ostrovski, Will Dabney, John Quan, and Remi Munos. Recurrent ex- perience replay in distributed reinforcement learning. In International Conference on Learning Representations, 2019. URL https / http:///ppppeeeee//// net/ Eorum?i httppparrrrr....eeee//rrrrrrrrrr
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima, 2016.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization (2014). arXiv preprint arXiv:1412.6980, 15, 2015.
Yann LeCun, John S Denker, and Sara Solla. Optimal brain damage. In Advances in neural information processing systems, pp. 598-605, 1990.
Vincent Liu, Raksha Kumaraswamy, Lei Le, and Martha White. The utility of sparse representa- tions for control in reinforcement learning. In Proceedings of the AAAI Conference Artificial Intelligence, volume 33, pp. 4384 4391, 2019.
Marlos C. Machado, Mare G Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61:523-562, 2018.
Dominik Meyer, Rémy Degenne, Ahmed Omrane, and Hao Shen. Accelerated gradient temporal difference learning algorithms. In 2014 IEEE Symposium Adaptive Dynamic Programming and Reinforcement Learning (ADPRL), pp- 1-8. IEEE, 2014.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602. 2013.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pp.
2016.
Pavlo Molchanov, Stephen Tyree, Tero Karras, Timo Aila, and Jan Kautz. Pruning convolutional neural networks for resource efficient inference. arXiv preprint 11Xvv:661.006400 2016.
Yu Nesterov. method of solving a convex programming problem with convergence rate o11\3²) (1/k2). In Sov. Math. Dokl, volume 27, 1983.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.
Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Song. Assessing generalization in deep reinforcement learonnnnn arXiv preprint 1eXivv/8800.200 2018.
Yangchen Pan, Adam White, and Martha White. Accelerated gradient temporal difference learning. arXiv preprint arXiv: 1611:09321.88 1.09288, 2016.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, pp. 8026-8037, 2019.
Boris Polyak. Some methods of speeding up the convergence of iteration methods. USSR Com- putational Mathematics and Mathematical Physics, 4(5):1-17, 1964.
10

---

Under review as a conference paper at ICLR 2021
Joshua Romoff, Peter Henderson, David Kanaa, Emmanuel Bengio, Ahmed Touati, Pierre-Luc Ba- con, and Joelle Pineau. Tdprop: Does jacobi preconditioning help temporal difference learning?, 2020.
M. Sarigüil and M. Avci. Performance comparison of different momentum techniques on deep inforcement learning. Journal of Information and Telecommunication, 2(2):205-216, 2018. doi: 10.1080/24551839 201...24888..4..1
Robert Schapire and Manfred K Warmuth. the worst-case analysis of cemporal-different learning algorithms. Machine Learning, 22(1-3):95-121, 1996.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 2017.
Satinder P Singh and Richard S Sutton. Reinforcement learning with replacing eligibility traces.
Machine learning, 22(1-3):123-155 1996.
Tao Sun, Han Shen, Tianyi Chen, and Dongsheng Li. Adaptive temporal difference learning with linear function approximation. arXiv preprint urXiv:2002.08537, 2020.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum in deep learning. In International conference on machine learning, pp- 1139-1147, 2013
Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning, (1):9-44, 1988.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Richard S Sutton, Csaba Szepesvári, and Hamid Reza Maei. A convergent (n) algorithm for off-policy temporal-different learning with linear function approximation. Advances in neural information processing systems, 21(11):600-1116, 2008.
Richard S Sutton, Hamid Reza Maei, Doina Precup. Shalabh Bhatnagar, David Silver, Csaba Szepesvári, and Eric Wiewiora. Fast gradient-descent methods for with linear function approximation. In Proceedings of the 26th Annual International Conference on Machine Learning, pp. 993-1000, 2009.
learning
Harm van Seijen and Rich Sutton. True online td(A). In International Conference Machine Learning, pp. 692-700, 2014.
Harm van Seijen and Rich Sutton. A deeper look at planning as learning from replay. International conference on machine learning, pp. 2314-2322 23222, 2015.
Nino Vieillard, Bruno Scherrer, Olivier Pietquin, and Matthieu Geist. Momentum in reinforcement learning. In International Conference on Artificial Intelligence and Statistics, pp. 2529-22 2538, 2020.
Randall Wilson and Tony R Martinez. The general inefficiency of batch training for gradient descent learning. Neural networks, 16(10):1429-453 2003.
Sam Witty, Jun Ki Lee, Emma Tosch, Akanksha Atrey, Michael Littman and David Jensen. Measuring and characterizing generalization deep reinforcement learning. arXiv preprint arXiv: 028.8....8.... 2018
Shangtong Zhang, Wendelin Bochmer, and Shimon Whiteson. Deep residual reinforcement leam- ing. arXiv preprint arXiv:1905.01172, 2019.
11

---

Under review as a conference paper at ICLR 2021
A RELATED WORK
To the best of our knowledge, no prior work attempts to derive corrected momentum-SGI update adapted to the Temporal Difference method. That being said, wealth of papers are looking to accelerate TD and related methods.
On momentum, traces, and gradient acceleration in TD
From an RL perspective, our work has some similarity to the so-called eligibility traces mechanism. In particular, in the True Online TD(A) method of van Seijen & Sutton (2014), the authors derive strict-online update (i.e. weights are updated at every MDP step, using only information from past steps, rather than future information as the X-return perspective) where the main mechanism of the derivation lies in finding update by assuming (at least analytically) that one "start over" and reuse all past data iteratively at each step of training, and then from this analytical assumption derive a recursive update (that doesn't 't require iterating through all past data). The extra values that have to be kept to compute the recursive updates are then called traces. This is akin to how we conceptualize / (6), and derive ja.
The conceptual similarities of the work of van Seijen & Sutton (2015) with our work are also in- teresting. There, the authors analyse what "retraining from scratch" means (i.e., again, iteratively restarting from 0g E "R"" but with some ideal target 0* (e.g. the current parameters) by redoing se- quentially all the TD(0) updates using 6* for all the n transitions a replay buffer, costing (nm) They derive an online update showing that one can continually learn at a cost of O(m²) rather than paying (mm) at each step. The proposed update also reminiscent of our method in that aims to perform an approximate batch update without computing the entire batch gradient, and also main- tains extra momentum-like vectors and matrices. We note that the derivation there only works in the linear TD case.
In a way, such an insight can be found in the original presentation of TD(A) of Sutton (1988), where the TD(A) parameter update is written as (equation (4) in the original paper, but with adapted notation):
~d, = afre + (o. (st)
x Vv. Vo,
k=1
Remark the use of 0, in the sum; in the linear case since Vq, Vo, (s+) d(sx), the sum does not depend on Og and thus can be computed recursively. A posteriori, if one can find a way to cheaply compute Ve, Ve, (sk) Vk, perhaps using the method we propose, it may be an interesting way to perform TD(A) using non-linear function approximator.
Our analysis is also conceptually related to the work of Schapire & Warmuth (1996), where a worst- case analysis of TD*(A is performed using best-case learner as the performance upper bound. This is similar to our momentum oracle; just as the momentum oracle is the "optimal" approxima- tion of the accumulation gradients coming from all past training examples, the best-case learner of Schapire& & Warmuth (1996) is the set parameters that is optimal when one is allowed to look at all past training examples (in contrast to online TD learner).
Before moving on from TD(A), let us remark that eligibility traces and momentum, while simi- lar, estimate different quantities. The usual (non-replacing) traces estimate the exponential moving average of the gradient of Vo. while momentum does so for the objective J (itself a function of V&):
C = (1-
Σ 1 \ Pe (1-11- 3) Σ g 3'
Our method also has similarities with residual gradient methods (Baird, 1995). A recent example of this is the work of Zhang et al. (2019), who adapt the residual gradient for deep neural networks. Residual methods learn by taking the gradient of the TD loss with respect to both the current value and the next state value (S'), but this comes at the cost of requiring two independent samples of S' (except in deterministic environments).
Similarly, our work is related to the "Gradient TD" family of methods (Sutton et al., 2008; 2009). These methods attempt to maintain expectation (over states) of the TD update, which allows to
12

---

Under review as a conference paper at ICLR 2021
directly optimize the Bellman objective. While the exact relationship between GTD and "momen- tum TD" is not known, they both attempt to maintain an "expected update" and adjust parameters according to it; the first approximates the one-step linear TD solution, while the latter approxi- mates the one-step batch TD update. Note that linear GTD methods can also be accelerated with momentum-style updates (Meyer et al., 2014), low-rank approximations for part of the Hessian (Pan et al., 2016), and adaptive learning rates (Gupta et al., 2019).
More directly related this work that of Sun al. (2020), who show convergence properties of rescaled momentum for linear TD(O). While most (if not every) deep reinforcement learning method implicitly uses some form of momentum and/or adaptive learning rate as part of the deep learning toolkit, Sun et al. (2020) properly analyse the use of momentum in a (linear) TD context. Gupta (2020) also analyses momentum in the context of linear TD(O) and TD(A,, with surprising nega- tive results suggesting naively applying momentum may hurt stability and convergence in minimal MDPs.
Another TD-aware adaptive method is that of Romoff et al. (2020), who derive per-parameter adap- tive learning rates, reminiscent of RMSProp (Hinton et al., 2012), by considering a (diagonal) Jacobi preconditioning that takes into account the bootstrap term in TD.
Finally, we note that, as far as we know, recent deep RL works all use some form of adaptive gradient method, Adam (Kingma & Ba, 2015) being optimizer of choice, closely followed by RMSProp (Hinton et al., 2012); notable examples of such works include those of Mnīh et al. (2013), Schulman et al. (2017), Hessel et al. (2018), and Kapturowski et al. (2019). We also note the work of Sarigül & Avei (2018), comparing various SGD variants on the game of Othello, showing significant differences based on the choice of optimizer.
Balduzzi et al. (2017) note that while theory suggests that Taylor expansions around parameters should not be useful because of the "non-convexity" of ReLU neural networks, there nonetheless exists local regions in parameter space where the Taylor expansion is consistent. Much earlier work by Engelbrecht (2000) also suggests that Taylor expansions of small sigmoid neural networks are easier to optimize. Using Taylor approximations around parameters to find how to prune neural networks also appears to be an effective approach with a long history (LeCun et al., 1990; Hassibi & Stork, 1993; Engelbrecht, 2001; Molchanov et al., 2016).
On Taylor approximations
On policy-gradient methods and others While not discussed this paper, another class of meth- ods used to solve RL. problems are PG methods. They consist in taking gradients of the objective wrt a directly parameterized policy (rather than inducing policies from value functions). We note in particular the work of Baxter & Bartlett (2001), who analyse the bias of momentum-lic like cumulated policy gradients (referred to as traces therein), showing that 3 the momentum parameter should be chosen such that 1/(1 3) exceeds the mixing time of the MDP.
Let us also note the method of Vieillard et al. (2020), Momentum Value Iteration, which uses the concept of an exponential moving average objective for decoupled (with its own parameters) action-value function from which the greedy policy being evaluated is induced. This moving av- erage is therein referred to momentum; even though it is properly speaking the optimizational acceleration of Polyak (1964), its form is similar.
B COMPLETE DERIVATIONS
B. DERIVATION OF MOMENTUM CORRECTION
In momentum, H4 be rewritten as:
14 = (1 3)
3 V
(21)
We argue that ideal unbiased momentum #4 would approximate the large batch case by only discounting past minibatches and using current parameters 0, rather than past parameters 0,:
1 del (1
33
3 Ve-ii Vo, (((O,)
(22)
13

---

Under review as a conference paper at ICLR 2021
Note that the only difference between (21) and (22) is the use of 8, vs 8t. The only way to compute 14 exactly is to recompute the entire sum after every parameter update. Alternatively, we could somehow approximate this sum. Below we will define o V Ji which we will then approx- imate with z. We will then show that this approximation has a recursive form which leads to an algorithm.
We want to correct accumulated past gradients 9, to their "ideal" form g. as above. We do so with their Taylor expansion around 0, and we write the correction of the gradient gi computed at learning time corrected at time / as:
g g: 20(1 1;1) = Vog. AO( 1; + ((1)001))
(23)
= 9 - A0
(24)
where (f(t; =; 0g. g: Vg Z, Ve.9t-
Here we are agnostic of the particular form of Z, which will depend on the loss and learning algo- rithm, and not necessarily the so-called Hessian. To see why this the case, and for the derivation of Z for the squared loss, cross-entropy and TD(0), see section B.2.
Let's now express J in terms of gi and updates μr, At each learning step. the current momentum μ4 is multiplied with the learning rate a to update the parameters, which allows us to more precisely write gf:
0: = 0:11
0/4 0g
(25)
i=1
A0(1;) = 0, - 02 = 0p a
1. 0g + a
Pk
-0
(26)
k=1
k=1
k=
Φ = + z 20 - 1;1) 9: aZ
1k
(27)
k=1
We can now write ius the approximated μi using gf:
1-1
(1- +(1 111))
Σ
8t-i0t
(28)
1-1
[-1
1-1
= (1 - 8) + (1 8) Σ
8r-1 9i 111-1-1 -11-0 3)
Σ Σ
3
-
07
ik
(29)
k=
11
14 - a01 - (8)
3 4 z jis extract
(30)
k=
1-
14: a(1 () Σ Σ
gr- 2T juu change the sum indices for convenience
(31)
H4 7t extract T/
(32)
Let's try find recursive form for ?t!
14

---

Under review as a conference paper at ICLR 2021
t-1
1-2 k 2
7 n-- = a(1 - 3)
gl z Ps a(1 - (8
3 10-1-:10 z--1-izf pr.
k=1 1m1
k= 1
a(1 - 8)
8 144 -3 2 (m +
Σ Br 2, 1h--1
i=1
1-2 k
t-1
a(1 11-)) ()
Σ
>
(5 1-8--11-2111 pr
Σ gt 
PP-1
t-1
(5 1)1-111 + a. - Σ
gr 1-- 777 M-1
=
1=1
Let
G = (1 3) Σ g
1=1
(5 1)1--11 1 + 0:BGe-1
1tt M--1 =
-(1 (1-8)m-1 + 0:357- (1-1
Th = Bm--1 - 11-1-11-11
We can now write the full update as:
A He - 1t
N BT--1 + a T 144-1
He
(1 - 3)
9: -1- -11-8) 334-1
i=1
Se (1 3) g1-1 21 (1 B)Z, + BG
with Yo
This corresponds to an algorithm where one maintains and Ç.
B.2 DERIVATION OF TAYLOR EXPANSIONS AND Z
For least squares regression, the expansion around of g(0) is simply the Hessian of the loss 82:
8² 1 7 fo(z))
(33)
g(0) V00²
(34)
g(0 + A0) ((0) + H A0 o(110]13)
(35)
This Hessian has the form
H, & Vof V
(36)
52 being small when neural network is trained and 2 being expensive to compute, common approximation to H, is to only rely on the outer product. Thus we can write:
Zreg Vol & Vol
(37)
For classification, or categorical crossentropy. Z has exactly the same form, but where f is the log-likelihood (i.e. output of log-softmax) of the correct class.
15

---

Under review as a conference paper at ICLR 2021
For TD(O), the expansion is more subtle. Since TD is semi-gradient method, when computing g(0), gradient of the TD loss 8², we ignore the bootstrap target's 'ss derivative, i.c. we hold (s') constant (unlike in GTD). On the other hand, when computing the Taylor expansion around 8, do care about how Vo(s') changes, and so its gradient comes into play:
9 = (Vo(x) Voo(') VoVo(2)
g(0 + A0); = 9:(0) + (VoVa(>)
VolV(e)) -
A0 40 40 20
9(0 + A0) = =g(() + ((VoVo(z) -VVoV(<')) @ VoVa(z)) AO + VV)V((>)
Similarly for TD we ignore the second order derivatives and write:
ZTD (VoVo(a) [oVo(<')) & VoVo(1)
For an n-step TD objective, the correction is very similar:
ZTD(M) = V - VoVo(z)4n)) VVV & VoVo(z,)
For forward-view TD(A) objective this is also similar, but more expensive:
ZTD(A)= (V)) (1
...)) VoVo(z))
Finally, to avoid maintaining values for Z, it is possible to only maintain the diagonal of Z or some block-diagonal approximation, at some performance cost.
B.3 THE LINEAR CASE
Much RL analysis is done in the linear case, as it is easier to make precise statements about convergence there. We write down such case below, but we were unfortunately unable to derive any interesting analyses from it.
Recall that the proposed method attempts to approximate μ²:
mi (1 B) g VV ((
1=1
which in the linear case Vo(x) 0 ((x)) simply:
= (1 - B)
3r 85,0
which depends on 6, through ò, the TD error. As such we can write as:
ZTD (VoVo(x) - VoV((z)) & VoVo(z) = (6- (0-70'00
which interestingly does not depend on 8. To the best of our knowledge, and linear algebra skills, this lack of dependence on f does not allow for simplification of the proposed correction mechanism (in contrast with linear eligibility traces) due to the pairwise multiplicative I, dependencies that emerge.
B.4 MINIBATCHING
It is possible only do 3 (or in the supervised case) "backward" calls (e.g. torch. autograd. grad in PyTorch) even when computing this correction for several examples. This is in contrast other adaptive gradient methods (Romoff et al., 2020) which require per-input
16

---

Under review as a conference paper at ICLR 2021
gradients (these can be computed efficiently with methods such Dangel et al. (2020), but remain more expensive).
Because the gradient care to correct is the gradient of the minibatch, i.c. the of gradients, then we only have to correct that sum, and not individual gradients. This allows us to compute Z out of the gradients of the minibatch instead of the individual gradients. When the gradient is:
9: VoJi(0) - V Je.i
Σ VoJ2..(0)
(38)
then the first derivative used in its Taylor expansion is, by linearity of the derivative:
0g:
Vola Voli + H
(39)
06
Vo fi
Σ f8d
(40)
Ignoring the Hessian we only have to use the gradients obtained for the sum (or the mean if J is mean) of losses of the minibatch, which requires only one backpropagation regardless of batch size. The derivation analogous for TD(0). Note that this would not work we used the second order derivative due the multiplicative 6τ,2 factors.
C ADDITIONAL FIGURES
100
Baseline: i
Baseline: /
Mean Squared Error, [Va Vaya
0.5
Mean Squared Error, IB(V Vaaa
Oracle: xe'
Oracle: j²'
Corrected:
Corrected: f
Corrected: Pding
Corrected: Adiag
10
0.1
10-2
.95
.975
.99
.995
.999
64
B
Minibatch size
Figure 11: TD(0) policy evaluation on Moun- tain Car with varying on a replay buffer. The MSE measured after 5k SGD steps against a pretrained V= Shaded areas 95% confidence runs (10 seeds per setting). We use a minibatch size of 4 to reveal interesting trends.
Figure 10: TD(O) policy evaluation on Moun- tain Car with varying minibatch size on a replay buffer. The MSE is measured after 5k SGD steps against pretrained VT, Shaded areas are bootstrapped 95% confidence runs (20 seeds per setting).
bootstrapped
17

---

Under review as a conference paper at ICLR 2021
0.40
Oracle: j
Mean Squared Error, E(VV Va)a
Other layers: Daing
Baseline: f
0.23
Other layerst ja
0.35
Corrected: A
Cosine Similarity
Corrected: Adias
0.30
0.22
0.25
0.21
0.20
0.15
0
1000
2000
3000
4000
5000
SGD steps
bottom used top used
Figure 12: TD(0) policy evaluation Moun- tain Car with RBF on variety of hyperparame- settings (10 seeds per setting). Here o²
Figure 13: TD(0) policy evaluation Moun- tain Car with an MLP. We vary the number of layers whose parameters are used for full f cor- rection (n² params); e.g. when "bottom used" is 3, the first layers, those closest the input, are used; when "top used" is 1, only the last layer, that predicts V from embeddings, is used. The parameters of other layers either corrected with the diagonal correction or use normal mo- mentum. Correcting "both ends" is not better than just the bottom (not shown here).
4.0
4.0
Mean Squared Error, (VV Vee
Oracle: fe
3.5
Baseline: J
3.5
Corrected: P
3.0
Corrected: Adias
Mean Squared Error, BB(V Va
2.5
2.5
2.0
2.0
Oracle: f²
Baseline:
1.5
1.5
Corrected: i
Corrected: Painn
1.0
1.0
0
1000
2000
3000
4000
5000
0
1000
2000
3000
4000
5000
SGD steps
SGD steps
Figure 14: TD(0) policy evaluation Ac- robot with varying hyperparameters on a replay buffer. The MSE is measured after 5k SGD steps against a pretrained !a, Shaded areas are bootstrapped 95% confidence runs (10 seeds per setting).
Figure 15: TD(0) policy evaluation Cartpole with varying hyperparameters replay buffer. The MSE is measured after 5k SGD steps against pretrained V Shaded areas are boot- strapped 95% confidence (10 seeds per set- ting).
18

---

Under review as a conference paper at ICLR 2021
0.40
30.0
Baseline (frozen targets):
Adam
0.35
Baseline: Y
27.5
Adam (frozen target)
Corrected: j
Corrected: Adiag/ V +
Mean Squared Error, [V -Ve2
Corrected: Paing
Mean Squared Error, P(VV Gm)2
25.0
0.30
Oracle:
22.5
0.25
20.0
0.20
17.5
0.15
15.0
0.10
12.5
0.05
10.0
0.00
1000
2000
3000
4000
5000
100
200
300
400
500
SGD steps
SGD steps (x 10
Figure 16: Replication of Figure 3 including the frozen targets baseline.
Figure 17: Replication of Figure 9 including the frozen targets baseline. Interestingly the models trained with frozen targets eventually become more precise than those without, but this only happens after very long time. This is explained by the stability required for bootstrapping when TD errors become increasingly small, which is casily addressed by keeping the target network fixed.
0.10
0.5
Oracle: je"
Oracle: st
Mean Squared Error, [B(V Vaja
Baseline: p
Baseline: J
0.08
Corrected: j
0.4
Corrected: ft
Corrected: Ading
MSE, E(V% Vxya
Corrected: Painn
0.06
0.3
0.04
0.2
0.02
0.00
0.0
10000
20000 30000 40000 50000
10-2
10-1
SGD steps
Learning rate
Figure 18: Replication of Figure 3 with 10x more training steps. Methods gradually con- verge to the value function.
Figure 19: Effect of the learning rate on Moun- tain Car, replay buffer policy evaluation, MSE after 5k training steps.
19

---

Under review as a conference paper at ICLR 2021
D HYPERPARAMETER SETTINGS AND COMMENTS
All experiments are implemented using PyTorch (Paszke et al., 2019). We use Leaky ReLUs throughout. All experimental code is available in supplementary materials.
we did experiment with ReLU, tanh, ELU, and SELU activation units. The latter 3 units have more stable Taylor expansions for randomly initialized neural networks, but in terms of experimental results, Leaky ReL.Us were always significantly better.
On Leaky ReLUs:
D.1 FIGURE
We use an MLP with 4 layers of width 12y-
We use the cross-product of ny (8,16,32), 3 e 0.99), {0.005, 0.013. Rmbb
{4, 16, 64).
D.2 FIGURE 2
We use a convolutional model with the following sequence of layers, following PyTorch convention: Conv2d(3, n4, 2, 1), Conv2d(n,, 2n, 3, 1), Conv2d(2n4, 2n,, 2, 1), Conv2d(2n,, nhe 3, 1, 1), Flatten() Linear(16n,, 4n Linear(4nn, 4m4), Linear(4n,, 10), with LeakyReLUs between each layer.
We use the cross-product of na E {8, 16), 8 E 0.99,,99,,.. a E f0.005, 0.01}, Romb {4, 16, 64).
D.3 FIGURE 3 AND 16
We use an MLP with 4 layers of width ma.-
We use the cross-product of ny {8, 12, 8 (0.9,0.99). a e 0,5,0..,00055, amb
{4, 16, 64).
D.4 FIGURE 4
We use an MLP with 4 layers of width nA.
We use the cross-product of n& [16), 8 0.99, a {0.005, 0.001, 0.0005), Ammb (1).
D.5 FIGURE AND 12
We linear layer top of grid RBF representation with each gaussian having variance of o² /griid.
For the RBF
Agrid 20, 0.1, Romb 16, B 0.99. For the MLP we n. 16,
0.1, Bmmh = 16, 3 = 0.99.
D.6 FIGURE 6 AND
We use RBFs with o² E (4,355,3 2,5, 1,0.5,, 0.5), E (10,20), a e (0.1,0001).
D.7
FIGURE 8
We use an MLP with 4 layers of width ny, = 16, a = 0.1, 7mm = 16, 3 0.95.
D.8 FIGURE 9 AND 17
We use the convolutional model of Mnih et al. (2013) with the same default hyperparameters, fol- lowing PyTorch convention: Conv2d(4, ny, 8, stride=4, padding=4), Conv2d(nn, 2n,, 4, stride=2, padding=2), Conv2d(2n,, 2n,, 3, padding=1) 1), Flatten(), Linear(2n, x 12 x 12, 16n,). Linear(16n,, ).cts).
20

---

Under review as a conference paper at ICLR 2021
We n4 64, Rmmb 32, for Adam use = 5 x 10-5 and 3 = 0.99, for method a = 10 4 and 3 = 0.9 (these choices are the result of a minor hyperparameter search which the best values were picked, equal amounts of compute went towards our method and the baseline as to avoid "poor baseline cherry picking"). For the frozen target baseline we update the target every 2.5k steps.
We were able to find similar differences between Adam and our method with a much smaller model, but using the full correction instead of the diagonal one. Although the full correction outperforms Adam when both use this small model, using so few parameters is not as accurate as the original model described above, and we omit these results.
This small model is: Conv2d(4, ThA 3, 2, 1), Conv2d(rиk. 742 3, 2. 1), Conv2d(ns, nA. 3, 2, D), Conv2d(ns, 74,, 3, 1), Conv2d(nu, Dacts. 6). For T4 16 (which we used for experiments), this model has little less than 16k parameters, making Z about 250M scalars. While this is large, this easily fits on a modern GPU, and the extra computation time required for the full correction mostly still comes from computing two extra backward passes, rather than from computing Z and the correction.
This is beyond the scope of this paper, but is worth of note: Interestingly this small model still works quite well for control (both with Adam and our method). We have not tested this extensively, but, perhaps contrary popular RL wisdom surrounding Deep Q-Learning, decent MsPacman agents with (1) no replay buffer, but rather 4 or more parallel environments (the more the better) in A2C (Mnih et al., 2016) (2) no frozen target network (3) Q network with only 16k parameters rather than the commonplace 4.8M-parameter Atari DQN model. The only "trick" required is use 5-step Sarsa instead of 1-step TD (as suggested by the results of Fedus et al. (2020), although in their case a replay buffer is used).
were able to train
D.9 FIGURE 10
use the same configuration than Figure 3 with nmb E (1,4,16,6)) 64).
D.10
FIGURE 11
We use the same configuration than Figure 3 with 3 E {.8,.9, .975, .99, -9, -995,
D.11
FIGURE 13
We use MIP with 6 layers of width = 16, = 0.1, 7mbb 16, , = 0.9. We vary which layers get used in the corrected momentum; see caption.
D.12 FIGURE 14 AND 15
We use the settings as Figure 3, with the exception that we do 5-step TD in Figure 14, for Acrobot, and 3-step TD for Figure 15, for Cartpole. Using n > appears necessary for convergence for both method and the baseline.
21


Reviews:
[{"review_id": "oqa_KHZ3qKR", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper studies the approach to 'correct the bias' of momentum in supervised learning and reinforcement learning (especially, in the temporal-difference algorithm). The proposed approach can maintain a moving average of the 'correct' gradient direction over past objective functions via extra updates. And the experiments show that it does speed up the convergence. (ps: I did not find the comparisons over the loss functions at the stationary points).\n\nAs we know TD learning uses a semi/partial/biased gradient with respect to the mean square Bellman error. Thus, the effect of applying momentum to training is an open question. I feel very glad to see such a paper studying it and adding a reasonable modification. The approach averages the TD update over different steps via a modified momentum update. Such an approach reduces the variance, which is usually very large in TD, but also uses additional computational power.\n\nI found the motivation not very convincing for me. The paper seems to claim that a good momentum should carry the information of past objective function and calculate their gradients only with respect to the parameter in the current iterate. However, momentum terms do have the reason to use the gradient wrt previous parameters, as it can speed up the convergence in many situations, e.g. ravines. Another concern is that TD is not typically an online learning problem. The objective function in the early phase is not informative to the update in the later phases as the values of the next states are inaccurate. Therefore, I feel the theoretical reason for 'correcting' the 'bias', which is defined by this paper, is not warranted very well. \n\nI suggest the paper to rigorously discuss more on why such 'bias' has to be corrected. Alternatively, the proposed method is designed simply for the purpose of reducing the variance of the stochastic gradient for the case of minibatch update. But the paper needs to argue how it is better than using a larger batch or compare itself with other variance reduction approaches.", "extra_scores": {}, "extra_text": {}}, {"review_id": "nsdUVsilkB6", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "#### Idea\nThe paper explains and studies an interesting issue with momentum in TD learning, which is its staleness while doing TD updates which is in contrary to the supervised learning. However, the paper could benefit from more investigations.\n\n#### Comments\n- In section 2.1: It is said that to compute the oracle update, one has to recompute the entire sum. This may take a while to compute and it would be good if they have had provided a computation-time or wall-time comparison will prior methods.\n- In section 2.1: To handle the problem of long summation, \"an effective horizon\" is used. However, there's no explanations given behind this particular choice. Is it considered as a geometric distribution and thus this formula? Or if it's not, where does this formula come from?\n- In section 3.2: When the proposed method is applied to the Atari games, the problem of scalability shows itself. As mentioned by the authors: \"Since the architecture required to train ...\", the scalability of this approach is under question. \n- In section 3.2: Regarding this paragraph: \"Note that we do not use frozen targets ...\". The most important question about this work is its scalability, whereas in a prior approaches like frozen targets there is not such an issue while dealing with staleness issue. As in experiments, authors have not used frozen targets, which seems to be misleading when it comes to the experiments and results. It would be useful to see a comparison between these two approaches since they try to solve a similar problem.\n\n#### Minor issues\n- In section 3.1: \"We task a ...\" needs to be changed to \"We take a ...\"\n\n- In section 2.2: It would be better if you could give a more in depth explanation for eq 11: “Here the term multiplying ∆θ is not exactly the Hessian…”\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "AXq1l9-agl6", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThis paper extends the idea of momentum, commonly used in optimization literature, to Temporal Difference (TD) learning which is a widely used algorithm for policy evaluation in Reinforcement learning literature. The main challenge in this work is to account for the 'optimization bias' introduced by using momentum based updates. The authors propose to do this via a Taylor approximation to the TD update and empirically show the merits of this idea on some toy data sets as well as on an Atari game. \n\nReason for score: \n\nThe idea of using momentum for TD learning seems quite interesting. However, in my opinion, the paper seems to a somewhat incremental contribution, given the work of Sun et. al. 2020, which gives a theoretical analysis of momentum based updates for TD with linear function approximation. \n\n1) The bias correction via Taylor approximation of the gradients seems new to me. The authors claim that the bias doesn't play a significant role in supervised learning tasks. Can the authors provide some more citations to support this besides the 1-d regression task. For TD learning, the authors don't present results for bias corrections with second order terms. How do we know that we can safely ignore these for large networks? On the other hand, incorporating these terms poses a computational challenge. \n\n2) Prior work of Sun et. al. 2020, which proposes a momentum for a linear function approximation, uses projection to bound the iterates. Is the momentum update even stable for deep neural networks for large problems, without some additional tricks? The authors only test toy examples (the Atari example has some combination of Adam and momentum which makes it ).  \n\n3) For the comparison on MsPacman, why is a combination with Adam required? The authors claim it is for speeding up learning. In this case, maybe the authors can have the MSE plot for just Adam updates for a better comparison? \n\n4) Many of the plots are truncated at a few thousand steps. What happens after that? It might be helpful to visualize the loss at limit points that different updates converge towards (roughly). Graphs with more steps might be more useful.\n\n5) The authors claim their method improves sample complexity as compared to say the more widely used trick of freezing network. But Figure 17 shows that models with frozen targets eventually become more precise. This might be a reason why this idea may not perform well in the control setting, as Figure 17 suggests that there is probably a residual bias in all momentum based updates (when used with neural networks). \n\nOverall, I am unsure whether the contributions here are strong enough to justify a paper publication.\n\n--------- Update after author response ---------\n\nI thank the authors for their response. I have updated my score. I am slightly leaning toward acceptance though I think the paper might benefit from a revision based on some of the points raised by other reviewers, including comments about motivation for correcting the bias as well as the scalability (the proposed method requires storing n^2 additional values which is expensive for large networks). ", "extra_scores": {}, "extra_text": {}}, {"review_id": "6TFTxdkcdc", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a modification to momentum. Extra terms are added to account for the drift in parameters through the gradient updates that contribute to the momentum. The effect of this modification is studied on supervised learning and TD learning and on different representations. The conclusion is that while the original momentum is good enough for supervised learning, this modification matters when there is bootstrapping, especially on representations with high interference.\n\nOverall, the empirical study highlights a direction for improving TD learning and the experiments are diverse and conclusive. Although the proposed approximations to the extra terms show limited improvement, experiments show that if the extra terms can be approximated more accurately, there will be a considerable gain. Computing the extra terms exactly is computationally expensive but it is possible that better approximations will be introduced later on. There are a few weaknesses that I describe below.\n\n1. There is little discussion on why we would want to account for the drift in parameters in momentum. Would this somehow reduce variance or improve the rate of convergence?\n\n2. Candidate step-sizes are too few and some differences could simply be due to the match between the chosen step-size and an algorithm. Consider sweeping over a finer grid.\n\n3. Section 3.1 says \"Note that this choice is purely illustrative and that our findings extend to similar simple functions.\" Is there evidence on other simple functions?\n\n4. A discussion on different optimizers somewhere in the paper or the appendix is needed to see how the proposed updates compare to existing work.\n\nMinor comments:\n \n1. Section 1.1, problem formulation and definition of V and Q: In a standard RL setting reward is a function of state and action.\n\n2. \\theta and V_\\theta are not defined. It is better to make it explicit that the function is parameterized by \\theta and outputs V_\\theta which estimates V^\\pi.\n\n---\nUpdate: I have the other reviews and rebuttal. I am still leaning towards acceptance, although I do agree with the other reviews that there is little motivation for the idea. While the experiments show improvement, a discussion on convergence rates or variance could show why one should try this idea in the first place, and if \\mu_* (rather than regular momentum) is indeed an ideal that the algorithm should approximate.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies the role of momentum in temporal difference (TD) learning algorithms, and how this can be systematically exploited to accelerate the TD type algorithms. More specifically, the authors point out that the momentum term could be quite biased, and propose a scheme to remedy this issue. However, the reviewers point out the lack of motivation about bias correction; it is unclear why bias correction is crucial to achieve acceleration. 
Decision: Reject