Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies

Introduction

Unrolled computation graphs arise in many scenarios in machine learning, including when training RNNs (Williams & Peng, 1990), tuning hyperparameters through unrolled computation graphs (Baydin et al., 2017; Domke, 2012; Maclaurin et al., 2015; Wu et al., 2018; Franceschi et al., 2017; Donini et al., 2019; Franceschi et al., 2018; Liu et al., 2018; Shaban et al., 2019), and training learned optimizers (Li & Malik, 2016, 2017; Andrychowicz et al., 2016; Wichrowska et al., 2017; Metz et al., 2018, 2019, 2020b, 2020a). Many methods exist for computing gradients in such computation graphs, including ones based on reverse-mode (Williams & Peng, 1990; Tallec & Ollivier, 2017b; Aicher et al., 2019; Grefenstette et al., 2019) and forward-mode (Williams & Zipser, 1989; Tallec & Ollivier, 2017a; Mujika et al., 2018; Benzing et al., 2019; Marschall et al., 2019; Menick et al., 2020) gradient accumulation. These methods have different tradeoffs with respect to compute, memory, and gradient variance.

Backpropagation through time involves backpropagating through a full unrolled sequence (e.g. of length $T$) for each parameter update. Unrolling a model over full sequences faces several difficulties: 1) the memory cost scales linearly with the unroll length, because we need to store intermediate activations for backprop (though this can be reduced at the cost of additional compute (Dauvergne & Hascoët, 2006; Chen et al., 2016)); 2) we only perform a single parameter update after each full unroll, which is computationally expensive and introduces large latency between parameter updates; 3) long unrolls can lead to exploding or vanishing gradients (Pascanu et al., 2013), and chaotic and poorly conditioned loss landscapes (Pearlmutter, 1996; Maclaurin et al., 2015; Parmas et al., 2018; Metz et al., 2019). This is especially true in meta-learning (Metz et al., 2019).

The most commonly-used technique to alleviate these issues is truncated backprop through time (TBPTT) (Werbos, 1990; Tallec & Ollivier, 2017b), which splits the full sequence into shorter sub-sequences and performs a backprop update after processing each sub-sequence. However, a critical drawback of TBPTT is that it yields biased gradients, that can severely impact training (e.g. only taking into account short-term dependencies). To address the poorly conditioned loss surfaces that often result from sequential computation, it can additionally 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. ES works well on pathological meta-optimization loss surfaces (Metz et al., 2019); however, due to the computational expense of running full unrolls, ES can only practically be applied in a truncated fashion, introducing bias.

An alternative to BPTT is real-time recurrent learning (RTRL), which performs forward gradient accumulation (Williams & 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 & Ollivier, 2017a; 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.

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 vanilla ES—which yields an unbiased estimate of the gradient even when using truncated sequences. PES is simple to implement, and because it is an evolution strategies-based approach, it retains desirable characteristics such as being trivially parallelizable, memory efficient, and broadly applicable to many different types of problems, including to non-differentiable target functions.

We introduce a method called Persistent Evolution Strategies (PES) to obtain unbiased gradient estimates for the parameters of an unrolled system from partial unrolls of the system.

We prove that PES is an unbiased gradient estimate for a smoothed version of the loss, and an unbiased estimate of the true gradient for quadratic losses.

We provide theoretical and empirical analyses of its variance. In addition, we describe a variance reduction technique for PES, that incorporates the analytic gradient (computed with standard backprop) of the most recent unroll of the dynamical system.

We demonstrate the applicability of PES in several illustrative scenarios: 1) we apply PES to tune hyperparameters including learning rates and momentums, by estimating hypergradients through partial unrolls of optimization algorithms; 2) we use PES to meta-train a learned optimizer; 3) we use PES to learn policy parameters for a continuous control task.

We provide a Colab notebook implementation of PES.

Background

We provide an overview of notation in Appendix A.

We consider unrolled computation graphs with state $\boldsymbol{s}_{t}$ updated based on parameters ${\boldsymbol{\theta}}$ via the recurrence:

where $\boldsymbol{x}_{t}$ is an optional input at step $t$. The objective function for optimizing ${\boldsymbol{\theta}}$ is the sum of per-timestep losses $L_{t}(\boldsymbol{s}_{t};{\boldsymbol{\theta}})$:

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 & Adams, 2019). For details on telescoping sums, see Appendix D. Instances of this problem setup include training RNNs, training learned optimizers, learning policies for control tasks, and unrolled optimization, as illustrated in Figure 1.

Optimization algorithms can be unrolled to yield computation graphs, in which the nodes are the model parameters at successive optimization steps. Estimating gradients through unrolled optimization has been used to tune hyperparameters (Domke, 2012; Maclaurin et al., 2015; Baydin et al., 2017; Donini et al., 2019; Franceschi et al., 2017) and train learned optimizers (Li & Malik, 2016, 2017; Andrychowicz et al., 2016; Wichrowska et al., 2017; Metz et al., 2019, 2020b, 2020a, 2018).

Truncation, or short horizon, bias poses a major challenge when unrolled optimization is decomposed into a sequence of short sequential unrolls of length $K\ll T$. These challenges have been demonstrated in gradient-based hyperparameter optimization (Wu et al., 2018) and in the training of learned optimizers (Metz et al., 2019). Approaches to mitigating short horizon bias are an area of active research (Micaelli & Storkey, 2020).

Evolution Strategies (ES) (Rechenberg, 1973; Nesterov & Spokoiny, 2017) refers to a family of methods for estimating a descent direction for arbitrary black-box functions using stochastic finite differences. Since ES only requires function evaluations and not gradients, it is a zeroth-order optimization method. The vanilla ES estimator is defined as:

where $\boldsymbol{\epsilon}^{(i)}\sim\mathcal{N}(0,\sigma^{2}I)$. ES is trivially parallelizable, and thus highly scalable—it has seen renewed interest in recent years as a viable optimization algorithm for reinforcement learning among other black-box problems (Salimans et al., 2017; Mania et al., 2018; Ha & Schmidhuber, 2018; Houthooft et al., 2018; Cui et al., 2018; Ha, 2020). The estimator in Eq. 3 has high variance, and thus many variance reduction techniques have been proposed, including control variates (Tang et al., 2020) and antithetic sampling (Owen, 2013). Antithetic sampling involves using pairs of function evaluations ${\boldsymbol{\theta}}+\boldsymbol{\epsilon}$ and ${\boldsymbol{\theta}}-\boldsymbol{\epsilon}$, yielding the following estimator:

where $N$ is even, and $\boldsymbol{\epsilon}^{(i)}\sim\mathcal{N}(0,\sigma^{2}I)$. Several methods have been proposed to improve the search space for ES, including covariance matrix adaptation ES (CMA-ES) (Hansen, 2016) and Guided ES (Maheswaranathan et al., 2018). A limitation of ES is that applying it to full unrolls is often computationally costly (as we only make one update to the system parameters every full unroll), while applying ES to partial unrolls suffers from truncation bias similarly to TBPTT. In contrast, PES allows computation of gradients from partial updates without incurring truncation bias.

Any approach that performs online parameter updates, including RTRL and its approximations, will suffer from hysteresis, which refers to the dependence of the state of a system on its history. This is due to the fact that if we update ${\boldsymbol{\theta}}$, then any accumulated state (e.g. in the case of RTRL, the accumulated Jacobian $\frac{d\boldsymbol{s}_{t}}{d{\boldsymbol{\theta}}}$) will be incorrect because it is computed from previous values of ${\boldsymbol{\theta}}$. To eliminate hysteresis completely, one would need to run the full sequence for a given problem for each parameter update, which is often prohibitively expensive. In practice, hysteresis can be mitigated by using sufficiently small learning rates; this introduces a tradeoff between training stability and training speed.

Related Work

In this section, we discuss additional related work on online learning algorithms, and on one special class of unrolled optimization problems: hyperparameter optimization (HO). Table 1 compares several approaches to gradient estimation in unrolled computation graphs, with respect to compute, memory, parallelization, unbiasedness, and smoothing. In addition, Table 4 in Appendix B provides a comparison of the HO algorithms mentioned in this section.

Real-time recurrent learning (RTRL) performs forward-mode gradient accumulation: it does not require storage of past states, but requires matrix-matrix products and storage of a matrix $G_{t}$ of size $\text{dim}(\boldsymbol{s}_{t})\times\text{dim}({\boldsymbol{\theta}})$. When $\text{dim}({\boldsymbol{\theta}})$ is large, as in RNN training, the cost of storing $G_{t}$ and the cost of computing the required matrix-matrix products is prohibitive. Several approaches propose efficient variants of RTRL based on cheaper, noisy approximations of $G_{t}$. Unbiased Online Recurrent Optimization (UORO) (Tallec & Ollivier, 2017a) uses an unbiased rank-1 approximation to the full matrix; Kronecker-Factored RTRL (KF-RTRL) (Mujika et al., 2018) uses a Kronecker product decomposition to approximate the RTRL update for a class of RNNs; and Optimal Kronecker Sum Approximation (OK) (Benzing et al., 2019) uses a similar approximation but with the lowest possible variance among methods within an approximation family. Cooijmans & Martens (2019) also draw a connection between UORO and REINFORCE applied to estimate the gradient of an RNN by injecting noise into the hidden states. In contrast, PES injects noise into the parameters.

There are three main approaches that can be categorized based on the types of problem-specific information used: 1) black-box approaches that do not consider the internal structure of the objective $L$; 2) gray-box approaches that make use of the fact that the objective is the result of an iterative optimization procedure (e.g. by using the validation performance of a model); and 3) gradient-based approaches that require access to the exact functional form of the objective $L$, and that require the objective to be differentiable in the hyperparameters. Black-box approaches include grid search, random search (Bergstra & Bengio, 2012), Bayesian optimization (BO) (Snoek et al., 2012), and ES (Salimans et al., 2017; Metz et al., 2019). Gray-box approaches include Freeze-Thaw BO (Swersky et al., 2014), successive halving (Jamieson & Talwalkar, 2016), Hyperband (Li et al., 2017), Population-Based Training (Jaderberg et al., 2017), and hypernetwork-based approaches to HO (Lorraine & Duvenaud, 2018; MacKay et al., 2019).

A key advantage of gradient-based approaches is that they scale to high-dimensional hyperparameters (e.g. millions of hyperparameters) (Lorraine et al., 2020). Maclaurin et al. (2015) differentiate through unrolled optimization to tune many hyperparameters including learning rates and weight decay coefficients. These methods can perform poorly, however, when the underlying meta-loss is not smooth. Additionally they cannot optimize non-differentiable objectives, for example accuracy rather than loss.

PES can be considered a gray-box approach as it does not require the objective to be differentiable like gradient-based approaches, but it does take into account the iterative optimization of the inner problem.

Persistent Evolution Strategies

In this section, we introduce a method to obtain unbiased gradient estimates from partial unrolls of a computation graph, called Persistent Evolution Strategies (PES). First, we derive the PES gradient estimator, prove that it is unbiased, and present a practical algorithm (Algorithm 2). Then we discuss the variance characteristics of PES, both theoretically and empirically.

where $\boldsymbol{\xi}_{t}=\sum_{\tau=1}^{t}\boldsymbol{\epsilon}_{\tau}$, Equation 4 relies on $L_{t}\left(\cdot\right)$ being independent of $\boldsymbol{\epsilon}_{\tau}$ for $\tau>t$, and Equation 6 similarly relies on $L_{t}\left(\cdot\right)$ only being a function of ${\boldsymbol{\theta}}_{\tau}$ for $\tau\leq t$. The PES estimator consists of Monte Carlo estimates of Equation 5,

where $\boldsymbol{\epsilon}^{(i)}$ are samples of $\boldsymbol{\epsilon}$, and $N$ is the number of Monte Carlo samples. Gradient estimates at each time step can be evaluated sequentially, and used to perform SGD.

In practice, we use antithetic sampling to reduce variance. The PES estimator with antithetic sampling, which we denote $\hat{\boldsymbol{g}}^{\text{PES-A}}$, is given by:

See Appendix F for a proof of the following Statement 4.1.

Based on Eq. 7, 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 $\boldsymbol{\xi}_{t}$ each particle has experienced over multiple unrolls. 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.

1 Variance Analysis

We use the total variance, $\text{tr}(\text{Var}(\hat{\boldsymbol{g}}^{\text{PES-A}}))$, to quantify the variance of the estimator. We provide a full derivation of the variance in Appendix G, and here we present some takeaways. The variance depends on the gradients of each loss term $L_{t}$ with respect to each of the per-timestep parameters ${\boldsymbol{\theta}}_{\tau}$. To gain insight into the structure of these gradients, we can arrange them in a matrix:

$M$ is upper-triangular due to the fact that $\nabla_{{\boldsymbol{\theta}}_{\tau}}L_{t}=0$ for all $\tau>t$. The variance of the PES estimator depends on the covariance between the gradients $\nabla_{{\boldsymbol{\theta}}_{\tau}}L_{t}$ in this matrix.

We consider two structures for $M$: 1) a diagonal structure, where the gradients $\nabla_{{\boldsymbol{\theta}}_{i}}L_{j}=0,\forall i\neq j$; and 2) an upper-triangular structure. For each of these two matrix structures, we consider two possibilities for the covariance between gradients: a) all gradients $\nabla_{{\boldsymbol{\theta}}_{i}}L_{j}$ are identical; b) all gradients are i.i.d. The total variance for each of the four resulting scenarios is shown in Table 2. It is possible for the $\boldsymbol{g}_{t}$ to have variance larger than any of these scenarios, though we do not observe this in practice.

To investigate the variance characteristics of PES empirically, we computed the variance in a toy setting. We used an LSTM with 5 hidden units and 5-dimensional embeddings, for character-level language modeling on the Penn Treebank corpus (Marcus et al., 1993) (with a vocabulary consisting of 50 unique tokens). We measured the variance of the $\hat{\boldsymbol{g}}^{\text{PES-A}}$ gradient estimate on a fixed sequence of $10^{4}$ characters. The ground-truth gradient of the smoothed objective was computed using vanilla ES with 5000 particles over the full sequence (without truncation). Figure 3 shows the variance of the PES gradient estimate using different numbers of unrolls, ranging from 1 (a single unroll for the full sequence) to $10^{4}$ (one unroll per input token). Note that we do not update the parameters of the RNN after each unroll; we simply accumulate the gradient estimates over all partial unrolls. We plot the variance normalized by the squared norm of the ground-truth gradient. We observe an initial drop in variance, and then a linear growth. Additional empirical variance measurements are presented in Figure 16 (Appendix G).

For functions $L$ that are differentiable, we can use the analytic gradient from the most recent partial unroll (e.g., backpropagating through the last $K$-step unroll) to reduce the variance of the PES gradient estimates. In Appendix H, we show how we can incorporate the analytic gradient in the ES estimate for $\frac{\partial L_{t}(\Theta)}{\partial{\boldsymbol{\theta}}}$, deriving the following estimator:

where $\boldsymbol{p}_{t}=\frac{\partial L_{t}(\Theta)}{\partial{\boldsymbol{\theta}}_{t}}$. We call the resulting estimator PES+Analytic. In Appendix H we describe the implementation of this estimator (Algorithm 4), which requires a few simple changes from the standard PES estimator. We also provide empirical variance measurements for PES+Analytic, using the same setup as was used for Figure 3; we found that it can reduce variance by 1-2 orders of magnitude, given the same number of particles as PES.

Experiments

First, we demonstrate via a toy experiment that PES does not suffer from truncation bias, allowing it to converge to correct solutions that are not found by TBPTT or truncated ES. Then, we apply PES to several illustrative scenarios: we use PES to meta-train a learned optimizer, learn a policy for continuous control, and optimize hyperparameters. All experiments used JAX (Bradbury et al., 2018). A simplified code snippet implementing PES is provided in Appendix M.

2 Learned Optimizer Meta-Optimization

In this section we demonstrate PES’s applicability for learned optimizer training. We meta-train an MLP-based learned optimizer as described in Metz et al. (2019). This optimizer is used to train a two hidden-layer, 128 unit, MLP on CIFAR-10 with a batch size of 128. Our meta-objective is the average training loss. We train with a total number of inner-steps of $T=1000$ and a truncation length of $K=4$, using both PES and truncated ES.

We outer-train with Adam, using a learning rate of $10^{-4}$ selected via grid search over half-orders of magnitude for each method independently. We use gradient clipping of $3$ applied to each gradient coordinate. We outer-train on 8 TPUv2 cores with asynchronous, batched updates of size $16$. To evaluate, we compute the meta-loss averaged over 20 inner initializations over the course of meta-training. Results can be found in Figure 5. Due to PES’s unbiased nature, PES achieves both lower losses, and is more consistent across random initializations of the learned optimizer.

3 Learning a Continuous Control Policy

Recent work (Salimans et al., 2017; Mania et al., 2018) has shown that ES-based algorithms can be a viable alternative to more complex RL algorithms. ES optimizes the parameters of a policy directly, by sampling parameters from a distribution, running an episode, and estimating the gradient; this is in contrast to standard RL algorithms that sample actions from a distribution output by a policy. Here, we demonstrate that PES can be used to train a policy for a continuous control problem using partial unrolls, improving on the efficiency of vanilla ES typically applied to full unrolls. We train a linear policy on the Swimmer-v2 MuJoCo environment, following Mania et al. (2018). For PES, the objective for each partial unroll is the sum of rewards over that unroll. We also applied vanilla ES to the partial unrolls to demonstrate that this naïve strategy does not work—truncation bias occurs for these control problems as well. Figure 7 compares vanilla ES applied to full episodes, ES applied to partial episodes, PES applied to partial episodes, and variants of ES from Mania et al. (2018) and Salimans et al. (2017). To evaluate policies, we computed the average full-episode reward over 50 random environment seeds. In Figure 7, we show the mean performance of each algorithm over 6 random seeds, with standard deviation shown by the shaded region. We see that PES reaches the same performance as full-unroll ES in slightly fewer total environment steps.

4 Hyperparameter Optimization

In this section we demonstrate that PES can be used for hyperparameter optimization across four different problems. We show that PES performs well when the meta-loss has many local minima, does not suffer from truncation bias, can be applied to non-differentiable objectives, and can be used to optimize many hyperparameters (both continuous and discrete) simultaneously.

First, we used PES to meta-optimize a learning rate schedule for a toy 2D regression problem that has one global minimum, but many local minima to which truncated gradient methods could converge. The inner optimization trajectories for different values of the outer-parameters are shown in Appendix C. We tuned a linear learning rate schedule parameterized by the initial and final log-learning rates, $\theta_{0}$ and $\theta_{1}$, respectively: $\alpha_{t}=\left(1-\frac{t}{T}\right)e^{\theta_{0}}+\frac{t}{T}e^{\theta_{1}}$. In Figure 6 we compare TBPTT, UORO, RTRL, ES, and PES applied to this meta-optimization task. We found that the gradient-based methods (TBPTT, UORO, and RTRL) got stuck in suboptimal regions due to high-frequency structure in the meta-loss landscape. ES makes more progress due to smoothing, but still suffers from truncation bias. PES smooths the meta-objective surface and is unbiased, converging to a substantially better solution.

Next, we used PES to meta-learn a learning rate schedule for an MLP classifier on MNIST. Following Wu et al. (2018), we used a two-layer MLP with 100 hidden units per layer and ReLU activations and the learning rate schedule parameterization $\alpha_{t}=\frac{\theta_{0}}{\left(1+\frac{t}{Q}\right)^{\theta_{1}}}$, where $\alpha_{t}$ is the learning rate at step $t$, $\theta_{0}$ is the initial learning rate, $\theta_{1}$ is the decay factor, and $Q$ is a constant fixed to 5000. This schedule is used for SGD with fixed momentum 0.9. The full unrolled inner problem consists of $T=5000$ optimization steps, and we apply vanilla ES and PES with truncation lengths $K\in\{10,100\}$, yielding 500 and 50 unrolls per inner problem, respectively. The meta-objective is the sum of training losses over the inner optimization trajectory. In Figure 8(a) we see that ES converges to a suboptimal region of the hyperparameter space due to truncation bias, while PES finds the correct solution.

Because PES only requires function evaluations and not gradients, it can optimize non-differentiable objectives such as accuracy rather than loss. We demonstrate this by tuning the same parameterization of learning rate schedule as before, but using the accuracy on the MNIST validation set as the meta-objective. Figure 8(b) compares the meta-optimization trajectories of ES and PES on the validation accuracy meta-objective; again ES is biased and fails to converge to the right solution, while PES works well.

Here, we show that PES can tune several hyperparameters simultaneously, and achieves better performance than random search with an uninformative search space, using less compute. We tuned both continuous and discrete hyperparameters: the number of units per hidden layer (discrete architectural hyperparameters) and per-parameter-block learning rates and momentum coefficients (continuous hyperparameters). We trained a 5-hidden-layer MLP (6 layers including the output layer mapping to logits) on FashionMNIST, yielding 29 hyperparameters in total. We set the maximum number of hidden units per layer to 100, and tuned sigmoid-transformed hyperparameters representing the fraction of hidden units to use. The meta-objective was the sum of validation losses over the inner optimization trajectory.

Figure 9 compares the best meta-objective values achieved by random search, vanilla ES, and PES, expressed in terms of the total number of inner iterations used (which accounts for the particles used in ES and PES). We ran each method with four random seeds, and plot the mean (solid lines) and the min/max (shaded region) performance. Each evaluation computes the mean meta-objective over 10 full inner problems using different random seeds for model initialization and data sampling. PES outperforms ES and random search, achieving lower loss using less compute.

Conclusion

We introduced a method for unbiased gradient estimation in unrolled computation graphs, called Persistent Evolution Strategies (PES). PES obtains gradients from truncated unrolls—which speeds up optimization by allowing for frequent parameter updates—while not suffering from truncation bias that affects many competing approaches. We show that PES is broadly applicable, with experiments demonstrating its application to an RNN-like task, hyperparameter optimization, reinforcement learning, and meta-training of learned optimizers.

Acknowledgements

We thank Sergey Ioffe and Niru Maheswaranathan for very helpful discussions and feedback on the paper.

References

Appendix A Notation

Table 3 summarizes the notation used in this paper.

Appendix B Hyperparameter Optimization Methods

Table 4 presents a comparison of several hyperparameter optimization approaches. We distinguish between black-box, gray-box, and gradient-based approaches, and focus our comparison on whether each method can tune optimization hyperparameters, regularization hyperparameters, and discrete hyperparameters, as well as whether the method requires multiple runs through the inner problem or is online (operating within the timespan of a single inner problem), and whether the method is unbiased, meaning that it will eventually converge to the optimal hyperparameters.

Appendix C Experiment Details

In this section, we provide details for the experiments from Section 5.

All experiments except for learned optimizer training were run on NVIDIA P100 GPUs (using only a single GPU per experiment). The learned optimizer experiment in Section 5.2 was trained on 8 TPUv2 cores; we used asynchronous multi-TPU training for convenience, not necessity (these experiments could be run on a single GPU if desired).

C.1 2D Toy Regression

The inner objective is a toy 2D function defined as:

This was manually designed to be a challenging problem for any meta-optimization method that suffers from truncation bias. In Figure 10 we visualize the outer loss surface (aka the meta-loss surface) and the inner loss surface for this task; we show the optimization trajectories on the inner loss surface corresponding to three different choices of optimization hyperparameters (shown by color-coded markers).

In our experiments, the total inner problem length was $T=100$, and we used truncated unrolls of length $K=10$. For ES and PES, we used perturbation variance $\sigma^{2}=1$, and 100 particles (50 antithetic pairs). We used Adam with learning rate 1e-2 as the outer optimizer for all methods (TBPTT, RTRL, UORO, ES, and PES).

C.2 Influence Balancing

where $A$ is a fixed $n\times n$ matrix with $A_{i,i}=0.5$, $A_{i,i+1}=0.5$ and 0 everywhere else. The vector on the right hand side consists of $\theta$ tiled $n$ times, with $p$ positive and $n-p$ negative copies. In our experiments, we used $n=23$ and $p=10$. The loss at each step is regression on the first index in the state vector $\boldsymbol{s}_{t}$:

For the influence balancing experiment, we used $n=23$ with 10 positive and 13 negative $\theta$’s. The state was initialized to a vector of ones, $\boldsymbol{s}_{0}=\mathbf{1}$, and $\theta$ was initialized to 0.5. We used gradient descent for optimization, with learning rate 1e-4. We did not use learning rate decay as was used in (Tallec & Ollivier, 2017a), as we did not find this to be necessary for convergence. For ES and PES we used perturbation scale $\sigma=0.1$ and $10^{3}$ particles.

C.3 MNIST Experiments

Following Wu et al. (2018), we used a two-layer MLP with 100 hidden units per layer and ReLU activations and the learning rate schedule parameterization $\alpha_{t}=\frac{\theta_{0}}{\left(1+\frac{t}{Q}\right)^{\theta_{1}}}$, where $\alpha_{t}$ is the learning rate at step $t$, $\theta_{0}$ is the initial learning rate, $\theta_{1}$ is the decay factor, and $Q$ is a constant fixed to 5000. This schedule is used for SGD with fixed momentum 0.9. We used mini-batches of size 100. The full unrolled inner problem consists of $T=5000$ optimization steps, and we used vanilla ES and PES with truncation lengths $K\in\{10,100\}$, yielding 500 and 50 unrolls per inner problem. The meta-objective is the sum of training softmax cross-entropy losses over the inner optimization trajectory. We used Adam as the outer-optimizer, and for each method (ES and PES), we performed a grid search over the outer-learning rates $\{0.01,0.03,0.1\}$ to find the most stable and fastest-converging setups. For both ES and PES, we used perturbation standard deviation $\sigma=0.1$, and 1000 particles (500 antithetic pairs).

We trained on FashionMNIST with minibatch size $100$ for $T=1000$ inner problem steps, using truncations of length $K=10$, yielding 100 unrolls per inner problem. For both ES and PES, we used $\sigma=0.3$ and used Adam with learning rate $1e-2$ as the outer optimizer. We used an MLP with ReLU activations and 5 hidden layers (6 layers including the output layer mapping the final hidden representation to logits). We tuned separate learning rates and momentum coefficients for SGD with momentum, for each weight matrix and bias vector in the network (this yields 24 hyperparameters, as we have 6 layers each with 2 parameter blocks and 2 hyperparameters tuned). We also tuned the number of units per hidden layer, by masking the output of each hidden layer, with a deterministic mask that zeros out part of the representation, effectively using only the first $n$ units. We tune the number of units in each of the 5 hidden layers, yielding 5 discrete hyperparameters, and 29 hyperparameters in total. Because we are effectively tuning the architecture of the MLP, we apply hidden unit masking at evaluation time in addition to training time. As the meta-objective, we used the sum of validation losses over the inner optimization trajectory.

To tune the number of hidden units, we used an unconstrained parameterization (in the real numbers) transformed by a sigmoid to the range $(0,1)$ which represents the fraction of units that are used, out of the maximum number of units per layer (set to be 100 in our experiments). The number of units per layer is determined by $\lfloor m_{i}*\text{sigmoid}(\theta_{i})\rfloor$ where $m_{i}$ is the maximum number of units for hidden layer $i$ and $\theta_{i}$ is the unconstrained parameterization for the fraction of units to be used.

For random search, we sampled learning rates uniformly at random in log-space, with range $(1\text{e-8},1\text{e1})$; we sampled momentum coefficients uniformly at random in logit-space corresponding to the sigmoid-transformed range $(0.01,0.999)$; and we sampled the number of hidden units per layer from the logit-space corresponding to the sigmoid-transformed range $(0.01,0.999)$ . For ES and PES, we initialized each learning rate uniformly at random in log space in the range $(1\text{e-4},1\text{e-2})$; we initialized each momentum coefficient uniformly at random in logit-space, to have the sigmoid-transformed range $(0.01,0.9)$; and we initialized the number of hidden units per layer in logit-space corresponding to the sigmoid-transformed range $(0.2,0.8)$. These ranges are slightly smaller than the ones used for random search in order to maintain meta-optimization stability; note from Figure 9 that the performance of both ES and PES is initially poor (prior to meta-optimization), indicating that these ranges for random initialization do not increase their performance compared to random search, and thus the improvement for PES is primarily due to its adaptation of the hyperparameters. For ES and PES, we used perturbation standard deviation $\sigma=0.3$, $N=10$ particles, and Adam with learning rate 0.01 for outer optimization. We ran each method four times with different random seeds, and plotted the mean performance, with the min and max shown by the shaded regions in Figure 9. We measured the best meta-objective value achieved so far during meta-optimization, as a function of total compute, which takes into account the number of inner iterations performed, as well as the number of parallel workers (or particles); total compute corresponds to the product of inner iterations and the number of workers.

In Figure 12(a), we tune hyperparameters for a 1.6M parameter ResNet on CIFAR-10 using ES and PES with $T=5000$, $K=20$, and $N=4$, targeting the sum of validation losses. In Figure 12(b), we train a learned optimizer on MNIST (similarly to Metz et al. (2019)). We use the same configuration as described in Section 5.2 but target a 2-hidden layer, 128 unit MLP trained on MNIST.

Here we show that truncation bias can also arise for regularization hyperparameters such as the $L_{2}$ regularization coefficient. We tune $L_{2}$ regularization for linear regression on the Yacht data from the UCI collection (Asuncion & Newman, 2007). We found the optimal $L_{2}$ coefficient using a fine-trained grid search. In Figure 13 we compare meta-optimization using ES and PES, starting from different initial $L_{2}$ coefficients; PES robustly converges to the correct solution in all cases. We used $\sigma=0.01$, $K=1$, and $N=4$ for both ES and PES.

C.4 Continuous Control Details

We used OpenAI Gymhttps://github.com/openai/gym to interface with MuJoCo. In our implementation, each antithetic pair shares a MuJoCo environment state, which is different between different antithetic pairs. The environment state is reset to the same point before running the partial unrolls of each particle in a pair, to control for randomness (e.g., the antithetic perturbations are evaluated starting from a common state). As is standard for MuJoCo environments, the length of a full episode is $T=1000$; we ran full-unroll ES with $K=1000$, and we used partial unrolls of length $K=100$ for truncated ES and PES. We used $10$ antithetic pairs for each of ES and PES. Following Mania et al. (2018), we used vanilla SGD to optimize the policy parameters. For each of ES and PES, we performed a grid search over learning rates and perturbation scales, both from the set $\{1.0,0.3,0.1,0.01\}$. To evaluate policies, we computed the average full-episode reward over 50 random environment seeds. In Figure 7, we show the mean performance of each algorithm over 6 random seeds, with standard deviation shown by the shaded region. Following Mania et al. (2018), we used a linear policy initialized as all 0s (the linear policy is a single weight matrix with no bias term). Also following Mania et al. (2018), we divided the rewards by their standard deviation (computed using the aggregated rewards from all antithetic pairs) before computing the ES/PES gradient estimates. We did not use state normalization, nor did we perform any heuristic selection of a subset of the best sampled perturbation directions (as used in the ARS V2 approach of Mania et al. (2018)).

Appendix D Telescoping Sums

If we wish to target the final loss $L_{T}$ as the meta-objective, we can define $p_{t}=L_{t}-L_{t-1}$, where $L_{-1}\equiv 0$. This yields the telescoping sum:

Targeting the final loss encourages different behavior than targeting the sum or average of the losses. Targeting the sum of losses encourages fast convergence (small $\sum_{t}L_{t}$), but not necessarily the smallest final loss $L_{T}$, while targeting the final loss encourages finding the smallest $L_{T}$ potentially at the expense of slower convergence (larger $\sum_{t}L_{t}$).

We performed an experiment using telescoping sums to target the final training loss, optimizing an exponential LR schedule for an MLP on FashionMNIST with $T=5000$, $K=20$, $N=100$ (Figure 14). Due to the computational expense of evaluating the loss on the full training set to obtain $L_{t}$ at each partial unroll, we selected a random minibatch at the start of each inner problem, which was kept fixed for the loss evaluations for that inner problem.

Appendix E Derivation of Persistent Evolution Strategies

Here we derive the PES estimator. The derivation here closely follows that in the text body, but shows additional intermediate steps in several places in the derivation. Also see Appendix I for an alternate derivation using stochastic computation graphs (Schulman et al., 2015).

E.2 PES is ES Over the Parameters at Each Unroll Step

where $\otimes$ denotes the Kronecker product, $\mathbf{I}$ has dimension $P\times P$, $\boldsymbol{1}^{\top}$ has dimension $1\times T$, and thus $\mathbf{I}\otimes\boldsymbol{1}^{\top}$ has dimension $P\times PT$. Note that because $\frac{\partial L(\Theta)}{\partial\operatorname{vec}\left(\Theta\right)}$ has dimension $PT\times 1$, the product $\left(\mathbf{I}\otimes\boldsymbol{1}^{\top}\right)\frac{\partial L(\Theta)}{\partial\operatorname{vec}\left(\Theta\right)}$ will be $P\times 1$. Next, we will apply ES to approximate the last RHS expression above:

We next show that $\boldsymbol{g}^{\text{PES}}$ decomposes into a sum of sequential gradient estimates,

where $\boldsymbol{\xi}_{t}=\sum_{\tau=1}^{t}\boldsymbol{\epsilon}_{\tau}$, Equation 15 relies on $L_{t}\left(\cdot\right)$ being independent of $\boldsymbol{\epsilon}_{\tau}$ for $\tau>t$, and Equation 18 similarly relies on $L_{t}\left(\cdot\right)$ only being a function of ${\boldsymbol{\theta}}_{\tau}$ for $\tau\leq t$.

The PES estimator consists of Monte Carlo estimates of Equation 17,

where $\boldsymbol{\epsilon}^{(i)}$ are samples of $\boldsymbol{\epsilon}$, and $N$ is the number of Monte Carlo samples. Gradient estimates at each time step can be evaluated sequentially, and used to perform SGD.

To illustrate how the expressions in the derivation above yield the desired gradient estimate, here we provide a concrete example using two-dimensional ${\boldsymbol{\theta}}$ with three steps of unrolling. The matrix $\Theta$ is:

The vectorized matrix $\operatorname{vec}\left(\Theta\right)$ and gradient $\frac{\partial L(\Theta)}{\partial\operatorname{vec}\left(\Theta\right)}$ are as follows:

The Kronecker product is: $\mathbf{I}\otimes\boldsymbol{1}^{\top}=\begin{bmatrix}1&0\\ 0&1\end{bmatrix}\otimes\begin{bmatrix}1&1&1\end{bmatrix}=\begin{bmatrix}1&1&1&0&0&0\\ 0&0&0&1&1&1\end{bmatrix}$. Thus, we have:

Similarly, to see how the PES derivation works, consider a matrix of perturbations $\boldsymbol{\epsilon}$ and its vectorization $\text{vec}(\boldsymbol{\epsilon})$ as follows:

This shows how the following statements are equivalent in our derivation:

Appendix F Proof that PES is Unbiased

Consider the PES estimator (using antithetic sampling) below:

Using the assumption that $L$ is quadratic and due to antithetic sampling, we can simplify this expression $L(\Theta+\boldsymbol{\epsilon})-L(\Theta-\boldsymbol{\epsilon})$ as follows:

Appendix G PES Variance

In this section, we derive the variance of PES. The antithetic PES estimator assuming quadratic $L$ is as follows:

For simplicity in the following derivation, we consider a Monte-Carlo estimate using a single particle pair. For $N$ particles, the variance will be scaled by a factor of $1\over N$. We use the total variance $\text{tr}(\text{Var}(\hat{\boldsymbol{g}}^{\text{PES}}))$ to quantify the variance of the estimator:

To derive term \raisebox{-.9pt}{1}⃝, we will expand out $\hat{\boldsymbol{g}}^{\text{PES-A}\top}\hat{\boldsymbol{g}}^{\text{PES-A}}$ into a sum of simple sub-expressions and use the linearity of expectation to combine them. To simplify notation, we use the shorthand $\boldsymbol{v}_{t}\equiv\operatorname{vec}\left(\boldsymbol{\epsilon}_{1\dots t}\right)$ and $\boldsymbol{g}_{t}\equiv\nabla_{\operatorname{vec}\left(\Theta_{1\dots t}\right)}L_{t}(\Theta)$. First, note that:

There are two types of terms in Eq. 33: terms of type \raisebox{-.9pt}{a}⃝, which have the form $\boldsymbol{g}_{i}^{\top}\boldsymbol{v}_{i}\boldsymbol{\xi}_{i}^{\top}\boldsymbol{\xi}_{i}\boldsymbol{v}_{i}^{\top}\boldsymbol{g}_{i}$, and terms of type \raisebox{-.9pt}{b}⃝, which have the form $\boldsymbol{g}_{i}^{\top}\boldsymbol{v}_{i}\boldsymbol{\xi}_{i}^{\top}\boldsymbol{\xi}_{j}\boldsymbol{v}_{j}^{\top}\boldsymbol{g}_{j}$ where $i\neq j$. We will derive the expectations of each of these two types of terms separately, and then combine the resulting sub-expressions.

As the first step in expanding out each term of type \raisebox{-.9pt}{a}⃝, note that:

Also, note that $\boldsymbol{\xi}_{i}^{\top}\boldsymbol{\xi}_{i}$ can be expanded as follows:

We see that there are two types of terms in \raisebox{-.9pt}{I}⃝ with non-zero expectation:

There are $i$ terms of this type, that make the following contribution to \raisebox{-.9pt}{I}⃝:

The second type of term in \raisebox{-.9pt}{I}⃝ with non-zero expectation has the form $\nabla_{{\boldsymbol{\theta}}_{m}}L_{i}(\Theta)^{\top}\boldsymbol{\epsilon}_{m}\boldsymbol{\epsilon}_{n}^{\top}\boldsymbol{\epsilon}_{n}\boldsymbol{\epsilon}_{m}^{\top}\nabla_{{\boldsymbol{\theta}}_{m}}L_{i}(\Theta)$ where $m\neq n$. Computing the expectation, we have:

The total contribution of terms of this type is:

Here, the terms with nonzero expectation have the form $\nabla_{{\boldsymbol{\theta}}_{m}}L_{i}(\Theta)^{\top}\boldsymbol{\epsilon}_{m}\boldsymbol{\epsilon}_{m}^{\top}\boldsymbol{\epsilon}_{n}\boldsymbol{\epsilon}_{n}^{\top}\nabla_{{\boldsymbol{\theta}}_{n}}L_{i}(\Theta)$ or $\nabla_{{\boldsymbol{\theta}}_{n}}L_{i}(\Theta)^{\top}\boldsymbol{\epsilon}_{n}\boldsymbol{\epsilon}_{m}^{\top}\boldsymbol{\epsilon}_{n}\boldsymbol{\epsilon}_{m}^{\top}\nabla_{{\boldsymbol{\theta}}_{m}}L_{i}(\Theta)$, both of which have expectation:

We have the following contribution from terms of this type, where the factor of 2 accounts for the two conditions ($\boldsymbol{\epsilon}_{n}\boldsymbol{\epsilon}_{m}^{\top}\boldsymbol{\epsilon}_{n}\boldsymbol{\epsilon}_{m}^{\top}$ and $\boldsymbol{\epsilon}_{m}\boldsymbol{\epsilon}_{m}^{\top}\boldsymbol{\epsilon}_{n}\boldsymbol{\epsilon}_{n}^{\top}$):

Next, we consider terms of type \raisebox{-.9pt}{b}⃝, which have the form $\boldsymbol{g}_{i}^{\top}\boldsymbol{v}_{i}\boldsymbol{\xi}_{i}^{\top}\boldsymbol{\xi}_{j}\boldsymbol{v}_{j}^{\top}\boldsymbol{g}_{j}$ where $i\neq j$. Note that we can expand $\boldsymbol{\xi}_{i}^{\top}\boldsymbol{\xi}_{j}$ as follows:

where we define $r=\min(i,j)$. Plugging in this expansion for $\boldsymbol{\xi}_{i}^{\top}\boldsymbol{\xi}_{j}$, we have:

Expanding \raisebox{-.9pt}{I}⃝, there are two types of terms of interest: ones of the form $\nabla_{{\boldsymbol{\theta}}_{m}}L_{i}(\Theta)^{\top}\boldsymbol{\epsilon}_{m}\boldsymbol{\epsilon}_{m}^{\top}\boldsymbol{\epsilon}_{m}\boldsymbol{\epsilon}_{m}^{\top}\nabla_{{\boldsymbol{\theta}}_{m}}L_{j}(\Theta)$, and ones of the form $\nabla_{{\boldsymbol{\theta}}_{n}}L_{i}(\Theta)^{\top}\boldsymbol{\epsilon}_{n}\boldsymbol{\epsilon}_{m}^{\top}\boldsymbol{\epsilon}_{m}\boldsymbol{\epsilon}_{n}^{\top}\nabla_{{\boldsymbol{\theta}}_{n}}L_{j}(\Theta)$. The expectation of the first type of term is:

The total contribution from terms like this is:

The expectation of the second type of term is:

The total contribution from terms of this type is:

Next, we look at terms in expression \raisebox{-.9pt}{II}⃝. We have two types of terms that have nonzero expectation:

Putting these components together, we have the following overall expression:

To obtain $\text{tr}(\text{Var}(\hat{\boldsymbol{g}}^{\text{PES-A}}))$, we subtract the following from the expression above:

G.1 Considering the Dependence on T𝑇T

The variance depends on the gradients of each loss term $L_{t}$ with respect to each of the per-timestep parameters ${\boldsymbol{\theta}}_{\tau}$. To gain insight into the structure of these gradients, we can arrange them in a matrix:

The RHS is upper-triangular due to the fact that $\nabla_{{\boldsymbol{\theta}}_{\tau}}L_{t}=0$ for all $\tau>t$. The variance of the PES estimator depends on the covariance between the gradients $\nabla_{{\boldsymbol{\theta}}_{\tau}}L_{t}$ in this matrix.

We consider two structures for the matrix: 1) a diagonal structure, where the gradients $\nabla_{{\boldsymbol{\theta}}_{i}}L_{j}=0,\forall i\neq j$; and 2) an upper-triangular structure as shown in the RHS of Eq. 74. For each of these two matrix structures, we will consider two scenarios for the covariance between gradients: a) all gradients $\nabla_{{\boldsymbol{\theta}}_{i}}L_{j}$ are identical; b) all gradients are i.i.d.

We denote the gradient of $L$ by $\boldsymbol{g}=\nabla_{{\boldsymbol{\theta}}}L(\Theta)=\sum_{t=1}^{T}\nabla_{{\boldsymbol{\theta}}}L_{t}(\Theta)=\sum_{t=1}^{T}\boldsymbol{g}_{t}$. (Note that $\nabla_{{\boldsymbol{\theta}}}L_{t}(\Theta)=\cancel{\nabla_{{\boldsymbol{\theta}}_{1}}L_{t}}+\cancel{\nabla_{{\boldsymbol{\theta}}_{2}}L_{t}}+\cdots+\nabla_{{\boldsymbol{\theta}}_{t}}L_{t}=\nabla_{{\boldsymbol{\theta}}_{t}}L_{t}$ due to the diagonal structure.)

To go from Eq. 78 to Eq. 79, we use the fact that $\nabla_{{\boldsymbol{\theta}}_{m}}L_{i}(\Theta)=0$ for $m\neq i$ and $\nabla_{{\boldsymbol{\theta}}_{n}}L_{j}(\Theta)=0$ for $n\neq j$. Next, note that when $M$ is diagonal, we have:

Recall our notation for the total gradient, $\boldsymbol{g}=\nabla_{{\boldsymbol{\theta}}}L(\Theta)=\sum_{t=1}^{T}\nabla_{{\boldsymbol{\theta}}}L_{t}(\Theta)=\sum_{t=1}^{T}\boldsymbol{g}_{t}$. If we assume that the gradients for each unroll are identical to each other, then:

If we assume that the gradients for each unroll are i.i.d., then:

G.1.2 Upper-Triangular Structure

Suppose all the terms in the matrix are equal, e.g., $\nabla_{{\boldsymbol{\theta}}_{i}}L_{j}=\boldsymbol{h},\forall i,j$. The total gradient $\boldsymbol{g}=\nabla_{{\boldsymbol{\theta}}}L(\Theta)$ is equal to the sum of the gradients in the upper-triangular matrix. Thus, $\boldsymbol{g}=\frac{T(T+1)}{2}\boldsymbol{h}$, so we can write:

Term \raisebox{-.9pt}{II}⃝ is as follows:

Next we derive each of the terms that arise in Eq. 111.

Combining all these terms, we obtain the following expression for the total variance:

We are interested in the scaling behavior as a function of the total gradient norm $\left|\left|\boldsymbol{g}\right|\right|^{2}$, where

Because the denominator in Eq. 127 is $\mathcal{O}(T^{4})$, we will have terms in the total variance that scale as:

In this case, by direct analogy to Equation 93, we have:

Here, the denominator is of order $O(T^{2})$, while the numerator is of order $O(T^{4})$, yielding variance that scales as $\mathcal{O}(T^{2})$:

Figure 16 shows the empirical variance for several potential scenarios. We performed an analysis similar to that in Section 4, measuring the variance of the PES gradient with respect to the number of unrolls for a small LSTM on the Penn TreeBank (PTB) dataset. We constructed synthetic data sequences to illustrate different scenarios: in Figure 16(a) we used a $10^{3}$ length sequence consisting of characters sampled uniformly at random from the PTB vocabulary, simulating the first scenario; in Figure 16(b) we used a $10^{3}$ length sequence consisting of a single repeated character, simulating the second scenario; Figure 16(c) shows the variance for real data—the first $10^{3}$ characters of PTB—which exhibits characteristics of both synthetic scenarios.

Appendix H Reducing Variance by Incorporating the Analytic Gradient

For functions $L$ that are differentiable, we can use the analytic gradient from the most recent partial unroll (e.g., backpropagating through the last $K$-step unroll) to reduce the variance of the PES gradient estimates. Below, we show how we can incorporate the analytic gradient in the ES estimate for $\frac{\partial L_{t}(\Theta)}{\partial{\boldsymbol{\theta}}}$:

We call the resulting estimator PES+Analytic. Algorithm 4 describes the implementation of this estimator, which requires a few simple changes from the standard PES estimator. We repeated the empirical variance measurement described in Section 4 and Appendix G using the PES+Analytic estimator, for each of the three scenarios from Appendix G, shown in Figure 18. Similarly to the other variance measurements, we report variance normalized by the squared norm of the true gradient. We found that variance increases with the number of unrolls, but the PES+Analytic variance is 1-2 orders of magnitude smaller than the standard PES variance.

Appendix I Connection to Gradient Estimation in Stochastic Computation Graphs

In this section, we show how PES can be derived using the framework for gradient estimation in stochastic computation graphs introduced in (Schulman et al., 2015). We follow their notation for this exposition: in Figure 19, squares represent deterministic nodes, which are functions of their parents; circles represent stochastic nodes which are distributed conditionally on their parents, and nodes not in squares or circles represent inputs. For notational simplicity, in the following exposition we consider 1-dimensional $\theta$. We represent the unrolled computation graph in terms of an input node $\theta$, that gives rise to a stochastic variable $\theta_{t}$ at each time step; the sampled $\theta_{t}$ is used to compute the state $s_{t}$, which is a deterministic function of the previous state $s_{t-1}$ and the current parameters $\theta_{t}$. The losses $L_{t}$ are designated as cost nodes, and our objective is $L=\sum_{t}L_{t}$.

Theorem 1 from (Schulman et al., 2015) gives the following general form for the gradient of the sum of cost nodes in such a stochastic computation graph. Here, $\mathcal{C}$ is the set of cost nodes; $\mathcal{S}$ is the set of stochastic nodes; $\text{DEPS}_{w}$ denotes the set of nodes that $w$ depends on; $a\prec^{D}b$ indicates that node $a$ depends deterministically on node $b$ (note that this relationship holds as long as there are no stochastic nodes along a path from $a$ to $b$; in our case, $\theta\prec^{D}\theta_{t}$ holds for all $t$); and $\hat{Q}_{w}$ is the sum of cost nodes downstream from node $w$.

For the computation graph in Figure 19, $\theta$ does not deterministically influence any of the cost nodes $L_{t}$, so the second term in the expectation in Eq. 136 will be 0. In addition, each stochastic node $\theta_{t}$, depends only on $\theta$, e.g. $\text{DEPS}_{\theta_{t}}=\{\theta\},\forall t$. Thus, our gradient estimate is:

$\hat{Q}_{\theta_{t}}$ is the sum of cost nodes downstream of $\theta_{t}$, thus $\hat{Q}_{\theta_{t}}=\sum_{i=t}^{T}L_{i}$. Now, each $\theta_{t}\sim\mathcal{N}(\theta,\sigma^{2})$, so we have:

where we used the reparameterization $\theta_{t}=\theta+\epsilon_{t}$ with $\epsilon_{t}\sim\mathcal{N}(0,\sigma^{2})$. Plugging this into Eq. 137, we have:

Appendix J Derivations and Compute/Memory Costs

Backpropagating through a full unroll of $T$ steps requires $T$ forward and backward passes, yielding compute $T(F+B)$; all $T$ states must be stored in memory to be available for gradient computation during backprop, yielding memory cost $TS$. Similarly, because TBPTT unrolls the computation graph for $K$ steps, it requires $K$ forward and backward passes, yielding computation $K(F+B)$, and requires storing $K$ states in memory, yielding memory cost $KS$. ARTBP is identical to TBPTT except that it randomly samples the truncation length in a theoretically-justified way to reduce or eliminate truncation bias. In theory, the sampled truncation lengths must allow for maximum length $T$, yielding worst-case compute $T(F+B)$ and memory cost $TS$. However, in practice this is often intractable, so truncation lengths may be sampled within a restricted range centered around $K$—this is no longer unbiased, but yields average case compute $K(F+B)$ and memory cost $KS$ (which is reported in Table 1).

and our objective is $L=\sum_{t=1}^{T}L_{t}$. In order to optimize this objective, we need the gradient $\nabla_{{\boldsymbol{\theta}}}L=\sum_{t=1}^{T}\frac{dL_{t}}{d{\boldsymbol{\theta}}}$. The loss at step $t$ is a function of $\boldsymbol{s}_{t}$, so we have:

Using Eq. 147 and the chain rule, we have:

Here, $G_{t}$ is $S\times P$, $H_{t}$ is $S\times S$, and $F_{t}$ is $S\times P$. RTRL maintains the Jacobian $G_{t}$, which requires memory $SP$; furthermore, instantiating the matrices $H_{t}$ and $F_{t}$ requires memory $S^{2}$ and $SP$, respectively, so the total memory cost of RTRL is $2SP+S^{2}$. The matrix multiplication $H_{t}G_{t-1}$ has computational complexity $S^{2}P$. The cost of computing the Jacobian $F_{t}$ is approximately $\text{min}\{S(F+B),P(F+B)\}$, depending on which of $S$ or $P$ is smaller-dimensional (and correspondingly whether we use forward-mode or reverse-mode automatic differentiation to compute the rows/columns of the Jacobian). Similarly, the cost of computing the Jacobian $H_{t}$ is approximately $S(F+B)$ (using either forward or reverse mode autodiff). Thus, the total computational cost of RTRL is: $S^{2}P+S(F+B)+\text{min}\{S(F+B)+P(F+B)\}$.

Note that, in general, it matters which of $\boldsymbol{s}_{t}$ or ${\boldsymbol{\theta}}$ is higher dimensional. In the case of unrolled optimization, $S$ is usually larger than $P$, causing RTRL to be particularly memory-intensive due to the $S\times S$ Jacobian $H_{t}$. The computation and memory costs we have derived here are expressed in a general form for state and parameter dimensions $S$ and $P$, respectively. In the case of RNN training, most prior work (such as (Tallec & Ollivier, 2017a; Mujika et al., 2018; Benzing et al., 2019)) assumes that the RNN parameters are of dimensionality $S^{2}$, where $S$ is the size of the hidden state. This is a simplification of the parameter count for RNNs, assuming that it is dominated by the hidden-to-hidden weight matrix.

Unbiased Online Recurrent Optimization (UORO) (Tallec & Ollivier, 2017a) approximates RTRL by maintaining a rank-1 estimate of the Jacobian $G_{t}$ as:

The reparameterization gradient estimator is $\hat{g}^{\text{reparam}}=\frac{1}{N}\sum_{i=1}^{N}\nabla_{\boldsymbol{\theta}}L({\boldsymbol{\theta}}+\sigma\boldsymbol{\epsilon}^{(i)})$, where $\boldsymbol{\epsilon}^{(i)}\sim\mathcal{N}(0,I)$. With respect to computational complexity, this is equivalent to BPTT: its compute cost is $T(F+B)$ and its memory cost is $TS$.

ES applied to an unroll of length $K$ requires performing $K$ forward passes—it does not require any backward passes, since ES is not gradient-based (e.g., it is a zeroth-order optimization algoritm). Because ES does not require backprop, it does not need to store the intermediate states in memory, only the most recent state, yielding memory cost $S$ that is independent of the unroll length. Using ES with $N$ particles yields total compute and memory costs $NKF$ and $NS$, respectively.

As PES is an evolutionary strategies-based method, it also does not require backward passes; applied to unrolls of length $K$, PES has compute cost $KF$. In addition to storing the current state of size $S$ as in ES, PES also maintains a perturbation accumulator for each particle; thus, the memory cost of a single PES chain is $S+P$. Using PES with $N$ particles yields total compute and memory costs $NKF$ and $N(S+P)$, respectively.

Similarly to standard PES, we need to maintain a collection of $N$ states, each of size $S$, and $N$ perturbation accumulators, each of size $P$, yielding memory cost $N(S+P)$; unrolling each state for $K$ steps requires computational cost $NKF$. To incorporate the analytic gradient, we need to maintain one additional particle that is unrolled using the mean ${\boldsymbol{\theta}}$ rather than a perturbed version ${\boldsymbol{\theta}}+\boldsymbol{\epsilon}$; this adds memory cost $S$. The main computational and memory overhead comes from the gradient computation through the partial unroll of length $K$: similarly to TBPTT, this requires storing $K$ intermediate states, yielding memory cost $KS$, and requires $K$ forward and $K$ backward operations, yielding computational cost $K(F+B)$. Combined with the memory and computational cost of standard PES, we have total compute cost $NKF+K(F+B)$ and total memory cost $N(S+P)+(K+1)S$.

Appendix K Diagrammatic Representation of Algorithms

Figure 20 provides diagrammatic representations of ES and PES. For each partial unroll, vanilla ES starts from a shared initial state $\boldsymbol{s}^{(0)}$ that is evolved in parallel using perturbed parameters ${\boldsymbol{\theta}}+\boldsymbol{\epsilon}^{(i)}$. After each truncated unroll, the mean parameters ${\boldsymbol{\theta}}$ are used to update the state, which then becomes the initial state for the next truncated unroll; no information is passed between truncated unrolls for vanilla ES. In contrast, PES maintains a set of states $\boldsymbol{s}^{(i)}$ that are evolved in parallel, each according to a different perturbation of the parameters ${\boldsymbol{\theta}}$ in each truncated unroll. Intuitively, these states maintain their history between truncated unrolls, since we accumulate the perturbations experienced by each state over the course of meta-optimization; when we reach the end of an inner problem, the states are reset to the same initialization, and the perturbation accumulators are reset to $\boldsymbol{0}$.

Appendix L Ablation Studies

In this section, we show an ablation study over the the number of particles $N$, and the truncation length $K$ (which controls the number of unrolls per inner-problem). In Figure 21 we show the sensitivity of PES to these meta-parameters for a version of the 2D regression problem (from Section 5.4) with total inner problem length $T=10,000$.

Appendix M Implementation

Code Listing 1 presents a simple JAX implementation of the toy 2D regression meta-learning problem from Section 5.4, in a self-contained, runnable example. PES is easy to implement efficiently in JAX by making use of the construct jax.vmap (or jax.pmap in settings with multiple workers) to parallelize the unrolling computations over $N$ particles.