今天小編分享的科技經驗： 英偉達nGPT重塑Transformer，AI訓練速度暴增20倍！文本越長，加速越快，歡迎閲讀。
【新智元導讀】LLM 訓練速度還可以再飙升 20 倍！英偉達團隊祭出全新架構歸一化 Transformer（nGPT），上下文越長，訓練速度越快，還能維持原有精度。
AI 的未來，或許就此改寫 ......
最近，英偉達團隊抛出的一枚重磅炸彈，提出了全新神經網絡架構——歸一化 Transformer（nGPT），基于超球面（hypersphere）進行表示學習。
相較于 Transformer 架構本身，nGPT 直接将 LLM 訓練速度提升至高 20 倍，而且還保持了原有精度。
也就意味着，原本需要一個月完成的訓練，在未來可能只需 1-2 天的時間就能搞定。
無疑為通向 AGI 終極目标，注入了一針強心劑！
論文地址：https://arxiv.org/pdf/2410.01131
在 nGPT 中，所有的向量（嵌入、MLP、注意力矩陣、隐藏狀态），都被歸一化為部門範數（unit norm）。
輸入後的 token 在超球面表面上移動，每一層都通過「位移」來貢獻最終的輸出預測，其中位移量是由 MLP 和注意力模塊進行定義的，其向量組件都位于同一個超球面上。
實驗表明，nGPT 達到相同精度所需的訓練步驟減少了 4-20 倍，具體取決于序列長度：
- 1k 上下文，訓練速度提高 4 倍
- 4k 上下文，訓練速度提高 10 倍
- 8k 上下文，訓練速度提高 20 倍
可以看出，上下文越長，訓練越快。
Reddit 網友表示，「我很好奇它還能擴展到多大程度。如果它能在更長的上下文中大幅擴展，這意味着像 o1 這樣的模型将會獲得顯著的訓練速度優勢」。
還有人表示，「下一代模型将會更高效、更智能」。
nGPT 全新架構，超球面上歸一化
毋庸置疑，Transformer 架構是現代大模型的基礎。
不過，當前基于 Transformer 搭建的大模型都是計算密集型的，需要耗費大量的資源和時間。
為了改進其訓練穩定性、推理成本、上下文長度、魯棒性等方面，AI 科學家已進行了大量的修改嘗試。
其中，最突出的發現是，歸一化技術對于 Transformer 性能改善起着重要作用，比如 LayerNorm 和 RMSNorm。
另一種模型歸一化方法是，通過權重衰減（weight decay）控制權重範數。
不過，最新研究又對權重衰減的作用進行評估，并且轉向更多地關注旋轉，而非僅僅關注向量範數。
越來越多的證據表明，在超球面上進行表示學習與更穩定的訓練、更大的嵌入空間可分離性以及在下遊任務上的更好性能相關。
而且，還有新研究表明，Transformer 隐式地執行梯度下降作為元優化器。
由此，英偉達團隊提出了，在歸一化 Transformer 新視角下，統一該領網域的各種發現和觀察。
這項研究的主要貢獻在于：
- 在超球面上優化網絡參數
建議将形成網絡矩陣嵌入維度的所有向量歸一化，使其位于部門範數超球面上。這種方法将矩陣 - 向量乘法轉化為餘弦相似度的計算，其範圍限定在 [ -1,1 ] 之間。而且歸一化消除了對權重衰減的需求。
- 歸一化 Transformer 作為超球面上的可變度量優化器
歸一化 Transformer 本身在超球面上執行多步優化（每層兩步），其中注意力和 MLP 更新的每一步，都由特征學習率控制——這些是可學習的可變度量矩陣的對角線元素。
對于輸入序列中的每個 token
，歸一化 Transformer 的優化路徑從超球面上對應于其輸入嵌入向量的點開始，移動到超球面上最能預測下一個
的嵌入向量的點。
- 更快的收斂
研究證明，歸一化 Transformer 将達到相同精度所需的訓練步驟減少了 4-20 倍。
Transformer 演變：從 GPT 到 nGPT
嵌入層歸一化
标準的 decoder-only Transformer 的訓練目标是根據輸入序列的前序 tokens 來預測後面的 token，在 token 預測時，模型會引入兩個可學習的嵌入矩陣 Einput 和 Eoutput，分别用來從輸入詞轉為詞嵌入，以及從詞嵌入轉為預測輸出。
在模型訓練期間，通常使用對應嵌入向量的點積來計算 token 相似度，但嵌入向量的範數（norms）不受限制的，可能會導致相似性計算存在偏差。
為了提高相似性估計的準确性，研究人員在新架構中提出，在訓練算法的每一步之後，對 Einput 和 Eoutput 中的嵌入向量進行歸一化。
智能體在預測文本中的下一個詞時，會使用因果掩碼（casual masking）來确保模型在預測 token 時不會「偷看」到之後的詞，造成信息泄露，從而讓模型能夠同時預測多個詞并計算預測誤差，提高訓練效率，同時保持了按順序預測詞的能力。
在輸入詞序列後，模型會在預測序列中的每個位置都生成一個輸出向量，然後計算出一個 logits 向量 zi 來表示詞匯表中每個詞出現的可能性，可以輔助模型理解不同詞在當前上下文中的重要性：
之後用 softmax 函數把 zi 轉為概率值，并選取概率最高的詞作為下一個詞的預測。
由于 nGPT 的嵌入矩陣已經歸一化了，所以 zi 的值範圍為 [ −1,1 ] ，也會限制 softmax 後得到的概率分布的置信度，也可以叫做温度。
為了在訓練過程中調整置信度，nGPT 又引入了一個可學習的縮放參數 sz，通過逐元素地縮放 logits，模型可以更靈活地預測的置信度，更好地學習到在不同情況下如何做出更準确的預測：
層 / 塊歸一
标準 Transformer 架構需要對隐藏層狀态 h 進行 L 層變換，包括一個自注意力（ATTN）和多層感知機（MLP）。
其中 RMSNorm 也可以替換成其他歸一化（normalization）函數。
隐藏層的參數更新，其實就是在一個超平面上（維度為隐藏層的向量長度）尋找兩個點（原參數和新參數）的最短距離。
1985 年，Shoemake 提出了球面線性插值（SLERP，Spherical Linear Interpolation），可以沿着球面上兩點之間的最短路徑找到中間點，研究人員發現該方法還可以通過更簡單的線性插值（LERP，linear interpolation）來得到近似解，從而降低計算量：
按最短路徑尋找來説，參數更新過程可以描述為：
其中 a 和 b 是球面上的兩個點，對應到 nGPT 上，a 也就是隐藏層狀态，b 是經過注意力機制或 MLP 塊後的狀态，梯度就是 g=a-b，B 為可變矩陣。
在拟牛頓方法中，B 可以近似于逆黑塞矩陣，當 B 是一個對角線元素非負的對角矩陣時，αB 就變成了一個向量，其元素對應于 B 的對角線元素乘以學習率 α，也可以稱之為特征學習率（eigen learning rates）。
eigen 源自德語詞，意為「自己的」（own），可以指代 Transformer 的内部結構。
所以 nGPT 中的參數更新方程可以寫為：
其中 αA 和 αM 是可學習的參數，分别用于注意力和多層感知機（MLP）模塊的歸一化輸出 hA 和 hM
與基礎 Transformer 相比，在 nGPT 的最終層之後不需要再進行額外的歸一化了。
自注意力塊
注意力機制可以説是 Transformer 中最重要的模塊，序列中的每個 token 都能夠關注到其他所有 token，從而讓模型具有捕捉長距離依賴關系的能力。
模型會把處理後的信息分解成三個部分：查詢（q，query）、鍵（k，key）和值（v，value），可以輔助确定哪些信息是重要的，以及信息之間是如何相互關聯的。
為了确保模型能夠理解每個詞在序列中的位置，模型中通常還會在 query 和 key 向量之間加入旋轉位置嵌入（Rotary Position Embeddings，RoPE）。
然後通過計算 query 向量和 key 向量的點積、縮放、應用 softmax 得到注意力權重，對 value 向量進行加權求和，得到注意力得分。
在實踐中，Transformer 一般都會用到多個注意力頭，其中每個頭的注意力機制都是獨立計算，最後再通過一個可學習的投影矩陣 Wo 合并所有頭輸出。
在計算注意力得分的過程中，權重矩陣沒有受到太多限制，可能會導致最終得分過大或過小。
在 nGPT 中，研究人員對 q 向量和 k 向量進行歸一化，還引入了一些可調整的參數（sqk），以确保權重矩陣在處理位置信息時不會失真，更準确地捕捉到句子中詞與詞之間的關系，從而做出更好的預測和決策。
MLP 塊
在标準 Transformer 中，隐藏層收入通過 RMSNorm 進行歸一化，然後經過兩個線性投影生成中間向量（暫不考慮偏置項）：
然後使用 SwiGLU 門控激活函數，以及一個線性變換得到最終門控激活。
在 nGPT 中，研究人員提出對線性投影的權重矩陣進行歸一化，并引入可學習的縮放因子，能夠更充分地利用處理信息時的非線性特性，在處理復雜信息時更加靈活。
多層感知機模塊的輸出不會因為縮放調整而發生變化。
Adam 高效學習率
Adam 優化算法通過動量和梯度幅度的估計來調整每次的學習步長，同時考慮了當前及過去的梯度信息。
在 nGPT 中，研究人員同樣引入了一個可訓練的縮放參數向量，對特定的參數進行更精細的控制，确保每個參數都能以最适合自己的速度進行學習，從而進一步提高學習效率。
在不影響全局學習率的情況下，對特定的參數進行調整，提供了更大的靈活性和控制力。
變化總結
和基礎 Transformer 相比，nGPT 主要做了七個改變：
1、移除所有歸一化層，比如 RMSNorm 或 LayerNorm；
2、在每個訓練步驟之後，沿着嵌入維度對所有矩陣，包括輸入輸出嵌入矩陣，以及各種權重矩陣進行歸一化處理；
3、修改了隐藏層參數更新方程；
4、調整注意力機制中的 softmax 縮放因子，對 q 和 k 進行重新縮放和歸一化；
5、對 MLP 塊的中間狀态進行重新縮放；
6、對 logits 進行重新縮放；
7、移除權重衰減和學習率預熱步驟。
上下文越長，訓練速度越快
接下來，研究人員在 OpenWebText 數據集上訓練了基礎基礎 Transformer（GPT）和歸一化 Transformer（nGPT），并在一系列标準下遊任務上對其進行評估。
實驗中，使用了 0.5B 和 1B（包括嵌入）兩種參數規模的模型。兩種參數規模的模型 0.5B 和 1B（包含嵌入）。
訓練加速
圖 1 顯示了，在訓練過程中，10 億參數且樣本長度為 4k token 的 GPT 和 nGPT 模型的驗證損失。
經過 2 萬次迭代後，nGPT 達到了與 GPT 在 20 萬次迭代（約 4000 億個 token）後，才能達到的相同驗證損失。
這表明，在迭代次數和使用 token 數量方面，nGPT 實現了 10 倍的加速。
再來看圖 2，展示了 nGPT 和 GPT 在三個方面的性能差距是如何變化的：總 token 數量、上下文長度、參數規模。
在 1k、4k 和 8k token 上下文中，訓練 0.5B 和 1B 的 nGPT 模型分别約快 4 倍、10 倍和 20 倍。
圖 3 在下遊任務中顯示了類似的性能，證實加速不僅反映在困惑度上，也反映在任務表現上。
研究人員觀察到，對于較長的訓練運行，nGPT 顯示出一些飽和現象，這暗示在當前可訓練參數數量下，模型容量可能已接近極限。
神經網絡參數檢查
圖 4 顯示，雖然 nGPT 保持固定的嵌入範數（這是設計使然），但 GPT 表現出明顯的變化。
從嵌入的協方差矩陣計算得出的特征值分布（已經由其中位數歸一化）顯示，GPT 的輸入嵌入具有更高的條件數，尤其是在 1B 模型中。
嵌入之間的成對點積分布表明，即使在 nGPT 中，嵌入也并非均勻分布在超球面上（在那裏點積會接近 0），而是形成簇——這可能反映了語言數據中的自然模式。
由于 GPT 的嵌入形成了一個超橢球體（hyper-ellipsoid ) ，如向量範數的分布所示，其點積往往具有更高的值。
GPT 輸入嵌入的病态性質（ill-conditioned nature）可能導致涉及這些嵌入的計算問題。
下圖 5 展示了，注意力和 MLP 矩陣在不同層深度上的中位數條件數（跨多個頭）—— 0.5B 模型有 24 層，1B 模型有 36 層。
與 nGPT 相比，GPT 模型的注意力矩陣呈現顯著更高的條件數。
對這些矩陣的進一步檢查，GPT 的注意力矩陣表現出退化為低秩矩陣的趨勢，可能減少了這些塊的學習容量。
下圖 6 展示了，（左圖）注意力模塊和 MLP 模塊的特征學習率，（中圖）應用于 MLP 中間狀态的縮放因子，（右圖）應用于 QK 點積之前的縮放因子。
