Low-Variance Gradient Estimation in Unrolled Computation Graphs with ES-Single

TL;DR

This paper introduces ES-Single, a simple and low-variance evolution strategies-based method for unbiased gradient estimation in unrolled computation graphs, improving over previous methods like PES especially for long inner problems.

Contribution

ES-Single is a novel, unbiased gradient estimator that is simpler to implement and has lower, constant variance compared to PES, enabling better performance on long unrolled problems.

Findings

ES-Single has lower variance than PES.

ES-Single outperforms PES on multiple tasks.

ES-Single is unbiased for quadratic problems.

Abstract

We propose an evolution strategies-based algorithm for estimating gradients in unrolled computation graphs, called ES-Single. Similarly to the recently-proposed Persistent Evolution Strategies (PES), ES-Single is unbiased, and overcomes chaos arising from recursive function applications by smoothing the meta-loss landscape. ES-Single samples a single perturbation per particle, that is kept fixed over the course of an inner problem (e.g., perturbations are not re-sampled for each partial unroll). Compared to PES, ES-Single is simpler to implement and has lower variance: the variance of ES-Single is constant with respect to the number of truncated unrolls, removing a key barrier in applying ES to long inner problems using short truncations. We show that ES-Single is unbiased for quadratic inner problems, and demonstrate empirically that its variance can be substantially lower than that of PES. ES-Single consistently outperforms PES on a variety of tasks, including a synthetic benchmark task, hyperparameter optimization, training recurrent neural networks, and training learned optimizers.