分享丨大規模神經網絡優化:神經網絡損失空間“長”什麼樣?

轉自 PaperWeekly 作者 鄭奘巍

©作者 |鄭奘巍

單位 |新加坡國立大學

研究方向 |高效機器學習、神經網絡優化

We should use the special structure properties of f(for example f is a given by a neural network)to optimize it faster, instead of purely relying on optimization algorithms.

― Yuanzhi Li

Adam 雖然稱作”通用“優化算法,雖然在大語言模型等 Transformer 架構上表現的好,但它也並不是萬能的,在視覺模型的泛化性,以及一些凸問題上收斂的表現都沒有 SGD 要好。正如本篇開頭所引,我們要關注特殊網絡結構的特殊性質(比如 Transformer)。

如何刻畫網絡的優化性質呢?在優化相關的論文中,通常通過分析 Hessian 矩陣及其特徵值,或者將損失函數進行一維或二維的可視化來分析網絡的優化性質。我們希望這些指標能夠幫助我們更好的理解網絡損失的 landscape,優化器優化軌跡的性質等等。我們希望將這些指標刻畫的性質與優化器的設計關聯起來。我們也希望找到合適的指標來反應隨着網絡參數、數據量、批量大小等參數的變化,網絡的優化性質如何變化「R1」。

本文討論的方向如果有新的動向,也會持續更新。如果有任何錯誤/問題/討論,歡迎在評論區留言。

Hessian 陣

大量對神經網絡損失空間的研究都是基於 Hessian 陣的。Hessian 陣是損失函數的二階導數。由於 Hessian 陣是實對稱矩陣,可以應用特徵分解。它的特徵值和特徵向量可以幫助我們理解損失函數的 landscape。這裡我們不太嚴謹的介紹一些優化文章中常見的概念。

1.1 常見概念

Curvature 曲率:數學定義爲一點處密切圓弧的半徑的倒數,它刻畫了函數與直線的差距。在梯度爲 0 的點時,Hessian 陣行列式爲曲率。

Eigenvalue 特徵值:Hessian 陣的特徵值越大,對應的特徵向量(Eigenvector)方向就越陡(凸性越強)。最大特徵值對應的方向爲最大麴率的方向。如果考慮以特徵向量對應方向的單位向量作爲基張成的空間,該空間稱爲特徵空間 Eigenspace,對應的單位向量爲 Eigenbasis。

Conditioning number 條件數:Hessian 陣的最大特徵值與最小特徵值的比值。條件數越大,系統越不穩定。

Axis-aligned 座標軸對齊:Hessian 陣的特徵向量是否與座標軸平行。

Smoothness 光滑(一般指 Lipschitz 梯度條件) : 。 如果是二階可導凸函數,則有 。即 Smoothness 限制了 Hessian 陣特徵值的大小(和變化幅度的上限),不能大於 。

Strong Convexity 強凸性:雖然神經網絡一般都是非凸優化,但這裡我們引入凸性用以和光滑對應,以便於理解。強凸性的定義爲 。即 Strong Convexity 限制了 Hessian 陣特徵值的大小(和變化幅度的下限),不能小於 m。

Sharpness / Flatness 平坦程度:對 Hessian 特徵譜分佈情況的描述。Sharp 指 Hessian 陣特徵值含有大量大值,Flat 指 Hessian 陣特徵值含有大量小值(絕大部分接近 0)。另一種定義方法是直接將 Sharpness 定義爲最大的特徵值 。Hessian 陣退化(degenerate)/ 奇異(singular)/ 低秩 / 線性相關是 Flat 的一種表現。通常收斂點平坦有利於泛化性能(Hessian 陣是實對稱矩陣,特徵值和奇異值相等)。

常見地形(terrain, landscape):

駐點(stationary point),臨界點(critical point):

鞍點:H 非正定的駐點

plateaus:平坦地,擁有小曲率的平面

basins:盆地,局部最優點的集合

wells:井,鞍點的集合

局部的泰勒展開到二階:

1.2 費雪信息

Fisher information matrix 費雪信息:

費雪信息總是對稱半正定陣。由於 ,所以 Fisher information (二階矩)衡量了不同樣本對應導數的方差( )。

,所以 Fisher information 同時是使用 log-likelihood 衡量的函數的 Hessian 陣的期望。證明可參考此處:https://mark.reid.name/blog/fisher-information-and-log-likelihood.html

分類模型建模的分佈是 。在訓練初期,網絡的輸出 與真實分佈的差距很大。如果根據網絡輸出的標籤採樣,可以得到對應分佈的 Hessian 陣(半正定)但並不是我們關心的真實分佈的 Hessian 陣。

當模型分佈與真實分佈的差距很小時,Fisher information 可以近似爲 Hessian 陣的期望。

通過鏈式法則即可證明。這裡的第一項被稱作 Gauss-Newton 陣,它衡量了由於網絡特徵 的變動貢獻的 Hessian 陣值。第二項則是由於輸入值的變動貢獻的 Hessian 陣值。由於第二項通常很小,所以我們可以近似認爲黑塞陣由 Gauss-Newton 陣決定。通過變形,我們可以得到 Gauss-Newton 陣的另一種形式,衡量了樣本的帶權二階矩。可見 Gauss-Newton 陣與 Fisher information 有着密切的聯繫。

臨界點的數量與性質

爲了更直觀的理解,不會介紹太多理論推導。討論網絡優化性質的文章往往觀點鮮明,不少文章得出的結論甚至是相互矛盾。從今天的視角往前看,我們主要介紹目前被廣泛接受的觀點。但也要注意,部分結論是在 CNN,MLP,LSTM 上用監督學習進行,對應的是 over-paramterized 網絡(參數量遠大於數據量)。目前的 LLM 則一定程度是 under-parameterized 網絡。

1. 局部最優點的數量隨着網絡參數的增加而指數級增加,但不同局部最優點的訓練損失相差不大

[AHW96] 用單神經元網絡說明了局部最優點的數量隨着網絡參數的增加而指數級增加。LeCun 引領的一系列工作 [CHM+14,CLA15,SGA15,K16] 使用 spin-glass 物理模型闡釋了簡化的神經網絡中差局部最優解隨着模型參數增大指數級減小,以及局部最優解附近相當平坦。[DY16] 則證明了多層神經網絡中相似的結論。

[SEG+18] 中提到,在全連接層中交換任意兩個神經元可以得到等價的神經網絡。這個性質讓隨機初始化的網絡可能收斂到非常不同的局部最優值。[DVS18] 則用過實驗發現,不同的局部最小值之間可以通過平坦的路徑相連。[KWL+20] 同樣支持了這點:任意兩個局部極小值之間都可以通過一條不會太”扭曲“的路徑相連,這條路徑上的損失最多比局部最優處的損失高一點。

2. 鞍點的數量比上局部最優點的數量隨着網絡參數的增加而指數級增加

[DPG+14] 理論分析了鞍點數量和局部最優點數量的比例隨着網絡參數的增加而指數級增加,同時鞍點通常被 plateaus 所包圍。[LSJ+16] 則通過理論分析補充,訓練時不會陷入鞍點,但是會減慢速度。[CHM+14] 同意該看法,並認爲訓練神經網絡的過程可以看作避開鞍點,並通過選擇鞍點的一端從而打破“對稱性”的過程。在鞍點處優化器的選擇會影響到最終收斂點的性質。

3. Hessian 的低秩性質

[SBL17,SEG+18] 通過計算 CNN 網絡訓練時的特徵譜,發現 Hessian 陣特徵值譜分佈特點:

在訓練初期,Hessian 特徵值以 0 爲中心對稱分佈。隨着訓練進行不斷向 0 集中,H 退化嚴重。訓練末期極少量小於 0。

接近收斂時 Hessian 特徵譜分爲兩部分,大量接近 0,離羣點遠離大部分特徵值,離羣點的數量接近於訓練數據中的類別數量(或聚類數量)。即 Hessian 陣此時是低秩的。

數據不變的情況下,增大參數僅僅使 Hessian 陣接近 0 的部分更寬;參數不變的情況下,數據量的變化會改變 Hessian 陣特徵譜的離羣點。

參數量越大,批量大小越小,訓練結束時 Hessian 陣特徵值的負值越多(但總體比例都很小。[BKX19] 補充,在訓練早期 Hessian 陣就會變得基本沒有負值。

[DVS18,BKX19] 同樣發現了離羣點數量(近似 Hessian 陣的秩)與分類任務中的類別數量相關。最近的一篇工作 [SHS23] 則說明,CNN 中 Hessian 陣的秩隨着參數增加而增大。[SKV20] 分析了每層的 Hessian 特徵值譜,發現與整個網絡的 Hessian 特徵值譜相似。

[GRD18] 則進一步描繪了整個訓練過程中的低秩現象:Hessian 最大的幾個特徵值對應的特徵向量張成的子空間儘管會變化,但變化程度不大。訓練的梯度會很快的進入到這個子空間中,但在這個子空間中並沒有什麼規律。故整個神經網絡的訓練似乎是發生在一個低維的子空間上的。

4. Gauss-Newton 陣可以近似 Hessian 陣

[SEG+18,SKV20] 都指出了這一點。實驗如圖所示。Sophia 等二階優化器正是利用了這點進行對角黑塞陣的計算。[LGZ+19,P19] 則驗證了梯度的二階矩的大特徵值與 Hessian 陣的大特徵值近似。

▲ [LGZ+19] 不同層以及整個網絡 Hessian 陣與 Gauss-Newton 陣特徵值分佈對比

5. 平坦的收斂點具有更好的泛化性,大的批量容易收斂到 sharp 的收斂點

[KMN+17] 指出 Flat minima 有利於泛化,大批量訓練容易收斂到 sharp minima,小批量訓練容易收斂到 flat minima。[SEG+18] 同意這一點,但反對 [KMN+17] 中認爲 flat minima 和 sharp minima 被 well 分割開。通過實驗證明能夠從 sharp minima 平坦的“走”到 flat minima。[YGL+18] 分析了不同批量大小收斂點的 Hessian 特徵譜,發現了相同的結論。

一般認爲平坦的損失空間,既有利於優化算法的梯度下降,得到的收斂點性質也會更好。故很多文章從該角度解釋網絡設計的成功原因。譬如殘差鏈接,BN,更寬的網絡(over-paramterized)都有助於讓損失空間更平滑。

優化路徑性質

這裡先介紹兩篇近期的工作,之後有機會再補充。

[LBZ+21] 沿着前人的工作,探討了初始權重與訓練完成後的權重,兩者插值得到的一系列權重的性質。可以看到不同初始化的網絡權重往往收斂到同一個盆地。雖然與訓練路徑不同,插值路徑也是單調遞減,沒有凸起的現象。

[GGG+21] 討論了好的優化路徑(通過模型設計如增加 BN 或者選擇好的優化超參)應該讓優化器快速進入到平坦的區域,從而能夠使用較大的學習率。本文選擇觀察 即最大的特徵值以衡量網絡的 sharpness。文章發現網絡初始化方法,Warmup,BatchNorm 等方法都有效的減小了 ,從而可以使用更高的學習率。作者通過實驗發現,SGD 算法能夠穩定收斂的條件是 和學習率的乘積應該小於 2。(由於 Hessian 最小的值接近 0,所以控制 也可以看做控制 Hessian 陣的條件數)

▲ [GGG+21] 左:不同學習率下不同方法的表現;右:最大特徵值和學習率的變化

[PL23] 工作則是將優化器每步進行分解。根據局部泰勒展開,作者定義 gradient correlation 和 directional sharpness 兩項:

▲ [PL23] 函數的單次更新的泰勒展開

作者發現,不同的優化器的 direction sharpness 的表現很不相同。Sharpness 刻畫的是一個點局部的平坦程度,而 directional sharpness 則刻畫了優化方向上的平坦程度。在語言模型上,作者發現 SGD 優化器的 direction sharpness 顯著高於 Adam 優化器。作者認爲這是 Adam 優化器收斂速度快的原因。作者還通過梯度的值裁剪進一步減小 Sharpness 程度,從而加速收斂。

▲ [PL23] 不同方法的 directional sharpness 對比

我們已經知道收斂點是需要越平坦泛化性越好的。但從這兩篇文章中可以看到,我們希望優化路徑始終能夠在平坦的區域探索。以往的文章會認爲這些區域可能會導致優化陷入鞍點,但目前看來,優化空間中存在大量這樣平坦的空間,優化器應該在這些較爲平坦的空間中“大踏步”的進行探索。

Transformer 優化性質

前面的分析都是比較泛化的,爲了更好的設計針對 Transformer 的優化器,我們還需要分析 Transformer 的特殊性質。比如文章 [LLG+20] 就指出了 Transformer 訓練時的一些困難。同樣的,Adam 在 Transformer 上表現顯著好於 SGD,也暗示了 Transformer 的優化性質與 CNN 等網絡不同。

Transformer 的分析文章很多都是對模型設計的分析,如果有機會,我會單獨寫一篇文章進行介紹。這裡只介紹一篇最近的工作 [LYB+23],以補完本節的完整性。

該工作用大量實驗揭示了非常直接但有趣的一個現象:Transformer 架構的激活層非常的稀疏。在 T5-Base 中,只有 3% 的神經元被激活了(ReLU 輸出值非 0)。並且增大模型會使其更加稀疏,以及和數據似乎無關。無論使用視覺還是語言,甚至是隨機數據,以及不重複的數據,都有該現象出現。作者簡單的討論了這種現象可能和優化動力學相關。筆者認爲這個現象可能會給 Transformer 的優化器設計一些啓示。

▲ [LYB+23] 激活值的稀疏性現象

參考文獻

[AHW96] Exponentially many local minima for single neurons

[CHM+14] The Loss Surfaces of Multilayer Networks[CLA15] Open Problem: The landscape of the loss surfaces of multilayer networks[SGA15] Explorations on High Dimensional Landscapes[DY16] No bad local minima: Data independent training error guarantees for multilayer neural networks[K16] Deep Learning without Poor Local Minima[SEG+18] Empirical Analysis of the Hessian of Over-Parametrized Neural Networks[DVS18] Essentially No Barriers in Neural Network Energy Landscape[DPG+14] Identifying and attacking the saddle point problem in high-dimensional non-convex optimization[LSJ+16] Gradient Descent Converges to Minimizers[SBL17] Eigenvalues of the Hessian in Deep Learning: Singularity and Beyond[KWL20] Explaining Landscape Connectivity of Low-cost Solutions for Multilayer Nets[GRD18] Gradient Descent Happens in a Tiny Subspace[SKV20] A Deeper Look at the Hessian Eigenspectrum of Deep Neural Networks and its Applications to Regularization[LGZ+19] Hessian based analysis of SGD for Deep Nets: Dynamics and Generalization[P19] Measurements of Three-Level Hierarchical Structure in the Outliers in the Spectrum of Deepnet Hessians[SHS23] The Hessian perspective into the Nature of Convolutional Neural Networks[KMN+17] On large-batch training for deep learning generalization gap and sharp minima[YGL+18] Hessian-based Analysis of Large Batch Training and Robustness to Adversaries[GGG+21] A Loss Curvature Perspective on Training Instability in Deep Learning[PL23] Toward Understanding Why Adam Converges Faster Than SGD for Transformers[LLG+20] Understanding the Difficulty of Training Transformers[LYB+23] The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers[LBZ+21] Analyzing Monotonic Linear Interpolation in Neural Network Loss Landscapes