Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies (Extended Abstract)

1 Introduction

Unrolled computation graphs arise in many scenarios in machine learning, including when training RNNs [Williams and Peng, 1990], tuning hyperparameters through unrolled computation graphs [Franceschi et al., 2018], and training learned optimizers [Wichrowska et al., 2017; Metz et al., 2018].

Classic algorithms for computing gradients in such unrolled computation graphs include reverse- and forwardmode gradient accumulation. Backpropagation through time (BPTT) involves backpropagating through a full unrolled sequence for each parameter update. This faces several difficulties: 1) the memory cost scales linearly with the unroll length, due to storage of intermediate activations for backprop 2) it only performs one parameter update after each full unroll, which is expensive and introduces large latency between parameter updates; 3) long unrolls can lead to exploding or vanishing gradients [Pascanu et al., 2013], and chaotic loss landscapes [Parmas et al., 2018; Metz et al., 2019]. A common technique to alleviate these issues is truncated backprop through time (TBPTT) [Werbos, 1990], which splits the full sequence into shorter subsequences and performs an update after processing each subsequence. However, TBPTT yields biased gradients that can severely impact training.

An alternative to BPTT is real-time recurrent learning (RTRL), which performs forward gradient accumulation [Williams and Zipser, 1989]. RTRL enables online parameter updates (after each partial unroll) and does not suffer from truncation bias; however, its memory and compute requirements render it intractable for large-scale problems. Many approximations to RTRL have been proposed [Tallec and Ollivier, 2017; Mujika et al., 2018; Benzing et al., 2019], but most have high variance, are complicated to implement, or are only applicable to a restricted class of models.

∗This is an extended abstract of [Vicol et al., 2021] that received the best paper award at ICML 2021.

To address the poorly conditioned loss surfaces that often result from sequential computation, it can be useful to minimize a smoothed version of the loss. Evolution strategies (ES) is a family of algorithms that estimate gradients using stochastic finite-differences, and which provide an unbiased estimate of the gradient of the objective smoothed with a Gaussian. We introduce an approach to unbiased gradient estimation using short, truncated unrolls, called Persistent Evolution Strategies (PES). In PES, we accumulate the perturbations experienced by the outer parameters in each partial unroll-rather than starting perturbations from scratch as in ES-which yields an unbiased estimate of the gradient even when using truncated sequences. PES is simple to implement, and because it is ES-based, it retains desirable characteristics such as being parallelizable, memory efficient, and broadly applicable to many problems, including to non-differentiable target functions.

Problem Setup. We consider unrolled computation graphs with state st updated based on parameters θ via the recurrence:

where xt is an optional input at step t. The objective function for optimizing θ is the sum of per-timestep losses Lt(st; θ):

This setup is general, even encompassing situations where we want to consider only the final loss at step T, which can be expressed using a telescoping sum of loss differences between successive steps [Beatson and Adams, 2019]. Instances of this problem setup include training RNNs, training learned optimizers, learning policies for control tasks, and hyperparameter optimization, as illustrated in Figure 1.

Smoothing & Evolution Strategies (ES). Unrolling optimization for many steps can lead to pathological meta-loss Objective:

where ⊗ denotes the Kronecker product, ϵ = (ϵ1, . . . , ϵT )⊤ is a matrix of perturbations ϵt to be added to the θt at each timestep, and the expectation is over entries in ϵ drawn from an i.i.d. Gaussian with variance σ2 . This ES approximation is an unbiased estimator of the gradient of the Gaussian-smoothed objective Eϵ[L(Θ + ϵ)]. We next show that gPES decomposes into a sum of sequential gradient estimates,

where ξt = ∑tτ=1 ϵτ , Equation 3 relies on Lt (·) being independent of ϵτ for τ > t, and Equation 5 similarly relies on Lt (·) only being a function of θτ for τ ≤ t. The PES estimator consists of Monte Carlo estimates of Equation 4,

where ϵ(i) are samples of ϵ, and N is the number of Monte Carlo samples. Gradient estimates at each time step can be evaluated sequentially, and used to perform SGD. When using antithetic sampling (e.g., positive and negative perturbation pairs in Eq. 6), we denote the estimator by ˆgPES-A . See the full paper for a proof of the following Statement 2.1.

Statement 2.1 (PES is unbiased). Let θ ∈ RP and L(θ) = ∑Tt=1 Lt(θ). Suppose that ∇θL(θ) exists, and assume that L is quadratic, so that it is equivalent to its second-order Taylor series expansion: L(Θ+ϵ) = L(Θ)+vec(ϵ)⊤∇vec(Θ)L(Θ)+ 12 vec(ϵ)⊤∇2vec(Θ)L(Θ) vec(ϵ). Then, bias(ˆgPES-A ) = Eϵ[ˆgPES-A ] - ∇θL(θ) = 0.

Algorithm. Based on Eq. 6, we see that we can obtain unbiased gradient estimates from partial unrolls by: 1) not resetting the particles between unrolls, and 2) accumulating the perturbations ξt each particle has experienced over multiple unrolls.

surfaces that exhibit near-discontinuities and chaotic structure [Parmas et al., 2018; Metz et al., 2019]. Optimization on such non-smooth landscapes fails due to exploding gradients or gets stuck in poor local minima. An effective method to address these pathologies is to smooth the meta-loss surface, e.g. descend the Gaussian-blurred objective L(θ) = E˜θ∼N(θ,σ2 I) [L(˜θ)] [Staines and Barber, 2012; Metz et al., 2019]. Evolution strategies provides an unbiased estimate of the gradient ∇θE˜θ∼N(θ,σ2 I) [L(˜θ)]. The ES estimator is defined as: ˆgES = 1 Nσ2 ∑Ni=1 ϵ(i)L(θ + ϵ(i) ), where ϵ(i) ∼ N(0, σ2 I). ES is trivially parallelizable, and thus highly scalable; it has seen renewed interest as a viable algorithm for reinforcement learning among other black-box problems [Ha and Schmidhuber, 2018; Cui et al., 2018; Ha, 2020]. However, applying ES to full unrolls is costly, while using partial unrolls leads to truncation bias, similarly to TBPTT.

2 Persistent Evolution Strategies

We introduce a method to obtain unbiased gradient estimates from partial unrolls of a computation graph, called Persistent Evolution Strategies (PES). We first derive the PES gradient estimator, and then present a practical algorithm (Algorithm 2).

Derivation. Unrolled computation graphs (as illustrated in Figure 1) depend on shared parameters θ at every timestep; in order to account for how these contribute to the overall gradient ∇θL(θ), we use subscripts θt to distinguish between applications of θ at different steps, where θt = θ, ∀t. We further define Θ = (θ1, . . . , θT )⊤, which is a matrix with the per-timestep θt as its rows. For notational simplicity in the following derivation, we drop the dependence on st and explicitly include the dependence on each θt, writing Lt(st; θ) as either Lt(θ1, . . . , θt) or simply Lt(Θ). We wish to compute the gradient ∇θL(θ) of the total loss over all unrolls. We begin by writing this gradient in terms of the full gradient ∂L(Θ) ∂ vec(Θ) ∈ RPT×1 , and then using ES to approximate ∂L(Θ) ∂ vec(Θ) ,

Algorithm 1 Truncated Evolution Strategies (ES) applied to partial unrolls of a computation graph.

Input: s0, initial state K, truncation length for partial unrolls N, number of particles σ, standard deviation of perturbations α, learning rate for ES optimization Initialize s = s0 s(i) = s0 Initialize ξ(i) ← 0 for i ∈ {1, . . . , N} while true do ˆgES ← 0 for i = 1, . . . ,N do ϵ(i) = { draw fromN(0, σ2 I) i odd -ϵ(i-1) i even ˆL(i) K ← unroll(s, θ + ϵ(i) , K) ξ(i) ← ξ(i) + ϵ(i) ˆgES ← ˆgES + ϵ(i) ˆL(i) K end for ˆgES ← 1 Nσ2 ˆgES s ← unroll(s, θ, K) θ ← θ - αˆgES end while Input: s0, initial state K, truncation length for partial unrolls N, number of particles σ, standard deviation of perturbations α, learning rate for PES optimization Initialize s(i) = s0 for i ∈ {1, . . . ,N} Initialize ξ(i) ← 0 for i ∈ {1, . . . , N} while true do ˆgPES ← 0 for i = 1, . . . ,N do ϵ(i) = { draw fromN(0, σ2 I) i odd -ϵ(i-1) i even s(i) , ˆL(i) K ← unroll(s(i) , θ + ϵ(i) , K) ξ(i) ← ξ(i) + ϵ(i) ˆgPES ← ˆgPES + ξ(i) ˆL(i) K end for ˆgPES ← 1 Nσ2 ˆgPES s ← unroll(s, θ, K) θ ← θ - αˆgPES end while

Algorithm 2 Persistent evolution strategies (PES). Differences from ES are highlighted in purple.

Figure 2: A comparison of vanilla ES and PES gradient estimators, applied to partial unrolls of a computation graph. The conditional statement for ϵ(i) is used to implement antithetic sampling. For clarity, we describe the meta-optimization updates to θ using SGD, but we typically use Adam in practice.

The resulting algorithm is simple to implement, requiring only minor modifications from vanilla ES. Algorithm 1 describes truncated ES applied to partial unrolls, where it suffers from short horizon bias. Algorithm 2 shows PES applied to the same problem, where it provides unbiased gradient estimates. Both algorithms (Fig. 2) are shown with antithetic sampling (perturbations are paired with their negations), which drastically reduces variance.

3 Experiments

In experiments, we use PES to meta-train a learned optimizer (Figure 3), and learn a policy for continuous control (Figure 4). We additionally demonstrate the use of PES for hyperparameter optimization, where we show that PES performs well when the meta-loss has many local minima (Figure 5), does not suffer from truncation bias (Figure 6a), can be applied to non-differentiable objectives (Figure 6b), and can be used to optimize many hyperparameters (both continuous and discrete) simultaneously (Figure 7). A detailed description of each experimental setup is given in the full paper, [Vicol et al., 2021]. A tutorial Colab notebook demonstrating the technique, and reproducing some of the experimental results, is available at this link.

4 Conclusion

We introduce a method called Persistent Evolution Strategies (PES), which divides a computation graph into a series of truncated unrolls, and performs an evolution strategies-based update after each unroll. PES eliminates bias from these truncations by accumulating correction terms over the entire sequence of unrolls. PES allows for rapid parameter updates, has low memory usage, is unbiased, and has reasonable variance characteristics. We show that PES is broadly applicable, with experiments demonstrating its use in hyperparameter optimization, reinforcement learning, and meta-training of learned optimizers.

Acknowledgements

This work was done while on internship at Google. We thank Sergey Ioffe and Niru Maheswaranathan for very helpful discussions and feedback on the paper.