何愷明LeCun改造Transformer！9行代碼替代歸一化性能不減還加速
入選CVPR 2025，代碼已開源
夢晨 克雷西 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI
何愷明LeCun聯(lián)手：Transformer不要歸一化了，論文已入選CVPR2025。
歸一化長期以來一直被認(rèn)為是必不可少的，在現(xiàn)代神經(jīng)網(wǎng)絡(luò)中無處不在。
但團(tuán)隊認(rèn)為可以換用一種非常簡單的技術(shù)，他們提出DyT（Dynamic Tanh)，直接替代Layer Norm或RMSNorm，性能達(dá)到或超過標(biāo)準(zhǔn)Transformer。
DyT模塊可以用幾行PyTorch代碼實現(xiàn)：
class DyT(nn.Module):
def __init__(self, num_features, alpha_init_value=0.5):
super().__init__()
self.alpha = nn.Parameter(torch.ones(1) * alpha_init_value)
self.weight = nn.Parameter(torch.ones(num_features))
self.bias = nn.Parameter(torch.zeros(num_features))
def forward(self, x):
x = torch.tanh(self.alpha * x)
return x * self.weight + self.bias從視覺的ViT/MAE，到語言模型的LLaMA，再到語音、DNA系列等模態(tài)都可以用，完整代碼庫已開源。
網(wǎng)友評價說，只要元素級運算就能實現(xiàn)和歸一化一樣的效果，這對于效率優(yōu)化來說簡直是免費的午餐。
前Salesforce首席科學(xué)家、搜索引擎You的CEO Richard Socher也表示，這項成果加強了他之前的假設(shè)——原始的Transformer只是眾多等效神經(jīng)結(jié)構(gòu)之中的一個。
Transformer不要歸一化了
團(tuán)隊首先實證研究了訓(xùn)練網(wǎng)絡(luò)中歸一化層的行為。
選取ViT、wav2vec 2.0和DiT三種訓(xùn)練好的網(wǎng)絡(luò)，對每個網(wǎng)絡(luò)采樣一個小批量樣本進(jìn)行前向傳播，測量LayerNorm層在可學(xué)習(xí)仿射變換前的輸入和輸出，建立輸入輸出元素的一一對應(yīng)關(guān)系，從而直接可視化兩者關(guān)系。
結(jié)果發(fā)現(xiàn)，LayerNorm傳統(tǒng)上被認(rèn)為是線性變換，但實際整體居然呈現(xiàn)出出類似tanh函數(shù)非線性變換效果。
受到這一相似性啟發(fā)，團(tuán)隊提出DyT作為歸一化層的直接替代品， DyT 層定義如下：
DyT(x) = γ * tanh?(αx) + β
其中α是可學(xué)習(xí)的標(biāo)量參數(shù)，負(fù)責(zé)縮放調(diào)整輸入；γ和β是可學(xué)習(xí)的per-channel矢量參數(shù)，將輸出縮放回任何尺度。
DyT適用于注意力塊、FFN塊和最終歸一化層，盡管它可能看起來像是激活函數(shù)，但在這篇研究中不會改變原始架構(gòu)中激活函數(shù)的任何部分，網(wǎng)絡(luò)的其他部分也保持不變。
同時團(tuán)隊還觀察到，幾乎不需要調(diào)整原始架構(gòu)使用的超參數(shù)即可使DyT表現(xiàn)良好。
實驗選用多種任務(wù)和模型架構(gòu)，DyT在大多數(shù)情況下能達(dá)到與歸一化層相當(dāng)甚至更好的性能。
視覺監(jiān)督學(xué)習(xí)，選用ViT和ConvNeXt，在ImageNet-1K分類任務(wù)中訓(xùn)練，DyT在兩種架構(gòu)的不同模型尺寸下，性能均略優(yōu)于LayerNorm，且模型收斂行為高度一致，表明二者學(xué)習(xí)動態(tài)相似。
視覺自監(jiān)督學(xué)習(xí)，選用兩種訓(xùn)練目標(biāo)不同的網(wǎng)絡(luò)MAE和DINO，DyT的表現(xiàn)與LayerNorm相當(dāng)。
擴散模型實驗中，訓(xùn)練了三個不同尺寸的DiT模型，用FID分?jǐn)?shù)評估生成圖像質(zhì)量。
僅用 tanh (αx) 函數(shù)替換DiT中LN 層的歸一化變換，保留其仿射參數(shù)（用于class conditionin），結(jié)果相差也不大。
語言模型實驗中，用DyT代替了LLaMA默認(rèn)的RMSNorm，在所有四種參數(shù)規(guī)模上的性能與RMSNorm相當(dāng)，在整個訓(xùn)練過程中，訓(xùn)練損失保持一致。
語音、DNA序列模型中的情況也類似。
但DyT的作用不止于此，在訓(xùn)練效率方面也有很大提升。
以LLaMA 7B為研究對象，分別采用RMSNorm和DyT進(jìn)行實驗，測量在使用單個長度為4096 tokens的序列時，100次前向傳遞（推理）和100次前向-后向傳遞（訓(xùn)練）所需的總時間。
在BF16精度下，DyT顯著縮短了計算時間，另外在FP32精度下觀察到類似的趨勢。
但DyT也有局限性，在非Transformer模型中，如替換ResNet的Batch Norm時效果不佳，是否以及如何適應(yīng)其他類型歸一化層的模型還需進(jìn)一步研究
作者團(tuán)隊
包括何愷明和LeCun在內(nèi)，本文的作者一共有五位，其余三位分別是第一作者、紐約大學(xué)博士生Jiachen Zhu，以及來自Meta FAIR實驗室的劉壯和陳鑫磊。
雖然這是五人第一次聚在一起發(fā)表的論文，但其中部分成員之間的合作已經(jīng)進(jìn)行過多次。
項目負(fù)責(zé)人、Meta FAIR實驗室科學(xué)家劉壯，就是何愷明的一位“老搭檔”。
和何愷明一樣，劉壯本科畢業(yè)自清華，并且也是CVPR最佳論文獎得主——他是CVPR2017最佳論文DenseNet的第一作者。
2017年，劉壯從清華姚班畢業(yè)，進(jìn)入加州大學(xué)伯克利分校攻讀博士學(xué)位，師從Trevor Darrell，是賈揚清的同門師弟。
博士畢業(yè)后，劉壯進(jìn)入Meta AI Research工作。在此之前，他已經(jīng)在Meta實習(xí)了一年多時間，期間和謝賽寧合作，發(fā)表了ConvNeXt。
還有浙大校友陳鑫磊， 目前是Meta FAIR實驗室的研究科學(xué)家，研究興趣集中于預(yù)訓(xùn)練，特別是有自監(jiān)督或是多模態(tài)視覺表示的預(yù)訓(xùn)練。
發(fā)表在CVPR上、目前谷歌學(xué)術(shù)引用量達(dá)8998次的MAE開山論文，陳鑫磊與何愷明是共同一作，謝賽寧也參與其中。
第一作者Jiachen Zhu，來自重慶，本科就讀于香港理工大學(xué)，取得了計算機和工商管理雙重學(xué)位。
本科畢業(yè)3年后，Jiachen Zhu重回校園，到紐約大學(xué)先后攻讀計算機碩士和博士學(xué)位，目前仍然在讀，博士生導(dǎo)師就是LeCun。
谷歌學(xué)術(shù)信息顯示，除了本次的新成果之外，Jiachen Zhu自2022年至今一共還發(fā)表過5篇論文，其中3篇為一作或共同一作，每篇均有LeCun的參與。
并且Jiachen Zhu去年以Meta實習(xí)生身份發(fā)表的一篇關(guān)于多模態(tài)理解與生成的論文，也與劉壯、陳鑫磊以及LeCun的參與。
論文地址：
https://arxiv.org/abs/2503.10622v1
GitHub地址：
https://github.com/jiachenzhu/DyT
項目主頁：
https://jiachenzhu.github.io/DyT/
參考鏈接：
[1]https://x.com/liuzhuang1234/status/1900370738588135805
相關(guān)閱讀
