Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies

在當前深度學習及強化學習等領域中，許多模型訓練過程涉及對「未展開計算圖」（unrolled computation graphs）進行反向梯度估計。然而，未展開計算圖的長期依賴特性常導致梯度估計效率與偏差問題，尤其當需要對多步過程或迭代方法進行優化時，傳統的反向傳播（Backpropagation through Time, BPTT）方法面臨記憶體瓶頸與偏差累積的困境。針對這個挑戰，Vicol 等人在其 2021 年 ICML 頂會上發表的論文《Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies》提出了一種創新的無偏差梯度估計方法，稱為 Persistent Evolution Strategies（持久演化策略，簡稱 PES）。該論文不僅獲得了 ICML 2021 的 Outstanding Paper 獎項，更為複雜時間依賴的優化問題提供了全新的解決視角。

研究背景與動機

在許多機器學習問題中，尤其是強化學習（Reinforcement Learning）、元學習（Meta-learning）以及隱狀態模型（如 RNN）訓練過程中，模型的計算流程往往需要「展開」或「展開計算圖」。例如，為了優化一個 RNN，我們往往需要對其在多個時間步的狀態進行展開，之後通過 BPTT 來更新參數。然而，展開長時間序列時，存儲中間狀態和計算梯度的開銷巨大，且梯度估計易受截斷（truncation）偏差影響，影響訓練效果。

為了克服這些限制，有研究嘗試使用無偏且低記憶體消耗的黑盒優化方法，例如進化策略（Evolution Strategies, ES），其通過直接擾動參數，估計目標函數的梯度。ES 方法可直接在黑盒環境下優化，避免了對整個展開計算圖反向傳播的需求。然而，ES 本身的估計方差相對較大，且在長期依賴問題上應用受限。

為此，作者致力於設計一種既能保持無偏差特性，又能有效利用展開計算圖結構以提高梯度估計效率的方法，解決長期依賴和記憶體限制的雙重挑戰，這成為本論文的研究核心動機。

核心方法與創新

論文提出的 Persistent Evolution Strategies（PES）是一種基於演化策略的無偏梯度估計方法，針對未展開計算圖中的長期時間步問題特別設計。其核心創新可概括為以下幾點：

1. 持久記憶的漸近梯度估計： PES 方法在每一步優化迭代中，不完全重新隨機抽樣梯度方向，而是將過去的隨機擾動持續保留，結合新擾動一起進行梯度估計。透過這種「持久化」機制，PES 能在多次迭代間累積信息，有效抑制估計方差並提升穩定性。
2. 無偏性保證： 不同於常見的截斷 BPTT 所產生的偏差，PES 能保證梯度估計在數學上無偏，即梯度估計的期望等於真實梯度。這對於保證訓練收斂和提升優化品質非常關鍵。
3. 高效利用展開計算圖結構： PES 利用展開計算圖中狀態轉移函數的可微分結構，將隨機擾動從參數空間映射到狀態空間，結合直接元素級梯度信息，實現更精準且穩健的梯度估計。
4. 降低記憶體需求： 傳統 BPTT 在長展開時間軸容易導致巨大的記憶體消耗。PES 通過隨機擾動追蹤與增量更新，使得在不用完整保存中間狀態的前提下，仍能取得有效梯度估計，顯著減少空間複雜度。

這些設計綜合起來，使 PES 成為一種在保持無偏的同時，又能高效穩定控制方差的梯度估計工具，特別適合在包含多步展開計算圖的複雜優化任務中使用。

主要實驗結果

為驗證 PES 的效能，作者在多項經典且具挑戰性的任務中進行實驗，包括：

• 循環神經網路（RNN）訓練： 對標準語言模型及序列預測任務進行訓練，PES 在長序列中能取得比截斷 BPTT 更優的梯度估計質量，從而達到更好的模型性能。
• 強化學習問題： 在須要長期時間依賴的策略優化任務中，PES 相較於傳統 ES 和梯度估計方法顯示出更快的收斂速度和更穩定的表現。
• 元學習任務： 元優化場景中，PES 的無偏估計有效降低元梯度的高方差問題，提升了元學習的泛化能力及訓練穩定性。

實驗結果一致證明，PES 方法在方差控制、數值穩定以及終極任務性能方面均有顯著優勢，且能有效應對長時間步未展開計算圖所帶來的挑戰。

對 AI 領域的深遠影響

PES 方法的提出，對深度學習和強化學習領域有多方面重大啟示：

• 解決長期依賴梯度消失與偏差問題： 長時間展開的模型訓練是一大瓶頸，PES 提供了理論和實踐上無偏且低方差的梯度估計方案，將助力未來更深層次時間序列模型及複雜動態系統的訓練。
• 促進黑盒優化與可微分模型融合： PES 將演化策略與可微分結構巧妙結合，激發了黑盒優化方法在深度學習中的創新應用潛力，未來可推動更多混合優化技術的誕生。
• 降低硬體資源需求，提升可擴展性： 傳統 BPTT 隨著序列長度擴大記憶體消耗大幅增加，PES 以更節省資源的無偏估計方法，適合部署於受限硬體環境與大規模分佈式訓練。
• 拓展元學習與強化學習可訓練性邊界： PES 有助於有效估計複雜動態系統中的梯度，從而可能推動元學習與強化學習算法向更深更大尺度範圍發展。

總結而言，Persistent Evolution Strategies 利用持久擾動策略，在展開計算圖中實現無偏且高效的梯度估計，這不僅解決了長序列優化中的核心難題，也為未來多領域複雜系統的訓練提供強大理論和方法支持，具備廣泛而深遠的學術及工業應用前景。

---

論文資訊
📄 Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies
👥 Vicol, Metz, Sohl-Dickstein
🏆 ICML 2021 · Outstanding Paper
🔗 arxiv.org/abs/2112.02434