©作者 | 陳醒濠
Transformer 已經成為了語言和視覺任務中常用的基礎架構之一。然而，由於 Transformer 結構高計算開銷的影響，其在端側等資源受限裝置中的應用依然面臨很大的挑戰。我們對 Transformer 結構中的標準化層和注意力機制兩個模組的最佳化策略進行了深入探索，從而構建一個高效的 Transformer 結構。其中，LayerNorm 作為 Transformer 結構中常用的標準化層，但模型推理時仍需計算資料的統計值，導致了推理的低效。
我們提出了漸進式的 LayerNorm 替換策略，並對標準的 BatchNorm 進行了改進以更好地取代 LayerNorm 層。同時，我們採用了一種簡單高效的線性注意力模組（Simplified Linear Attention），來獲得更強的模型效能。我們將這兩種策略的結合簡稱為 SLAB。
我們在影像分類、目標檢測以及語言任務上都進行了大量的實驗，獲得了很好的效果。例如，我們的 SLAB-Swin-S 在 ImageNet1k 資料集上獲得了 83.6% 的分類精度，相對 Flatten-Swin-S 在精度提升 0.1% 的情況下，時延減少了 2.4ms。
論文標題：
SLAB: Efficient Transformers with Simplified Linear Attention and Progressive Re-parameterized Batch Normalization
論文連結：
https://arxiv.org/abs/2405.11582
PyTorch連結：
https://github.com/xinghaochen/SLAB
Mindspore程式碼：
https://github.com/mindspore-lab/models/tree/master/research/huawei-noah/SLAB
方法
1.1 漸進式重引數化BatchNorm
LN 作為 Transformer 中常用的標準化層結構，由於其在訓練和推理兩階段均存在均值和方差的計算，影響了 Transformer 的執行速度。與之相對，BN 僅在訓練階段存在均值和方差的計算，且在推理階段可與相鄰的線性層融合，可以去除標準化層對模型推理速度的影響。但是，在 Transformer 結構中將 LN 簡單替換為 BN 訓練會導致模型精度下降以及訓練崩潰等問題。為解決這個問題，我們對 BN 進行了最佳化，並提出了漸進式重引數化批標準化策略。
首先，重引數化批標準化的定義如下：
式中， 是一個可學習引數。其中， RepBN 可以透過調節 BN 的權值和偏移量，是特定層 BN 操作被跳過；當 為 0 時，RepBN 等效為純 BN 結構。同時，RepBN 能重引數化為 BN 的表現形式，並實現與相鄰線性層的融合。
其次，為增強 BN 在 Transformer 結構中的訓練穩定性，我們引入了漸進式替換策略。其表示形式如下：
式中， 是一個超引數，用於控制 LN 和 RepBN 的輸出比例。在訓練開始階段， 一般設定為 1，此時 LN 在模型中發揮主導作用；在訓練結束階段， 將衰減至 0，此時模型將轉變為純 BN 組成的結構。在實際應用中，我們採用了簡單的線性替換策略， 的值輸出如下：
其中， 為訓練中包含 LN 的總訓練步數， 為模型當前的訓練步數。相對於其他衰減策略，我們發現線性策略更為簡單且高效。因此，後續實驗中我們均採用了線性衰減的策略。
1.2 簡單線性注意力
Attention 是 Transformer 網路中重要的模組之一。為進一步壓縮模型計算量，我們引入了線性注意力模組。在該模組中，我們僅使用了硬體親和的 ReLU 運算元作為相似度函式，並增加了一個深度可分離模組增強區域性特徵提取。該簡單線性注意力模組（simplified linear attention, SLA）形式如下：
式中，DWC 表示深度可分離卷積。
實驗結果
2.1 分類任務
我們在 ImageNet1k 資料集上進行了實驗，實驗結果證明在多個 backbone 上，我們的 PRepBN 均獲得了與 LN 相當甚至更好的效能。從實驗結果看，相當基於 LN 的模型，PRepBN 模型的分類精度有 0.1%~1.4% 的提升。而基於我們 SLAB 的模型，能在精度與 Flatten Transformer 相當的情況下，減少模型的推理的時延。
2.2 檢測任務
此外，我們驗證了不同 backbone 在 COCO 資料集上的效果。從實驗結果可以看出，我們的方法實現了與原 Backbone 模型相當的效能，但擁有更低的模型推理時延。
2.3 語言任務
我們基於 Adaptive inputs 方法在 Wikitext-103 資料集上評測了 PRepBN 在語言任務的能力。同時，我們也將 PRepBN 應用在了 LlaMA-350M 模型中，並評測了模型在下游任務的效能。從實驗結果可以看出，我們的 PRepBN 方法在語言任務上也表現出了不錯的效能，精度無損的情況下將 LLaMA-350M 速度從 44 tokens/s 提升到了 50.4 tokens/s。
總結
我們對 Transformer 結構中的標準化層和注意力機制兩個模組的最佳化策略進行了深入探索，提出了漸進式的 LayerNorm 替換策略，同時採用一種簡單高效的線性注意力模組，來獲得更加高效的 Transformer 模型架構。這個方法在影像分類、目標檢測以及語言任務上進行了大量的實驗驗證，在精度無損的情況下，大幅提升 Transformer 的推理效率。
更多閱讀
#投 稿 通 道#
讓你的文字被更多人看到
如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。
總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋樑，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。
PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學術熱點剖析、科研心得或競賽經驗講解等。我們的目的只有一個，讓知識真正流動起來。
📝 稿件基本要求：
• 文章確係個人原創作品，未曾在公開渠道發表，如為其他平臺已發表或待發表的文章，請明確標註
• 稿件建議以 markdown 格式撰寫，文中配圖以附件形式傳送，要求圖片清晰，無版權問題
• PaperWeekly 尊重原作者署名權，並將為每篇被採納的原創首發稿件，提供業內具有競爭力稿酬，具體依據文章閱讀量和文章質量階梯制結算
📬 投稿通道：
• 投稿郵箱：[email protected]
• 來稿請備註即時聯絡方式（微信），以便我們在稿件選用的第一時間聯絡作者
• 您也可以直接新增小編微信（pwbot02）快速投稿，備註：姓名-投稿
△長按新增PaperWeekly小編
🔍
現在，在「知乎」也能找到我們了
進入知乎首頁搜尋「PaperWeekly」
點選「關注」訂閱我們的專欄吧
·
·
·
·