[論文レビュー] 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は、各部分アンラップごとに摂動を蓄積することで、再起動を避け、切り捨てバイアスを排除しながら、低メモリ使用量、低分散、高速な更新を維持する。これにより、ハイパーパramータチューニングや学習された最適化手法の訓練が可能になる。

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.

研究の動機と目的

• アンラップド計算グラフにおける切り捨て誤差を伴う逆誤差伝搬によって生じる高バイアスを解消すること。これは、メタラーニングやハイパーパramータ最適化の性能を制限する要因である。
• TBPTT（バイアス付き勾配）やRTRL（高メモリ/計算コスト）および標準的ES（切り捨て設定下で高分散）といった既存手法の限界を克服すること。
• 不偏で、メモリ効率が良く、順序付きモデルにおけるリアルタイムかつオンラインでのパラメータ更新に適した勾配推定手法を開発すること。
• 安定的かつスケーラブルな勾配推定メカニズムを提供することで、学習された最適化手法の効果的な訓練とハイパーパramータチューニングを可能にすること。
• 計算およびメモリのオーバーヘッドを低減したバックプロパゲーションの理論的裏付けが明確で実用的な代替手段を提供すること。

提案手法

• PESは、アンラップド計算グラフを短く固定長のセグメント（切り捨てアンラップ）に分割し、各セグメント後に進化戦略（ES）を適用する。
• 各アンラップで摂動を再起動するのではなく、PESは時間経過とともに摂動を蓄積することで、切り捨てにもかかわらず不偏な勾配推定を実現する。
• 勾配の分散を低減するために反対称サンプリングを用い、対称な摂動（正と負）を生成することで統計的効率を向上させる。
• PESは、全粒子における摂動重み付き損失の平均を勾配推定値として計算し、時間経過とともに補正項を蓄積する。
• 最近のアンラップからの解析的勾配と組み合わせることで、収束速度を向上させる分散低減技術を導入する。
• JAXのJITコンパイルおよびベクトル化を効果的に活用することで、高い並列性と低メモリフットプリントを実現し、効率的な実装を実現する。

実験結果

リサーチクエスチョン

• RQ1 完全なバックプロパゲーションのメモリおよび計算コストを負担せずに、アンラップド計算グラフにおける不偏勾配推定を達成できるか？
• RQ2 切り捨てアンラップにわたって摂動を蓄積することで、TBPTTのような標準的切り捨てアンラップ手法に内在するバイアスが解消されるか？
• RQ3 PESは、順序付きモデルにおけるリアルタイムかつオンラインでのパラメータ更新を可能にしながら、低分散の勾配を維持できるか？
• RQ4 ハイパーパramータチューニングや学習された最適化手法の訓練において、PESはTBPTT、RTRL、標準的ESといった既存手法と比較して収束速度、安定性、メモリ使用量の点で優れているか？
• RQ5 PESは、メタ最適化やハイパーパramータ探索で一般的に見られる非微分可能または病理的損失関数の領域に対しても、効果的に適用可能か？

主な発見

• PESは、損失関数の平滑化版に対する不偏勾配推定を提供し、2次損失関数に対しては正確に不偏である。
• 特に最近のアンラップからの解析的勾配と組み合わせた場合、切り捨て設定下で標準的ESよりも顕著に低い勾配分散を達成する。
• ハイパーパramータチューニングのタスクにおいて、PESはTBPTTや標準的ESよりも速くかつより安定して収束し、最終的な性能は同等またはそれ以上である。
• PESは、学習された最適化手法のメタトレーニングを効果的に可能にし、内側ループのハイパーパラメータの安定的かつ効率的な最適化を示している。
• 連続的制御タスクへのスケーラビリティが高く、アンラップド最適化目的関数を用いたポリシー学習においても、頑健な性能を示している。
• 実験的結果から、PESは低メモリ使用量を維持するとともに、頻繁なパラメータ更新を可能にし、リアルタイムおよびオンライン学習のシナリオに適していることが確認された。