何愷明聯手LeCun改進Transformer！成功去掉歸一化層！中稿CVPR2025

點選下方卡片，關注“CVer”公眾號

AI/CV重磅乾貨，第一時間送達

點選進入—>【頂會/頂刊】投稿交流群

新增微訊號：CVer2233，小助手會拉你進群！

掃描下方二維碼，加入CVer學術星球！可以獲得最新頂會/頂刊上的論文idea和CV從入門到精通資料，及最前沿應用！發論文/搞科研/漲薪，強烈推薦！

轉載自：機器之心

何愷明又雙叒叕發新作了，這次還是與圖靈獎得主 Yann LeCun 合作。

這項研究的主題是沒有歸一化層的 Transformer（Transformers without Normalization），並已被 CVPR 2025 會議接收。

Meta FAIR 研究科學家劉壯的推文

過去十年，歸一化層已經鞏固了其作為現代神經網路最基本元件之一的地位。這一切可以追溯到 2015 年批歸一化（batch normalization）的發明，它使視覺識別模型的收斂速度變得更快、更好，並在隨後幾年中獲得迅速發展。從那時起，研究人員針對不同的網路架構或領域提出了許多歸一化層的變體。

如今，幾乎所有現代網路都在使用歸一化層，其中層歸一化（Layer Norm，LN）是最受歡迎之一，特別是在占主導地位的 Transformer 架構中。

歸一化層的廣泛應用很大程度上得益於它們在最佳化方面的實證優勢。除了實現更好的結果之外，歸一化層還有助於加速和穩定收斂。隨著神經網路變得越來越寬、越來越深，歸一化層的必要性變得越來越重要。因此，研究人員普遍認為歸一化層對於有效訓練深度網路至關重要，甚至是必不可少的。這一觀點事實上得到了微妙證明：近年來，新架構經常尋求取代注意力層或卷積層，但幾乎總是保留歸一化層。

本文中，研究者提出了 Transformer 中歸一化層的一種簡單平替。他們的探索始於以下觀察：LN 層使用類 tanh 的 S 形曲線將其輸入對映到輸出，同時縮放輸入啟用並壓縮極值。

受此啟發，研究者提出了一種元素級運算，稱為 Dynamic Tanh（DyT），定義為：DyT (x) = tanh (αx)，其中 α 是一個可學習引數。此運算旨在透過 α 學習適當的縮放因子並透過有界 tanh 函式壓縮極值來模擬 LN 的行為。值得注意的是，與歸一化層不同，DyT 可以實現這兩種效果，而無需計算啟用資料。

論文一作 Jiachen Zhu 為紐約大學四年級博士生、二作陳鑫磊（Xinlei Chen）為 FAIR 研究科學家，專案負責人為劉壯。

論文標題：Transformers without Normalization
論文地址：https://arxiv.org/pdf/2503.10622
專案主頁：https://jiachenzhu.github.io/DyT/
GitHub 地址：https://github.com/jiachenzhu/DyT

DyT 使用起來非常簡單，如下圖 1 所示，研究者直接用 DyT 替換視覺和語言 Transformer 等架構中的現有歸一化層。實證結果表明，使用 DyT 的模型可以在各種設定中穩定訓練並獲得較高的最終效能。同時，DyT 通常不需要在原始架構上調整訓練超引數。

DyT 模組可以透過短短幾行 PyTorch 程式碼來實現。

該工作挑戰了「歸一化層對訓練現代神經網路必不可少」這一觀念，並提供了有關歸一化層屬性的實證見解。此外，初步結果表明，DyT 可以提升訓練和推理速度，從而成為以效率為導向的網路設計的候選方案。

劉壯發推稱，對他而言，歸一化層一直是深度學習中比較神秘的內容。這項工作讓他對歸一化層的作用有了更深的理解。另外，考慮到模型訓練和推理需要數千萬的算力需求，DyT 有潛力助力成本降低。他很期待接下來 DyT 的應用。

歸一化層有什麼作用？

要去掉 Transformer 中的歸一化層，首先要做的當然是瞭解歸一化層有什麼用。

該團隊透過實證研究對此進行了分析。為此，他們使用了三個不同的經過訓練的 Transformer 模型：一個 Vision Transformer（ViT-B）、一個 wav2vec 2.0 Large Transformer 和一個 Diffusion Transformer（DiT-XL）。

他們使用這三個模型取樣了一小批樣本，並讓其前向透過整個網路。然後，他們監測了其中歸一化層的輸入和輸出，即歸一化操作前後的張量。

由於 LN 會保留輸入張量的維度，因此可以在輸入和輸出張量元素之間建立一一對應關係，從而可以直接視覺化它們的關係。這個對映關係見圖 2。

具有層歸一化的類 tanh 對映。對於這三個模型，該團隊發現，它們的早期 LN 層（圖 2 第 1 列）的輸入 – 輸出關係基本上是線性的。但是，更深的 LN 層卻有更有趣的表現。

可以觀察到，這些曲線的形狀大多與 tanh 函式表示的完整或部分 S 形曲線非常相似（見圖 3）。

人們可能預期 LN 層會對輸入張量進行線性變換，因為減去平均值和除以標準差都是線性運算。LN 以每個 token 的方式進行歸一化，僅對每個 token 的啟用進行線性變換。

由於 token 具有不同的平均值和標準差，因此這種線性並不對輸入張量的所有啟用都成立。儘管如此，該團隊表示依然很驚訝：實際的非線性變換竟然與某個經過縮放的 tanh 函式高度相似！

對於這樣一個 S 型曲線，可以看到其中心部分（x 值接近零的部分）仍然主要呈線性形狀。大多數點（約 99%）都屬於這個線性範圍。但是，仍有許多點明顯超出此範圍，這些點被認為具有「極端」值，例如 ViT 模型中 x 大於 50 或小於 -50 的點。

歸一化層對這些值的主要作用是將它們壓縮為不太極端的值，從而與大多數點更加一致。這是歸一化層無法透過簡單的仿射變換層近似的地方。

該團隊假設，這種對極端值的非線性和不成比例的壓縮效應正是歸一化層的關鍵之處。

前段時間的一篇論文《On the Nonlinearity of Layer Normalization》同樣重點指出了 LN 層引入的強非線性，並且表明這種非線性可以增強模型的表徵能力。

此外，這種壓縮行為還反映了生物神經元對大輸入的飽和（saturation）特性，這種現象大約一個世紀前就已經被觀察到。

token 和通道的歸一化。LN 層如何對每個 token 執行線性變換，同時以這種非線性方式壓縮極端值呢？

為了理解這一點，該團隊分別按 token 和通道對這些點進行視覺化。圖 4 給出了 ViT 的第二和第三個子圖的情況，但為了更清晰，圖中使用了取樣的點子集。

在圖 4 左邊兩個小圖中，使用了同一顏色標記每個 token 的啟用。可以觀察到，任何單個 token 的所有點確實都會形成一條直線。但是，由於每個 token 都有不同的方差，因此斜率也不同。輸入 x 範圍較小的 token 往往具有較小的方差，並且歸一化層將使用較小的標準偏差來除它們的啟用，從而讓直線有較大的斜率。

總的來說，它們形成了一條類似於 tanh 函式的 S 形曲線。在右側的兩個小圖中，同樣使用相同的顏色標記各個通道的啟用。可以看到，不同通道的輸入範圍往往存在巨大差異，只有少數通道（例如紅色、綠色和粉色）會表現出較大的極端值 —— 而這些通道正是被歸一化層壓縮得最厲害的。

Dynamic Tanh（DyT）

既知根本，正當創新。基於歸一化層和擴充套件版 tanh 函式的相似性，該團隊提出了 Dynamic Tanh（DyT），並且這可以作為歸一化層的直接替代。

給定一個輸入張量 x，DyT 層的定義如下：

其中 α 是一個可學習的標量引數，允許根據輸入的範圍以不同的方式縮放輸入，並會考慮不同的 x 尺度（圖 2）。也因此，他們將整個操作命名為「動態」tanh。γ 和 β 是可學習的每通道向量引數，與所有歸一化層中使用的引數相同 —— 它們允許輸出縮放到任何尺度。這有時被視為單獨的仿射層；這裡，該團隊將它們視為 DyT 層的一部分，就像歸一化層也包括它們一樣。演算法 1 給出了用類 PyTorch 的虛擬碼實現的 DyT。

要想將 DyT 層整合到現有架構中，方法很簡單：直接用一個 DyT 層替換一個歸一化層（見圖 1）。這適用於注意力塊、FFN 塊和最終歸一化層內的歸一化層。

儘管 DyT 可能看起來像或可被視為啟用函式，但本研究僅使用它來替換歸一化層，而不會改變原始架構中啟用函式的任何部分，例如 GELU 或 ReLU。網路的其他部分也保持不變。該團隊還觀察到，幾乎不需要調整原始架構使用的超引數即可使 DyT 表現良好。

關於縮放參數。在這裡，總是簡單地將 γ 初始化為全一向量，將 β 初始化為全零向量，後接歸一化層。對於 scaler 引數 α，除了 LLM 訓練外，預設初始化為 0.5 通常就足夠了。除非另有明確說明，否則在後續的實驗中，α 均被初始化為 0.5。

說明。DyT 並非一種新型的歸一化層，因為它在前向傳遞過程中會獨立地對張量中的每個輸入元素進行操作，而無需計算統計資料或其他型別的聚合。但它確實保留了歸一化層的效果，即以非線性方式壓縮極端值，同時對輸入的中心部分執行近乎線性的變換。

DyT 在實驗中的表現

為了驗證 DyT 的效果，研究團隊在不同任務和領域中測試了 Transformer 及其他架構，將原始架構中的 LN 或 RMSNorm 替換為 DyT 層，並按照官方開源方案進行訓練和測試。

視覺監督學習

研究團隊在 ImageNet-1K 分類任務上訓練了 Base 和 Large 兩種規模的 Vision Transformer（ViT）和 ConvNeXt 模型。

選擇 ViT 和 ConvNeXt 是因為它們既具代表性，又分別採用不同機制：ViT 基於注意力機制，ConvNeXt 基於卷積操作。從表 1 的 Top-1 分類準確率來看，DyT 在兩種架構和不同規模模型上均優於 LN。圖 5 中展示的 ViT-B 和 ConvNeXt-B 的訓練損失曲線。

表 1：ImageNet-1K 上的監督分類準確率。DyT 在兩種架構和不同模型規模上均實現了優於或等同於 LN 的效能表現。

視覺自監督學習

研究團隊測試了兩種流行的視覺自監督學習方法：何愷明的 MAE 和 DINO。

這兩種方法都預設使用 Vision Transformer 作為骨幹網路，但訓練目標不同。MAE 使用重建損失進行訓練，而 DINO 則使用聯合嵌入損失。研究團隊先在 ImageNet-1K 資料集上進行無標籤預訓練，然後新增分類層並用標籤資料微調來測試預訓練模型。表 2 展示了微調的結果。在自監督學習任務中，DyT 和 LN 的表現基本持平。

擴散模型

研究者在 ImageNet-1K 上訓練了三個尺寸分別為 B、L 和 XL 的 DiT 模型。需要注意的是，在 DiT 中，LN 層的仿射引數用於類調節，DyT 實驗中也保留了這一引數，只是用 tanh (αx) 函式替換了歸一化遷移。訓練結束，如表 3 所示，與 LN 相比，DyT 的 FID 值相當或有所提高。

LLM

這些模型是按照 LLaMA 中概述的原始配方在帶有 200B tokens 的 The Pile 資料集上進行訓練的。在帶有 DyT 的 LLaMA 中，研究者在初始嵌入層之後添加了一個可學習的標量引數，並調整了 α 的初始值（第 7 節）。下表 4 報告了訓練後的損失值，並按照 OpenLLaMA 的方法，在 lm-eval 的 15 個零樣本任務上對模型進行了基準測試。如表 4 所示，在所有四種規模的模型中，DyT 的表現與 RMSNorm 相當。

圖 6 展示了損失曲線，顯示了所有模型大小的相似趨勢，訓練損失在整個訓練過程中都非常接近。

語音自監督學習。研究者在 LibriSpeech 資料集上預訓練了兩個 wav2vec 2.0 Transformer 模型。表 5 報告了最終的驗證損失。在兩種模型規模下，DyT 的表現都與 LN 相當。

DNA 序列建模

在長程 DNA 序列建模任務中，研究者對 HyenaDNA 模型和 Caduceus 模型進行了預訓練。結果如表 6，在這項任務中，DyT 保持了與 LN 相當的效能。

α 初始化

非 LLM 模型的 α 初始化

非 LLM 模型對 α_0 相對不敏感。圖 9 展示了在不同任務中改變 α_0 對驗證效能的影響。

α_0 越小，訓練越穩定。圖 10 展示了使用 ImageNet-1K 資料集對有監督 ViT 訓練穩定性的消減。

將 α_0 = 0.5 設為預設值。根據研究結果，研究者將 α_0 = 0.5 設定為所有非 LLM 模型的預設值。這種設定既能提供與 LN 相當的訓練穩定性，又能保持強大的效能。

LLM 模型的 α 初始化

調整 α_0 可以提高 LLM 效能。如前所述，預設設定 α_0 = 0.5 在大多數任務中表現良好。然而，研究者發現調整 α_0 可以大幅提高 LLM 效能。他們對每個 LLaMA 模型都進行了 30B tokens 的預訓練，並比較了它們的訓練損失，從而調整了它們的 α_0。

表 11 總結了每個模型的調整後 α_0 值，其中有兩個重要發現：

1. 較大的模型需要較小的 α_0 值。一旦確定了較小模型的最佳 α_0 值，就可以相應地縮小較大模型的搜尋空間；

2. 注意力塊的 α_0 值越高，效能越好。對注意力塊中的 DyT 層初始化較高的 α 值，而對其他位置（即 FFN 區塊內或最終線性投影之前）的 DyT 層初始化較低的 α 值，可以提高效能。

為了進一步說明 α_0 調整的影響，圖 11 展示了兩個 LLaMA 模型損失值的熱圖。這兩個模型都受益於注意力塊中較高的 α_0，從而減少了訓練損失。

模型寬度主要決定了 α_0 的選擇。我們還研究了模型寬度和深度對最優 α_0 的影響。研究者發現，模型寬度對確定最優 α_0 至關重要，而模型深度的影響則微乎其微。表 12 顯示了不同寬度和深度下的最佳 α_0 值，表明較寬的網路可以從較小的 α_0 值中獲益，從而獲得最佳效能。另一方面，模型深度對 α_0 的選擇影響微乎其微。

從表 12 中可以看出，網路越寬，「注意力」和「其他」所需的初始化就越不均衡。研究者假設，LLM 的 α 初始化的敏感度與其他模型相比過大的寬度有關。

更多研究細節，可參考原論文。

何愷明在MIT授課的課件PPT下載

在CVer公眾號後臺回覆：何愷明，即可下載本課程的所有566頁課件PPT！趕緊學起來！

CVPR 2025 論文和程式碼下載

在CVer公眾號後臺回覆：CVPR2025，即可下載CVPR 2025論文和程式碼開源的論文合集

ECCV 2024 論文和程式碼下載

在CVer公眾號後臺回覆：ECCV2024，即可下載ECCV 2024論文和程式碼開源的論文合集

CV垂直方向和論文投稿交流群成立

掃描下方二維碼，或者新增微訊號：CVer2233，即可新增CVer小助手微信，便可申請加入CVer-垂直方向和論文投稿微信交流群。另外其他垂直方向已涵蓋：目標檢測、影像分割、目標跟蹤、人臉檢測&識別、OCR、姿態估計、超解析度、SLAM、醫療影像、Re-ID、GAN、NAS、深度估計、自動駕駛、強化學習、車道線檢測、模型剪枝&壓縮、去噪、去霧、去雨、風格遷移、遙感影像、行為識別、影片理解、影像融合、影像檢索、論文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。

一定要備註：研究方向+地點+學校/公司+暱稱（如Mamba、多模態學習或者論文投稿+上海+上交+卡卡），根據格式備註，可更快被透過且邀請進群