SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\operatorname {t r} \left(\operatorname {C o v} \left[ \operatorname {G P E S} _ {K = c} (\theta) \right]\right) = \frac {(d + 2)}{T} \sum_ {t = 1} ^ {T} \left(\| g ^ {t} \| _ {2} ^ {2}\right) - \left\| \frac {1}{T} \sum_ {t = 1} ^ {T} g ^ {t} \right\| _ {2} ^ {2} + \frac {1}{T} \sum_ {t = 1} ^ {T} \left(\frac {d}{2} \sum_ {j = 1, j ^ {\prime} = 1} ^ {\lceil t / c \rceil} \left\| g _ {c, j} ^ {t} - g _ {c, j ^ {\prime}} ^ {t} \right\| _ {2} ^ {2}\right). \tag {9}
$$
$^7 [x]$  is smallest integer  $\geq x$  .  $\lfloor x\rfloor$  is the largest integer  $\leq x$
<sup>8</sup>Proofs for all the Lemmas and Theorems are in Appendix D.
Figure 3: Total variance of  $\mathrm{GPES}_K$  vs. noise-sharing period  $K$  for different  $\theta_{i}$  's from the learned trajectory of PES.  $\mathrm{GPES}_{K = T}$  (NRES) has the lowest total variance among estimators of its class (including PES) for each  $\theta_{i}$ .
To understand how the value of  $K = c$  changes the total variance, we notice that only the nonnegative third term in (9) depends on it. This term measures the pairwise squared distance between nonoverlapping partial sums  $g_{c,j}^{t}$  for all  $j$ . When  $c = T$ , for every  $t \in [T]$ , there is only a single such partial sum as  $\lceil t / c \rceil = 1$ . In this case, this third term reduces to its smallest value of 0. Thus:
Corollary 6. Under Assumption 2, when  $W = 1$ , the gradient estimator  $\mathrm{GPES}_{K = T}(\theta)$  has the smallest total variance among all  $\{\mathrm{GPES}_K: K \in [T]\}$  estimators.
Remark 7. To understand Corollary 6 intuitively, notice that at a given time step  $t$  (i.e., a length-1 truncation window), any  $\mathrm{GPES}_{K = c}$  gradient estimator  $(c\in [T])$  aims to unbiasedly estimate the total derivative of the smoothed loss at this step with respect to  $\theta$ , which we have denoted by  $g^{t}$ . By applying a new Gaussian noise perturbation every  $c < T$  steps, the  $\mathrm{GPES}_{K = c}$  estimators indirectly estimate  $g^{t}$  by first unbiasedly estimating the gradients inside each size-  $c$  noise-sharing window:  $\{g_{c,j}^{t}\}_{j = 1}^{\lceil t / c\rceil}$  and then summing up the result (notice  $g^{t} = \sum_{j = 1}^{\lceil t / c\rceil}g_{c,j}^{t}$ ). To obtain this extra (yet unused) information about the intermediate partial derivatives, these estimators require more randomness and thus suffer from a larger total variance than the  $\mathrm{GPES}_{K = T}$  estimator which directly estimates  $g^{t}$ .
Experimental Verification of Corollary 6. We empirically verify Corollary 6 on a meta-training learned optimizer task ( $T = 1000$ ; see additional details in Section 5.2). Here we save a trajectory of  $\theta_{i}$  learned by PES ( $i$  denotes training iteration) and compute the total variance of the estimated gradients (without hysteresis) by GPES with different values of  $K$  in Figure 3(a) ( $W = 1$ ). In agreement with theory,  $K = T$  has the lowest variance.
4 Noise-Reuse Evolution Strategies
NRES has lower variance than PES. As variance reduction is desirable in stochastic optimization [18], by Corollary 6, the gradient estimator  $\mathrm{GPES}_{K = T}$  is particularly attractive and can serve as a variance-reduced replacement for PES. When  $K = T$ , we only need to sample a single  $\epsilon$  once at the beginning of an episode (when  $\tau = 0$ ) and reuse the same noise for the entirety of that episode before it resets. This removes the need to keep track of the cumulative applied noise ( $\xi$ ) (Figure 2(b)), making the algorithm simpler and more memory efficient than PES. Due to its noise-reuse property, we name this gradient estimator  $\mathrm{GPES}_{K = T}$  the Noise-Reuse Evolution Strategies (NRES) (pseudocode in Algorithm 2). Concurrent with our work, Vicol [19] independently proposes a similar algorithm with different analyses. We discuss in detail how our work differs from [19] in Appendix B. Despite Theorem 5 assuming  $W = 1$ , one can relax this assumption to any  $W$  that divides the horizon length  $T$ . By defining a "mega" UCG whose single transition step is equivalent to  $W$  steps in the original UCG, we can apply Corollary 6 to this mega UCG and arrive at the following result.
Corollary 8. Under Assumption 2, when  $W$  divides  $T$ , the NRES gradient estimator has the smallest total variance among all  $\mathrm{GPES}_{K = cW}$  estimators  $c \in [T / W]$ .
NRES is a replacement for FullES. By sharing the same noise over the entire horizon, the smoothing objective of NRES is the same as FullES's. Thus, we can think of NRES as the online counterpart to the offline algorithm FullES. Hence NRES can act as a drop-in replacement to FullES in UCGs. A single FullES worker runs  $2T$  unroll steps for each gradient estimate, while a single NRES runs only  $2W$  steps. Motivated by this, we compare the average of  $T / W$  i.i.d. NRES gradient estimates with 1 FullES gradient estimate as they require the same amount of compute.
NRES is more parallelizable than FullES. Because the  $T / W$  NRES gradient estimators are independent of each other, we can run them in parallel. Under perfect parallelization, the entire NRES gradient estimation would require  $O(W)$  time to complete. In contrast, the single FullES gradient estimate has to traverse the UCG from start to finish, thus requiring  $O(T)$  time. Hence, NRES is  $T / W$  times more parallelizable than FullES under the same compute budget (Figure 4(a)).
(a) Gradient estimation of FullES vs NRES under the same per-update budget of  $2T$
Figure 4: (a) Comparison of FullES and NRES gradient estimation under the same unroll budget. Unlike FullES which can only use a single noise perturbation  $\epsilon$  to unroll sequentially for an entire episode of length  $T$ , NRES can use  $T / W$  parallel step-unlocked workers each unrolling inside its random truncation windows of length  $W$  with independent perturbations  $\epsilon^{(i)}$ . This results in a  $T / W \times$  speed-up and variance reduction (Theorem 9) over FullES. (b) The total variance of NRES and FullES estimators under the same compute budget at the same set of  $\theta_{i}$  checkpoints in Figure 3(a). NRES achieves significantly lower total covariance.
(b)  $T = 1000, W = 1$  (±95% ci)
iter. j of θj (×104)	total variance tr(Cov[·])
W/T ∑i=1T/W NRESi(θj)	FullES(θj)
1	1.47 ± 0.38	864.34 ± 385.01
2	15.74 ± 4.90	513.50 ± 74.71
3	4.28 ± 0.78	201.38 ± 30.18
4	3.20 ± 0.50	684.85 ± 86.90
5	4.47 ± 0.88	181.06 ± 24.97
6	4.42 ± 0.68	288.61 ± 46.77
7	5.33 ± 1.08	448.30 ± 62.44
8	6.62 ± 0.89	154.86 ± 28.21
NRES can often have lower variance than FullES. We next compare the variance of the average of  $T / W$  i.i.d. NRES gradient estimates with the variance of 1 FullES gradient estimate:
Theorem 9. Under Assumption 2, for any  $W$  that divides  $T$ , if
$$
\left. \sum_ {k = 1} ^ {T / W} \left\| \sum_ {t = W \cdot (k - 1) + 1} ^ {W \cdot k} g ^ {t} \right\| _ {2} ^ {2} \leq \frac {d + 1}{d + 2} \left\| \sum_ {j = 1} ^ {T / W} \sum_ {t = W \cdot (k - 1) + 1} ^ {W \cdot k} g ^ {t} \right\| _ {2} ^ {2}, \right. \tag {10}
$$
then for iid  $\{\mathrm{NRES}_i(\theta)\}_{i = 1}^{T / W}$  estimators,  $\operatorname {tr}(\operatorname {Cov}(\frac{1}{T / W}\sum_{i = 1}^{T / W}\operatorname {NRES}_i(\theta))\leq \operatorname {tr}(\operatorname {Cov}(\operatorname {FullES}(\theta)).$
Remark 10. To understand the inequality assumption in (10), we notice that it relates the sum of the squared 2-norm of vectors  $\{\sum_{t = W\cdot (k - 1) + 1}^{W\cdot k}g^t\}_{k = 1}^{T / W}$  with the squared 2-norm of their sum. When these vectors are pointing in similar directions, this inequality would hold (to see this intuitively, consider the more extreme case when all these vectors are exactly in the same direction). Because each term  $\sum_{t = W\cdot (k - 1) + 1}^{W\cdot k}g^{t}$  can be understood as the total derivative of the sum of smoothed losses in the  $k$  th truncation window with respect to  $\theta$ , we see that inequality (10) is satisfied when, roughly speaking, different truncation windows' gradient contributions are pointing in similar directions. This is often the case for real-world applications because if we can decrease the losses within a truncation window by changing the parameter  $\theta$ , we likely will also decrease other truncation windows' losses. At a high-level, Theorem 9 shows that for many practical unrolled computation graphs, NRES is not only better than FullES due to its better parallelizability but also better due to its lower variance given the same computation budget.
Empirical Verification of Theorem 9. We empirically verify Theorem 9 in Figure 4(b) using the same set up of the meta-training learned optimizer task used in Figure 3(a). Here we compare the total variance of averaging  $T / W = 1000$  i.i.d. NRES estimators versus using 1 FullES gradient estimator (same total amount of compute). We see that NRES has a significantly lower total variance than FullES while also allowing  $T / W = 1000$  times wall-clock speed up due to its parallelizability.
NRES is particularly suitable for optimization in UCGs in two scenarios: 1) when the loss surface exhibits extreme local sensitivity; 2) when automatic differentiation of the loss is not possible/gives noninformative (e.g., zero) gradients. In this section, we focus on three applications exhibiting these properties: a) learning Lorenz system's parameters (sensitive), b) meta-training learned optimizers (sensitive), and c) reinforcement learning (nondifferentiable), and show that NRES outperforms existing AD and ES methods for these applications. When comparing online gradient estimation methods, we keep the number of workers  $N$  used by all methods the same for a fair comparison. For the offline method FullES, we choose its number of workers to be  $W / T \times$  the number of NRES workers on all tasks (in order to keep the same number of unroll steps per-update) except for the learned optimizer task in Section 5.2 where we show that NRES can solve the task faster while using much fewer per-update steps than FullES.
Figure 5: (a) The pathological training loss surface of the Lorenz system problem (left) and the optimization trajectory of different  $\mathrm{GPES}_K$  gradient estimators (right). NRES's trajectory is the smoothest because of its lowest variance. (b) Different ES methods' loss convergence on the same problem. NRES converges the fastest.
5.1 Learning dynamical system parameters
In this application, we consider learning the parameters of a Lorenz system, a canonical chaotic dynamical system. Here the state  $s_t = (x_t, y_t, z_t) \in \mathbb{R}^3$  is unrolled with two learnable parameters  $a, r^9$  with the discretized transitions  $(dt = 0.005)$  starting at  $s_0 = (x_0, y_0, z_0) = (1.2, 1.3, 1.6)$ :
$$
x _ {t + 1} = x _ {t} + a \left(y _ {t} - x _ {t}\right) d t; y _ {t + 1} = y _ {t} + \left[ x _ {t} \cdot \left(r - z _ {t}\right) - y _ {t} \right] d t; z _ {t + 1} = z _ {t} + \left[ x _ {t} \cdot y _ {t} - 8 / 3 \cdot z _ {t} \right] d t.
$$
Due to the positive constraint on  $r > 0$  and  $a > 0$ , we parameterize them as  $\theta = (\ln (r),\ln (a))\in \mathbb{R}^2$  and exponentiate the values in each application. We assume we observe the ground truth  $z$ -coordinate  $z_{t}^{\mathrm{gt}}$  for  $t\in [T], T = 2000$  steps unrolled by the default parameters  $(r^{\mathrm{gt}},a^{\mathrm{gt}}) = (28,10)$ . For each step  $t$ , we measure the squared loss  $L_{t}^{s}(s_{t})\coloneqq (z_{t} - z_{t}^{\mathrm{gt}})^{2}$ . Our goal is to recover the ground truth parameters  $\theta_{\mathrm{gt}} = (\ln (28),\ln (10))$  by optimizing the average loss over all time steps using vanilla SGD. We first visualize the training loss surface in the left panel of Figure 5(a) (also see Figure 8 in the Appendix) and notice that it has extreme sensitivity to small changes in the parameter  $\theta$ .
To illustrate the superior variance of NRES over other  $\mathrm{GPES}_K$  estimators, we plot in the right panels of Figure 5(a) the optimization trajectory of  $\theta$  using gradient estimator  $\mathrm{GPES}_K$  with different values of  $K$  under the same SGD learning rate. We see that NRES's trajectory exhibits the least amount of oscillation due to its lowest variance. In contrast, we notice that PES's trajectory is highly unstable, thus requiring a smaller learning rate than NRES to achieve a possibly slower convergence. Hence, we take extra care in tuning each method's constant learning rate and additionally allow PES to have a decay schedule. We plot the convergence of different ES gradient estimators in wall-clock time using the same hardware in Figure 5(b). (We additionally compare against automatic differentiation methods in Figure 9 in the Appendix; they all perform worse than the ES methods shown here.)
In terms of the result, we see that NRES outperforms 1) TES, as NRES is unbiased and can better capture long-term dependencies; 2) PES, as NRES has provably lower variance, which aids convergence in stochastic optimization; 3) FullES, as NRES can produce more gradient updates in the same amount of wall clock time than FullES (with parallelization, each NRES update takes  $O(W)$  time instead of FullES's  $O(T)$  time). Additionally, we plot the asymptotically converged loss value when we train with a significantly larger number of particles ( $N = 20000$ ) for PES and NRES. We see that by only using  $N = 200$  particles, NRES can already converge around its asymptotic limit, while PES is still far from reaching its limit within our experiment time.
5.2 Meta-training learned optimizers
In this application [3], the meta-parameters  $\theta$  of a learned optimizer control the gradient-based updates of an inner model's parameters. The inner state  $s_t$  is the optimizer state which consists of both the inner model's parameters and its current gradient momentum statistics. The transition function  $f_t$  computes an additive update vector to the inner parameters using  $\theta$  and a random training batch and outputs the next optimizer state  $s_{t + 1}$ . Each time step  $t$ 's meta-loss  $L_{t}^{s}$  evaluates the updated inner parameters' generalization performance using a sampled validation batch.
We consider meta-training the learned optimizer model given in [3] ( $d = 1762$ ) to optimize a 3-layer MLP on the Fashion MNIST dataset for  $T = 1000$  steps. (We show results on the same task
<sup>9</sup>We don't learn the third parameter, fixed at  $8/3$ , so that we can easily visualize a 2-d loss surface.
(a)
Figure 6: (a) Different ES gradient estimators' training loss convergence on the learned optimizer task in wall-clock time. NRES reaches the lowest loss fastest. (b) The loss convergence of  $\mathrm{GPES}_K$  gradient estimators with difference  $K$  values on the same task. NRES converges the fastest due to its reduced variance.
with higher-dimension and longer horizon in Appendix E.1.2.) This task is used in the training task distribution of the state of the art learned optimizer VeLO  $[20]^{10}$ . The loss surface for this problem has high sharpness and many suboptimal minima as previously shown in Figure 1(a). We meta-train with Adam using different gradient estimation methods with the same hardware and tune each gradient estimation method's meta learning rate individually. Because AD methods all perform worse than the ES methods, we defer their results to Figure 10 in the Appendix and only plot the convergence of the ES methods in wall-clock time in Figure 6(a).
Here we see that NRES reaches the lowest loss value in the same amount of time. In fact, PES and FullES would require 5 and  $9 \times$  (respectively) longer than NRES to reach a loss NRES reaches early on during its training, while TES couldn't even reach that loss within our experiment time. It is worth noting that, for this task, NRES only requires  $2 \cdot N \cdot W = 200$  unrolls to produce an unbiased, low-variance gradient estimate, which is even smaller than the length of a single episode ( $T = 1000$ ). In addition, we situate NRES's performance within our proposed class of  $\mathrm{GPES}_K$  estimators in Figure 6(b). In accordance with Corollary 8, NRES converges fastest due to its reduced variance.
5.3 Reinforcement Learning
It has been shown that ES is a scalable alternative to policy gradient and value function methods for solving reinforcement learning tasks [14]. In this application, we learn a linear policy  ${}^{11}$  (following [21]) using different ES methods on the Mujoco [22] Swimmer ( $d = 16$ ) and Half Cheetah task ( $d = 102$ ). We minimize the average of negative per-step rewards over the horizon length of  $T = 1000$ , which is equivalent to maximizing the undiscounted sum rewards. Unlike [14, 21, 15], we don't use additional heuristic tricks such as 1) rank conversion of rewards, 2) scaling by loss standard deviation, or 3) state normalization. Instead, we aim to compare the pure performance of different ES methods assuming perfect parallel implementations. To do this, we measure a method's performance as a function of the number of sequential environment steps it used. Sequential environment steps are steps that have to happen one after another (e.g., the environment steps within the same truncation window). However, steps that are parallelizable don't count additionally in the sequential steps. Hence, the wall-clock time under perfect parallel implementation is linear with respect to the number of sequential environment steps used. As all the methods we compare are iterative update methods, we additionally require that each method use the same number of environment steps per update when measuring each method's required number of sequential steps to solve a task. We tune the SGD learning rate individually for each method and plot their total rewards progression on both Mujoco tasks in Figure 7.
Here we see that TES fails to solve both tasks due to the short horizon bias [23], making it unable to capture the long term dependencies necessary to solve the tasks. On the other hand, PES, despite being unbiased, suffers from high variance, making it take longer (or unable in the case of Half Cheetah) to solve the task than NRES. As for FullES, despite using the same amount of compute per gradient update as NRES, it's much less parallelizable as discussed in Section 4 - it takes much longer time  $(10\times$  and  $60\times)$  than NRES assuming perfect parallelization. In addition to the number
<sup>10</sup>We show the performance of ES methods on another task from this distribution in Appendix E.1.2.
We additionally compare the ES methods on learning a non-linear ( $d = 726$ ) policy on the Half-Cheetah task in Appendix E.1.3.
(a) Swimmer,  $T = 1000$
Figure 7: ES methods' performance vs. the number of sequential environment steps used in solving the Mujoco (a) Swimmer task and (b) Half Cheetah task. NRES solves the tasks fastest under perfect parallelization.
(b) Half Cheetah,  $T = 1000$
of sequential steps, we additionally show the total number of environment steps used by each method in Table 3 in the Appendix - NRES also uses the least total number of steps to solve both tasks, making it the most sample efficient.
Beyond the most related work in Section 2, in this section, we further position NRES relative to existing zeroth-order gradient estimation methods. We also provide additional related work on automatic differentiation (AD) methods for unrolled computation graphs in Appendix B.
Zeroth-Order Gradient Estimators. In this work, we focus on zeroth-order methods that can estimate continuous parameters' gradients to be plugged into any first-order optimizers, unlike other zeroth-order optimization methods such as Bayesian Optimization [24], random search [25], or Trust Region methods [26, 27]. We also don't compare against policy gradient methods [9], because they assume internal stochasticity of the unrolling dynamics, which may not hold for deterministic policy learning [e.g., 22]. Within the space of evolution strategies methods, many works have focused on improving the vanilla ES method's variance by changing the perturbation distribution [28-31], considering the covariance structure [32], and using control variates [33]. However, these works do not consider the unrolled structure of UCGs and are offline methods. In contrast, we reduce the variance by incorporating this unrolled aspect through online estimation and noise-reuse. As the aforementioned variance reduction methods work orthogonally to NRES, it is conceivable that these techniques can be used in conjunction with NRES to further reduce the variance.
7 Discussion, Limitations, and Future Work
In this work, we improve online evolution strategies for unbiased gradient estimation in unrolled computation graphs by analyzing the best noise-sharing strategies. By generalizing an existing unbiased method, Persistent Evolution Strategies, to a broader class, we analytically and empirically identify the best estimator with the smallest variance and name this method Noise-Reuse Evolution Strategies (NRES). We demonstrate the convergence benefits of NRES over other automatic differentiation and evolution strategies methods on a variety of applications.
Limitations. As NRES is both an online method and an ES method, it naturally inherits some limitations shared by all methods of these two classes, such as hysteresis and variance's linear dependence on the dimension  $d$ . We provide a detailed discussion of these limitations in Appendix F.
Future Work. There are some natural open questions: 1) choosing a better sampling distribution for NRES. Currently the isotropic Gaussian's variance  $\sigma^2$  is tuned as a hyperparameter. Whether there are better ways to leverage the sequential structure in unrolled computation graphs to automate the selection of this distribution is an open question. 2) Incorporating hysteresis. Our analysis assumes no hysteresis in the gradient estimates and we haven't observed much impact of it in our experiments. However, understanding when and how to correct for hysteresis is an interesting direction.
We thank Kevin Kuo, Jingnan Ye, Tian Li, and the anonymous reviewers for their helpful feedback.