summaries: The authors introduce a theoretical model for delayed gradients in asynchronous training. It is a very nice model and solving the corresponding differential equation allows to study its stability. Authors derive stability bounds for pure SGD (learning rate needs to decrease with delay) and for SGD with momentum, where they introduce a nice momentum formulation that improves stability. These are nice insights and good results and they are validated by experiments. More experiments and practical analysis would be welcome though. Some example questions: would introducing some sychronization help? Is the lower learning rate hurting training speed when measures as wall-clock time to accuracy?

summaries: This paper studies how asynchrony affects model training by investigating dynamic stability of minimum points that A-SGD can access. They point out that not all local minimum points are accessible, and asynchrony can affect which minimum points can be accessed, and thus helps to explain why models trained by A-SGD have higher generalization gap. The authors also propose shifted-momentum that utilize momentum for asynchronous training.

summaries: The authors model A-SGD as a dynamical system, where parameters are updated with delayed gradients. The authors analyze the stability of this system, and they first derive that the learning rate must scale linearly with the inverse of the delay around a minimum to remain stable. Using a similar analysis they show that the standard way of incorporating momentum into A-SGD requires small learning rates for high momentum values, and they propose "shifted momentum," which allows for stability under higher momentum values. Experimentally, the authors show that around minima the learning rate needed to retain stability scales linearly with the inverse of the delay, that there appears to be an analogous threshold when training models from scratch, that shifted momentum allows for higher momentum values, and finally that on several datasets A-SGD with an appropriate learning rate is able to generalize at least as well as large batch synchronous training.

abstractText: Background: Recent developments have made it possible to accelerate neural networks training significantly using large batch sizes and data parallelism. Training in an asynchronous fashion, where delay occurs, can make training even more scalable. However, asynchronous training has its pitfalls, mainly a degradation in generalization, even after convergence of the algorithm. This gap remains not well understood, as theoretical analysis so far mainly focused on the convergence rate of asynchronous methods. Contributions: We examine asynchronous training from the perspective of dynamical stability. We find that the degree of delay interacts with the learning rate, to change the set of minima accessible by an asynchronous stochastic gradient descent algorithm. We derive closed-form rules on how the learning rate could be changed, while keeping the accessible set the same. Specifically, for high delay values, we find that the learning rate should be kept inversely proportional to the delay. We then extend this analysis to include momentum. We find momentum should be either turned off, or modified to improve training stability. We provide empirical experiments to validate our theoretical findings.

authors:
Yossi Arjevani
Ohad Shamir
Nathan Srebro

title: A Tight Convergence Analysis for Stochastic Gradient Descent with Delayed Updates. jun 2018

authors:
Mahmoud Assran
Nicolas Loizou
Nicolas Ballas
Michael Rabbat

title: Stochastic Gradient Push for Distributed Deep Learning

authors:
Jianmin Chen
Xinghao Pan
Rajat Monga
Samy Bengio
Rafal Jozefowicz

title: Revisiting distributed synchronous sgd

venue: arXiv preprint arXiv:1604.00981,

authors:
Wei Dai
Yi Zhou
Nanqing Dong
Hao Zhang
Eric Xing

title: Toward Understanding the Impact of Staleness in Distributed Machine Learning, sep 2018

authors:
Jeffrey Dean
Greg Corrado
Rajat Monga
Kai Chen
Matthieu Devin
Mark Mao
Andrew Senior
Paul Tucker
Ke Yang
Quoc V Le

title: Large scale distributed deep networks

authors:
Jia Deng
Wei Dong
Richard Socher
Li-Jia Li
Kai Li
Li Fei-Fei

title: Imagenet: A large-scale hierarchical image database

venue: In Computer Vision and Pattern Recognition,

authors:
Sanghamitra Dutta
Gauri Joshi
Soumyadip Ghosh
Parijat Dube
Priya Nagpurkar

title: Slow and Stale Gradients Can Win the Race: Error-Runtime Trade-offs in Distributed SGD

authors:
Priya Goyal
Piotr Dollár
Ross Girshick
Pieter Noordhuis
Lukasz Wesolowski
Aapo Kyrola
Andrew Tulloch
Yangqing Jia
Kaiming He

title: Accurate, large minibatch sgd: training imagenet in 1 hour

venue: arXiv preprint arXiv:1706.02677,

authors:
Ido Hakimi
Saar Barkai
Moshe Gabel
Assaf Schuster

title: Taming Momentum in a Distributed Asynchronous Environment

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

authors:
Elad Hoffer
Itay Hubara
Daniel Soudry

title: Train longer, generalize better: closing the generalization gap in large batch training of neural networks

venue: In NIPS,

authors:
Nitish Shirish Keskar
Dheevatsa Mudigere
Jorge Nocedal
Mikhail Smelyanskiy
Ping Tak Peter Tang

title: On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima

title: Learning multiple layers of features from tiny images

authors:
Yann Lecun
Leon Bottou
Yoshua Bengio
Patrick Ha

title: Gradient-Based Learning Applied to Document Recognition

authors:
Yann A LeCun
Léon Bottou
Genevieve B Orr
Klaus-Robert Müller

title: Efficient backprop

venue: In Neural networks: Tricks of the trade,

title: Scaling Distributed Machine Learning with the Parameter Server

venue: Proceedings of the 2014 International Conference on Big Data Science and Computing - BigDataScience

authors:
Xiangru Lian
Yijun Huang
Yuncheng Li
Ji Liu

title: Asynchronous parallel stochastic gradient for nonconvex optimization

venue: In Advances in Neural Information Processing Systems,

authors:
Xiangru Lian
Huan Zhang
Cho-Jui Hsieh
Yijun Huang
Ji Liu

title: A comprehensive linear speedup analysis for asynchronous stochastic parallel optimization from zeroth-order to first-order

authors:
Xiangru Lian
Wei Zhang
Ce Zhang
Ji Liu

title: Asynchronous decentralized parallel stochastic gradient descent

venue: arXiv preprint arXiv:1710.06952,

authors:
Tianyi Liu
Shiyang Li
Jianping Shi
Enlu Zhou
Tuo Zhao

title: Towards understanding acceleration tradeoff between momentum and asynchrony in nonconvex stochastic optimization

authors:
Ioannis Mitliagkas
Ce Zhang
Stefan Hadjis
Christopher Re

title: Asynchrony begets momentum, with an application to deep learning

venue: 54th Annual Allerton Conference on Communication,

authors:
Kamil Nar
S Shankar Sastry

title: Step Size Matters in Deep Learning

title: Discrete-time signal processing

venue: Pearson Education India,

authors:
Benjamin Recht
Christopher Re
Stephen Wright
Feng Niu

title: Hogwild: A lock-free approach to parallelizing stochastic gradient descent

authors:
Christopher J Shallue
Jaehoon Lee
Joe Antognini
Jascha Sohl-Dickstein
Roy Frostig
George E Dahl

title: Measuring the effects of data parallelism on neural network training

venue: arXiv preprint arXiv:1811.03600,

authors:
Karen Simonyan
Andrew Zisserman

title: Very deep convolutional networks for large-scale image recognition

venue: arXiv preprint arXiv:1409.1556,

authors:
Lei Wu
Chao Ma
Weinan E

title: How SGD Selects the Global Minima in Over-parameterized Learning: A Dynamical Stability Perspective

authors:
Yan Yan
Tianbao Yang
Zhe Li
Qihang Lin
Yi Yang

title: A unified analysis of stochastic momentum methods for deep learning

venue: arXiv preprint arXiv:1808.10396,

authors:
Sergey Zagoruyko
Nikos Komodakis

title: Wide residual networks

venue: arXiv preprint arXiv:1605.07146,

authors:
Wei Zhang
Suyog Gupta
Xiangru Lian
Ji Liu

title: Staleness-aware async-sgd for distributed deep learning

venue: arXiv preprint arXiv:1511.05950,

heading: 1 INTRODUCTION

text:
Training deep neural networks (DNNs) requires large amounts of computational resources, often using many devices operating over days and weeks. Furthermore, with ever-increasing model size and available data, the amount of compute used was noted to increase exponentially over the past few years (Amodei & Hernandez, 2018). Fortunately, the operations made in DNN training are highly suitable for parallelization over many devices. Most commonly, the parallelization is done by "data-parallelism" in which the data is divided into separate batches which are distributed between different devices during training. By using many devices, with each device highly utilized, one could train on a large number of samples without paying in run-time. However, this training is done synchronously, with synchronization between the different devices done on every iteration (S-SGD). With the growth in the number of devices participating in the training, the synchronization becomes a prominent bottleneck.
To overcome the synchronization bottleneck, asynchronous training (A-SGD) has been proposed (Dean et al., 2012; Recht et al., 2011). The basic idea behind such training is that when a device finishes calculating its batch gradients, an update to the DNN parameters is immediately performed, without waiting for other devices. This approach is known as asynchronous centralized training with a parameter server (PS) that updates the parameters (Li, 2014). The main problem in such training is that, since the PS updates the parameters whenever a device communicates with him, the parameters that are being used in other devices calculation are no longer up-to-date. This phenomena is called gradient staleness or delay as we will refer to it, and it causes a deterioration in generalization performance when using A-SGD.
∗Equal contribution
The generalization deterioration can be seen in Fig. 1. Unless the learning rate is significantly decreased, we can see a generalization gap: a deterioration in the generalization of A-SGD (with delay τ = 32) from the baseline (S-SGD, τ = 0) value near the steady state — i.e., after we are near full convergence, following many training epochs (2000). Due to the generalization gap demonstrated in the figure, A-SGD with a parameter server and large delays is not commonly used — despite it is relatively simple to implement and despite its vast potential to accelerate training. Therefore, our main goal in this work is to shed some theoretical light on the generalization gap problem and use it to improve A-SGD training.
Previous theoretical analysis of A-SGD (Liu et al., 2018; Lian et al., 2016; 2015; Dai et al., 2018; Dutta et al., 2018; Arjevani et al., 2018) focused on analyzing the convergence rate, i.e., the time it takes to reach steady state, and not on the properties of the obtained solution. In contrast, in our paper we focus on understanding how the delay affects the selection of the solution we converge to, and changes in this selection can impact generalization.
We tackle these questions from the perspective of dynamical stability. The dynamical stability approach was used in (Nar & Sastry, 2018; Wu et al., 2018), to study which minima are accessible under a specific choice of optimization algorithm and hyperparameters. In (Nar & Sastry, 2018), the authors analyzed gradient descent (GD) algorithm as a discrete-time nonlinear dynamical system. By analyzing the Lyapunov stability of this system for different minima, they showed a relation between the learning rate choice and the accessible minima set, i.e., the subset of the local optima that GD can converge to. In (Wu et al., 2018), the authors focused on stochastic gradient descent (SGD), defined a criterion to evaluate the stability of a specific minimum and used this criterion to show how the learning rate and batch-size play a role in SGD minima selection process.
Here, we use dynamical stability to analyze the dynamics of A-SGD. Using this approach, the main question we try to tackle is:
How do learning rate, delay and momentum interact and affect the minima selection process?
Contributions
We start by modelling A-SGD as a dynamical system with delay. By analyzing the stability properties of that system, we find:
• There exist an inverse linear relation between the delay of A-SGD, and the threshold learning rate i.e., the critical learning rate in which a minimum loses stability. • This implies that in order to keep a specific minimum stability as we change the delay, we
need to change the learning rate inversely proportional to the delay. • Momentum has a crucial role in determining stability properties for A-SGD. Depending on
the specific training algorithm, it can either hurt the stability or improve it. We show how to modify momentum to stabilize A-SGD training.
With our theoretical findings, we derive closed form rules when training asynchronously. We suggest to scale the learning rate inversely to the delay in order to reach a better generalization performance. In section 3 we provide experiments on popular classification tasks to support our theoretical analysis. It is interesting to contrast this simple linear relation with the more complicated picture observed for the scaling of the learning rate with the batch size. Specifically, it has been observed by Shallue et al. (2018) that the optimal learning rate scales differently with the batch size when changing the models and datasets, i.e., there is no general "rule of thumb" that applies in all cases.
A similar scaling rule for the learning rate was proposed before (Zhang et al., 2015), however, there it lacked a theoretical basis or a demonstration of its effectiveness with large effective batch sizes, i.e., the sum of all the minibatch sizes of all the workers, as we do here.

heading: 2 THEORETICAL ANALYSIS

heading: 2.1 PRELIMINARIES AND PROBLEM SETUP

text:
Consider the problem of minimizing the empirical loss
f(x) = N∑ i=1 fi(x) , (1)
where N is the number of samples, x ∈ Rd and ∀i = 1, ..., N : fi : Rd → R is twice continuously differentiable function. The update rule for minimizing f(x) in eq. 1 using asynchronous stochastic gradient descent (A-SGD) is given by
xt+1 = xt − η∇fn(t−τ)(xt−τ ) , (2)
where η is the learning rate, n(t) is a random selection process of a sample from {1, 2, . . . , N} at iteration t, and τ is some delay due the asynchronous nature of our training. For simplicity, we focus on a fixed delay τ . The extension to stochastic delay is discussed in section 2.2.1.
Our main goal is understanding how different factors such as the gradient staleness τ and the learning rate η affect the minima selection process in neural networks. In other words, we would like to understand how the different hyperparameters interact to divide the minima of the loss to two sets: those we can converge to and those we cannot converge to.
To analyze this we suppose we are in a vicinity of some minimum point x∗, which implies∇f(x∗) = 0. Given some values of η and τ , we ask whether or not we can converge to x∗, using the theory of dynamical stability. As we shall see, if the learning rate η or delay τ are too high, then this minimum loses stability, i.e., A-SGD cannot converge to x∗ , in expectation.
Specifically, we examine the expectation of the linearized dynamics around x∗. We show in appendix A that, to ensure stability, it is sufficient that the following one-dimensional equation is stable:
xt+1 − xt + ηaxt−τ = 0 , (3)
where we defined the Hessian H = 1N N∑ i=1 ∇2fi(x∗), the sharpness term a = λmax (H) as the maximal eigenvalue of the Hessian, and xt is the projection of the expectation of the perturbation xt − x∗ on the eigenvector which corresponds to the maximal eigenvalue a. Additionally, we show in appendix A that the characteristic equation of eq. 3 is
zτ+1 − zτ + ηa = 0 . (4)

heading: 2.2 STABILITY ANALYSIS

text:
For given (η, τ), we would like to determine if the minimum point x∗ is asymptotically stable (Oppenheim, 1999) in expectation. This will enable us to divide the minima into two sets: those we might converge to since they are stable, and those we cannot possibly converge (except for a measure zero set of initialization) since they are unstable. We use the characteristic equation (eq. 4) to define the stability criterion:
Definition 1 (Stability criterion) For given (η, τ), we say that the dynamics in eq. 3 are stable if all the roots of the corresponding characteristic equation (eq. 4) are inside the unit circle. This condition is equivalent to xt → 0 for any initialization.
We would like to find conditions on when and how can we change the hyperparameters (η, τ) to maintain the same stability criterion. This, for example, will enable us to understand how to compensate for the delay introduced by asynchrony using the learning rate. In order to analyze stability we examine the characteristic equation in eq. 4 and ask: when does the maximal root of this polynomial have unit magnitude?

heading: 2.2.1 THE INTERACTION BETWEEN LEARNING RATE AND DELAY

text:
Recall the characteristic equation zτ+1 − zτ + aη = 0 . (5)
This polynomial has τ +1 roots. In order to ensure stability we require that the root with the maximal magnitude, i.e., the maximal root, will be inside the unit circle. We want to find the threshold learning rate, i.e., the learning rate in which the maximal root is exactly on the unit circle, so we are at the threshold of stability. In appendix B, we show that this requires that
aη = 2 sin
( π
4τ + 2
) . (6)
Using Taylor approximation we get
aη = 2
( π
4τ + 2 +O
( 1
τ3
)) ⇒ 1
aη =
1 π (2τ + 1) +O
( 1
τ
) . (7)
We observe numerically that the error between the exact solution (eq. 6) and the linear approximation (eq. 7) is smaller than 0.05 for τ ≥ 1. Additionally, in the appendix Fig. 6 we demonstrate the high accuracy of the analytic approximation of the relation between 1aη and τ , compared to the numerical evaluation of the value of 1aη for which the maximal root of eq. 5 is on the unit circle.
Implications: We can see that to maintain stability for a given minimum point (or a given sharpness value) the learning rate should be kept inversely proportional to the delay. This implies that given some minimum x∗ for τ = 0 and a corresponding learning rate η0 for which this minimum is stable, we can evaluate the learning rate which will ensure this minimum remain stable for larger delay values. To do this, we first estimate a using eq. 6: aη0 = 2 sin ( π
4·0+2 ) = 2⇒ a = 2η0 . Note that
we do not use eq. 7 to evaluate a since we are interested in τ = 0, while the relation in eq. 7 is only a good approximation for τ ≥ 1. Next, given some delay τ , we substitute a into eq. 7 in order to evaluate a learning rate value for which the stability of all minima is maintained as in τ = 0
η ≈ π a(2τ + 1) = πη0 4(τ + 0.5) . (8)
Related empirical results: In section 3.1 we show empirically the existence of a stability threshold for different values of delay and compare it with our theoretical threshold.
Stochastic delay: In appendix C we extend our theoretical analysis to derive the characteristic equation for A-SGD with stochastic delay, i.e., when τ is drawn from some discrete distribution. We analyze two cases:
• Discrete uniform distribution: τ ∼ unif{a, b}. • Gaussian distribution discrete approximation: ∀µ,∀k = 1, ..., 2µ : Pr(τ = k) = Pr(k − 0.5 ≤ z ≤ k + 0.5) where z ∼ N(µ, 1). In Zhang et al. (2015), the authors showed empirically a delay distribution that resembles to Gaussian distribution.
For both cases, we show numerically that the inverse relation between the learning rate η and the expected delay Eτ still applies.

heading: 2.2.2 THE EFFECTS OF MOMENTUM ON STABILITY

text:
The optimization step of A-SGD with momentum is defined as follows: vt+1 = mvt − η (1−m)∇fn(t−τ)(xt−τ ) xt+1 = xt + vt+1 , (9)
where xt are the weights, vt and ∇fn(t)(xt) are the velocity and gradients at time t, respectively, m is the momentum parameter, η is the step size and τ is the delay. The constant (1−m) term which multiplies the gradients is called "dampening".
In appendix D we extend our theoretical analysis to derive the characteristic equation for A-SGD with momentum. In addition, in section D.2.1 we show numerically that an inverse relation between the learning rate η and the delay τ also applies when using momentum. Particularly,
1. We need to keep the learning rate inversely proportional to the delay. 2. For a given delay, larger momentum values require smaller learning rate for stability.
Therefore, the range of stable learning rates decreases when the momentum parameter (’m’) increases which can make tuning the learning rate difficult, especially with large delays. These results suggest that, to ensure stability when training with large delay, it is recommended to work without momentum.
Relation to previous results: The conclusion that when using A-SGD with high delay values it’s recommended to turn off momentum is consistent with the results of Mitliagkas et al. (2017); Liu et al. (2018). In Liu et al. (2018), for streaming PCA, the authors showed it is necessary to reduce momentum in order to ensure convergence and acceleration through asynchrony. In Mitliagkas et al. (2017) the authors suggested that asynchrony adds implicit momentum to the training process and therefore it is necessary to reduce momentum when working asynchronously.
Next, we discuss a new method for introducing momentum into asynchronous training. This method enables asynchronous training where increasing the momentum parameter improves the stability.
Shifted momentum As discussed, training asynchronously is less stable when momentum is used. We propose to utilize momentum in an alternative method, called shifted momentum, for training asynchronously where the momentum benefits the training process and particularly its stability. We observe this method tends to improve the convergence stability compared to training without momentum (or with the original momentum), as is demonstrated in appendix Fig. 11. We suggest that instead of exchanging the gradients terms, as in eq. 9, we exchange the entire velocity term. This way, each worker calculates the velocity term based in its own gradients and momentum buffer. The formulation of this method is given in the following equation:
vt+1 = mvt−τ − η (1−m)∇fn(t−τ)(xt−τ ) xt+1 = xt + vt+1
or xt+1 = xt +m (xt−τ − xt−τ−1)− η (1−m)∇fn(t−τ)(xt−τ ) (10)
In appendix D we show that the characteristic equation of eq. 10 is zτ+2 − zτ+1 + (η (1−m) a−m) z +m = 0 , (11)
where we use the sharpness term a = λmax (H) as before.
We would like to characterize how the learning rate, delay and momentum affect minima stability. Thus, for each momentum value, we evaluate numerically the value of 1aτ for which the maximal root of eq. 11 is on the unit circle, i.e., when we are on the threshold of stability. In Fig. 4 we see that the inverse relation between the learning rate η and the delay τ still applies when using shifted momentum. Moreover, we can see that in contrast to regular momentum, in shifted momentum, a larger momentum value m enables us to work with larger step-size. Therefore, increasing the momentum improves the stability.
Relation to previous results: The analysis of delayed gradients and velocity might also shed light on the stability of recent asynchronous training approaches, where the workers communicating through weight averaging (Lian et al., 2017; Assran et al., 2019). This method is similar to ours in the sense that each worker holds its own momentum buffer. In parallel to us, (Hakimi et al., 2019) used a variation of shifted momentum with Nesterov momentum (denoted as Multi-ASGD), and demonstrated empirically its generalization benefits, but without a theoretical analysis of this benefit.

heading: 3 EXPERIMENTS

text: In this section, we provide experiments to support our theoretical findings. In the following subsections we: (1) demonstrate the relationship between hyperparameters and minima selection, and (2) show our findings can help improve asynchronous training. Details about the experiments and the implementation can be found in appendix E. 1

heading: 3.1 HOW HYPERPARAMETERS AFFECT MINIMA SELECTION

text:
In this section, we demonstrate how the delay and learning rate affect the accessible minima set. We train with a fixed learning rate.
Minima stability. To demonstrate how delay and learning rate change the stability of a minimum, we start from a model that converged to some minimum and train with different values of delay and learning rate to see if the algorithm leaves that minimum, i.e., the minimum becomes unstable. We do so by training a VGG-11 on CIFAR10 for 10,000 epochs – until we reach a steady state. This training is done without delay, momentum and weight decay. Next, introduce a delay τ , change the learning rate, and continue to train the model. Fig. 2 shows the number of epochs it takes to leave the steady state for a given delay τ and learning rate η. We observe that for certain (τ, η) pairs, the algorithm stays in the minimum (below the black circle) while for others it leaves that minimum after some number of epochs (above the black circle). Importantly, we can see in the right panel in Fig. 2 that, as predicted by the theory, the inverse learning rate, 1/η, where η is the maximal learning rate in which we did not diverge, scales linearly with the delay τ . Additional details are given in appendix F.
Generalization with delay. With delay, we need to adjust the hyperparameters in order to maintain the accessible minima set. However, even if a certain pair of (τ ,η) changes the accessible set, it is possible that there are still accessible minima that generalize well. We perform experiments to investigate what type of generalization we can expect from training with different (τ ,η) pairs. We examine this by training a VGG-11 on CIFAR10 with different (τ ,η) pairs until we reach a steady state (around 6,000 epochs). We use plain SGD without momentum or weight decay. We compare the generalization performance achieved by each pair of learning rate and delay. We repeat the experiment for each pair four times and report mean and standard deviation values. The results are presented in Fig. 3. We can see a linear scaling between the delay values and the learning rate empirical stability threshold, i.e., the learning rate in which the error diverge. In addition, we see that the smallest error (with small variance) is obtained at a learning rate which is somewhat smaller then the learning rate at the empirical threshold of stability, as expected (LeCun et al., 2012).
1Code will be available at https://github.com/paper-submissions/delay_stability
Shifted momentum stability. In subsection 2.2.2 we introduced shifted momentum as a modification of the standard momentum update, in which increasing the momentum value (’m’) improves stability. We validate the stability properties of this method by training a fully connected model with MNIST. The model has 3 layers, 1024 neurons in each layer with ReLU activations. Again, training is done synchronously with a large batch and no weight decay until it reaches a steady state. Then, after reaching a minimum, we continued to train the model, each time with a different triplet of (τ, η,m). We found the maximum learning rate for a given m and τ for which training does not diverge, i.e, this maximum learning rate is at the edge of stability. Fig. 4 depicts the empirical results. The value of a = λmax(H) is estimated at the minimum point using power iteration. Two important results emerge from the graph. First, in order to maintain the stability threshold, the learning rate has to be inversely proportional to the delay, as described in subsection 2.2. Second, with increased momentum values, training becomes more stable. Meaning, we can use larger learning rates with larger momentum value. This stands in contrast to the analysis of the original momentum: as we show numerically in appendix D, increasing the momentum value m in the standard momentum algorithm decreases the stability and hence we need to decrease the learning rate as the momentum value grows larger.

heading: 3.2 IMPROVING ASYNCHRONOUS TRAINING

text:
So far we focused on the steady state behavior of A-SGD, i.e. after many iterations. However, one of the prominent motivations for training asynchronously is to speed-up training. In this subsection, we examine whether our findings are also relevant for improving the accuracy of models trained asynchronously with little to no excess budget of epochs, compared to equivalent large batch regime. In large batch training, the learning rate is scaled as a function of the ratio between the large batch size and the small one. As Shallue et al. (2018) pointed out, we can expect a different learning rate scaling for different models/datasets. We start with a set of models and datasets that follow a square root scaling of the learning rate with the batch size (Hoffer et al., 2017). Recall that our inversely linear scaling rule (i.e., η ∝ 1/τ ) applies to the large batch learning rate. We demonstrate the validity of our inversely linear learning rate scaling by comparing our method with one that does not account for the delay τ . The results are presented in Table 1. We can see that using the small batch regime with delay (ASGD) does not converge at all. With our inversely linear learning rate scaling (+LR), there is a small gap from the equivalent large batch. If we also increase the budget of epochs by 30%, our optimization regime (+30%) generalizes better than the large batch (LB). Because A-SGD has better run-time performance compared to S-SGD (Dutta et al., 2018; Assran et al., 2019), such excess in epochs may still result in improved overall wall clock time (this depends on the hardware details).
ImageNet. We next experiment with ResNet50 and ImageNet, which is a popular benchmark for large batch and asynchronous training because of its complexity and scale. Our focus is asynchronous centralized training with high delay. To the best of our knowledge, the generalization gap is not closed yet in this setting. This is an important setting because of its simplicity and its scalability potential. Optimization details can be found in appendix E.1
In Fig. 5, we compare three A-SGD optimizations trained with ImageNet. As can be seen in the figure, the generalization gap decreases from 4.33% to 1.69%, just by turning off momentum. This corresponds with our analysis and previous results about the relation between momentum and delay. Shifted momentum similarly improves generalization, as can be seen in the right panel. This validates our analysis on the beneficial role the momentum parameter has, using shifted momentum. The learning rate is first scaled linearly with the total batch size of all the workers, as in Goyal et al. (2017). Our inverse linear scaling with delay cancels this increase, and we get a similar learning rate as in regular small batch training. In contrast, when we scaled the learning rate with 1/ √ τ , training did not converge at all.

heading: 4 DISCUSSION

text:
Motivation. A fundamental problem with asynchronous training is that it seems to cause a degradation in generalization, even after training has converged to a steady state. Identifying the root of this problem and resolving it is essential in order to scale parallelism. When approaching this problem, we need to address two issues. First, what is the origin of this generalization gap? Second, how can we decrease this generalization gap?
Conclusions. Our work tackles both questions by studying the behaviour of A-SGD and the relation between the delay and the hyperparameters, e.g., the learning rate and momentum. We analyze A-SGD from the theoretical perspective of dynamical stability. We find that a linear connection between the inverse learning rate and delay is necessary in order to keep the set of accessible minima
the same. In addition, we find how momentum affects the accessible minima set. Our analysis suggests that, when increasing delay, momentum degrades stability. Therefore, it should be turned off in order to maintain performance, as was previously observed empirically (Mitliagkas et al., 2017; Dai et al., 2018; Liu et al., 2018). However, we observe that momentum can have some benefits in improving convergence stability without delay. To maintain these benefits, we propose to modify momentum into a method called shifted momentum. We supply a theoretical analysis of A-SGD with shifted momentum and conclude that, in contrast to regular momentum, this method improves stability when delay is increased. In addition, we supply empirical evidence to support these findings. Lastly, we demonstrate, these findings can improve the generalization performance of distributed training algorithms and bring us one step closer to closing the generalization gap.
Future Directions. It is interesting to examine how the dynamical stability analysis holds during training. Specifically, what implications can we derive from the analysis, before we reach the steady state (e.g., a minimum of the loss), given a tight epochs budget. Our empirical findings in this case (Fig. 13 ) suggest that the optimal learning rate chosen for the steady state may remain optimal during the convergence process. This might be because the hyperparameters at the stability threshold are closely related to the hyperparameters that achieve the optimal convergence speed to that minimum. For example, in GD it easy to show that ηthreshold = 2ηoptimal. Therefore, both stability and convergence rate may be optimally maintained using similar scaling relations in the hyperparameters. This is an interesting question for future study.
Another interesting question for future study is to understand how the stability conditions change when some synchronization is introduced. Several practical A-SGD methods use synchronization to some extent, e.g. (Assran et al., 2019; Chen et al., 2016). For synchronization methods that reduce the delay in expectation, we can intuitively expect an improvement in dynamical stability, from our analysis of stochastic delay (appendix C). However, further analysis of the stability method for each method is required for an exact answer.
Answering these questions may lead to a more precise understanding of the trade-offs between S-SGD and the various A-SGD methods.

heading: A STABILITY ANALYSIS FOR EQ. 2

text:
In this section, we would like to analyze the stability of eq. 2. In particular, we would like to examine if a given minimum point x∗ is stable. To analyze this we suppose we are in a vicinity of some minimum point x∗, which implies∇f(x∗) = 0, and ask whether or not we can converge to x∗ using the theory of dynamical stability. Formally, consider the linearized dynamics around x∗: xt+1 = xt − η (G(x∗;nt−τ ) +∇G (x∗;nt−τ ) (xt−τ − x∗)) , where we denoted G(xt;nt) = ∇fnt(xt) and changed the notation of the time index in n(t) to the subscript, i.e., nt.
Examining the expectation of this equation we obtain
Ext+1 (1) = Ext − η
( 1
N N∑ i=1 G(x∗; i) + 1 N N∑ i=1 ∇G (x∗; i) (Ext−τ − x∗) ) (2) = Ext − ηH (Ext−τ − x∗) ,
where in (1) we used ∀k = 1, ..., N : Pr(nt = k) = 1N , and in (2) we used the fact that x ∗ is a minimum point and thus N∑ i=1 G(x∗; i) = ∇f(x∗) = 0 and denoted H = 1N N∑ i=1 ∇G(x∗; i) =
1 N N∑ i=1 ∇2fi(x∗) .
Using st = Ext − x∗ variable change we obtain st+1 = st − ηHst−τ . (12) From the Spectral Factorization Theorem we can write H = UAU> where U is an orthogonal matrix and A = diag(a1, ..., an) is a diagonal matrix. Substituting this into eq. 12 we obtain: st+1 = st − ηUAU>st−τ . Multiplying the last equation with U> from the left, denoting s̃t = U>st, and using the fact that U is an orthogonal matrix, i.e., U>U = I we get:
s̃t+1 = s̃t − ηAs̃t−τ . Recall that A = diag(a1, ..., aN ) is a diagonal matrix. Therefore, analyzing the dynamics of the last equation is equivalent to analyzing the dynamics of the following N one-dimensional equations:
∀i = 1, ...N : s̃t+1(i) = s̃t(i)− ηais̃t−τ (i) , where st(i) denotes the ith component of the vector st.
In order to ensure stability of the N -dimensional dynamical system, it’s sufficient to require that s̃t+1(1) = s̃t(1)− ηa1s̃t−τ (1) , (13)
where we assume without loss of generality a1 ≥ a2 ≥ · · · ≥ aN . This is true since the dynamics for i = 1 is the first one that loses stability.
To simplify notations we define xt = s̃t(1) and the sharpness term a = a1 = λmax (H) , the maximal singular value of H. Using these notations, we can re-write eq. 13 as:
xt+1 = xt − ηaxt−τ ⇔ xt+1 − xt + ηaxt−τ = 0 . (14)
Since this is a linear time invariant difference equation, of order τ + 1, its solution is, generically:
xt = τ+1∑ i=1 ciz t i ,
where ci are determined by the initial conditions, and zi are the roots of the characteristic equation of eq. 14. To find this characteristic equation, we substitute xt = zt into eq. 14, and obtain zt+1 − zt + ηazt−τ = 0⇔ zτ+1 − zτ + ηa = 0 Note that for xt → 0, it is required that ∀i, |zi| < 1.

heading: B FINDING THE CONDITION FOR THE MAXIMAL ROOT OF THE CHARACTERISTIC EQUATION (EQ. 4) TO BE ON THE UNIT CIRCLE

text:
To simplify the notation, we denote α = aη. Recall the characteristic eq. (eq. 5)
zτ+1 − zτ + α = 0 .
Since we are looking for solution on the unit circle we substitute z = eiθ into the equation where i is the unit imaginary number and θ is some angle. We obtain:
eiθ(τ+1) − eiθτ + α = 0 .
Using eiθ = cos(θ) + i sin(θ) we get
cos (θ(τ + 1)) + i sin (θ(τ + 1))− (cos (θτ) + i sin (θτ)) + α = 0⇒{ cos (θ(τ + 1))− cos (θτ) + α = 0 sin (θ(τ + 1))− sin (θτ) = 0
Using trigonometric product-sum identities we obtain−2 sin ( θ(2τ+1) 2 ) sin ( θ 2 ) + α = 0 2 cos ( θ(2τ+1)
2 ) sin ( θ 2 ) = 0
(15)
If sin ( θ 2 ) = 0 then we get from the first equation α = 0. This is a contradiction since we assume
that α > 0. Thus, sin ( θ 2 ) 6= 0. This implies from the second equation in eq. 15 that
cos
( θ(2τ + 1)
2
) = 0⇒ θ(2τ + 1)
2 = π 2 + πk , for k = 0, 1, ..., 2τ .
Substituting this result into the first equation in eq. 15 and using α > 0 we obtain that the possible solutions must satisfy
α = 2 sin
( (1 + 4k)π
4τ + 2
) , for k = 0, 1, ..., τ .
Also, since sin(x) = sin(π − x) we can eliminate symmetric solution and get
α = 2 sin
( (1 + 4k)π
4τ + 2
) , for k = 0, 1, ..., bτ + 1
2 c .
Note that for k = 0, 1, ..., b τ+12 c: 0 < (1+4k)π 4τ+2 < π 2 and thus sin ( (1+4k)π 4τ+2 ) > 0 is monotonically increasing with k. Since the magnitude of the characteristic equation increases with α and we are interested in the maximal root we are looking for the minimal α value. This value corresponds to k = 0. Thus, the maximal root satisfies:
α = sin
( π
4τ + 2
) .

heading: C STABILITY ANALYSIS FOR EQ. 2 WITH STOCHASTIC DELAY

text:
We assume that τ(t) is drawn from some distribution. In this section, we would like to analyze the stability of eq. 2. In particular, we would like to examine if a given minimum point x∗ is stable. To analyze this we suppose we are in a vicinity of some minimum point x∗, which implies∇f(x∗) = 0, and ask whether or not we can converge to x∗ using the theory of dynamical stability. Formally, consider the linearized dynamics around x∗:
xt+1 = xt − η ( G(x∗;nt−τ(t)) +∇G ( x∗;nt−τ(t) ) ( xt−τ(t) − x∗ )) ,
where we denoted G(xt;nt) = ∇fnt(xt) and changed the notation of the time index in n(t) to the subscript, i.e., nt.
Examining the expectation of this equation we obtain
Ext+1 (1) = Ext − η
( 1
N N∑ i=1 G(x∗; i) + 1 N N∑ i=1 ∇G (x∗; i) (∑ k Ext−k Pr(τ(t) = k)− x∗ ))
(2) = Ext − ηH ∑ k Pr(τ(t) = k) (Ext−k − x∗) ,
where in (1) we used ∀t and ∀k = 1, ..., N : Pr(nt = k) = 1N , and in (2) we used the fact that x ∗ is a minimum point and thus N∑ i=1 G(x∗; i) = ∇f(x∗) = 0 and denoted H = 1N N∑ i=1 ∇G(x∗; i) =
1 N N∑ i=1 ∇2fi(x∗) .
Using st = Ext − x∗ variable change we obtain st+1 = st − ηH ∑ k Pr(τ(t) = k)st−k . (16)
From the Spectral Factorization Theorem we can write H = UAU> where U is an orthogonal matrix and A = diag(a1, ..., an) is a diagonal matrix. Substituting this into eq. 16 we obtain:
st+1 = st − ηUA ∑ k Pr(τ(t) = k)U>st−k .
Multiplying the last equation with U> from the left, denoting s̃t = U>st, and using the fact that U is an orthogonal matrix, i.e., U>U = I we get:
s̃t+1 = s̃t − ηA ∑ k Pr(τ(t) = k)s̃t−k .
Recall that A = diag(a1, ..., aN ) is a diagonal matrix. Therefore, analyzing the dynamics of the last equation is equivalent to analyzing the dynamics of the following N one-dimensional equations:
∀i = 1, ...N : s̃t+1(i) = s̃t(i)− ηai ∑ k Pr(τ(t) = k)s̃t−k(i) ,
where st(i) denotes the ith component of the vector st.
In order to ensure stability of the N -dimensional dynamical system, it’s sufficient to require that s̃t+1(1) = s̃t(1)− ηa1 ∑ k Pr(τ(t) = k)s̃t−k(1) , (17)
where we assume without loss of generality a1 ≥ a2 ≥ · · · ≥ aN . This is true since the dynamics for i = 1 is the first one that loses stability.
To simplify notations we define xt = s̃t(1) and the sharpness term a = a1 = λmax (H) , the maximal singular value of H. Using these notations, we can re-write eq. 13 as:
xt+1 = xt − ηa ∑ k Pr(τ(t) = k)xt−k ⇔ xt+1 − xt + ηa ∑ k Pr(τ(t) = k)xt−k = 0 . (18)
The characteristic equation of eq. 18 is z − 1 + ηa ∑ k Pr(τ(t) = k)z−k = 0 (19)
C.1 STABILIY ANALYSIS
We would like to find conditions on when and how can we change the hyperparameters (η,Eτ) to maintain the same stability criterion (definition 1). In order to analyze stability we examine the characteristic equation in eq. 18 and ask: when does the maximal root of this polynomial have unit magnitude?
1. τ ∼ unif{a, b}:
z − 1 + ηa b∑
k=a
1 n z−k = 0 (20)
2. Gaussian distribution discrete approximation: ∀µ,∀k = 1, ..., 2µ : Pr(τ = k) = Pr(k − 0.5 ≤ z ≤ k + 0.5) where z ∼ N(µ, 1):

heading: D THEORETICAL ANALYSIS WITH MOMENTUM

text:
In this section, we repeat the steps of the analysis in section 2 for A-SGD with momentum.
D.1 PRELIMINARIES AND PROBLEM SETUP
Consider the problem of minimizing the empirical loss
f(x) = N∑ i=1 fi(x) , (21)
where N is the number of samples, x ∈ Rd and ∀i = 1, ..., N : fi : Rd → R is twice continuously differentiable function. The update rule for minimizing f(x) in eq. 21 using asynchronous stochastic gradient descent with momentum (A-MSGD) is given by
vt+1 = mvt − η (1−m)G(xt−τ ;nt−τ ) xt+1 = xt + vt+1
or
xt+1 = xt +m (xt − xt−1)− η (1−m)G(xt−τ ;nt−τ ) , (22)
where m is the momentum, η is the learning rate, G(xt;nt) = ∇fnt(xt), nt is a random selection process of a sample from {1, 2, . . . , N} at iteration t, and τ is some delay due the asynchronous nature of our training.
Our main goal is understanding how different factors such as the gradient staleness τ , and hyperparameters such as the momentum m and learning rate η affect the minima selection process in neural networks. In other words, we would like to understand how the different hyperparameters interact to divide the global minima of our network to two sets: those we can converge to and those we cannot converge to.
To analyze this we suppose we are in a vicinity of a minimum point x∗, which implies∇f(x∗) = 0, and ask whether or not we can converge to x∗ using the theory of dynamical stability. Formally, consider the linearized dynamics around some minimum point x∗:
xt+1 = xt +m (xt − xt−1)− η (1−m) (G(x∗;nt−τ ) +∇G (x∗;nt−τ ) (xt−τ − x∗)) .
Examining the first moment of this equation we obtain
Ext+1 = Ext +m (Ext − Ext−1)− η (1−m)
( 1
N N∑ i=1 G(x∗; i) + 1 N N∑ i=1 ∇G (x∗; i) (Ext−τ − x∗) ) = Ext +m (Ext − Ext−1)− η (1−m)H (Ext−τ − x∗) ,
where 1N N∑ i=1 G(x∗; i) = 1N∇f(x ∗) = 0 since x∗ is a minimum point and H = 1N N∑ i=1 ∇G(x∗; i) =
1 N N∑ i=1 ∇2fi(x∗) .
Using st = Ext − x∗ variable change we obtain
st+1 = st +m (st − st−1)− η (1−m)Hst−τ . (23)
From the Spectral Factorization Theorem we can write H = UDU> where U is an orthogonal matrix and A = diag(a1, ..., an) is a diagonal matrix. Substituting this into eq. 23 we obtain:
st+1 = st +m (st − st−1)− η (1−m)UAU>st−τ .
Multiplying the last equation with U> from the left, denoting s̃t = U>st, and using the fact that U is an orthogonal matrix, i.e., U>U = I we get:
s̃t+1 = s̃t +m (s̃t − s̃t−1)− η (1−m)As̃t−τ .
Recall that A = diag(a1, ..., aN ) is a diagonal matrix. Therefore, analyzing the dynamics of the last equation is equivalent to analyzing the dynamics of the following N one-dimensional equations:
∀i = 1, ...N : s̃t+1(i) = s̃t(i) +m (s̃t(i)− s̃t−1(i))− η (1−m) ais̃t−τ (i) , where st(i) denotes the ith component of the vector st.
In order to ensure stability of the N -dimensional dynamical system, it’s sufficient to require that
s̃t+1(1) = s̃t(1) +m (s̃t(1)− s̃t−1(1))− η (1−m) a1s̃t−τ (1) , (24) where we assume without loss of generality a1 ≥ a2 ≥ · · · ≥ aN . This is true since the dynamics for a1 is the first one that loses stability.
To simplify notations we define xt = s̃t(1) and the sharpness term a = a1 = λmax (H), the maximal singular value of H. Using these notations, we can re-write eq. 24 as:
xt+1 = xt +m (xt − xt−1)− η (1−m) axt−τ ⇔ xt+1 − (m+ 1)xt +mxt−1 + η (1−m) axt−τ = 0 . (25)
The characteristic equation of this difference equation is
zτ+1 − (1 +m) zτ +mzτ−1 + η (1−m) a = 0 . (26)
D.2 STABILITY ANALYSIS
We would like to find conditions on when and how can we change the hyperparameters (m, η, τ) to maintain the same stability criterion (definition 1). In order to analyze stability we examine the characteristic equation in eq. 26 and ask: when does the maximal root of this polynomial have unit magnitude?
D.2.1 THE INTERACTION BETWEEN LEARNING RATE, DELAY, AND MOMENTUM
For the general case and different values of momentum, we evaluate numerically the value of 1aτ for which the maximal root of the general characteristic equation (eq. 26) is on the unit circle, i.e., when we are on the threshold of stability. We show the results in Fig. 10. We observe that there is still an inverse relation between the learning rate and delay when using momentum. Specifically, we observe that:
1. We need to decrease the learning rate as the delay increases. 2. The lines becomes more steep for larger values of momentum. This implies that, for a given
delay, maintaining stability for larger values of momentum requires smaller learning rate. Therefore, it is easier to maintain stability with m = 0.
D.3 SHIFTED MOMENTUM
In this section, we repeat the steps of the analysis in section D.1 for A-SGD with shifted momentum. The update rule for minimizing f(x) in eq. 21 using asynchronous stochastic gradient descent with shifted momentum is given by
vt+1 = mvt−τ − η (1−m)∇fn(t−τ)(xt−τ ) xt+1 = xt + vt+1
or
xt+1 = xt +mvt−τ − η (1−m)∇fn(t−τ)(xt−τ ) = xt +m (xt−τ − xt−τ−1)− η (1−m)∇fn(t−τ)(xt−τ )
Following the same steps as in section D.1 we obtain the following one-dimensional equation:
xt+1 = xt +m (xt−τ − xt−τ−1)− η (1−m) axt−τ ⇔ xt+1 − xt + (η(1−m)a−m)xt−τ +mxt−τ−1 = 0 .
The last equation characteristic equation is:
zτ+2 − zτ+1 + (η(1−m)a−m)z +m = 0 .

heading: E EXPERIMENTS DETAILS

text:
We experimented with a set of popular classification tasks: MNIST (Lecun et al., 1998), CIFAR10, CIFAR100 (Krizhevsky, 2009) and ImageNet (Deng et al., 2009). In order to incorporate delay in the experiments, we keep replicas of the model parameters and perform SGD step with different replica at a time according to a round robin scheduling. The implementation is done with PyTorch framework. Section C in the appendix shows that constant delay, i.e., round robin, acts similarly in terms of stability to several other delay distributions. In addition, Zhang et al. (2015) shows empirically with realistic distribution that constant delay distribution is a good model for delay distribution. Fig. 9 shows empirically that round robin scheduling works in a similar way to discrete Gaussian distribtion.
Although stochastic gradient descent with momentum was originally introduced with the dampening term, many drop the latter when using momentum SGD. This is true for most DNN frameworks like Caffe, PyTorch, and TensorFlow where the dampening term is set to zero by default. As Shallue et al.
(2018); Yan et al. (2018) pointed out, without dampening, the momentum scales the learning rate so that the effective learning rate ηeff is equal to ηeff = η1−m after sufficiently many iterations T (i.e., when T →∞). Because we alter the momentum value in our experiments, we use dampening and scale the learning rate accordingly to keep the effective learning rate the same as without dampening.
E.1 IMAGENET OPTIMIZATION DETAILS
Our baseline is Goyal et al. (2017) large batch training. They trained ResNet50 for 90 epochs with a large batch of 8192 samples and reached the same accuracy level of 23.8% as the small batch training. They used linear scaling with the batch size (scaling the learning rate by the large batch size divided by the baseline batch size) and a warm-up of the learning rate at the first 5 epochs. We use 32 workers, each with a minibatch size of 256.

heading: F MINIMA STABILITY EXPERIMENTS

text:
In this section, we provide additional details about the minima stability experiment presented in section 3. As discussed in section 3, we are interested to examine for which pairs of (τ, η) the minima stability remains the same.
In Fig. 12 we show the validation error of such (τ, η) pairs as a function of epochs. We note that these graphs are the same experiment as in Fig. 2. As can be seen, for larger learning rates, it takes less epochs to leave the minimum. It is interesting to see that after leaving the minimum, the A-SGD algorithm converges again to a minimum with generalization as good as the baseline (at τ = 0). This suggests that the minima selection process of A-SGD is affected by the whole optimization path. In other words, suppose we start the optimization from a minimum with good generalization (since it was selected using optimization with τ = 0), and then it becomes unstable due to a change in the values of (η, τ), as in this experiment. These results in Fig. 12 suggest we typically converge to a stable minimum with similar generalization properties, possibly nearby the original minimum. In contrast, if we start to train from scratch using the same (η, τ) pair which lost stability in our experiment, we typically get a generalization gap (as observed in our experiments), which suggests the optimization path might have taken a very different path from the start, leading to other regions with worse generalization than the original minimum.

heading: G ADDITIONAL EXPERIMENTS

text: We examine what generalization performance we can expect during training, before reaching a minimum. We do so by using the same experiment in Fig. 1. We sample the validation error of each learning rate at different epochs. Results are presented in Fig. 13. Using our suggested learning rate modification (dividing the baseline large batch learning rate by the delay) gives the near optimal learning rate that can be seen in the figure.

title: AT STABILITY’S EDGE: HOW TO ADJUST HYPER- PARAMETERS TO PRESERVE MINIMA SELECTION IN ASYNCHRONOUS TRAINING OF NEURAL NETWORKS?


summaries: This paper presents a new model-based reinforcement learning method, termed hindsight modelling. The method works by training a value function which, in addition to depending on information available at the present time is conditioned on some learned embedding of a partial future trajectory. A model is then trained to predict the learned embedding based on information available at the current time-step. This predicted value is fed in place of the actual embedding to the same value model, to generate a value prediction for the current time-step. So instead of just learning a value function based on future returns, the method uses a two-step process of learning an embedding of value relevant information from the future and then learns to predict that embedding.

summaries: The paper proposes a way to learn better representation for RL by employing a hindsight model-based approach. The reasoning is that during training, we can observe the future trajectory and use features from it to better predict past values/returns. However to make this practical, the proposed approach fits an approximator to predict these features of the future trajectory from the current state and then subsequently, use them to predict the value. The authors claim that this extra information can be used to learn a better representation in some problems and lead to faster learning of good policies (or optimal value functions)

summaries: Value-Driven Hindsight Modelling proposes a method to improve value function learning. The paper introduces the hindsight value function which estimates the expected return at a state conditioned on the future trajectory of the agent. How use this hindsight value function is not obvious, since an agent does not have access to the future states needed in order to take actions (for Q-Learning) and the hindsight value function is a biased gradient estimator for training policy gradient methods.

abstractText: Value estimation is a critical component of the reinforcement learning (RL) paradigm. The question of how to effectively learn predictors for value from data is one of the major problems studied by the RL community, and different approaches exploit structure in the problem domain in different ways. Model learning can make use of the rich transition structure present in sequences of observations, but this approach is usually not sensitive to the reward function. In contrast, model-free methods directly leverage the quantity of interest from the future but have to compose with a potentially weak scalar signal (an estimate of the return). In this paper we develop an approach for representation learning in RL that sits in between these two extremes: we propose to learn what to model in a way that can directly help value prediction. To this end we determine which features of the future trajectory provide useful information to predict the associated return. This provides us with tractable prediction targets that are directly relevant for a task, and can thus accelerate learning of the value function. The idea can be understood as reasoning, in hindsight, about which aspects of the future observations could help past value prediction. We show how this can help dramatically even in simple policy evaluation settings. We then test our approach at scale in challenging domains, including on 57 Atari 2600 games.

authors:
Marc G Bellemare
Yavar Naddaf
Joel Veness
Michael Bowling

title: The arcade learning environment: An evaluation platform for general agents

venue: Journal of Artificial Intelligence Research,

authors:
Lars Buesing
Theophane Weber
Yori Zwols
Sebastien Racaniere
Arthur Guez
Jean-Baptiste Lespiau
Nicolas Heess

title: Woulda, coulda, shoulda: Counterfactually-guided policy search

venue: arXiv preprint arXiv:1811.06272,

authors:
Lasse Espeholt
Hubert Soyer
Remi Munos
Karen Simonyan
Volodymir Mnih
Tom Ward
Yotam Doron
Vlad Firoiu
Tim Harley
Iain Dunning

title: Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures

venue: arXiv preprint arXiv:1802.01561,

authors:
Carles Gelada
Saurabh Kumar
Jacob Buckman
Ofir Nachum
Marc G Bellemare

title: Deepmdp: Learning continuous latent space models for representation learning

authors:
Arthur Guez
Mehdi Mirza
Karol Gregor
Rishabh Kabra
Sebastien Racaniere
Theophane Weber
David Raposo
Adam Santoro
Laurent Orseau
Tom Eccles

title: An investigation of model-free planning

venue: In International Conference on Machine Learning,

authors:
Tuomas Haarnoja
Aurick Zhou
Pieter Abbeel
Sergey Levine

title: Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor

venue: arXiv preprint arXiv:1801.01290,

authors:
Max Jaderberg
Volodymyr Mnih
Wojciech Marian Czarnecki
Tom Schaul
Joel Z Leibo
David Silver
Koray Kavukcuoglu

title: Reinforcement learning with unsupervised auxiliary tasks

venue: arXiv preprint arXiv:1611.05397,

authors:
Steven Kapturowski
Georg Ostrovski
Will Dabney
John Quan
Remi Munos

title: Recurrent experience replay in distributed reinforcement learning

venue: In International Conference on Learning Representations,

authors:
Volodymyr Mnih
Koray Kavukcuoglu
David Silver
Andrei A Rusu
Joel Veness
Marc G Bellemare
Alex Graves
Martin Riedmiller
Andreas K Fidjeland
Georg Ostrovski

title: Human-level control through deep reinforcement learning

authors:
Lerrel Pinto
Marcin Andrychowicz
Peter Welinder
Wojciech Zaremba
Pieter Abbeel

title: Asymmetric actor critic for image-based robot learning

venue: arXiv preprint arXiv:1710.06542,

authors:
John Schulman
Filip Wolski
Prafulla Dhariwal
Alec Radford
Oleg Klimov

title: Proximal policy optimization algorithms

venue: arXiv preprint arXiv:1707.06347,

authors:
Richard S Sutton
Andrew G Barto

title: Reinforcement learning: An introduction

authors:
Richard S Sutton
Joseph Modayil
Michael Delp
Thomas Degris
Patrick M Pilarski
Adam White
Doina Precup

title: Horde: A scalable real-time architecture for learning knowledge from unsupervised sensorimotor interaction

venue: In The 10th International Conference on Autonomous Agents and Multiagent Systems-Volume

title: Model regularization for stable sample rollouts

venue: In Proceedings of the Thirtieth Conference on Uncertainty in Artificial Intelligence,

authors:
Hado Van Hasselt
Arthur Guez
David Silver

title: Deep reinforcement learning with double qlearning

venue: In Thirtieth AAAI conference on artificial intelligence,

authors:
Vladimir Vapnik
Rauf Izmailov

title: Learning using privileged information: similarity control and knowledge transfer

venue: Journal of machine learning research,

authors:
Yuke Zhu
Ziyu Wang
Josh Merel
Andrei A. Rusu
Tom Erez
Serkan Cabi
Saran Tunyasuvunakool
János Kramár
Raia Hadsell
Nando de Freitas
Nicolas Heess

title: Reinforcement and imitation learning for diverse visuomotor

venue: skills. CoRR,

heading: 1 INTRODUCTION

text:
Consider a baseball player trying to perfect their pitch. The player performs an arm motion and releases the ball towards the batter, but suppose that instead of observing where the ball lands and the reaction of the batter, the player only gets told the result of the play in terms of points or, worse, only gets told the final result of the game. Improving their pitch from this experience appears hard and inefficient, yet this is essentially the paradigm we employ when optimizing policies in model-free reinforcement learning. The scalar feedback that estimates the return from a state (and action), encoding how well things went, drives the learning while the accompanying observations that may explain that result (e.g. flight path of the ball or the way the batter anticipated and struck the incoming baseball) are ignored. To intuitively understand how such information could help value prediction, consider a simple discrete Markov chain X → Y → Z, where Z is the scalar return and X is the observation from which we are trying to predict Z. If the space of possible values of Y is smaller than X , then it may be more efficient to estimate both P (Y |X) and P (Z|Y ) rather than directly estimating P (Z|X).1 In other words observing and then predicting Y can be advantageous to directly estimating the signal of interestZ. Model-based RL approaches would duly exploit the observed Y (by modeling the transition Y |X), but Y would, in general scenarios, contain information that is irrelevant to Z and hard to predict. Building a full high-dimensional predictive model to indiscriminately estimate all possible future observations, including potentially chaotic details of the ball trajectory and the spectators’ response, is a challenge that may not pay off if the task-relevant predictions (e.g., was the throw accepted, was the batter surprised) are error-ridden. Model-free RL methods directly consider the relation X to Z, and focus solely upon predicting and optimising this goal, rather than attempting to learn the full dynamics. These methods have recently dominated the literature, and have attained the best performance in a wide array of complex problems with high-dimensional observations (Mnih et al., 2015; Schulman et al., 2017; Haarnoja et al., 2018; Guez et al., 2019).
1In the discrete case, this follows from a counting argument from the size of the probability tables involved.
In this paper, we propose to augment model-free methods with a lightweight model of future quantities of interest. The motivation is to model only those parts of the future observations (Y ) that are needed to obtain better value predictions. The major research challenge is to learn, from observational data, which aspects of the future are important to model (i.e. what Y should be). To this end, we propose to learn a special value function in hindsight that receives future observations as an additional input. This learning process reveals features of the future observations that would be most useful for value prediction (e.g. flight path of the ball or the reaction of the batter), if provided by an oracle. These important features are then predicted, in advance, using only information available at test time (at the time of releasing the baseball, we knew the identity of the batter, the type of throw and spin given to the ball). Learning these value-relevant features can help representation learning for an agent and provide an additional useful input to its value and policy. Experimentally, hindsight value functions surpassed the performance of model-free RL methods in a challenging association task (Portal Choice). When hindsight value functions were added to the prior state-of-the-art RL method for Atari games, they significantly increased median performance from 833% to 965%.

heading: 2 BACKGROUND AND NOTATION

text:
We consider a reinforcement learning setting whereby an agent learns from interaction in a sequential decision-making environment (Sutton & Barto, 2011). An agent’s policy π, mapping states to an action distribution, is executed to obtain a sequence of rewards and observations as follows. At each step t, after observing state st, the policy outputs an action at, sampled from π(A|st), and obtains a scalar reward rt and the next-state st+1 from the environment. The sum of discounted rewards from state s is the return denoted by G = ∑∞ t=0 γ
tRt, with γ < 1 denoting the discount factor. Its expectation, as a function of the starting state, is called the value function, vπ(s) = Eπ[G|S0 = s]. An important related quantity is the action-value, or Q-value, which corresponds to the same expectation with a particular action executed first: qπ(s, a) = Eπ[G|S0 = s,A0 = a]. The learning problem consists in adapting the policy π in order to achieve a higher value vπ . This usually entails learning an estimate of vπ for the current policy π, this is the problem we focus on in this paper.
Note that in practice we are interested in partially-observed environments where the state of the world is not directly accessible. For this case, we can think of replacing the observed state s in the case of the fully-observed case by a learned function that depends on past observations.

heading: 3 VALUE LEARNING

heading: 3.1 DIRECT LEARNING

text:
A common approach to estimate v (or q) is to represent it as a parametric function vθ (or qθ) and directly update its parameters based on sample returns of the policy of interest. Value-based RL algorithms vary in how they construct a value target Y from a single trajectory. They may regress vθ towards the Monte-Carlo return (Yt = Gt), or exploit sequentiality in the reward process by relying on a form of temporal-difference learning to reduce variance (e.g. the TD(0) target Yt = Rt + γvθ(St+1)). For a given target definition Y , the value loss Lv to derive an update for θ is: Lv(θ) = 12Es[(vθ(s) − Y )
2]. In constructing a target Yt based on a trajectory of observations and rewards from time t, the observations are either unused (for a Monte Carlo return) or only indirectly exploited (when bootstrapping to obtain a value estimate). In all cases, the trajectory is distilled into a scalar signal that estimates the return of a policy, and other relevant aspects of future observations are discarded. In particular in partially observed domains or domains with highdimensional observation spaces it can be difficult to discover correlations with this noisy signal.

heading: 3.2 MODEL-BASED APPROACH

text:
An indirect way of estimating the value is to first learn a model of the dynamics. For example a 1-step observation model mθ learns to predict the conditional distribution st+1, rt|st, at. Then a value estimate v(s) for state s can be obtained by autoregressively rolling out the model (until the end of the episode or to a fixed depth with a parametric value bootstrap).
The model is trained on potentially much richer data than the return signal since it exploits all information in the trajectory. Indeed, the observed transitions between states can reveal the structure behind a sparse reward signal. A drawback of classic model-based approaches is that they predict a high-dimensional signal, a task which may be costly and harder than directly predicting the scalar value . As a result, the approximation of the dynamics mθ may contain errors where it matters most for predicting the value (Talvitie, 2014). Although the observations carry all the data from the environment, most of it is not essential to the task (Gelada et al., 2019). The concern that modeling all observations is expensive also applies when the model is not used for actual rollouts but merely for representation learning. So while classic model-based methods fully use this highdimensional signal at some cost, model-free methods take the other extreme to focus only on the most relevant low-dimensional signal (the scalar return). Below we propose a method that strikes a balance between these paradigms.

heading: 3.3 HINDSIGHT VALUE AND MODEL

text:
We introduce a new value function estimate that can only be computed at training time, the hindsight value function v+. This value still estimates the expected return from a state st but it is further conditioned on k additional observations τ+t = st+1, st+2, . . . st+k occurring after time t: 2
v+(st, τ + t ) ≈ E[G|S0 = st, . . . , Sk = st+k]. (1)
Furthermore, we require v+ to follow this particular parametric structure:
v+(st, τ + t ; θ) = ψθ1(f(st),φθ2(τ + t )), (2)
where θ = (θ1, θ2), which forces information about the future trajectory through some vector-valued function φ ∈ Rd. Intuitively, v+ is estimating the expected return from a past time point using privileged access to future observations. Note that if k is large enough, then v+ simply estimates the empirical return from time t given access to the state trajectory. However, if k is small and φ is lowdimensional, then φ becomes a bottleneck representation of the future trajectory τ+t . By learning in hindsight, we identify features that are maximally useful to predict the return on the trajectory from time t. The hindsight value function is not a useful quantity by itself, since – because of its use of privileged future observations – we cannot readily use it at test time. Furthermore, it cannot be used as a baseline either, as when computing the policy gradient it will yield a biased gradient estimator. . Instead, the idea is to learn a model φ̂ of φ, that can be used at test time. We conjecture that if privileged features φ are useful for estimating the value, then the model of those features will also be useful for estimating the value function. We propose to learn the approximate expectation model φ̂η2(s) conditioned on the current state s and parametrized by η2, minimizing the following squared loss:
Lmodel(η2) = Es,τ+ [‖φθ2(s, τ+)− φ̂η2(s)‖22] (3)
where the expectation is taken over the distribution of states and partial trajectories τ+ resulting from that state.
The approximate model φ̂ can then be leveraged to obtain a better model-based value estimate vm(s; η) = ψη1(f(s), φ̂η2(s)). Although φ̂(s) cannot contain more information than included already in the state s, it can still benefit from having being trained using a richer signal before the value converges. Figure 3 summarizes the relation between the different quantities.

heading: 3.4 ILLUSTRATIVE EXAMPLE

text:
We consider the following example to illustrate how the approaches to estimating the value function can differ. There are no actions in this example3 and each episode consists of a single transition from initial state s to terminal state s′, with a reward r(s, s′) on the way.
Each instance of this example is parametrized by a square matrix W and a vector b sampled from a unit normal distribution, which determine the uncontrolled MDP. Initial states s are of dimension D and sampled from a multivariate unit normal distribution (si ∼ N(0, 1) for all state dimension i).
2In general τ+t can be defined to include full future transitions, actions and rewards. 3This can be understood as a Markov Reward Process or a policy evaluation setting
Given s = ( s1s2 ), where s1 and s2 are of dimension D1 and D2 (D = D1 + D2), the next state s′ = ( s′1 s′2 ) is determined according to the transition function: s′1 = MLP(s)+ and s ′ 2 = σ(Ws2+ b) where σ is the Heaviside function. s′1 acts as a distractor here, with additive noise ∼ N(0, 1). The reward obtained is r(s, s′) = ∑ i s (i) 1 ∑ i s ′(i) 2 / √ D. The true value in the start state is also v(s) = r(s, s′).
The key aspect of this domain is that s′ reveals structure that helps predict the value function in the start state s. This is made visually obvious in the trajectories sampled in this domain shown in Figure 1.
Let us consider how the different approaches to learning values presented above fare in this problem. For direct learning, the value from v(s′) is 0 since s′ is terminal, so any n-step return is identical to the Monte-Carlo return, that is, the information present in observation s′ is not leveraged. Results from learning v from s given the return is presented in Figure 2 (blue curve). A model-based approach first predicts s′ from s, then attempts to predict the value given s and the estimated next state. When increasing the input dimension, given a fixed capacity, the model does not focus its attention on the reward-relevant structure in s′ and makes error where it matters most. As a result, it can struggle to learn v faster than a model-free estimate (cf. red curve in Figure 2). When learning in hindsight, v+ can directly exploit the revealed structure in the observation of τ+, and as a result the hindsight value learns faster than the regular causal model-free estimate (cf. dotted yellow curve in Figure 2). This drives the learning of φ and its model φ̂, which directly gets trained to predict these useful features for the value. As a result, vm also benefits and learns faster than the regular v estimate on this problem (cf. green curve in Figure 2).

heading: 3.5 WHEN IS IT ADVANTAGEOUS TO MODEL IN HINDSIGHT?

text:
To understand the circumstances in which hindsight modelling provides a better value estimate, we first consider an analysis that relies on the following assumptions. Suppose that vmθ is sharing the same function ψ as v+ (i.e., θ1 = η1), and let ψ be linear. If we write ψθ1(f,φ) = ( ω1 ω2 ) > ( f φ ) + b,
where θ1 = (ω1, ω2), then we have for fixed values of the parameters:
E[(vm(s; η)− v+(s, τ+; θ)2] = E[‖ω>2 (φ(τ+; θ2)− φ̂(s; η2))‖2] (4) ≤ E[‖ω2‖2‖φ(τ+)− φ̂(s)‖2] (5) = ‖ω2‖2Lmodel(η2), (6)
using the Cauchy-Schwarz inequality. Let L define the value error for a particular value function v: L(v) = E[(v(s)−G)2] and L(v+) = E[(v+(s, τ+)−G)2]. Then we have:
L(vm) = E[(vm(s)− v+(s, τ+) + v+(s, τ+)−G)2] (7) ≤ 2(‖ω2‖2Lmodel(η2) + L(v+)), (8)
using the fact that E[(X +Y )2] ≤ 2(E[X2] +E[Y 2]) for random variables X and Y . If we assume L(v+) = CL(v) with C < 0.5 (i.e., estimating the value in hindsight with more information is an easier learning problem), then the following holds:
Lmodel(η2) < (1− 2C)L(v)
2‖ω2‖2 =⇒ L(vm) < L(v). (9)
In other words, this relates how small the modeling error needs to be to guarantee that the value error for vm is smaller than the value error for the direct estimate v. The modeling error can be large for different reasons. If the environment or the policy is stochastic, then there is some irreducible modeling error for the deterministic model. Even in these cases, a small C can make hindsight modeling advantageous. The modeling error could also be high because predicting φ is hard. For example, it could be that φ essentially encodes the empirical return, which means predicting φ is at least as hard as predicting the value function (Lmodel(η2) ≥ L(v)). Or it could be that φ is highdimensional, this could cause both a hard prediction problem but also would cause the acceptable threshold for Lmodel to decrease (since ‖θ2‖2 will grow). We address some of these concerns with specific architectural choices like v+ having a limited view on future observations and having low dimensional φ (see next section). Note that the analysis above ignores any advantage that could be obtained from representation learning when training φ̂ (if the state encoding function f shares parameters with φ̂).

heading: 4 ARCHITECTURE

text:
The architecture for Hindsight Modelling (HiMo) we found to work at scale and tested in the experimental section of the paper is described here. To deal with partial observability, we employ a recurrent neural network, the state-RNN, which replaces the state st with a learned internal state ht, a function of the current observation ot and past observations through ht−1: ht = f(ot, ht−1; η3), where we have extended the parameter description of vm as η = (η2, η1, η3). The model-based value function vm and the hindsight value function v+ share the same internal state representation h, but the learning of v+ assumes h is fixed (we do not backpropagate through the state-RNN in hindsight). In addition, we force φ̂ to only be learned through Lmodel, so that vm uses it as an additional input.
To summarize:
v+(ht, ht+k; θ) = ψθ1(ht,φθ2(ht+k)), (10)
vm(ht; η) = ψη1(ht, φ̂η2(ht)), (11)
with the bar notation denoting quantities treated as non-differentiable (i.e. where the gradient is stopped). The different losses in the HiMo architecture are combined in the following way:
L(θ, η) = Lv(η) + αLv+(θ) + βLmodel(η). (12)
A diagram of the architecture is presented in Figure 3, and further implementation details can be found in the appendix.
This architecture can be straightforwardly generalized to cases where we also output a policy πη for an actor-critic setup, providing h and φ̂ as inputs to a policy network.4 For a Q-value based algorithm like Q-learning, we predict a vector of values qm and q+ instead of vm and v+. Computing v+ and training φ̂ can be done in an online fashion by simply delaying the updates by k steps (just like the computation of an n-step return).

heading: 5 EXPERIMENTS

text: The illustrative example in Section 3.4 demonstrated the positive effect of hindsight modeling in a simple policy evaluation setting. In this section, we now explore these benefits in the context of policy optimization in challenging domains, a custom navigation task called Portal Choice, and Atari 2600. To demonstrate the generality and scalability of our approach we test hindsight value functions in the context of two high-performance RL algorithms, IMPALA (Espeholt et al., 2018) and R2D2 (Kapturowski et al., 2019).

heading: 5.1 PORTAL CHOICE TASK

text:
The Portal Choice (Fig. 4) is a two-phase navigation task where, in phase one an agent is presented with a contextual choice between two portals, whose positions vary between episodes. The position of the portal determines its destination in phase two, one of two different goal rooms (green and red rooms). Critically, the reward when terminating the episode in the goal room depends on both the color of the goal room in phase two and a visually indicated combinatorial context shown in the first phase. If the context matches the goal room color, then a reward of 2 is given, otherwise the reward is 0 when terminating the episode (see appendix for the detailed description).
An easy suboptimal solution is to select the portal at random and finish the episode in the resulting goal room by reaching the goal pixel, which will result in a positive reward of 1 on average. A more
4In this case, the total loss also contains an actor loss to update πη and a negative entropy loss.
difficult strategy is to be selective about which portal to take depending on the context, in order to get the reward of 2 at each every episode. A model-free agent has to learn the joint mapping from contexts and portal positions to rewards. Even if the task is not visually complex, the context is combinatorial in nature (the agent needs to count randomly placed pixels) and the joint configuration space of context and portal is fairly large (around 250M). Since the mapping from portal position to rooms does not depend on context, learning the portal-room mapping independently is more efficient in this scenario.
For this domain, we implemented the HiMo architecture within a distributed actor-critic agent, named IMPALA proposed by Espeholt et al. (2018). In this case, the target Yt to train vm (used as a critic in this context) and v+ is the V-trace target (Espeholt et al., 2018) to account for offpolicy corrections between the behavior policy and the learner policy. The actor shares the same network as the critic and receives h and φ̂ as inputs.
We found that HiMo+IMPALA learned reliably faster to reach the optimal behavior, compared to the vanilla IMPALA baseline that shared the same network capacity (see Figure 5a). The hindsight value v+ rapidly learns to predict whether the portal-context association is rewarding based on seeing the goal room color in the future. Then φ learns to predict the new information from the future that helps that prediction: the identity of the room (see analysis Fig 5c). The prediction of φ becomes effectively a model of the mapping from portal to room identity (since the context does not correlate with the room identity). Having access to such mapping through φ̂ helped the value prediction (Fig 5b), which led to better action selection. Note that if the two rooms were visually indistinguishable, for example with no red/green rooms separation, HiMo would not be able to offer any advantage over its model-free counterpart.

heading: 5.2 ATARI

text:
We tested our approach in Atari 2600 videogames using the Arcade Learning Environment (Bellemare et al., 2013). We added HiMo on top of Recurrent Replay Distributed DQN (R2D2), a DQN-
based distributed architecture introduced by Kapturowski et al. (2019) which achieved state-of-theart scores in Atari games.
In this value-based setting, HiMo trains qm(·, ·; η) and q+(·, ·; θ) based on n-step return targets:
Yt = g ( n−1∑ m=0 γmRt+m + γ ng−1 ( qm(St+n, A ∗; η−) )) , (13)
where g is an invertible function, η− are the periodically updated target network parameters (as in DQN by Mnih et al. (2015)), andA∗ = argmaxa q
m(St+n, a; η) (the Double DQN update proposed by Van Hasselt et al. (2016)). The details of the architecture and hyperparameters are described in the appendix.
We ran our approach on 57 Atari games for 200k gradient steps (around 1 day of training), with 3 seeds for each game. The evaluation averages the score between 200 episodes across seeds, each lasting a maximum of 30 minutes each and starting a random number (up to 30) of no-op actions. In order to compare scores between different games and aggregated results, we computed normalized scores for each game based on random and human performance so that 0% corresponds to random performance and 100% corresponds to human. We observed an increase of 132.5% in the me-
dian human normalized score compared to the R2D2 baseline with the same network capacity, aggregate results are reported in Table 1. Figure 6 details the difference in normalized score between HiMo and our R2D2 baseline for all games individually. We note that the original R2D2 results reported by Kapturowski et al. (2019), which used a similar hardware configuration but a different network architecture, were around 750% median human normalized score after a day of training.
In our experimental evaluation we observed that HiMo typically either offers improved data efficiency or has no overwhelming adverse effects in training performance. In Figure 6 we show training curves for a selection of representative Atari environments where at evaluation time HiMo both under-performed (left) and out-performed R2D2 (right); these seem to indicate that in the worst case scenario HiMo’s training performance reduces to R2D2’s.5
Bowling is one of the Atari games where rewards are delayed with relevant information being communicated through intermediate observations (the ball hitting the pins), just like the baseball example
5We will include a performance analysis over longer training regimes in a future version of the paper.
we have used in the introduction. We found HiMo to perform better than the R2D2 baseline in this particular game. We also ran HiMo with the actor-critic setup (IMPALA) described in the previous section, finding similar performance gain with respect to the model-free baseline. Learning curves for these experiments are presented in Figure 7.

heading: 6 RELATED WORK

text:
Recent work have used auxiliary predictions successfully in RL as a mean to obtain a richer signal for representation learning (Jaderberg et al., 2016; Sutton et al., 2011). However these additional prediction tasks are hard-coded and so they cannot adapt to the task demand when needed. We see them as a complementary approach to more efficient learning in RL.
Buesing et al. (2018) have considered using observations in an episode trajectory in hindsight to infer variables in a structural causal model of the dynamics, allowing to reason more efficiently in a model-based way about counterfactual actions. However this approach requires learning an accurate generative model of the environment.
In supervised learning, the learning using privileged information (LUPI) framework introduced by (Vapnik & Izmailov, 2015) considers ways of leveraging privileged information at train time. Although the techniques developed in that work do not apply directly in the RL setting, some of our approach can be understood in that setting as considering the future trajectory as the privileged information for a value prediction problem.
Privileged information coming from full state observation has been leveraged in RL to learn better critic in asymmetric actor-critic architectures (Pinto et al., 2017; Zhu et al., 2018). However this does not use future information and only applies to settings where special side-information (full state) is available at train time.

heading: 7 CONCLUSION

text: High-dimensional observations in the intermediate future often contain task-relevant features that can facilitate the prediction of an RL agent’s final return. We introduced a reinforcement learning algorithm, HiMo, that leverages this insight by the following two-stage approach. First, by reasoning in hindsight, the algorithm learns to extract relevant features of future observations that would be been most helpful for estimating the final value. Then, a forward model is learned to predict these features, which in turn is used as input to an improved value function, yielding better policy evaluation and training at test time. We demonstrated that this approach can help tame complexity in environments with rich dynamics at scale, yielding increased data efficiency and improving the performance of state-of-the-art model-free architectures.

heading: A APPENDIX

text:
A.1 GENERAL ARCHITECTURE DETAILS
To compute v+ and train φ̂ in an online fashion, we process fixed-length unrolls of the state-RNN and compute the hindsight value and corresponding updates at time t if t+k is also within that same unroll. Also, we update v+ at a slower rate (i.e., α < β) to give enough time for the model φ̂ to adapt to the changing hindsight features φ. In our experiments we found that even a low-dimensional φ (in the order of d = 3) and a relatively short hindsight horizon k (in the order of 5) are sufficient to yield significant performance boosts, whilst keeping the extra model computational costs modest.
A.2 PORTAL CHOICE
Environment The observation is a 7× 23 RGB frame (see Figure 4). There are 3 possible spawning points for the agent in the center and 42 possible portal positions (half of which lead to the green room, the other half leading to the red room). At the start of an episode, two portals, each leading to a different room, are chosen are random. They are both displayed as cyan pixels. Included in the observation in the first phase is the context, a random permutation in a 5 × 5 grid of N pixels, where is uniformly sampled at the start of each episode: N ∼ U{1, 10}. A fixed map f : {1, . . . , 10} → {0, 1} determines which contexts are rewarding with the green room, the rest being rewarding with the red room. The reward when reaching the goal is determined according to:
R = 2(f(N)G+ (1− f(N))(1−G)), (14)
where G ∈ {0, 1} is whether the reached room is green.
Network architecture The policy and value network takes in the observation and passes it to a ConvNet encoder (with filter channels [32, 32, 32], kernel shapes [4, 3, 3] applied with strides [2, 1, 1]) before being passed to a ConvLSTM network with 32 channels and 3x3 filters. The output of the ConvLSTM is the internal state h. The φ̂ network is a ConvNet with [32, 32, 32, 1] filter channels with kernels of size 3 except for a final 1x1 filter, whose output is flatten and passed to an MLP with 256 hidden units with ReLu activation, before a linear layer with dimension d = 3. The φ network is a similarly configured network with one less convolution layer and 128 hidden units in the MLP. The ψη network is an MLP with 256 hidden units followed by a linear layer that takes h and φ̂ as input and outputs the policy πm and the value vm. v+ is obtained similarly with a similar MLP that has a single scalar output. We used a future observation window of k = 5 steps in this domain and loss weights α = 0.25, β = 0.5. Unroll length was 20, and γ = 0.99. Optimization was done with the Adam optimizer (learning rate of 5e − 4), with batch size 32. The model-free baseline is obtained by using the same code and network architecture, and setting the modeling loss and hindsight value loss to 0 (α = β = 0).
A.3 ATARI
Hyper-parameters and infrastructure are the same as reported in Kapturowski et al. (2019), with deviations as listed in table 2. For our value target, we also average different n-step returns with exponential averaging as in Q(λ) (with the return being truncated at the end of unrolls). The Q network is composed of a convolution network (cf. Vision ConvNet in table) which is followed by an LSTM with 512 hidden units. What we refer to in the main text as the internal state h is the output of the LSTM. The φ and φ̂ networks are MLPs with a single hidden layer of 256 units and ReLu activation function, followed by a linear which outputs a vector of dimension d. The ψθ1 function concatenates h and φ as inputs to an MLP with 256 hidden units with ReLu activation function, followed by a linear which outputs q+ (a vector of dimension 18, the size of the Atari action set). qm is obtained by passing h and φ̂ to a dueling network as described by Kapturowski et al. (2019).
Other HiMo parameters are described in table 3. The R2D2 baseline with the same capacity is obtained by running the same architecture with α = β = 0.