Layer Normalization/zh
| Article | |
|---|---|
| Topic area | Deep Learning |
| Prerequisites | Backpropagation, Batch Normalization, Stochastic Gradient Descent |
概述
層歸一化是一種通過在單個訓練樣本的特徵維度上(而非跨樣本批次)計算均值和方差來標準化神經網絡層激活值的技術。該方法由 Jimmy Lei Ba、Jamie Ryan Kiros 和 Geoffrey Hinton 於 2016 年提出,旨在解決 Batch Normalization 在批次統計量不可靠或不可用場景下的局限性,例如循環網絡與在線學習。[1] 如今它已成為幾乎所有現代 Transformer 架構的標準組件,支撐着大型語言模型、視覺transformer以及許多序列模型。
其動機是通過降低每一層輸出對輸入幅度變化的敏感度來穩定訓練並加速收斂。通過將激活值重新歸一化為零均值和單位方差,層歸一化將信號在深層網絡中傳播的尺度保持在受控範圍內,進而在 Backpropagation 過程中維持梯度的良好行為。
直覺
深層網絡難以訓練,部分原因在於隨着前面各層權重的更新,每一層激活值的分佈會發生偏移。每一層必須不斷適應新的輸入統計量,這會減慢學習速度,並可能以病態的方式放大或衰減梯度。歸一化技術通過顯式固定激活值的某些統計量來應對這一問題。
Batch Normalization 在小批次中的樣本之間計算統計量,而層歸一化則在單個樣本的特徵之間計算統計量。這一差異雖屬概念性,卻影響深遠:統計量不再依賴於批次中的其他樣本,因此訓練和推理時操作完全相同,行為與批次大小無關,甚至可以在每次只處理一個樣本時應用。這種與批次組成無關的特性使其非常適合循環計算和自回歸生成,因為這些場景中每個 token 都是按序處理的。
一種有用的可視化方式是:想像網絡中某一位置上的激活值向量,對該向量進行重新中心化和重新縮放,使其各分量具有零均值和單位方差。然後,一個可學習的增益和偏置恢復了模型表示任意所需尺度和偏移的自由度。
公式化
設 $ x \in \mathbb{R}^d $ 是某一特定層在單個樣本上的激活值向量。層歸一化首先在 $ d $ 個分量上計算均值和方差:
$ {\displaystyle \mu = \frac{1}{d} \sum_{i=1}^{d} x_i, \qquad \sigma^2 = \frac{1}{d} \sum_{i=1}^{d} (x_i - \mu)^2.} $
接着對每個分量進行標準化,並應用一個由可學習向量 $ \gamma, \beta \in \mathbb{R}^d $ 參數化的逐元素仿射變換:
$ {\displaystyle \hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}, \qquad y_i = \gamma_i \hat{x}_i + \beta_i.} $
常數 $ \epsilon $(通常為 $ 10^{-5} $ 或 $ 10^{-6} $)可在激活值方差極小時防止除零。可學習的 $ \gamma $ 和 $ \beta $ 使網絡能夠恢復標準化後激活值的任意仿射變換(包括恆等變換),因此該操作不會限制該層的表示容量。
關鍵的是,相同的過程在訓練和推理時都運行。無需跟蹤任何滑動平均,批次大小為一時也不會出現行為變化。
在現代架構中的位置
在 Transformer 塊中,層歸一化應用於每個子層(多頭注意力和位置式前饋網絡)的兩側。常見兩種約定。
原始的表述方式,常稱為 Post-LN,將歸一化應用於殘差相加之後:$ \mathrm{LN}(x + \mathrm{Sublayer}(x)) $。這正是原始Transformer論文中所採用的方式。[2] 與之相對,Pre-LN 則在殘差相加之前對每個子層的輸入應用歸一化:$ x + \mathrm{Sublayer}(\mathrm{LN}(x)) $。研究表明,Pre-LN 在初始化時可產生條件更好的梯度,並且無需在深層模型中使用學習率預熱,因此它已成為當今大多數大型語言模型訓練時的默認選擇。[3]
除 Transformers 之外,層歸一化也是循環單元內部(在每個時間步穩定循環過程)、圖神經網絡內部(由於圖大小不一,批次統計量可能產生誤導)以及強化學習代理(數據以非獨立同分佈流的形式到來)的標準選擇。
訓練與推理
由於歸一化統計量僅依賴於當前樣本的激活值,層歸一化不會引入訓練與測試之間的差異。在典型實現中,前向傳播會計算 $ \mu $ 和 $ \sigma^2 $,對 $ x $ 進行標準化,應用仿射變換,並存儲中間張量以用於反向傳播。反向傳播通過一個閉式表達式將梯度傳過歸一化,該表達式考慮了每個輸出分量通過共享的均值和方差對每個輸入分量的依賴關係。
開銷適中。對於具有 $ d $ 個特徵的激活張量,前向和反向均需 $ O(d) $ 的工作量以及一次跨特徵維度的同步。在現代加速器上,這通常受限於內存帶寬而非計算能力,因此將標準化與仿射變換融合為單個內核的實現被廣泛使用。
初始化約定很簡單。增益 $ \gamma $ 初始化為一,偏置 $ \beta $ 初始化為零,因此該層初始時(在標準化之後)為恆等映射。此後,參數會在訓練過程中漂移到模型認為有用的尺度和偏移上。
變體
若干改進對原始表述的某些部分進行了刪減或重構。
RMSNorm(均方根層歸一化)放棄了均值中心化步驟和偏置 $ \beta $,僅按激活值的均方根進行歸一化:$ \hat{x}_i = x_i / \sqrt{\tfrac{1}{d}\sum_j x_j^2 + \epsilon} $。[4] 它的計算開銷更小,在許多語言建模任務上實證表現達到或超過層歸一化;LLaMA、PaLM 以及若干其他大型模型都採用了它。
ScaleNorm 用單個標量替代了逐元素增益,將每個向量歸一化到一個可學習的範數上。FixNorm 及類似變體進一步約束歸一化,使其位於超球面上。DeepNorm 重新縮放 Post-LN Transformer中的殘差分支,從而支持訓練數千層的網絡。[5]
Group Normalization 介於批歸一化和層歸一化之間:它將特徵維度劃分為若干組,並在每個組內進行歸一化。當只有一個組時,它退化為層歸一化;當組數等於通道數時,它退化為實例歸一化。當小批次大小使有效批歸一化變得不可行時,它在計算機視覺中被廣泛使用。
與其他歸一化方法的比較
這四種經典歸一化方案的差異僅在於計算均值和方差所沿的軸。
Batch Normalization 在每個特徵通道上沿批次維度計算統計量。它在批次較大的卷積視覺模型中表現極佳,但當批次較小、訓練與推理之間統計量發生漂移或序列長度可變時,性能會下降。
層歸一化對每個樣本沿特徵維度計算統計量。它不受批次大小影響,並且在訓練和推理時行為完全相同,這使其成為序列模型和Transformer的主導選擇。
實例歸一化按樣本和通道分別計算統計量,僅在空間維度上進行歸一化。它常見於風格遷移和生成式圖像模型。前面提到的組歸一化通過參數化歸一化前通道的分組方式,將這一族進行了推廣。
一條實用的經驗法則:對於具有大批次和穩定輸入分佈的視覺模型,優先使用Batch Normalization;對於序列模型、Transformer,以及任何批次組成不可靠的場景,優先使用層歸一化(或 RMSNorm)。
為什麼有效
原始論文將層歸一化的有效性歸因於與Batch Normalization相同的假設,即減少了內部協變量偏移。後續研究對該解釋提出了質疑。一項有影響力的分析表明,批歸一化通過改善損失及其梯度的 Lipschitz 性質來平滑優化景觀,而該效應並不需要協變量偏移的減少。[6] 類似的論證也被用於層歸一化,並附有額外證據表明該操作充當了梯度範數的隱式正則化器,並將激活值的幅度和方向解耦。
無論具體機制為何,實證效應都很穩健:層歸一化使訓練更加穩定,對學習率選擇不那麼敏感,對深層架構也更具容忍度。
局限性
層歸一化並非沒有代價。沿特徵維度的歸約引入了一種比純逐點操作更難並行化的串行依賴,在帶寬受限的硬件上,它可能在 Transformer 塊的開銷中佔據相當可觀的一部分。RMSNorm 及類似的簡化方案部分地緩解了這一問題。
該技術還假設被歸一化的特徵共享一種應被標準化的共同尺度。在某些架構中,某些特徵攜帶着歸一化會抹去的有意義的幅度信息;可學習的 $ \gamma $ 和 $ \beta $ 原則上可以恢復它,但優化器可能難以恢復那些從未以標準化形式出現過的尺度。
最後,對方差極低的 token 進行層歸一化即使加入了 $ \epsilon $ 項也可能產生數值不穩定的輸出,尤其是在低精度訓練中。混合精度實現通常將歸一化本身保留在更高精度下以避免這種失敗模式。
另請參閱
- Batch Normalization
- Group Normalization
- RMSNorm
- Transformer
- Backpropagation
- Stochastic Gradient Descent
參考文獻
- ↑ Template:Cite arxiv
- ↑ Template:Cite arxiv
- ↑ Xiong, Yang, et al. On Layer Normalization in the Transformer Architecture, 2020.
- ↑ Template:Cite arxiv
- ↑ Wang et al., DeepNet: Scaling Transformers to 1,000 Layers, 2022.
- ↑ Santurkar et al., How Does Batch Normalization Help Optimization?, NeurIPS 2018.