計算機視覺研究院專欄
作者：Edison_G
來自複旦大學、薩裡大學和華為諾亞方舟實驗室的研究者首次提出一種無 softmax Transformer。
視覺 Transformer (ViT) 借助 patch-wise 圖像标記化和自注意力機制已經在各種視覺識别任務上實作了 SOTA。然而，自注意力子產品的使用使得 Transformer 類模型的空間和時間複雜度都是 O(n^2)。自然語言處理領域的研究者們已經進行了各種讓 self-attention 計算逼近線性複雜度的嘗試。
近日，來自複旦大學、薩裡大學和華為諾亞方舟實驗室的研究者在一項研究中經過深入分析表明，這些嘗試要麼在理論上存在缺陷，要麼在實驗中對視覺識别無效，并進一步發現這些方法的局限性在于在近似過程中仍然保持 softmax 自注意力。具體來說，傳統的自注意力是通過對标記特征向量之間的縮放點積（scaled dot-product）進行歸一化來計算的。保持這種 softmax 操作阻礙了線性化 Transformer 的複雜度。基于此，該研究首次提出了一種無 softmax Transformer（softmax-free transformer，SOFT）。
為了去除 self-attention 中的 softmax，使用高斯核函數（Gaussian kernel function）代替點積相似度，無需進一步歸一化。這使得可以通過低秩矩陣分解來近似一個完整的自注意力矩陣。通過使用 Newton-Raphson 方法計算其 Moore-Penrose 逆來實作近似的穩健性。ImageNet 上的大量實驗表明，SOFT 顯着提高了現有 ViT 變體的計算效率。至關重要的是，對于線性複雜性，SOFT 中允許更長的 token 序列，進而在準确性和複雜性之間實作卓越的權衡。
- 論文位址：https://arxiv.org/abs/2110.11945
- 項目位址：https://github.com/fudan-zvg/SOFT
Transformer 模型存在一個瓶頸，即計算和記憶體使用的二次複雜度。這是自注意力機制的内在特征：給定一系列 token（例如，單詞或圖像塊）作為輸入，自注意力子產品通過将一個 token 與所有其他 token 相關聯來疊代地學習特征表示。這導緻計算（時間）和記憶體（空間）中 token 序列長度為 n 的二次複雜度 O(n 2 )，因為在推理過程中需要計算和儲存 n × n 大小的注意力矩陣。這個問題在視覺中尤為嚴重：即使空間分辨率适中，在 tokenization 的 2D 圖像也會産生比 NLP 中的序列長得多的序列。是以，這種二次複雜性阻止了 ViT 模型以高空間分辨率對圖像進行模組化，這對于視覺識别任務通常是至關重要的。
一種自然的解決方案是通過近似來降低自注意力計算的複雜性。事實上，在 NLP 中已經有很多嘗試 [33, 5, 18, 38]。例如，[33] 采取了一種天真的方法，通過可學習的預測來縮短 Key 和 Value 的長度。這種粗略的近似将不可避免地導緻性能下降。相比之下，[5, 17] 都利用核心機制來近似 softmax 歸一化，以線性化自注意力中的計算。[18] 取而代之的是采用散列政策來選擇性地計算最相似的對。最近，[38] 使用 Nyström 矩陣分解通過多項式疊代重建完整的注意力矩陣，以逼近地标矩陣的僞逆。
盡管如此，softmax 歸一化在矩陣分解過程中隻是簡單地重複，這在理論上是不可靠的。該研究通過實驗發現，當應用于視覺時，這些方法都不是有效的（參見第 4.2 節）。該研究發現現有高效 Transformer 的局限性是由使用 softmax self-attention 引起的，并首次提出了一種無 softmax 的 Transformer。更具體地說，在所有現有的 Transformer（有或沒有線性化）中，在 token 特征向量之間的縮放點積之上需要一個 softmax 歸一化。保持這種 softmax 操作挑戰任何後續的線性化工作。
為了克服這個障礙，該研究提出了一種新的無 softmax 的自注意力機制，命名為 SOFT，在空間和時間上具有線性複雜度 O(n)。具體來說，SOFT 使用 Gaussian kernel 來定義相似度（self-attention）函數，不需要後續的 softmax 歸一化。有了這個 softmax-free 注意力矩陣，該研究進一步引入了一種新的低秩矩陣分解算法來逼近。通過采用 Newton-Raphson 方法可靠地計算矩陣的 Moore-Penrose 逆，理論上可以保證近似的穩健性。
該研究的主要貢獻包括：
- 提出了一種具有線性空間和時間複雜度的新型 softmax-free Transformer；
- 該研究的注意力矩陣近似是通過一種具有理論保證的新型矩陣分解算法來實作的；
- 為了評估該方法在視覺識别任務上的性能，該研究使用 SOFT 作為核心自注意力元件設計了一系列具有不同能力的通用骨幹架構。大量實驗表明，具有線性複雜性（圖 1b），SOFT 模型可以将更長的圖像 token 序列作為輸入。是以，在模型大小相同的情況下，SOFT 在準确度 / 複雜度權衡方面優于 ImageNet [9] 分類上最先進的 CNN 和 ViT 變體（圖 1a）。
下圖 2 給出了該模型的示意圖。
圖 2：所提出的無 softmax 自注意力 (SOFT) 方法的示意圖。P.E.：位置嵌入。虛線：線性投影。dh：每個注意力頭的隐藏暗淡。◦ 表示矩陣點積。
作者采用了兩個實驗設定。在第一個設定下，對于所有方法，該研究使用相同的 Tiny（表 2）架構進行公平比較。也就是說，用每個基線自己的注意力塊替換 SOFT 中的核心自注意力塊，而架構的其餘部分保持不變。請注意，[35] 的空間縮減子產品是 Linformer [34] 的特例。研究者将減速比設定為與該方法相同。使用相同的統一采樣思想，該研究将 Nyströmformer（用于 NLP 任務）的 1D 視窗平均替換為 2D 平均池化（用于圖像）。下采樣率與該研究的方法的保持一緻。還值得一提的是，Reformer [19] 沒有官方代碼釋出，本地敏感哈希（LSH）子產品對輸入 token 的長度有嚴格的要求，是以該研究的比較中不包括這種方法。
從下表 1 可以觀察到：
- 與 Tiny 架構上的 Transformer 相比，Linear Transformer 方法大大減少了記憶體和 FLOP，同時保持了相似的參數大小；
- SOFT 方法在所有線性化方法中實作了最好的分類精度；
- 該方法的推理速度與其他線性 Transformer 相當，訓練速度比 Nystromformer 稍慢，并且都比 Performer 和 Linformer 慢。
研究者指出：該模型的訓練速度緩慢主要是由于 Newton-Raphson 疊代，它隻能按順序應用以確定 Moore-Penrose 逆的準确性。總之，由于同等的推理速度，研究者認為訓練成本的增加是值得為卓越的準确性付出的代價。
該研究與最先進的替代方案進行比較，并報告 ImageNet-1K 驗證集上的 top-1 準确率。FLOP 的計算批大小為 1024。從圖 1a 和表 3 中得出以下觀察結果：(i) 總體而言，ViT 及其變體比 CNN 産生更好的分類準确度。(ii) 該研究在最近基于純視覺 Transformer 的方法中取得了最佳性能，包括 ViT [11] 和 DeiT [31]，以及最先進的 CNN RegNet [26]。(iii)SOFT 在所有變體中都優于最相似的（在架構配置中）Transformer 對應物 PVT [35]。由于注意力子產品是主要差別，這直接驗證了該模型的有效性。(iv) 該方法還擊敗了旨在解決 ViT 效率限制的最新 ViT 變體 Twins，并且所需的參數和浮點計算都更少。
為了深入了解如何使用 SOFT 及替代方法學習注意力，圖 3 顯示了各種比較模型的注意力掩碼。對于每個模型，論文中給出了前兩個注意力頭的輸出。很明顯，SOFT 在捕捉像素之間的局部和長距離關系方面表現出魯棒性和多功能性。有趣的是，盡管 SOFT 在 ImageNet [9] 中的對象分類資料集上進行了訓練，但它似乎能夠學習同一類别中的執行個體之間共享的語義概念和執行個體特定的特征。
感興趣的讀者可以閱讀論文原文，了解更多研究細節。
© THE END
轉載請聯系本公衆号獲得授權
計算機視覺研究院學習群等你加入！
計算機視覺研究院主要涉及深度學習領域，主要緻力于人臉檢測、人臉識别，多目标檢測、目标跟蹤、圖像分割等研究方向。研究院接下來會不斷分享最新的論文算法新架構，我們這次改革不同點就是，我們要着重”研究“。之後我們會針對相應領域分享實踐過程，讓大家真正體會擺脫理論的真實場景，培養愛動手程式設計愛動腦思考的習慣！
掃碼關注
計算機視覺研究院