LLM訓練速度還可以再飆升20倍！輝達團隊祭出全新架構歸一化Transformer（nGPT），上下文越長，訓練速度越快，還能維持原有精度。

最近，輝達團隊拋出的一枚重磅炸彈，提出了全新神經網路架構——歸一化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倍，具體取決於序列長度：

可以看出，上下文越長，訓練越快。

Reddit網友表示，「我很好奇它還能擴展到多大程度。如果它能在更長的上下文中大幅擴展，這意味著像o1這樣的模型將會獲得顯著的訓練速度優勢」。

還有人表示，「下一代模型將會更高效、更智能」。

毋庸置疑，Transformer架構是現代大模型的基礎。

不過，當前基於Transformer搭建的大模型都是計算密集型的，需要耗費大量的資源和時間。

為了改進其訓練穩定性、推理成本、上下文長度、魯棒性等方面，AI科學家已進行了大量的修改嘗試。

其中，最突出的發現是，歸一化技術對於Transformer性能改善起著重要作用，比如LayerNorm和RMSNorm。

另一種模型歸一化方法是，通過權重衰減（weight decay）控制權重范數。

不過，最新研究又對權重衰減的作用進行評估，並且轉向更多地關注旋轉，而非僅僅關注向量范數。

越來越多的證據表明，在超球面上進行表示學習與更穩定的訓練、更大的嵌入空間可分離性以及在下游任務上的更好性能相關。

而且，還有新研究表明，Transformer隱式地執行梯度下降作為元最佳化器。

由此，輝達團隊提出了，在歸一化Transformer新視角下，統一該領域的各種發現和觀察。

建議將形成網路矩陣嵌入維度的所有向量歸一化，使其位於單位范數超球面上。這種方法將矩陣-向量乘法轉化為餘弦相似度的計算，其範圍限定在 [-1,1] 之間。而且歸一化消除了對權重衰減的需求。

- 歸一化Transformer作為超球面上的可變度量最佳化器

歸一化Transformer本身在超球面上執行多步最佳化（每層兩步），其中注意力和MLP更新的每一步，都由特徵學習率控制——這些是可學習的可變度量矩陣的對角線元素。

研究證明，歸一化Transformer將達到相同精度所需的訓練步驟減少了4-20倍。

標準的decoder-only Transformer的訓練目標是根據輸入序列的前序tokens來預測後面的token，在token預測時，模型會引入兩個可學習的嵌入矩陣Einput和Eoutput，分別用來從輸入詞轉為詞嵌入，以及從詞嵌入轉為預測輸出。

在模型訓練期間，通常使用對應嵌入向量的點積來計算token相似度，但嵌入向量的范數（norms）不受限制的，可能會導致相似性計算存在偏差。

為了提高相似性估計的精準性，研究人員在新架構中提出，在訓練演算法的每一步之後，對Einput和Eoutput中的嵌入向量進行歸一化。

智能體在預測文字中的下一個詞時，會使用因果掩碼（casual masking）來確保模型在預測token時不會「偷看」到之後的詞，造成資訊洩露，從而讓模型能夠同時預測多個詞並計算預測誤差，提高訓練效率，同時保持了按順序預測詞的能力。

在輸入詞序列後，模型會在預測序列中的每個位置都生成一個輸出向量，然後計算出一個logits向量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），可以輔助確定哪些資訊是重要的，以及資訊之間是如何相互關聯的。

彙整鉅亨號中所有使用「nGPT」關鍵字，所發布的內容 (依照時間排序，最新在前，舊的在後)`