JOVANA
Explore Library Glossary Getting Started Three Levels Fields How it works Mission
Join the mission
All guides

大規模最佳化:大批次、雜訊與變異

當你跨數百個加速器訓練時,批次極大、梯度是一個平均。學會臨界批次大小、用於大批次訓練的逐層縮放,以及為何變異縮減在深度學習中罕見。

為何更大的批次會停止帶來幫助

小批次梯度是真實梯度的一個帶雜訊估計,而這個雜訊有幫助——它讓 SGD 逃離陡峭區域並進行探索。當你跨多個裝置擴展分散式訓練時,自然的做法是放大批次讓每個裝置都忙碌。但有一個天花板:超過某個批次大小後,加倍批次不再讓收斂步數減半,因為梯度估計已經夠準、多出來的樣本大多是冗餘的。你為每步付出兩倍計算,換來的卻是趨近於零的步數減少。

這個轉折由臨界批次大小(critical batch size)主宰:在它以下,每步更多資料能買到成比例的更快收斂;在它以上,報酬急遽遞減。大致知道你的問題的臨界點落在哪裡,正是「叢集被善用」與「叢集在冗餘樣本上燒電」之間的差別。

正如神经网络的缩放定律,增大批量带来的加速也会递减:一旦超过临界批量,每步投入更多数据几乎换不来收益。

损失对规模的对数-对数缩放定律曲线,随规模增大而趋于平缓。

梯度雜訊尺度:量測臨界批次

梯度雜訊尺度給你一個實際的數字。想法是:估計每樣本梯度的變異,並與平均梯度的範數平方相比。當這個訊雜比很差(樣本之間的梯度彼此分歧很大)時,大批次有幫助,因為平均能抵消分歧;當訊雜比很好時,你已接近臨界批次,再增長就是浪費。雜訊尺度正是那個比值,而它在各種任務上都出奇地能良好預測臨界批次大小。

B_{\text{noise}}=\frac{\operatorname{tr}(\Sigma)}{\lvert G\rvert^{2}}

梯度噪声尺度:单样本梯度方差与平均梯度范数平方之比,决定了有效批量大小。

# estimate by comparing gradient norms at two batch sizes b_small, b_big
# B_noise ~ tr(Sigma) / |g|^2  (variance of grads over squared mean-grad norm)
g_small = grad(loss, batch=b_small)
g_big   = grad(loss, batch=b_big)
# fit |g_est|^2 = |g_true|^2 + tr(Sigma)/B  across the two B values, solve for B_noise
雜訊尺度在訓練過程中上升,因此可用的批次大小會隨模型學習而增長。

LAMB:真的會收斂的大批次訓練

把批次大小推到數萬會搞垮天真的最佳化器:適合某一層的單一全域學習率會在另一層上衝過頭,因為各層在「權重範數對更新範數之比」上差異極大。LAMB(用於批次訓練的逐層自適應動量)以逐層信任比(trust ratio)修正這點:它先算出 Adam 式的更新,再重新縮放每一層的更新,使其大小與該層的權重範數成比例。沒有任何一層被允許走出相對於自身權重而言過於巨大的一步。

x_{t+1}^{(i)} = x_t^{(i)} - \eta_t\,\frac{\lVert x_t^{(i)}\rVert}{\lVert r_t^{(i)}\rVert}\, r_t^{(i)},\qquad r_t^{(i)}=\frac{m_t^{(i)}}{\sqrt{v_t^{(i)}}+\epsilon}+\lambda\,x_t^{(i)}

LAMB 的逐层信任比 ‖x‖/‖r‖ 重新缩放每一层的更新,使单一全局学习率适配所有层。

這種逐層正規化,正是讓 LAMB 能在大批次區制(每步數萬樣本)下訓練 BERT 而不損失最終準確度的關鍵,在硬體充足時大幅縮短實際預訓練時間。比起這個特定最佳化器,更普遍的教訓更為長久:在大規模下,正確的自適應單位往往是層,而非單一參數,因為層才是幾何各異的自然區塊。

變異縮減 SGD:漂亮的理論、尷尬的契合

如果梯度雜訊是障礙,為何不直接消除它?變異縮減 SGD——SVRG 與 SAGA 家族——做的正是這件事。SVRG 偶爾計算一次全資料集的梯度作為錨點,再用當前梯度與錨定梯度在某個控制點上的差來修正每個便宜的隨機梯度。在有限和的問題上,它可被證明達到線性收斂,遠勝過單純 SGD 的次線性率。

g_t = \nabla f_{i}(x_t) - \nabla f_{i}(\tilde{x}) + \nabla F(\tilde{x}),\qquad \nabla F(\tilde{x})=\frac{1}{n}\sum_{j=1}^{n}\nabla f_{j}(\tilde{x})

SVRG 的控制变量减去快照梯度并加上全量锚梯度,在不引入偏差的情况下消除噪声。

那為何它在深度學習中罕見?三個原因。在巨大資料集上算錨點梯度很昂貴;損失曲面是非凸的,保證因而減弱;而且關鍵在於,那個梯度雜訊本來正在幫助泛化,所以移除它可能在加速最佳化的同時傷害最終模型。變異縮減是一記犀利的提醒:更快的最佳化不等於更好的學習——這是值得帶進每一次最佳化器選擇的主題。

記憶體壓力與輕量化的轉向

在大規模下,最佳化器的狀態成了預算項目。Adam 為每個參數儲存兩個額外張量——一階與二階動量——因此光是最佳化器就能讓權重的記憶體占用翻倍甚至三倍。當模型剛好塞得下時,這份額外開銷決定了你能不能訓練。這份壓力催生了 Lion(演化出的符號動量),它只保留單一個動量緩衝,並以一個插值動量的符號(sign)來更新,相對於 Adam 把最佳化器記憶體減半。

因為每個 Lion 步都是一個符號——在每個座標上大小都是一——它的有效步幾何更接近正規化更新,而非曲率縮放更新,因此它通常需要比 Adam 更小的學習率與更大的權重衰減。Lion 是通往第 5 篇的橋梁,我們將在那裡問更大的問題:如果一條由自動搜尋找出的符號規則能與手工設計的 Adam 匹敵,那麼好的最佳化器究竟是如何被發現的?