Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies
This addresses a bottleneck in machine learning for scenarios like training learned optimizers and tuning hyperparameters, offering an incremental improvement over existing gradient estimation methods.
The paper tackles the problem of high variance, bias, slow updates, or large memory usage in optimizing parameters for unrolled computation graphs, such as in training RNNs or learned optimizers, by introducing Persistent Evolution Strategies (PES), which divides the graph into truncated unrolls and uses evolution strategies with correction terms to eliminate bias, resulting in unbiased, low-memory, and rapid updates with reasonable variance.
Unrolled computation graphs arise in many scenarios, including training RNNs, tuning hyperparameters through unrolled optimization, and training learned optimizers. Current approaches to optimizing parameters in such computation graphs suffer from high variance gradients, bias, slow updates, or large memory usage. We introduce a method called Persistent Evolution Strategies (PES), which divides the computation graph into a series of truncated unrolls, and performs an evolution strategies-based update step 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 experimentally demonstrate the advantages of PES compared to several other methods for gradient estimation on synthetic tasks, and show its applicability to training learned optimizers and tuning hyperparameters.