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

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

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

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

歸一化長期以來一直被認為是必不可少的，在現(xiàn)代神經(jīng)網(wǎng)絡(luò)中無處不在。

但團隊認為可以換用一種非常簡單的技術(shù)，他們提出DyT（Dynamic Tanh)，直接替代Layer Norm或RMSNorm，性能達到或超過標準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不要歸一化了

團隊首先實證研究了訓(xùn)練網(wǎng)絡(luò)中歸一化層的行為。

選取ViT、wav2vec 2.0和DiT三種訓(xùn)練好的網(wǎng)絡(luò)，對每個網(wǎng)絡(luò)采樣一個小批量樣本進行前向傳播，測量LayerNorm層在可學(xué)習(xí)仿射變換前的輸入和輸出，建立輸入輸出元素的一一對應(yīng)關(guān)系，從而直接可視化兩者關(guān)系。

結(jié)果發(fā)現(xiàn)，LayerNorm傳統(tǒng)上被認為是線性變換，但實際整體居然呈現(xiàn)出出類似tanh函數(shù)非線性變換效果。

受到這一相似性啟發(fā)，團隊提出DyT作為歸一化層的直接替代品， DyT 層定義如下：

DyT(x) = γ * tanh?(αx) + β

其中α是可學(xué)習(xí)的標量參數(shù)，負責(zé)縮放調(diào)整輸入；γ和β是可學(xué)習(xí)的per-channel矢量參數(shù)，將輸出縮放回任何尺度。

DyT適用于注意力塊、FFN塊和最終歸一化層，盡管它可能看起來像是激活函數(shù)，但在這篇研究中不會改變原始架構(gòu)中激活函數(shù)的任何部分，網(wǎng)絡(luò)的其他部分也保持不變。

同時團隊還觀察到，幾乎不需要調(diào)整原始架構(gòu)使用的超參數(shù)即可使DyT表現(xiàn)良好。

實驗選用多種任務(wù)和模型架構(gòu)，DyT在大多數(shù)情況下能達到與歸一化層相當甚至更好的性能。

視覺監(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)練目標不同的網(wǎng)絡(luò)MAE和DINO，DyT的表現(xiàn)與LayerNorm相當。

擴散模型實驗中，訓(xùn)練了三個不同尺寸的DiT模型，用FID分數(shù)評估生成圖像質(zhì)量。

僅用 tanh (αx) 函數(shù)替換DiT中LN 層的歸一化變換，保留其仿射參數(shù)（用于class conditionin），結(jié)果相差也不大。

語言模型實驗中，用DyT代替了LLaMA默認的RMSNorm，在所有四種參數(shù)規(guī)模上的性能與RMSNorm相當，在整個訓(xùn)練過程中，訓(xùn)練損失保持一致。

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

但DyT的作用不止于此，在訓(xùn)練效率方面也有很大提升。

以LLaMA 7B為研究對象，分別采用RMSNorm和DyT進行實驗，測量在使用單個長度為4096 tokens的序列時，100次前向傳遞（推理）和100次前向-后向傳遞（訓(xùn)練）所需的總時間。

在BF16精度下，DyT顯著縮短了計算時間，另外在FP32精度下觀察到類似的趨勢。

但DyT也有局限性，在非Transformer模型中，如替換ResNet的Batch Norm時效果不佳，是否以及如何適應(yīng)其他類型歸一化層的模型還需進一步研究

作者團隊

包括何愷明和LeCun在內(nèi)，本文的作者一共有五位，其余三位分別是第一作者、紐約大學(xué)博士生Jiachen Zhu，以及來自Meta FAIR實驗室的劉壯和陳鑫磊。

雖然這是五人第一次聚在一起發(fā)表的論文，但其中部分成員之間的合作已經(jīng)進行過多次。

項目負責(zé)人、Meta FAIR實驗室科學(xué)家劉壯，就是何愷明的一位“老搭檔”。

和何愷明一樣，劉壯本科畢業(yè)自清華，并且也是CVPR最佳論文獎得主——他是CVPR2017最佳論文DenseNet的第一作者。

2017年，劉壯從清華姚班畢業(yè)，進入加州大學(xué)伯克利分校攻讀博士學(xué)位，師從Trevor Darrell，是賈揚清的同門師弟。

博士畢業(yè)后，劉壯進入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ù)引用量達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