何愷明LeCun改造Transformer！9行代碼替代歸一化性能不減還加速

克雷西 2025-03-14 15:39:46 來源：量子位

夢晨 克雷西 發自 凹非寺
量子位 | 公眾號 QbitAI

何愷明LeCun聯手：Transformer不要歸一化了，論文已入選CVPR2025。

歸一化長期以來一直被認為是必不可少的，在現代神經網絡中無處不在。

但團隊認為可以換用一種非常簡單的技術，他們提出DyT（Dynamic Tanh)，直接替代Layer Norm或RMSNorm，性能達到或超過標準Transformer。

DyT模塊可以用幾行PyTorch代碼實現：

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系列等模態都可以用，完整代碼庫已開源。

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

前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