[[437909]]

視覺 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] 中的對象分類數據集上進行了訓練，但它似乎能夠學習同一類別中的實例之間共享的語義概念和實例特定的特征。

感興趣的讀者可以閱讀論文原文，了解更多研究細節。