AI的未來(lái)，或許就此改寫(xiě)......

最近，英偉達(dá)團(tuán)隊(duì)拋出的一枚重磅炸彈，提出了全新神經(jīng)網(wǎng)絡(luò)架構(gòu)——?dú)w一化Transformer（nGPT），基于超球面（hypersphere）進(jìn)行表示學(xué)習(xí)。

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

也就意味著，原本需要一個(gè)月完成的訓(xùn)練，在未來(lái)可能只需1-2天的時(shí)間就能搞定。

無(wú)疑為通向AGI終極目標(biāo)，注入了一針強(qiáng)心劑！

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

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

輸入后的token在超球面表面上移動(dòng)，每一層都通過(guò)「位移」來(lái)貢獻(xiàn)最終的輸出預(yù)測(cè)，其中位移量是由MLP和注意力模塊進(jìn)行定義的，其向量組件都位于同一個(gè)超球面上。

實(shí)驗(yàn)表明，nGPT達(dá)到相同精度所需的訓(xùn)練步驟減少了4-20倍，具體取決于序列長(zhǎng)度：

- 1k上下文，訓(xùn)練速度提高4倍

- 4k上下文，訓(xùn)練速度提高10倍

- 8k上下文，訓(xùn)練速度提高20倍

可以看出，上下文越長(zhǎng)，訓(xùn)練越快。

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

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

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

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

不過(guò)，當(dāng)前基于Transformer搭建的大模型都是計(jì)算密集型的，需要耗費(fèi)大量的資源和時(shí)間。

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

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

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

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

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

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

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

這項(xiàng)研究的主要貢獻(xiàn)在于：

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

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

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

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

對(duì)于輸入序列中的每個(gè)token ，歸一化Transformer的優(yōu)化路徑從超球面上對(duì)應(yīng)于其輸入嵌入向量的點(diǎn)開(kāi)始，移動(dòng)到超球面上最能預(yù)測(cè)下一個(gè)的嵌入向量的點(diǎn)。

- 更快的收斂

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

Transformer演變：從GPT到nGPT

嵌入層歸一化

標(biāo)準(zhǔn)的decoder-only Transformer的訓(xùn)練目標(biāo)是根據(jù)輸入序列的前序tokens來(lái)預(yù)測(cè)后面的token，在token預(yù)測(cè)時(shí)，模型會(huì)引入兩個(gè)可學(xué)習(xí)的嵌入矩陣Einput和Eoutput，分別用來(lái)從輸入詞轉(zhuǎn)為詞嵌入，以及從詞嵌入轉(zhuǎn)為預(yù)測(cè)輸出。

在模型訓(xùn)練期間，通常使用對(duì)應(yīng)嵌入向量的點(diǎn)積來(lái)計(jì)算token相似度，但嵌入向量的范數(shù)（norms）不受限制的，可能會(huì)導(dǎo)致相似性計(jì)算存在偏差。

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

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

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

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

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

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

層/塊歸一

標(biāo)準(zhǔn)Transformer架構(gòu)需要對(duì)隱藏層狀態(tài)h進(jìn)行L層變換，包括一個(gè)自注意力（ATTN）和多層感知機(jī)（MLP）。

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

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

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

按最短路徑尋找來(lái)說(shuō)，參數(shù)更新過(guò)程可以描述為：

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

在擬牛頓方法中，B可以近似于逆黑塞矩陣，當(dāng) B是一個(gè)對(duì)角線元素非負(fù)的對(duì)角矩陣時(shí)，αB就變成了一個(gè)向量，其元素對(duì)應(yīng)于B的對(duì)角線元素乘以學(xué)習(xí)率α，也可以稱(chēng)之為特征學(xué)習(xí)率（eigen learning rates）。

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

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

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

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

自注意力塊

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

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

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

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

在實(shí)踐中，Transformer一般都會(huì)用到多個(gè)注意力頭，其中每個(gè)頭的注意力機(jī)制都是獨(dú)立計(jì)算，最后再通過(guò)一個(gè)可學(xué)習(xí)的投影矩陣Wo合并所有頭輸出。

在計(jì)算注意力得分的過(guò)程中，權(quán)重矩陣沒(méi)有受到太多限制，可能會(huì)導(dǎo)致最終得分過(guò)大或過(guò)小。

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

MLP塊

在標(biāo)準(zhǔn)Transformer中，隱藏層收入通過(guò)RMSNorm進(jìn)行歸一化，然后經(jīng)過(guò)兩個(gè)線性投影生成中間向量（暫不考慮偏置項(xiàng)）：

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

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

多層感知機(jī)模塊的輸出不會(huì)因?yàn)榭s放調(diào)整而發(fā)生變化。

Adam高效學(xué)習(xí)率

Adam優(yōu)化算法通過(guò)動(dòng)量和梯度幅度的估計(jì)來(lái)調(diào)整每次的學(xué)習(xí)步長(zhǎng)，同時(shí)考慮了當(dāng)前及過(guò)去的梯度信息。

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

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

變化總結(jié)

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

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

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

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

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

5、對(duì)MLP塊的中間狀態(tài)進(jìn)行重新縮放；

6、對(duì)logits進(jìn)行重新縮放；

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

上下文越長(zhǎng)，訓(xùn)練速度越快

接下來(lái)，研究人員在OpenWebText數(shù)據(jù)集上訓(xùn)練了基礎(chǔ)基礎(chǔ)Transformer（GPT）和歸一化Transformer（nGPT），并在一系列標(biāo)準(zhǔn)下游任務(wù)上對(duì)其進(jìn)行評(píng)估。

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

訓(xùn)練加速

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

經(jīng)過(guò)2萬(wàn)次迭代后，nGPT達(dá)到了與GPT在20萬(wàn)次迭代（約4000億個(gè)token）后，才能達(dá)到的相同驗(yàn)證損失。

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

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

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

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

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

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

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

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

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

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

GPT輸入嵌入的病態(tài)性質(zhì)（ill-conditioned nature）可能導(dǎo)致涉及這些嵌入的計(jì)算問(wèn)題。

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

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

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

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