英偉達nGPT重塑Transformer，AI訓練速度暴增20倍！文本越長，加速越快

　　【新智元導讀】LLM訓練速度還可以再飆升20倍!英偉達團隊祭出全新架構(gòu)歸一化Transformer(nGPT)，上下文越長，訓練速度越快，還能維持原有精度。

　　AI的未來，或許就此改寫......

　　最近，英偉達團隊拋出的一枚重磅炸彈，提出了全新神經(jīng)網(wǎng)絡(luò)架構(gòu)——歸一化Transformer(nGPT)，基于超球面(hypersphere)進行表示學習。

　　相較于Transformer架構(gòu)本身，nGPT直接將LLM訓練速度提升至高20倍，而且還保持了原有精度。

　　也就意味著，原本需要一個月完成的訓練，在未來可能只需1-2天的時間就能搞定。

　　無疑為通向AGI終極目標，注入了一針強心劑!

　　論文地址：https://arxiv.org/pdf/2410.01131

　　在nGPT中，所有的向量(嵌入、MLP、注意力矩陣、隱藏狀態(tài))，都被歸一化為單位范數(shù)(unit norm)。

　　輸入后的token在超球面表面上移動，每一層都通過「位移」來貢獻最終的輸出預(yù)測，其中位移量是由MLP和注意力模塊進行定義的，其向量組件都位于同一個超球面上。

　　實驗表明，nGPT達到相同精度所需的訓練步驟減少了4-20倍，具體取決于序列長度：

　　- 1k上下文，訓練速度提高4倍

　　- 4k上下文，訓練速度提高10倍

　　- 8k上下文，訓練速度提高20倍

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

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

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

　　nGPT全新架構(gòu)，超球面上歸一化

　　毋庸置疑，Transformer架構(gòu)是現(xiàn)代大模型的基礎(chǔ)。

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

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

　　其中，最突出的發(fā)現(xiàn)是，歸一化技術(shù)對于Transformer性能改善起著重要作用，比如LayerNorm和RMSNorm。

　　另一種模型歸一化方法是，通過權(quán)重衰減(weight decay)控制權(quán)重范數(shù)。

　　不過，最新研究又對權(quán)重衰減的作用進行評估，并且轉(zhuǎn)向更多地關(guān)注旋轉(zhuǎn)，而非僅僅關(guān)注向量范數(shù)。

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

　　而且，還有新研究表明，Transformer隱式地執(zhí)行梯度下降作為元優(yōu)化器。

　　由此，英偉達團隊提出了，在歸一化Transformer新視角下，統(tǒng)一該領(lǐng)域的各種發(fā)現(xiàn)和觀察。

　　這項研究的主要貢獻在于：

　　- 在超球面上優(yōu)化網(wǎng)絡(luò)參數(shù)

　　建議將形成網(wǎng)絡(luò)矩陣嵌入維度的所有向量歸一化，使其位于單位范數(shù)超球面上。這種方法將矩陣-向量乘法轉(zhuǎn)化為余弦相似度的計算，其范圍限定在 [-1,1] 之間。而且歸一化消除了對權(quán)重衰減的需求。

　　- 歸一化Transformer作為超球面上的可變度量優(yōu)化器

　　歸一化Transformer本身在超球面上執(zhí)行多步優(yōu)化(每層兩步)，其中注意力和MLP更新的每一步，都由特征學習率控制——這些是可學習的可變度量矩陣的對角線元素。

　　對于輸入序列中的每個token

　　，歸一化Transformer的優(yōu)化路徑從超球面上對應(yīng)于其輸入嵌入向量的點開始，移動到超球面上最能預(yù)測下一個

　　的嵌入向量的點。

　　- 更快的收斂

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

　　Transformer演變：從GPT到nGPT

　　嵌入層歸一化

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

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

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

　　智能體在預(yù)測文本中的下一個詞時，會使用因果掩碼(casual masking)來確保模型在預(yù)測token時不會「偷看」到之后的詞，造成信息泄露，從而讓模型能夠同時預(yù)測多個詞并計算預(yù)測誤差，提高訓練效率，同時保持了按順序預(yù)測詞的能力。

　　在輸入詞序列后，模型會在預(yù)測序列中的每個位置都生成一個輸出向量，然后計算出一個logits向量zi來表示詞匯表中每個詞出現(xiàn)的可能性，可以輔助模型理解不同詞在當前上下文中的重要性：

　　之后用softmax函數(shù)把zi轉(zhuǎn)為概率值，并選取概率最高的詞作為下一個詞的預(yù)測。

　　由于nGPT的嵌入矩陣已經(jīng)歸一化了，所以zi的值范圍為[−1,1]，也會限制softmax后得到的概率分布的置信度，也可以叫做溫度。

　　為了在訓練過程中調(diào)整置信度，nGPT又引入了一個可學習的縮放參數(shù)sz，通過逐元素地縮放logits，模型可以更靈活地預(yù)測的置信度，更好地學習到在不同情況下如何做出更準確的預(yù)測：

　　層/塊歸一

　　標準Transformer架構(gòu)需要對隱藏層狀態(tài)h進行L層變換，包括一個自注意力(ATTN)和多層感知機(MLP)。

　　其中RMSNorm也可以替換成其他歸一化(normalization)函數(shù)。

　　隱藏層的參數(shù)更新，其實就是在一個超平面上(維度為隱藏層的向量長度)尋找兩個點(原參數(shù)和新參數(shù))的最短距離。

　　1985年，Shoemake提出了球面線性插值(SLERP，Spherical Linear Interpolation)，可以沿著球面上兩點之間的最短路徑找到中間點，研究人員發(fā)現(xiàn)該方法還可以通過更簡單的線性插值(LERP，linear interpolation)來得到近似解，從而降低計算量：

　　按最短路徑尋找來說，參數(shù)更新過程可以描述為：

　　其中a和b是球面上的兩個點，對應(yīng)到nGPT上，a也就是隱藏層狀態(tài)，b是經(jīng)過注意力機制或MLP塊后的狀態(tài)，梯度就是g=a-b，B為可變矩陣。

　　在擬牛頓方法中，B可以近似于逆黑塞矩陣，當 B是一個對角線元素非負的對角矩陣時，αB就變成了一個向量，其元素對應(yīng)于B的對角線元素乘以學習率α，也可以稱之為特征學習率(eigen learning rates)。

　　eigen源自德語詞，意為「自己的」(own)，可以指代Transformer 的內(nèi)部結(jié)構(gòu)。

　　所以nGPT中的參數(shù)更新方程可以寫為：

　　其中αA 和 αM是可學習的參數(shù)，分別用于注意力和多層感知機(MLP)模塊的歸一化輸出 hA和 hM

　　與基礎(chǔ) Transformer 相比，在nGPT的最終層之后不需要再進行額外的歸一化了。

　　自注意力塊

　　注意力機制可以說是Transformer中最重要的模塊，序列中的每個token都能夠關(guān)注到其他所有token，從而讓模型具有捕捉長距離依賴關(guān)系的能力。

　　模型會把處理后的信息分解成三個部分：查詢(q，query)、鍵(k，key)和值(v，value)，可以輔助確定哪些信息是重要的，以及信息之間是如何相互關(guān)聯(lián)的。

　　為了確保模型能夠理解每個詞在序列中的位置，模型中通常還會在query和key向量之間加入旋轉(zhuǎn)位置嵌入(Rotary Position Embeddings，RoPE)。

　　然后通過計算query向量和key向量的點積、縮放、應(yīng)用softmax得到注意力權(quán)重，對value向量進行加權(quán)求和，得到注意力得分。

　　在實踐中，Transformer一般都會用到多個注意力頭，其中每個頭的注意力機制都是獨立計算，最后再通過一個可學習的投影矩陣Wo合并所有頭輸出。

　　在計算注意力得分的過程中，權(quán)重矩陣沒有受到太多限制，可能會導致最終得分過大或過小。

　　在nGPT中，研究人員對q向量和k向量進行歸一化，還引入了一些可調(diào)整的參數(shù)(sqk)，以確保權(quán)重矩陣在處理位置信息時不會失真，更準確地捕捉到句子中詞與詞之間的關(guān)系，從而做出更好的預(yù)測和決策。

　　MLP塊

　　在標準Transformer中，隱藏層收入通過RMSNorm進行歸一化，然后經(jīng)過兩個線性投影生成中間向量(暫不考慮偏置項)：

　　然后使用SwiGLU 門控激活函數(shù)，以及一個線性變換得到最終門控激活。

　　在nGPT中，研究人員提出對線性投影的權(quán)重矩陣進行歸一化，并引入可學習的縮放因子，能夠更充分地利用處理信息時的非線性特性，在處理復(fù)雜信息時更加靈活。

　　多層感知機模塊的輸出不會因為縮放調(diào)整而發(fā)生變化。

　　Adam高效學習率

　　Adam優(yōu)化算法通過動量和梯度幅度的估計來調(diào)整每次的學習步長，同時考慮了當前及過去的梯度信息。

　　在nGPT中，研究人員同樣引入了一個可訓練的縮放參數(shù)向量，對特定的參數(shù)進行更精細的控制，確保每個參數(shù)都能以最適合自己的速度進行學習，從而進一步提高學習效率。

　　在不影響全局學習率的情況下，對特定的參數(shù)進行調(diào)整，提供了更大的靈活性和控制力。

　　變化總結(jié)

　　和基礎(chǔ)Transformer相比，nGPT主要做了七個改變：

　　1、移除所有歸一化層，比如RMSNorm或LayerNorm;

　　2、在每個訓練步驟之后，沿著嵌入維度對所有矩陣，包括輸入輸出嵌入矩陣，以及各種權(quán)重矩陣進行歸一化處理;

　　3、修改了隱藏層參數(shù)更新方程;

　　4、調(diào)整注意力機制中的softmax縮放因子，對q和k進行重新縮放和歸一化;

　　5、對MLP塊的中間狀態(tài)進行重新縮放;

　　6、對logits進行重新縮放;

　　7、移除權(quán)重衰減和學習率預(yù)熱步驟。

　　上下文越長，訓練速度越快

　　接下來，研究人員在OpenWebText數(shù)據(jù)集上訓練了基礎(chǔ)基礎(chǔ)Transformer(GPT)和歸一化Transformer(nGPT)，并在一系列標準下游任務(wù)上對其進行評估。

　　實驗中，使用了0.5B和1B(包括嵌入)兩種參數(shù)規(guī)模的模型。兩種參數(shù)規(guī)模的模型0.5B和1B(包含嵌入)。

　　訓練加速

　　圖1顯示了，在訓練過程中，10億參數(shù)且樣本長度為4k token的GPT和nGPT模型的驗證損失。

　　經(jīng)過2萬次迭代后，nGPT達到了與GPT在20萬次迭代(約4000億個token)后，才能達到的相同驗證損失。

　　這表明，在迭代次數(shù)和使用token數(shù)量方面，nGPT實現(xiàn)了10倍的加速。

　　再來看圖2，展示了nGPT和GPT在三個方面的性能差距是如何變化的：總token數(shù)量、上下文長度、參數(shù)規(guī)模。

　　在1k、4k和8k token上下文中，訓練0.5B和1B的nGPT模型分別約快4倍、10倍和20倍。

　　圖3在下游任務(wù)中顯示了類似的性能，證實加速不僅反映在困惑度上，也反映在任務(wù)表現(xiàn)上。

　　研究人員觀察到，對于較長的訓練運行，nGPT顯示出一些飽和現(xiàn)象，這暗示在當前可訓練參數(shù)數(shù)量下，模型容量可能已接近極限。

　　神經(jīng)網(wǎng)絡(luò)參數(shù)檢查

　　圖4顯示，雖然nGPT保持固定的嵌入范數(shù)(這是設(shè)計使然)，但GPT表現(xiàn)出明顯的變化。

　　從嵌入的協(xié)方差矩陣計算得出的特征值分布(已經(jīng)由其中位數(shù)歸一化)顯示，GPT的輸入嵌入具有更高的條件數(shù)，尤其是在1B模型中。

　　嵌入之間的成對點積分布表明，即使在nGPT中，嵌入也并非均勻分布在超球面上(在那里點積會接近0)，而是形成簇——這可能反映了語言數(shù)據(jù)中的自然模式。

　　由于GPT的嵌入形成了一個超橢球體(hyper-ellipsoid)，如向量范數(shù)的分布所示，其點積往往具有更高的值。

　　GPT輸入嵌入的病態(tài)性質(zhì)(ill-conditioned nature)可能導致涉及這些嵌入的計算問題。

　　下圖5展示了，注意力和MLP矩陣在不同層深度上的中位數(shù)條件數(shù)(跨多個頭)——0.5B模型有24層，1B模型有36層。

　　與nGPT相比，GPT模型的注意力矩陣呈現(xiàn)顯著更高的條件數(shù)。

　　對這些矩陣的進一步檢查，GPT的注意力矩陣表現(xiàn)出退化為低秩矩陣的趨勢，可能減少了這些塊的學習容量。

　　下圖6展示了，(左圖)注意力模塊和MLP模塊的特征學習率，(中圖)應(yīng)用于MLP中間狀態(tài)的縮放因子，(右圖)應(yīng)用于QK點積之前的縮放因子。