[论文解读] Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies

Paul Vicol, Luke Metz+1 more|arXiv (Cornell University)|Dec 27, 2021

Machine Learning and Data Classification被引用 7

一句话总结

本文提出持久演化策略（PES），一种在展开计算图中使用截断反向传播进行无偏梯度估计的方法。通过在部分展开过程中累积扰动而非每次重新开始，PES 消除了截断偏差，同时保持低内存使用、低方差和快速更新——从而实现高效的超参数调优和学习优化器训练。

ABSTRACT

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.

研究动机与目标

• 解决在展开计算图中因截断反向传播而引入的高偏差问题，该问题限制了元学习和超参数优化的性能。
• 克服现有方法（如 TBPTT 的梯度偏差、RTRL 的高内存/计算开销，以及标准 ES 在截断设置下的高方差）的局限性。
• 开发一种无偏、内存高效且适用于序列模型中实时、在线参数更新的梯度估计方法。
• 通过提供稳定、可扩展的梯度估计机制，实现学习优化器的有效训练和超参数调优。
• 为反向传播提供一个理论基础扎实、实用性强的替代方案，在展开系统中实现更低的计算和内存开销。

提出的方法

• PES 将展开计算图划分为短的、固定长度的片段（截断展开），并在每个片段后应用演化策略（ES）。
• 与每次展开后重置扰动不同，PES 在时间上累积扰动，从而在存在截断的情况下实现无偏梯度估计。
• 该方法使用对称抽样（antithetic sampling）以降低梯度方差，通过生成对称的扰动（正负对）提高统计效率。
• PES 将梯度估计计算为所有粒子中扰动加权损失的平均值，并随时间累积校正项。
• 引入一种方差减少技术，将 PES 与最近一次展开的解析梯度相结合，从而加快收敛速度。
• 该方法通过 JAX 的 JIT 编译和向量化高效实现，支持高并行性和低内存占用。

实验结果

研究问题

• RQ1 我们能否在不承担完整反向传播的内存和计算成本的前提下，实现展开计算图中无偏的梯度估计？
• RQ2 在截断展开中累积扰动是否能消除标准截断反向传播方法（如 TBPTT）固有的偏差？
• RQ3 PES 是否能在保持低梯度方差的同时，实现在序列模型中快速、在线的参数更新？
• RQ4 在超参数调优和学习优化器训练方面，PES 与现有方法（如 TBPTT、RTRL、标准 ES）相比，在收敛速度、稳定性和内存使用方面表现如何？
• RQ5 PES 是否能有效应用于元优化和超参数搜索中常见的不可微或病态损失曲面？

主要发现

• PES 对损失的平滑版本提供了无偏梯度估计，并且在二次损失下是精确无偏的。
• 在截断设置中，PES 的梯度方差显著低于标准演化策略，尤其在结合最近一次展开的解析梯度时更为明显。
• 在超参数调优任务中，PES 的收敛速度和稳定性优于 TBPTT 和标准 ES，最终性能相当或更优。
• PES 能够有效实现学习优化器的元训练，展示了对内层超参数的稳定且高效的优化。
• 该方法在连续控制任务中表现出良好的可扩展性，在具有展开优化目标的策略学习中展现出稳健性能。
• 实证结果证实，PES 保持了低内存使用，并支持频繁的参数更新，适用于实时和在线学习场景。