論文筆記 — Symbolic Discovery of Optimization Algorithms

12 min readFeb 27, 2023

Paper Link: https://arxiv.org/abs/2302.06675

Code Link: https://github.com/lucidrains/lion-pytorch

Introduction

Optimizer 在機器學習扮演相當重要的角色，在近幾年出現了許多不一樣的自適應優化器(Adaptive optimizer)，像是有著Decoupled weight decay的 Adam (AdamW)就是其中一種標準的優化器，用來訓練大多數深度神經網絡。

其中一個比較新的研究方向是希望可以自動去找到其中最佳的optimization Algorithm， Learning to optimizer (L2O) 就是這種方法，建議通過訓練參數化模型（例如神經網絡）來發現優化器以輸出更新。然而，那些通常在有限數量的小任務上訓練的黑盒優化器很難學習到最好的設置，而且更大的模型需要更多的Training steps 來訓練。

另一種方法則是應用reinforcement learning 或是 Monte Carlo Sampling 來發現新的優化器，其中搜索空間由由預定義操作數（例如梯度和動量）組成的樹定義和運算符（例如，一元和二進制數學運算）。然而，為了使搜索易於管理，他們通常通過使用固定操作數和限制樹的大小來限制搜索空間，因此可能會去限制發現的可能性。

在本文中，我們提出了一種將算法發現表述為程序搜索的方法，並將其應用於發現優化算法。有兩個主要挑戰。第一個是在無限稀疏的程序空間中尋找高質量的算法。第二個是進一步選擇可以從小型代理任務推廣到更大的、最先進的任務的算法。為了應對這些挑戰，我們採用了一系列技術，包括:

1. Warm-start
2. Restart
3. Abstract execution
4. Funnel selection
5. Program simplification.

此方法稱作：Lion，evoLved sIgn mOmeNtum 的縮寫（很牽強，但作者最大）。該算法與各種自適應算法的不同之處在於，它僅跟踪momentum並利用sign operation 來計算更新，從而降低內存開銷並在所有維度上統一更新幅度。儘管簡單，Lion 在一系列模型（Transformer、MLP、ResNet、U-Net 和 Hybrid）和任務（圖像分類、視覺語言對比學習、擴散、語言建模和微調）中展示了出色的性能。值得注意的是，通過在 BASIC 中用 Lion 替換 Adafactor，我們在ImageNet 上實現了 88.3% 的零樣本和 91.1% 的fine-tune accuracy，分別超過之前的最佳結果 2% 和 0.1%。

Symbolic Discovery of Algorithms

我們使用program format 的symbolic representation 具有以下優點：

1. 它符合算法必須作為Program 來執行的事實
2. 與神經網絡等參數化模型相比，program等Symbloic representation 更易於分析、理解和轉移到新任務
3. Program Length可用於估計不同程序的複雜性，從而更容易選擇更簡單、通常更具通用性的program。這項工作側重於深度神經網絡訓練的優化器，但該方法通常適用於其他任務。

Program Search Space

我們堅持以下三個標準：

搜索空間應該足夠靈活，以便能夠發現新算法
program 應該易於分析並融入機器學習工作流程
program 應著重於高級算法設計而不是low-level 實現細節。

Input/output signature

該程序定義了一個訓練函數，它對正在搜索的優化算法進行編碼，其中主要輸入是模型權重 (w)、梯度 (g) 和當前訓練步驟的學習率 (lr)。主要輸出是對權重的更新。該程序還包含初始化為零的額外變量，以在訓練期間收集歷史信息。

Building blocks

train 函數由一系列 assignment statements 組成，對statements 或local variable 的數量沒有限制。每個statements 調用一個使用常量或現有變量作為輸入的函數，結果值存儲在新的或現有的變量中。我們選擇了 45 個常見的數學函數，其中大部分對應於 NumPy 中的函數或線性代數中的運算，例如線性插值函數interp(x, y, a)，它等價於(1 — a) * x + a * y。

Mutations and redundant statements

進化搜索中使用的突變設計與program的representation 緊密相關。我們包括三種類型的突變：

在隨機位置插入一個帶有隨機選擇的函數和參數的新語句
刪除一個隨機選擇的語句
通過隨機改變其函數參數之一來修改隨機語句，它可以是變量或常量。

為了改變參數，我們將其替換為現有變量或通過從正態分佈 X ∼ N (0 1) 中採樣獲得的新生成常量。此外我們通過將現有常數乘以隨機因子 2^a 來改變它，其中 a ∼ N(0 1)。這些常量在優化算法中用作可調超參數，像是AdamW 中的峰值學習率和權重衰減。

Infinite and sparse search space

給定無限數量的語句和局部變量，以及可變常量的存在，程序搜索空間是無限的。即使我們忽略常量並限製程序長度和變量數量，潛在程序的數量仍然非常大。粗略估計可能的程序數量為：

其中 n 是可能函數的數量，n_v 是局部變量的數量，n_a 是每個語句的平均參數數量，l 是程序長度。更重要的是，挑戰來自搜索空間中高性能程序的稀疏性。為了說明這一點，我們進行了隨機搜索，評估了超過 2 Million 的程序。其中最好的程序仍然明顯不如 AdamW。

Efficient Search Techniques

我們應用正則化進化(Regularized Evolution) 來解決無限稀疏搜索空間帶來的挑戰。帶有warm-start和restart的進化，因為它簡單、可擴展，保留了 P 種算法，這些算法通過循環逐漸改進。每個循環隨機選擇 T<P 算法，並選擇表現最好的算法作為父算法，即 tournament selection（Goldberg 和 Deb，1991）。然後copy 並mutate 該父算法以產生一個子算法，該子算法被添加到種群中，同時刪除最舊的算法(註：有點像是基因演算法)。

通常，進化搜索從隨機候選開始，但我們將初始種群作為 AdamW warm-start 以加速搜索。默認情況下，我們使用的初始Tournament 規模為 2，參與規模為 1K。為了進一步提高搜索效率，我們採用兩種類型的重啟：

從初始程序重啟，由於進化的隨機性，這會導致不同的局部最優，並鼓勵探索。這可以通過並行運行多個搜索來完成
從迄今發現的最佳算法重新開始，進一步優化它，鼓勵利用。

下圖顯示了五個進化搜索實驗的平均值和標準誤差。我們通過僅允許進化中的常數突變來運行基於 AdamW 的超參數調整，並通過對隨機程序進行採樣來運行隨機搜索，兩者的計算量都增加了 4 倍。我們的搜索明顯優於兩個基線取得的最佳結果，如圖中的兩條虛線所示。

Pruning through abstract execution

我們建議從三個來源修剪程序空間中的冗餘(redundancies)：具有語法或類型/形狀錯誤的程序、功能等效的程序以及程序中的冗餘語句。在程序實際執行之前，我們執行一個抽象的執行步驟:

推斷變量類型和形狀以檢測有錯誤的程序，並不斷變異父程序，直到生成有效的子程序
生成一個 Hash Value，該Hash Value 唯一標識輸出是如何從輸入中計算出來的，允許我們緩存和查找語義重複的程序
識別在實際執行和分析過程中可以忽略的冗餘語句。例如，程序 4 是在去除程序 8 中的所有冗餘語句後得到的(下圖)。

Generalization: Program Selection and Simplification

Large generalization gap

我們希望優化器在不同的體系結構、數據集甚至不同的領域上表現良好，因此發現的算法需要表現出很強的分佈外泛化能力。進化過程中的稀疏搜索空間和固有噪聲進一步加劇了這一挑戰，導致不同運行之間的泛化特性不一致。

Funnel selection

為了縮小泛化差距，我們根據搜索適應度收集效果較好的程序，並使用一系列驗證任務添加額外的選擇步驟來選擇那些泛化更好的程序。為了節省計算，我們應用了一個漏斗選擇過程，逐漸增加元驗證任務的規模。

例如，從代理任務 A 開始，我們通過增加模型大小和訓練步驟來創建一個大 10 倍的任務 B。只有在任務 B 上超過基線的算法才會在任務 C 上進行評估，該任務 C 大 100 倍。這種方法使我們能夠逐漸過濾掉表現出較差泛化性能的算法，最終導致選擇能夠很好地泛化到更大任務的算法。

Simplification

更簡單的程序更容易理解，我們的直覺是它們更可能泛化，因此我們通過以下步驟簡化程序：

刪除了對通過抽象執行確定的最終輸出沒有貢獻的冗餘語句
刪除了非冗餘但在刪除時產生最小差異的語句
最後，手動重新排列語句，為變量分配清晰的描述性名稱，並將程序轉換為更簡單的數學等效形式。

Derivation and Analysis of Lion

Derivation

Search 和 Funnel Selection過程通過從原始程序 8自動刪除冗餘語句獲得程序 4。我們進一步簡化它得到最終算法（Lion）。在簡化過程中從程序4中刪除了幾個不必要的元素

cosh 函數被刪除，因為 m 將在下一次迭代中重新分配（第 3 行）
使用 arcsin 和 clip 的語句也被刪除，因為我們觀察到沒有它們質量不會下降。三個紅色語句轉換為一個符號函數。
程序 4 中同時使用了 m 和 v，但 v 僅改變了動量的更新方式，並且不需要單獨跟踪。

Analysis

Sign update and regularization

Lion 算法通過Sign operation在所有維度上產生具有統一幅度的更新，這在原理上不同於各種自適應優化器。直覺上，符號操作為更新添加了噪聲，作為一種正則化形式並有助於泛化，Lion 在 ImageNet 上訓練的 ViT-B/16 與 AdamW 相比具有更高的訓練誤差，但在驗證集上的準確率提高了 2%。

Momentum tracking

與 AdamW 和動量 SGD 中常用的 0.9 相比，Lion 中用於跟踪動量的默認 EMA 因子為 0.99 (β2)。在應用符號運算之前，當前梯度和動量以 0.9 (β1) 的因子進行插值。 EMA 因子和插值的這種選擇允許 Lion 在記住動量梯度和在更新中對當前梯度施加更多權重之間取得平衡。

Hyperparameter and batch size choices

與 AdamW 和 Adafactor 相比，Lion 更簡單並且超參數更少，因為它不需要 ε 和因式分解相關(factorization-related)的參數。 Lion 相對於 AdamW 的優勢隨著Batch Size的增加而擴大，這符合通過數據並行擴展模型訓練的常見做法。

Memory and runtime benefits

Lion 只保存動量，因此比 AdamW 等流行的自適應優化器佔用的內存更小，這在訓練大型模型和/或使用大批量時很有用。

例如，AdamW 需要至少 16 個 TPU V4 芯片來訓練圖像分辨率為 224 且批量大小為 4,096 的 ViT-B/16，而 Lion 只需要 8 個（均具有 bfloat16 動量）。另一個實際的好處是，由於 Lion 的簡單性，Lion 在我們的實驗中具有更快的運行時間（步數/秒），通常比 AdamW 和 Adafactor 提速 2–15%。