Transformer架構迎來歷史性突破!剛剛,何愷明LeCun、清華姚班劉壯聯手,用9行程式碼砍掉了Transformer「標配」歸一化層,創造了性能不減反增的奇蹟。
就在剛剛,何愷明和LeCun兩大巨頭聯手,徹底改變了Tranformer架構!
在Meta、NYU、MIT、普林斯頓等聯手的這項工作中,研究者們共同發現了一種沒有歸一化層的Transformer。
在深度學習中,歸一化層(Norm layer,下圖Transformer架構中的綠色區塊)無所不在,並且長期以來一直被認為是不可或缺的。
然而,這次何愷明、Yann LeCun等人卻證明:即使不使用歸一化層的Transformer模型,也就能達到相同甚至更好的性能!
而且方法也非常簡單,只需要9行程式碼。
簡單來說,他們發現了一種簡單得出人意料的方法,來取代歸一化層——縮放的tanh函數(沒錯,彷彿回到了80年代)。
他們稱之為動態Tanh(Dynamic Tanh,簡稱DyT)。
這個方法的核心,源自於一個非常簡單的觀察:LayerNorm將輸入轉換為輸出時,呈現出類似tanh的S形曲線。
它能夠壓縮極端值,同時在中心區域保持線性形態。
因此,研究者以提出的動態Tanh(DyT)層取代了歸一化層,過程很簡單。
具體來說,他們用DyT取代了LayerNorm/RSMNorm,並在以下Transformer結構上進行了測試:
結果表明,在所有測試場景中,DyT的效果都與歸一化層相當或更優。
而且,DyT比RMSNorm在H100s上更快!
這項研究的意義就在於,鑑於模型訓練和推理可能需要數千萬次的計算,因此DyT有極大潛力幫助降低成本,接下來就是找到應用。
網友們對此紛紛表示:非常精彩的研究!而它改變應用transformer歸一化基本假設的方式,尤其令人喜歡。
或許,這項研究能從根本解決深度學習問題?
也有人提出,tanh一直因為梯度消失問題而臭名昭著。如果用tanh取代歸一化層,是否會重新引入這個問題,進而可能減緩甚至阻礙深度模型的學習呢?
在過去的十年裡,歸一化層已經穩固地成為現代神經網路最基礎的元件之一。
如今,幾乎所有現代神經網路都採用歸一化層,其中層歸一化(Layer Normalization,LN)尤其流行,特別是在主導的Transformer 架構中。
原因主要在於它們對最佳化過程的實證優勢。因此,歸一化層被普遍認為是深度網路高效訓練的關鍵甚至不可或缺的組成部分。
新的神經網路架構往往會嘗試替換注意力層或摺積層,但幾乎都會保留歸一化層。
然而這項工作,卻提出了一種簡單的歸一化層替代方案,並應用於Transformer結構。
研究始於一個關鍵觀察:LN層將輸入對應為類似tanh的S形曲線,既縮放了輸入活化值,又抑制了極端值。
因此,研究者提出了逐元素(element-wise)操作,稱為動態tanh(Dynamic Tanh, DyT),定義如下:
其中,α 是可學習參數。
DyT的設計目標是:
這也是整個操作命名為「動態」Tanh的原因。
值得注意的是,與歸一化層不同,DyT無需計算啟動統計訊息,卻能夠同時實現縮放和極值抑制。
γ和β是可學習的、逐通道的向量參數,與所有歸一化層中使用的參數相同——它們允許輸出縮放到任意尺度。
有時這被認為是一個獨立的仿射層;但在新設計中,它們被視為DyT層的一部分,就像歸一化層也包含它們一樣。
有關DyT在類似Pytorch偽程式碼中的實現,請參閱演算法1。
將DyT層整合到現有架構中非常簡單:DyT層可以直接替換歸一化層(見圖1)。
這種替換適用於注意力模組、前饋網路(FFN)模組中的歸一化層,以及最終的歸一化層。
儘管DyT可能看起來像是啟動函數,或被認為是一種啟動函數,但本研究僅將其用於替換歸一化層,而不會改變原始架構中的任何啟動函數(如GELU或ReLU)。
網路的其他部分也保持不變。
此外,DyT在應用於原始架構時,幾乎不需要調整超參數就能達到良好的效能。
在所有三個模型中,可以觀察到較淺的LayerNorm(LN)層(圖2第一列)的輸入-輸出關係大多是線性的,在xy平面上接近一條直線。
然而,在更深的LN層中,可以發現了一些更具啟發性的現象。
在大多數深層LN層中,輸入-輸出曲線高度類似tanh函數所呈現的完整或部分S形曲線(見圖3)。
依常理,可能會認為LN層只是對輸入張量進行線性變換,因為減去平均值和除以標準差都是線性運算。
LN的歸一化方式是逐token進行的,因此它對每個token施加的變換是線性的。
然而,由於不同token具有不同的平均值和標準差,LN無法整體上保持輸入張量所有活化值的線性關係。
儘管如此,讓人仍然感到驚訝的是,實際的非線性變換竟然高度類似於縮放後的tanh函數。
在這種S形曲線中,可以觀察到:
然而,仍有一些點超出了這個範圍,即那些被認為是“極端”值的點。例如,在ViT模型中,x>50或x<-50的資料點。
歸一化層的主要作用之一,便是將這些極端值壓縮到相對較溫和的範圍,使其更符合大部分資料點的分佈。
正是在這個方面,歸一化層的行為無法被簡單的仿射變換所逼近。
研究假設,歸一化層對極端值的這種非線性和非均勻壓縮作用,正是使其至關重要且難以被取代的核心原因。
為了理解這一點,團隊分別按token和通道對這些點進行視覺化。
可以看到,在左側兩幅圖中,相同token(相同顏色)在不同通道上的資料點排列成直線,這是因為LN在每個token內按通道進行線性變換。
然而,當所有token的資料點整體繪製在一張圖上時,這些直線卻形成了一個非線性的tanh形曲線。
右側兩幅圖中,每個通道的輸入值分佈在x軸上的不同範圍,共同構成了整體的tanh形曲線。
某些通道(如紅色、綠色和粉紅色)包含更極端的x值,這些極端值被LN壓縮到較溫和的範圍。
為了證明DyT的有效性,在Transformer及其他幾種現代架構上進行了實驗,涵蓋了多個不同的任務和領域。
在每次實驗中,用DyT層取代原始架構中的LN(Layer Normalization)或RMSNorm,並遵循官方開源協定來訓練和測試兩種版本的模型。
值得注意的是,為了突顯DyT 適配的簡單性,所使用的超參數與原始歸一化模型完全相同。
在ImageNet-1K分類任務上,訓練了「Base」和「Large」規模的Vision Transformer(ViT)和ConvNeXt。
選擇這些模型是因為它們的廣泛流行以及不同的操作特性:ViT中的注意力機制和ConvNeXt中的捲積操作。
表1報告了Top-1分類準確率。 DyT在兩種架構和模型規模上均略優於LN。
進一步繪製了ViT-B和ConvNeXt-B的訓練損失曲線,如圖5所示。曲線表明,DyT和基於LN的模型的收斂行為高度一致。
使用兩種流行的視覺自監督學習方法進行基準測試:掩碼自編碼器(MAE)和DINO。
兩者預設均使用Vision Transformers作為骨幹網絡,但具有不同的訓練目標:MAE透過重建損失進行訓練,而DINO使用聯合嵌入損失。
依照標準的自我監督學習流程,首先在不使用任何標籤的情況下在ImageNet-1K上對模型進行預訓練,然後透過附加分類層並使用標籤對預訓練模型進行微調來測試其效能。
微調結果如表2所示。 DyT在自監督學習任務中始終表現與LN相當。
在ImageNet-1K上訓練了三種規模的Diffusion Transformer(DiT)模型,分別為B、L和XL。其patch大小分別為4、4和2。
需要注意的是,在DiT中,LN層的仿射參數用於類別條件化(class conditioning),在DyT實驗中也保留了這一點,僅將歸一化變換替換為tanh(αx)函數。
訓練完成後,使用標準的ImageNet「參考批次」評估Fréchet Inception Distance(FID)分數,結果如表3所示。 DyT在FID上取得了與LN相當或更好的性能。
研究人員預先訓練了LLaMA 7B、13B、34B和70B模型,以評估DyT相對於RMSNorm的效能,RMSNorm是LLaMA中預設使用的歸一化層。
如表4所示,DyT在所有四種模型規模上的表現均與RMSNorm相當。
圖6展示了損失曲線,顯示所有模型規模的訓練趨勢相似,訓練損失在整個訓練過程中高度一致。
在長程DNA序列建模任務中,預先訓練了HyenaDNA模型和Caduceus模型。
預訓練使用了來自人類參考基因組數據,並在GenomicBenchmarks上進行評估。
結果如表6所示。 DyT在該任務中保持了與LN相當的性能。
在LibriSpeech資料集上,預先訓練了兩個wav2vec 2.0 Transformer模型。
表5中報告了最終的驗證損失。
可以觀察到,DyT在兩種模型規模上的表現均與LN相當。
研究團隊進行了額外的實驗,以評估超參數調優的影響,特別是針對所有非大語言模型(non-LLM)的學習率和α初始化。
學習率調優:表13總結了使用原始學習率與調優後學習率訓練的模型之間的表現對比。
結果表明,調優學習率對DyT模型的效能提升較為有限。
這表明,最初為LN模型優化的原始超參數已經非常適合DyT模型。這項觀察結果進一步強調了DyT與LN模型之間的內在相似性。
此外,本次也研究了優化DyT模型初始值α₀ 的影響,相關結果如表14。
研究發現,在某些模型上調整α₀ 僅帶來輕微的效能提升,這表明預設初始值(α₀ = 0.5)通常已經能夠實現接近最優的效能。
研究者對DyT的重要特性進行了多方面分析。
首先,評估其計算效率,隨後進行了兩項研究,分別探討了tanh函數的作用以及可學習縮放參數α的影響。
最後,他們與其他試圖去除歸一化層的方法進行了比較分析。
研究在LLaMA 7B模型上,對RMSNorm和DyT進行了基準測試,測量了100次前向傳播(推理) 和100次前向-反向傳播(訓練) 的總時間。
實驗採用4096個token的單一序列,在Nvidia H100(BF16 精度) 上運行。
表7報告了所有RMSNorm或DyT層以及整個模型的運行時間。
結果顯示:
由此可見,DyT可能是面向效率優化的網路設計中的一個有前景的選擇。
訓練過程中的α
分析表明,α在整個訓練過程中始終緊密追蹤啟動值的1/std(標準差的倒數)。
如圖8(左)所示,α在訓練早期先下降後上升,但始終與輸入活化值的標準差保持一致的波動趨勢。這現象顯示:
訓練後的α
對訓練完成後的α值進行進一步分析,我們發現它與輸入活化的1/std有強相關性。
這一趨勢與深度殘差網路的特性一致,與先前研究所揭示的一致。
上述分析表明,α在一定程度上發揮了類似歸一化的功能,其學習到的值大致接近輸入活化的1/std。但要注意的是:
研究也發現,對α(記作α₀) 的初始化進行調整,很少會顯著提升模型表現。
唯一的例外是-LLM訓練,在該場景下,對α₀進行精細調整可以帶來明顯的效能提升。
對於非LLM模型,其表現對α₀的敏感度相對較低
如圖9所示,研究人員在不同任務上對α₀進行調節後,觀察其對驗證集表現的影響。所有實驗均遵循各自任務的原始實驗設定和超參數。
實驗結果表明,在較寬範圍的α₀取值(0.5-1.2 之間)下,模型性能基本上保持穩定,通常能取得較好的結果。
此外,研究也發現,調整α₀通常僅影響訓練曲線的早期階段。
唯一的例外是監督學習ViT-L(Vision Transformer Large) 相關實驗,在該場景下,當α₀超過0.6時,訓練會變得不穩定並發生發散。
在這種情況下,降低學習率可以恢復訓練穩定性。
較小的α₀能夠提升訓練穩定性
在先前觀察的基礎上,研究進一步分析了導致訓練不穩定的因素。
研究表明,隨著模型規模或學習率的增加,需要降低α₀以確保訓練穩定。相反,較高的α₀需要搭配較低的學習率,以緩解訓練不穩定性。
如圖10展示了,在ImageNet-1K資料集上,監督學習ViT訓練穩定性的消融實驗。
在實驗中,研究人員調整了學習率、模型規模和α₀值。結果表明,訓練較大的模型更容易失敗,因此需要使用較小的α₀或更低的學習率來確保穩定性。
此外,研究人員在基於LN(LayerNorm)的模型上也觀察到了類似的不穩定性模式,在相同條件下,將α₀設為0.5可獲得與LN類似的穩定性表現。
將α₀=0.5設為預設值
基於上述研究結果,研究者將α₀=0.5設為所有非LLM模型的預設值。此設定能夠在保證訓練穩定性的同時,維持良好的性能,並與LN的穩定性模式相當。
調整α₀可提升LLM性能
如前文所述,預設設定α₀=0.5在大多數任務中表現良好。然而,研究人員發現,對α₀進行調優可以顯著提高LLM的性能。
他們對LLaMA系列模型進行了α₀調優,具體方法是在30B token規模的資料上進行預訓練,並比較其訓練損失。
表11總結了每個模型的最優α₀值。研究中,得出了兩個關鍵結論:
提高注意力層(attention blocks)的α₀值有助於提升性能
研究發現,在注意力塊(attention blocks)中的DyT層設定較高的α₀,而在其他位置(如FFN塊或最終線性投影之前)設定較低的α₀,能夠提升模型性能。
為了進一步說明α₀調優的影響,圖11展示了兩種LLaMA模型的訓練損失熱圖。實驗顯示,在注意力區塊(attention blocks)中使用較高的α₀,能夠有效降低訓練損失,從而提升模型表現。
模型寬度主要決定α₀的選擇
研究人員也進一步研究了模型寬度(width)和深度(depth) 對最適α₀的影響。研究結果顯示:
表12總結了不同寬度和深度下的最優α₀值,進一步證明了寬度越大的模型,適合使用較小的α₀,而深度對α₀的影響可以忽略不計。
作者介紹 Jiachen Zhu
Jiachen Zhu是紐約大學柯朗數學研究所(CIMS)的四年級電腦科學博士生,導師是Yann LeCun。
在此之前,他於2015年在香港理工大學獲得學士學位,並於2020年在紐約大學獲得碩士學位。
他主要對影像或影片的自監督學習,以及各類神經網路架構的設計的研究感興趣。
Xinlei Chen(陳鑫磊)
Xinlei Chen目前是Meta FAIR的研究科學家。
他曾獲得了CMU語言技術研究所的博士學位,本科畢業於中國浙江大學電腦科學專業。目前,他的研究興趣集中在視覺表示的預訓練與理解。
何愷明
何愷明目前是麻省理工學院(MIT)電機工程與電腦科學系(EECS)的副教授,於2024年2月加入。
他本人的研究重點是,建構能夠從複雜世界中學習表示並行展智慧的電腦模型,研究的長期目標是用更強大的AI來增強人類智慧。
Yann LeCun
Yann LeCun圖靈獎三巨頭之一,目前擔任Meta首席科學家,紐約大學終身教授。
Zhuang Liu(劉壯)
劉壯是此計畫負責人,目前在Meta FAIR實驗室擔任研究科學家。
他曾獲得了UC柏克萊分校的電子工程與電腦科學系(EECS) 電腦科學博士學位,導師是Trevor Darrell教授。
本科畢業於清華姚班,此外,他曾在康乃爾大學、英特爾實驗室、Adobe Research、FAIR擔任訪問研究員或實習生。
劉壯的主要研究領域為深度學習與電腦視覺,聚焦於深度學習模型架構、訓練方法、運算效率及模型理解。
他曾主導開發了DenseNet(獲得CVPR最佳論文獎)和ConvNeXt。這兩個演算法皆成為深度學習和電腦視覺領域最廣泛應用的神經網路架構之一。
參考資料:
https://arxiv.org/abs/2503.10622
https://jiachenzhu.github.io/DyT/
https://x.com/liuzhuang1234/status/1900370738588135805 (新智元)