Nvidia發展全新nGPT 有望使訓練速度大增20倍

Nvidia (NVDA-US) 研究人員提出了一種新穎的神經網路架構，稱為歸一化 Transformer (nGPT)。研究人員表示，在 nGPT 中，單位範數對所有向量進行歸一化，包括嵌入、多層感知器 (MLP)、注意力矩陣和隱藏狀態。

在這一架構中，輸入標記在超球體 (hypersphere) 的表面移動，每一層都為輸出預測貢獻一個位移。實驗表明，nGPT 加速了學習，將訓練步驟減少了 4 到 20 倍，具體取決於序列長度。

例如，在 1k 上下文中，訓練速度提高了 4 倍；在 4k 上下文中，提高了 10 倍；而在 8k 上下文中，更是提高了 20 倍。

報導指出，如此能夠大幅提升大型語言模型 (LLM) 的訓練速度，最高可達 20 倍，同時維持了模型的精確度。

另一項研究透過超球面上的原型嵌入來統一分類和回歸，實現大範圍的分離。然而，在確保不同數據分佈和嵌入配置之間的性能一致方面仍然存在挑戰。

從 GPT 到 nGPT 的演進涉及修改基線 Transformer 以創建歸一化版本，重點是 Transformer 解碼器和自注意力 (self-attention)。這些變化可以擴展到編碼器 - 解碼器和交叉注意力設定。關鍵的修改涉及限制所有嵌入向量具有單位範數並在訓練期間標準化權重。

它確保透過計算歸一化向量之間的點積來進行預測，從而獲得更準確的相似性估計值。令牌嵌入和輸出邏輯透過可學習矩陣進行處理，並透過引入可訓練的縮放參數來調整預測置信度。

研究人員稱，nGPT 利用特徵學習率，這是一項關鍵創新，其中學習率應用於特定的模型維度，優化注意力和 MLP 區塊更新，以獲得更好的準確性和穩定性。