
來源 | 量子位
何愷明LeCun聯手:Transformer不要歸一化了,論文已入選CVPR2025。

歸一化長期以來一直被認為是必不可少的,在現代神經網路中無處不在。
但團隊認為可以換用一種非常簡單的技術,他們提出DyT(Dynamic Tanh),直接替代Layer Norm或RMSNorm,效能達到或超過標準Transformer。

DyT模組可以用幾行PyTorch程式碼實現:
classDyT(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))
defforward(self, x):
x = torch.tanh(self.alpha * x)
return x * self.weight + self.bias
從視覺的ViT/MAE,到語言模型的LLaMA,再到語音、DNA系列等模態都可以用,完整程式碼庫已開源。

網友評價說,只要元素級運算就能實現和歸一化一樣的效果,這對於效率最佳化來說簡直是免費的午餐。

前Salesforce首席科學家、搜尋引擎You的CEO Richard Socher也表示,這項成果加強了他之前的假設——原始的Transformer只是眾多等效神經結構之中的一個。

Transformer不要歸一化了
團隊首先實證研究了訓練網路中歸一化層的行為。
選取ViT、wav2vec 2.0和DiT三種訓練好的網路,對每個網路取樣一個小批次樣本進行前向傳播,測量LayerNorm層在可學習仿射變換前的輸入和輸出,建立輸入輸出元素的一一對應關係,從而直接視覺化兩者關係。

結果發現,LayerNorm傳統上被認為是線性變換,但實際整體居然呈現出出類似tanh函式非線性變換效果。

受到這一相似性啟發,團隊提出DyT作為歸一化層的直接替代品, DyT 層定義如下:
DyT(x) = γ * tanh(αx) + β
其中α是可學習的標量引數,負責縮放調整輸入;γ和β是可學習的per-channel向量引數,將輸出縮放回任何尺度。
DyT適用於注意力塊、FFN塊和最終歸一化層,儘管它可能看起來像是啟用函式,但在這篇研究中不會改變原始架構中啟用函式的任何部分,網路的其他部分也保持不變。
同時團隊還觀察到,幾乎不需要調整原始架構使用的超引數即可使DyT表現良好。
實驗選用多種任務和模型架構,DyT在大多數情況下能達到與歸一化層相當甚至更好的效能。
視覺監督學習,選用ViT和ConvNeXt,在ImageNet-1K分類任務中訓練,DyT在兩種架構的不同模型尺寸下,效能均略優於LayerNorm,且模型收斂行為高度一致,表明二者學習動態相似。

視覺自監督學習,選用兩種訓練目標不同的網路MAE和DINO,DyT的表現與LayerNorm相當。

擴散模型實驗中,訓練了三個不同尺寸的DiT模型,用FID分數評估生成影像質量。
僅用 tanh (αx) 函式替換DiT中LN 層的歸一化變換,保留其仿射引數(用於class conditionin),結果相差也不大。

語言模型實驗中,用DyT代替了LLaMA預設的RMSNorm,在所有四種引數規模上的效能與RMSNorm相當,在整個訓練過程中,訓練損失保持一致。

語音、DNA序列模型中的情況也類似。

但DyT的作用不止於此,在訓練效率方面也有很大提升。
以LLaMA 7B為研究物件,分別採用RMSNorm和DyT進行實驗,測量在使用單個長度為4096 tokens的序列時,100次前向傳遞(推理)和100次前向-後向傳遞(訓練)所需的總時間。
在BF16精度下,DyT顯著縮短了計算時間,另外在FP32精度下觀察到類似的趨勢。

但DyT也有侷限性,在非Transformer模型中,如替換ResNet的Batch Norm時效果不佳,是否以及如何適應其他型別歸一化層的模型還需進一步研究
作者團隊
包括何愷明和LeCun在內,本文的作者一共有五位,其餘三位分別是第一作者、紐約大學博士生Jiachen Zhu,以及來自Meta FAIR實驗室的劉壯和陳鑫磊。
雖然這是五人第一次聚在一起發表的論文,但其中部分成員之間的合作已經進行過多次。
專案負責人、Meta FAIR實驗室科學家劉壯,就是何愷明的一位“老搭檔”。
和何愷明一樣,劉壯本科畢業自清華,並且也是CVPR最佳論文獎得主——他是CVPR2017最佳論文DenseNet的第一作者。
2017年,劉壯從清華姚班畢業,進入加州大學伯克利分校攻讀博士學位,師從Trevor Darrell,是賈揚清的同門師弟。
博士畢業後,劉壯進入Meta AI Research工作。在此之前,他已經在Meta實習了一年多時間,期間和謝賽寧合作,發表了ConvNeXt。

還有浙大校友陳鑫磊, 目前是Meta FAIR實驗室的研究科學家,研究興趣集中於預訓練,特別是有自監督或是多模態視覺表示的預訓練。
發表在CVPR上、目前谷歌學術引用量達8998次的MAE開山論文,陳鑫磊與何愷明是共同一作,謝賽寧也參與其中。

第一作者Jiachen Zhu,來自重慶,本科就讀於香港理工大學,取得了計算機和工商管理雙重學位。
本科畢業3年後,Jiachen Zhu重回校園,到紐約大學先後攻讀計算機碩士和博士學位,目前仍然在讀,博士生導師就是LeCun。
谷歌學術資訊顯示,除了本次的新成果之外,Jiachen Zhu自2022年至今一共還發表過5篇論文,其中3篇為一作或共同一作,每篇均有LeCun的參與。

並且Jiachen Zhu去年以Meta實習生身份發表的一篇關於多模態理解與生成的論文,也與劉壯、陳鑫磊以及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
技術交流群邀請函
△長按新增小助手
掃描二維碼新增小助手微信
請備註:姓名-學校/公司-研究方向
(如:小張-哈工大-對話系統)
即可申請加入自然語言處理/Pytorch等技術交流群
關於我們
MLNLP 社群是由國內外機器學習與自然語言處理學者聯合構建的民間學術社群,目前已經發展為國內外知名的機器學習與自然語言處理社群,旨在促進機器學習,自然語言處理學術界、產業界和廣大愛好者之間的進步。
社群可以為相關從業者的深造、就業及研究等方面提供開放交流平臺。歡迎大家關注和加入我們。

掃描二維碼新增小助手微信
關於我們
