作者丨happy??編輯丨極市平臺

導讀

?

本文介紹了復旦大學&華為諾亞提出的一種新穎的softmax-free的Transformer—SOFT。所提SOFT顯著改善了現有ViT方案的計算效率，更為關鍵的是：SOFT的線性復雜度可以允許更長的token序列，進而取得更佳的精度-復雜度均衡。

論文鏈接：https://arxiv.org/pdf/2110.11945.pdf

代碼鏈接：https://github.com/fudan-zvg/SOFT

項目鏈接：https://fudan-zvg.github.io/SOFT/

本文是復旦大學&華為諾亞關于Transformer中自注意力機制復雜度的深度思考，首次提出了一種新穎的softmax-free 的Transformer 。本文從softmax self-attention局限性出發，分析了其存在的挑戰；然后由此提出了線性復雜度的SOFT；再針對線性SOFT存在的訓練問題，提出了一種具有理論保證的近似方案。所提SOFT在ImageNet分類任務上取得了比已有CNN、Transformer更佳的精度-復雜度均衡。

Abstract

ViT通過圖像塊序列化+自注意力機制將不同CV任務性能往前推了一把。然而，自注意力機制會帶來更高的計算復雜度與內存占用。在NLP領域已有不同的方案嘗試采用線性復雜度對自注意力進行近似。然而，本文的深入分析表明：NLP中的近似方案在CV中缺乏理論支撐或者無效。

我們進一步分析了其局限性根因：softmax self-attention 。具體來說，傳統自注意力通過計算token之間的點乘并歸一化得到自注意力。softmax操作會對后續的線性近似帶來極大挑戰。基于該發現，本文首次提出了SOFT(softmax-free transformer )。

為移除自注意力中的softmax，我們采用高斯核函數替代點乘相似性且無需進一步的歸一化。這就使得自注意力矩陣可以通過低秩矩陣分析近似 。近似的魯棒性可以通過計算其MP逆(Moore-Penrose Inverse)得到。

ImageNet數據集上的實驗結果表明：所提SOFT顯著改善了現有ViT方案的計算效率 。更為關鍵的是：SOFT的線性復雜度可以允許更長的token序列，進而取得更佳的精度-復雜度均衡。

Contributation

本文的貢獻主要包含以下幾點：

• 提出一種新穎的線性空間、時間復雜度 的softmax-free Transformer ；

• 所提注意力矩陣近似可以通過具有理論保證的矩陣分解算法 計算得到；

• 所提SOFT在ImageNet圖像分類任務上取得了比其他ViT方案更佳的精度-復雜度均衡 (見下圖a)。

Method

Softmax-free self-attention formulation

上圖給出了本文所提SOFT架構示意圖。我們首先來看一下該注意力模塊的設計。給定包含n個token的輸入序列，自注意力旨在挖掘所有token對之間的相關性 。

具體來說，X首先線性投影為三個維的query、key以及values：

自注意力可以表示為如下廣義形式：

自注意力的關鍵函數包含一個非線性函數與一個相關函數。自注意力的常規配置定義如下：

雖然該softmax自注意力已成為首選且很少受到質疑，但是它并不適合進行線性化。為構建線性自注意力設計，我們引入了一種sfotmax-free自注意力函數：通過高斯核替換點乘操作。定義如下：

為保持注意力矩陣的對稱性，我們設置投影矩陣相同，即。所提自注意力矩陣定義如下：

為描述的簡單性，我們定義為矩陣形式：。所提自注意力矩陣S具有三個重要屬性：

• 對稱性

• 所有元素均在[0,1]范圍內；

• 所有對角元素具有最大值1。

我們發現：當采用無線性化的核自注意力矩陣時，transformer的訓練難以收斂 。這也就解釋了：為何softmax自注意力在transformer中如此流行。

Low-rank regularization via matrix decomposition with linear complxity

為解決收斂于二次復雜度問題，我們利用了矩陣分解作為帶低秩正則的統一解，這就使得模型復雜度大幅下降，且無需計算全部的自注意力矩陣。

作出上述選擇因為在于：S為半正定矩陣，且無需后接歸一化。我們將S表示為塊矩陣形式：

其中，。通過上述分解，注意力矩陣可以近似表示為：

其中，表示A的MP逆。更多關于MP逆的信息建議查看原文，這里略過。

在上述公式，A和B是S通過隨機采樣m個token得到的子矩陣，可表示為(我們將其稱之為bottleneck token )。然而，我們發現：隨機采樣對于m非常敏感。因此，我們通過利用結構先驗探索兩種額外的方案：

• 采用一個核尺寸為k、stride為k的卷積學習；

• 采用一個核尺寸為k、stride為k的均值池化生成。

通過實驗對比發現：卷積層學習 具有更好的精度 。由于K與Q相等，因此。給定m個token，我們計算A和P：

最終，我們得到了SOFT的正則化后的自注意力矩陣：

上圖Algorithm1給出所提SOFT流程，它涉及到了MP逆計算。一種精確且常用的計算MP逆的方法是SVD，然而SVD對于GPU訓練不友好。為解決該問題，我們采用了Newton-Raphson方法，見上圖Algorithm2：一種迭代算法。與此同時，作者還給出了上述迭代可以最終收斂到MP逆的證明。對該證明感興趣的同時強烈建議查看原文公式，哈哈。

Instantiations

上面主要聚焦于softmax-free self-attention 模塊的介紹，接下來我們將介紹如何利用SOFT模塊構建Transformer模型。我們以圖像分類任務為切入點，以PVT作為基礎架構并引入所提SOFT模塊構建最終的SOFT模型，同時還在stem部分進行了微小改動。下表給出了本文所提方案在不同容量大小下的配置信息。

Experiments

上表對比了所提方案與現有線性Transformer模型的性能，從中可以看到：

• 相比基線Transformer，線性Transformer能夠大幅降低內存占用與FLOPs，同時保持相當參數量；

• 所提SOFT在所有線性方案中取得了最佳分類精度；

• 所提SOFT與其他線性方案的推理速度相當，訓練速度稍慢。

上圖給出了不同方案的序列長度與內存占用之間的關系，從中可以看到：所提SOFT確實具有線性復雜度的內存占用 。

上表給出了所提方案與其他CNN、Transformer的性能對比，從中可以看到：

• 總體來說，ViT及其變種具有比CNN更高的分類精度；

• 相比ViT、DeiT等Transformer方法以及RegNet等CNN方法，所提SOFT取得了最佳性能；

• 相比PVT，所提方案具有更高的分類精度，直接驗證了所提SOFT模塊的有效性；

• 相比Twins與Swin，所提SOFT具有相當的精度，甚至更優性能。

此外，作者還在NLP任務上進行了對比，見上表，很明顯：SOFT又一次勝出 。

往期精彩回顧適合初學者入門人工智能的路線及資料下載機器學習及深度學習筆記等資料打印機器學習在線手冊深度學習筆記專輯《統計學習方法》的代碼復現專輯 AI基礎下載黃海廣老師《機器學習課程》視頻課黃海廣老師《機器學習課程》711頁完整版課件

本站qq群554839127，加入微信群請掃碼：

總結

以上是生活随笔為你收集整理的【深度学习】去掉softmax后Transformer会更好吗？复旦华为诺亚提出SOFT：轻松搞定线性近似...的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。