用無超參自動梯度下降訓(xùn)練ImageNet數(shù)據(jù)集

作者：朱先忠 2023-05-23 08:00:00

最近，我們團(tuán)隊推出了一種稱為自動梯度下降（AGD：automatic gradient descent）的優(yōu)化器，它可以在不使用超參數(shù)的情況下訓(xùn)練ImageNet數(shù)據(jù)集。基于這一優(yōu)化器算法，可以消除在模型訓(xùn)練時對昂貴且耗時的學(xué)習(xí)速率調(diào)整、學(xué)習(xí)速率衰減調(diào)度器等選項(xiàng)的需求。

譯者 | 朱先忠

審校 | 重樓

簡介

我與Jeremy Bernstein、Kevin Huang、Navid Azizan和Yisong Yue一起參與了這個項(xiàng)目。您可以參閱Jeremy先生的GitHub代碼庫以獲得一個干凈的Pytorch實(shí)現(xiàn)，也可以參閱我的GitHub代碼庫以獲得具有更多功能的實(shí)驗(yàn)版本。下面的圖1總結(jié)了我們提出的AGD算法與Adam和SGD兩種優(yōu)化算法之間的對比情況。

圖1：實(shí)線顯示訓(xùn)練精度，虛線顯示測試精度

其中，左圖展示的是與我們的方法相反的具有默認(rèn)超參數(shù)的Adam和SGD算法在CIFAR-10數(shù)據(jù)集上的深度全連接網(wǎng)絡(luò)（FCN）上不理想的表現(xiàn)。中間圖展示的是Adam和SGD算法的學(xué)習(xí)率網(wǎng)格搜索對比情況。我們所開發(fā)的優(yōu)化器的性能與經(jīng)過充分調(diào)參后的Adam和SGD性能差不多。右圖展示了我們所開發(fā)的AGD算法將ImageNet數(shù)據(jù)集訓(xùn)練到相當(dāng)高的測試精度。

動機(jī)

任何訓(xùn)練過深度神經(jīng)網(wǎng)絡(luò)的人都可能有過不得不調(diào)整模型的學(xué)習(xí)速率的經(jīng)歷。這樣做的目的無處乎：（1）確保訓(xùn)練的最大效率；（2）找到正確的學(xué)習(xí)率可以顯著提高整體模型的泛化應(yīng)用能力。當(dāng)然，這個過程也會伴隨著開發(fā)人員的巨大代價的付出。

圖2：為什么學(xué)習(xí)率對模型優(yōu)化非常重要

為了最大限度地提高收斂速度，你總會想盡可能找到最佳學(xué)習(xí)率：這是很大的一個值！而且，目標(biāo)函數(shù)中的非線性項(xiàng)可能會把你折磨得“死去活來”。

然而，對于SGD算法來說，最佳學(xué)習(xí)率在很大程度上取決于正在訓(xùn)練的體系結(jié)構(gòu)。找到這個值通常需要一個代價高昂的網(wǎng)格搜索程序，需要涵蓋許多數(shù)量級內(nèi)容。此外，還需要選擇和調(diào)整其他超參數(shù)，如動量和學(xué)習(xí)率衰減調(diào)度器。

為此，我們提出了一種稱為自動梯度下降（AGD）的優(yōu)化器，它不需要使用學(xué)習(xí)率就可以訓(xùn)練大范圍的架構(gòu)和數(shù)據(jù)集，可以在ImageNet數(shù)據(jù)集上從基礎(chǔ)模型一直擴(kuò)展訓(xùn)練到ResNet-50這樣的復(fù)雜模型。這種算法消除了對任何超參數(shù)調(diào)整的需要（因?yàn)榇怂惴ú恍枰治鲇行W(xué)習(xí)率和學(xué)習(xí)率衰減值），從而節(jié)省了計算成本，并大大加快了模型訓(xùn)練的過程。

為什么我們需要超參數(shù)呢？

究其原因，一個深度學(xué)習(xí)系統(tǒng)往往都是由許多相互關(guān)聯(lián)的組件組成：架構(gòu)、數(shù)據(jù)、損失函數(shù)和梯度等。這些組件的交互方式都是圍繞著一種結(jié)構(gòu)進(jìn)行的；但是，到目前為止，還沒有人確切地弄清楚這個結(jié)構(gòu)。所以，我們需要進(jìn)行大量的調(diào)整工作（例如學(xué)習(xí)率、初始化、調(diào)度器等），以確保快速收斂，并避免過度擬合。

然而，盡管完美地描述這些相互作用可以消除優(yōu)化過程中的所有自由度；但是，目前這些自由度都是通過手動超參數(shù)調(diào)整來實(shí)現(xiàn)的。二階方法中目前使用Hessian矩陣來表征目標(biāo)對權(quán)重擾動的敏感性，并以這種方式去除自由度——然而，這種方法可能計算密集，因此對于大型模型來說并不實(shí)用。

我們團(tuán)隊通過分析這些相互作用的特征，最終找到了AGD優(yōu)化方法：

對于給定的數(shù)據(jù)和架構(gòu)，我們根據(jù)權(quán)重的變化來約束神經(jīng)網(wǎng)絡(luò)輸出的變化。
我們將目標(biāo)的變化（一次批運(yùn)算中所有輸入的總損失）與神經(jīng)網(wǎng)絡(luò)輸出的變化聯(lián)系起來。
我們將這些結(jié)果結(jié)合在一種所謂的多數(shù)最小化方法（majorise-minimise approach）中。我們對目標(biāo)進(jìn)行多數(shù)化——也就是說，我們推導(dǎo)出與目標(biāo)相切的目標(biāo)的上界。然后我們可以將這個上限降到最低，因?yàn)槲覀冎肋@會讓我們的算法加快收斂。如圖3所示，紅色曲線表示目標(biāo)函數(shù)的多數(shù)化，藍(lán)色曲線表示目標(biāo)函數(shù)。

圖3：左側(cè)面板顯示了多數(shù)最小化背后的基本思想

最小化目標(biāo)函數(shù)（藍(lán)色）是通過最小化一系列上限或多數(shù)（紅色）來實(shí)現(xiàn)的。右側(cè)面板顯示了權(quán)重的變化如何導(dǎo)致函數(shù)的變化，進(jìn)而導(dǎo)致單個數(shù)據(jù)點(diǎn)的損失的變化，從而導(dǎo)致目標(biāo)的變化。我們用?W來約束?L，并用它來構(gòu)建我們的多數(shù)化。

基于Pytorch框架的AGD算法實(shí)現(xiàn)

在本節(jié)中，我們將介紹我們所開發(fā)的AGD優(yōu)化算法的所有關(guān)鍵部分基于Pytorch框架的實(shí)現(xiàn)代碼。相應(yīng)的草圖推導(dǎo)可參考文后的附錄A。

參數(shù)操作

我們所使用的參數(shù)化與傳統(tǒng)PyTorch默認(rèn)值略有不同。AGD算法可以在不借助參數(shù)化支持的情況下即可完成；不過，使用參數(shù)化可以簡化分析過程。對于完全連接的層l，我們使用正交初始化，并進(jìn)行適當(dāng)?shù)目s放以便使奇異值具有大小：sqrt((l的輸入維度)/(l的輸出維度))。

import math
import torch
from torch.nn.init import orthogonal_

def singular_value(p):
 sv = math.sqrt(p.shape[0] / p.shape[1])
 if p.dim() == 4:
 sv /= math.sqrt(p.shape[2] * p.shape[3])
 return sv

class AGD:
 @torch.no_grad()
 def __init__(self, net, gain=1.0):

 self.net = net
 self.depth = len(list(net.parameters()))
 self.gain = gain #用于加速訓(xùn)練的可選用的超參數(shù)——在后面的任何實(shí)驗(yàn)中并不使用

 for p in net.parameters():
 if p.dim() == 1: raise Exception("Biases are not supported.")
 if p.dim() == 2: orthogonal_(p)
 if p.dim() == 4:
 for kx in range(p.shape[2]):
 for ky in range(p.shape[3]):
 orthogonal_(p[:,:,kx,ky])
 p *= singular_value(p)

在此，我們之所以使用這種規(guī)范化，是因?yàn)樗哂蠵yTorch默認(rèn)參數(shù)化所沒有的優(yōu)秀特征，包括寬度穩(wěn)定性、對激活函數(shù)中的抗爆炸式增長能力以及促進(jìn)特征學(xué)習(xí)。這與Greg Yang和Edward Hu二人開發(fā)的muP庫相似。

更新操作

這一步可以分成兩個單獨(dú)的部分。第一個是計算eta（η），即“自動學(xué)習(xí)率”，它實(shí)現(xiàn)對所有層的更新進(jìn)行縮放。注意，Eta對梯度范數(shù)有對數(shù)依賴性——當(dāng)梯度很小時，Eta近似為線性（就像標(biāo)準(zhǔn)優(yōu)化器一樣）；但是，當(dāng)它們非常大時，對數(shù)會自動執(zhí)行一種梯度修剪。

@torch.no_grad()
 def step(self):
 G = 0
 for p in self.net.parameters():
 G += singular_value(p) * p.grad.norm(dim=(0,1)).sum()
 G /= self.depth
 eta = math.log(0.5 * (1 + math.sqrt(1 + 4*G)))

 for p in self.net.parameters():
 factor = singular_value(p) / p.grad.norm(dim=(0,1), keepdim=True)
 p -= self.gain * eta / self.depth * factor * p.grad

 return log

使用eta乘以層的權(quán)重范數(shù)、乘以歸一化梯度并除以深度來更新每個層。除以深度實(shí)現(xiàn)了按深度縮放。有趣的是，分析中沒有使用梯度歸一化，因?yàn)橄馎dam這樣的其他優(yōu)化器也啟發(fā)式地引入了類似的想法。

實(shí)驗(yàn)

進(jìn)行這些實(shí)驗(yàn)的目標(biāo)是為了測試AGD優(yōu)化算法的能力：（1）在廣泛的體系結(jié)構(gòu)和數(shù)據(jù)集上的收斂性能；以及（2）實(shí)現(xiàn)與調(diào)參后的Adam和SGD算法相當(dāng)?shù)臏y試精度。

接下來的圖4顯示了從完全連接網(wǎng)絡(luò)（FCN）到ResNet-50網(wǎng)絡(luò)的四種架構(gòu)在從CIFAR-10到ImageNet的數(shù)據(jù)集上的學(xué)習(xí)曲線。我們將實(shí)線所示的AGD優(yōu)化器與虛線所示的標(biāo)準(zhǔn)優(yōu)化器進(jìn)行比較（SGD基于ImageNet數(shù)據(jù)集，調(diào)優(yōu)后的Adam對應(yīng)于其他三種情況）。最上面一行顯示了訓(xùn)練目標(biāo)（損失）和自動學(xué)習(xí)率η。最下面一行顯示了訓(xùn)練和測試的準(zhǔn)確性。圖5展示了在8層FCN上進(jìn)行比較的AGD、調(diào)優(yōu)后的Adam與調(diào)優(yōu)后的SGD。我們看到所有三種算法的性能非常相似，達(dá)到了幾乎相同的測試精度。

圖6顯示，AGD在較寬的深度（2至32）和寬度（64至2048）范圍內(nèi)訓(xùn)練FCN。圖7顯示了AGD對批次大小（從32到4096）和4層FCN的依賴性。無論批量大小，它似乎都會收斂到一個很好的最優(yōu)值！

圖4——四種架構(gòu)上的AGD與Adam算法比較：CIFAR-10上的16層深的FCN、CIFAR-10的ResNet-18、CIFAR-100上的VGG-16和ImageNet-1k上的ResNet-50。AGD與超參數(shù)調(diào)優(yōu)后的Adam保持合理的速度（這需要網(wǎng)格搜索幾個數(shù)量級）！

這些實(shí)線表示AGD，虛線表示Adam（除了ImageNet，我們使用SGD）。最上面一行顯示了訓(xùn)練目標(biāo)（即損失）和訓(xùn)練期間自動學(xué)習(xí)率η的值。最下面一行顯示了訓(xùn)練和測試的準(zhǔn)確性。

圖5：具有均方誤差損失的深度為8的FCN網(wǎng)絡(luò)上的AGD、Adam與SGD算法比較

值得注意的是，Adam和SGD都調(diào)整了它們的學(xué)習(xí)率。在左邊圖形中，我們繪制了訓(xùn)練和測試目標(biāo)函數(shù)（即損失）。中間圖形顯示了訓(xùn)練和測試的準(zhǔn)確性。右邊圖形顯示了每個訓(xùn)練周期期間權(quán)重的平均值、最小值和最大值變化。

圖6：AGD算法在很大的深度和寬度范圍內(nèi)能夠自如地收斂。不過，較小的架構(gòu)往往缺乏實(shí)現(xiàn)低損耗的能力，但AGD仍可以訓(xùn)練它們！

圖7：為了檢查AGD算法不僅適用于批量大小128，這里還提供了一個深度為4的FCN模型的批量大小選擇

結(jié)論

總的來說，本文實(shí)現(xiàn)了一個“架構(gòu)感知”類型優(yōu)化器：自動梯度下降（AGD），它能夠在指定的批量范圍大小下訓(xùn)練成功從CIFAR-10數(shù)據(jù)集上的FCN等小型模型系統(tǒng)直到大型數(shù)據(jù)集ImageNet上的ResNet-50等大型模型系統(tǒng)，而無需手動進(jìn)行超參數(shù)調(diào)整。

雖然使用AGD并沒有從機(jī)器學(xué)習(xí)中完全刪除所有的超參數(shù)，但剩下的超參數(shù)——批量大小和體系結(jié)構(gòu)——通常能夠?qū)崿F(xiàn)均衡時間/計算方面的預(yù)算代價。

盡管如此，仍然有許多工作要做。例如，我們沒有明確考慮由于批次大小而引入梯度的隨機(jī)性。此外，我們也沒有研究像權(quán)重衰減這樣的規(guī)則化。雖然我們在增加對仿射參數(shù)（在批處理范數(shù)層中）和偏差項(xiàng)的支持方面做了一些工作，但我們還沒有對其進(jìn)行廣泛的測試，也沒有像本文引用的其他結(jié)果那樣從理論上證明它是合理的。

也許最重要的是，我們?nèi)匀恍枰M(jìn)行轉(zhuǎn)換器（Transformer）所需的分析，并在NLP任務(wù)中測試AGD。在OpenWebText2上對GPT-2進(jìn)行的初步實(shí)驗(yàn)表明，AGD在這種情況下也非常有效！

最后，如果你想嘗試使用AGD算法的話，請查看Jeremy的GitHub代碼以獲得一個干凈的版本，或者查看我的GitHub以獲得支持偏置和仿射參數(shù)的開發(fā)版本！我們希望你會覺得這個算法有用。

附錄A

最后，我們將在這里簡要介紹一下我們論文證明過程中的一些重要步驟。這是為那些想看看此算法的主要思想是如何結(jié)合在一起的人設(shè)計的，而不需要經(jīng)過我們論文中的完整證明。

其中，方程（1）明確規(guī)定了如何將數(shù)據(jù)集S上的總體目標(biāo)分解為各個數(shù)據(jù)點(diǎn)。L表示損失，x表示輸入，y表示目標(biāo)，w表示權(quán)重。方程（2）顯示了目標(biāo)線性化誤差的分解——在權(quán)重Δw發(fā)生變化的情況下，高階項(xiàng)對損失ΔL（w）變化的貢獻(xiàn)。目標(biāo)的線性化誤差很重要，因?yàn)樗扔谠跈?quán)重w-bounding下擴(kuò)展的損失中的高階項(xiàng)的貢獻(xiàn)，這將告訴我們在高階項(xiàng)變得重要之前我們可以移動多遠(yuǎn)，并確保我們正在采取合理大小的步長實(shí)現(xiàn)收斂。

方程（2）中的RHS上的第一項(xiàng)是兩個高維向量、模型的線性化誤差和損耗相對于f（x）的導(dǎo)數(shù)之間的內(nèi)積。由于沒有明確的理由說明為什么這兩個向量應(yīng)該對齊，我們假設(shè)它們的內(nèi)積為零。

將L（W+ΔW）加到方程（2）的每一側(cè)，并注意到損失的線性化誤差恰好是布雷格曼散度（Bregman Divergence），我們可以進(jìn)一步簡化符號表示：

Bregman散度是兩點(diǎn)之間距離的度量（在這種情況下，是神經(jīng)網(wǎng)絡(luò)兩個不同參數(shù)選擇的輸出），用嚴(yán)格凸函數(shù)定義——在這種情況中，即損失函數(shù)。

對于均方誤差損失來說，計算Bregman散度實(shí)際上是非常簡單的。并且，我們?nèi)菀椎贸鋈缦卤磉_(dá)：

其中，d是網(wǎng)絡(luò)的輸出維度。我們現(xiàn)在可以斷言存在以下縮放表達(dá)式。盡管所有這些都有點(diǎn)武斷；但是，采用這種形式會使問題分析變得更加簡單。

我們對網(wǎng)絡(luò)輸出的大小使用以下兩個界限。其中，方程（5）限制了網(wǎng)絡(luò)輸出的大小，并且僅將（輸入縮放）和（權(quán)重縮放）應(yīng)用于完全連接的網(wǎng)絡(luò)。方程（6）限制了f（x）的最大變化與權(quán)重W的變化。（6）中的第二個不等式在大深度下最緊，但在任何深度下都成立。

現(xiàn)在，我們將方程（6）代入方程（4），并將所有項(xiàng)顯式展開，得到方程（7）。

在關(guān)于梯度條件的附加假設(shè)下，我們可以用方程（8）中定義的G代替方程（7）中的和，這在本文中進(jìn)行了詳細(xì)討論。最后，我們得到方程（9）——這是多數(shù)化——圖3中的紅線。我們通過對η進(jìn)行微分來最小化多數(shù)化，并在exp（η）中求解由此產(chǎn)生的二次方，保留正解。于是，我們有了以下新的表達(dá)形式：