解決Transformer固有缺陷：復(fù)旦大學(xué)等提出線性復(fù)雜度SOFT

來自復(fù)旦大學(xué)、薩里大學(xué)和華為諾亞方舟實(shí)驗(yàn)室的研究者首次提出一種無 softmax Transformer。

視覺 Transformer (ViT) 借助 patch-wise 圖像標(biāo)記化和自注意力機(jī)制已經(jīng)在各種視覺識別任務(wù)上實(shí)現(xiàn)了 SOTA。然而，自注意力模塊的使用使得 Transformer 類模型的空間和時間復(fù)雜度都是 O(n^2)。自然語言處理領(lǐng)域的研究者們已經(jīng)進(jìn)行了各種讓 self-attention 計算逼近線性復(fù)雜度的嘗試。

近日，來自復(fù)旦大學(xué)、薩里大學(xué)和華為諾亞方舟實(shí)驗(yàn)室的研究者在一項研究中經(jīng)過深入分析表明，這些嘗試要么在理論上存在缺陷，要么在實(shí)驗(yàn)中對視覺識別無效，并進(jìn)一步發(fā)現(xiàn)這些方法的局限性在于在近似過程中仍然保持 softmax 自注意力。具體來說，傳統(tǒng)的自注意力是通過對標(biāo)記特征向量之間的縮放點(diǎn)積（scaled dot-product）進(jìn)行歸一化來計算的。保持這種 softmax 操作阻礙了線性化 Transformer 的復(fù)雜度?；诖?，該研究首次提出了一種無 softmax Transformer（softmax-free transformer，SOFT）。

為了去除 self-attention 中的 softmax，使用高斯核函數(shù)（Gaussian kernel function）代替點(diǎn)積相似度，無需進(jìn)一步歸一化。這使得可以通過低秩矩陣分解來近似一個完整的自注意力矩陣。通過使用 Newton-Raphson 方法計算其 Moore-Penrose 逆來實(shí)現(xiàn)近似的穩(wěn)健性。ImageNet 上的大量實(shí)驗(yàn)表明，SOFT 顯著提高了現(xiàn)有 ViT 變體的計算效率。至關(guān)重要的是，對于線性復(fù)雜性，SOFT 中允許更長的 token 序列，從而在準(zhǔn)確性和復(fù)雜性之間實(shí)現(xiàn)卓越的權(quán)衡。

論文地址：https://arxiv.org/abs/2110.11945

項目地址：https://github.com/fudan-zvg/SOFT

Transformer 模型存在一個瓶頸，即計算和內(nèi)存使用的二次復(fù)雜度。這是自注意力機(jī)制的內(nèi)在特征：給定一系列 token（例如，單詞或圖像塊）作為輸入，自注意力模塊通過將一個 token 與所有其他 token 相關(guān)聯(lián)來迭代地學(xué)習(xí)特征表示。這導(dǎo)致計算（時間）和內(nèi)存（空間）中 token 序列長度為 n 的二次復(fù)雜度 O(n 2 )，因?yàn)樵谕评磉^程中需要計算和保存 n × n 大小的注意力矩陣。這個問題在視覺中尤為嚴(yán)重：即使空間分辨率適中，在 tokenization 的 2D 圖像也會產(chǎn)生比 NLP 中的序列長得多的序列。因此，這種二次復(fù)雜性阻止了 ViT 模型以高空間分辨率對圖像進(jìn)行建模，這對于視覺識別任務(wù)通常是至關(guān)重要的。

一種自然的解決方案是通過近似來降低自注意力計算的復(fù)雜性。事實(shí)上，在 NLP 中已經(jīng)有很多嘗試 [33, 5, 18, 38]。例如，[33] 采取了一種天真的方法，通過可學(xué)習(xí)的預(yù)測來縮短 Key 和 Value 的長度。這種粗略的近似將不可避免地導(dǎo)致性能下降。相比之下，[5, 17] 都利用內(nèi)核機(jī)制來近似 softmax 歸一化，以線性化自注意力中的計算。[18] 取而代之的是采用散列策略來選擇性地計算最相似的對。最近，[38] 使用 Nystr?m 矩陣分解通過多項式迭代重建完整的注意力矩陣，以逼近地標(biāo)矩陣的偽逆。

盡管如此，softmax 歸一化在矩陣分解過程中只是簡單地重復(fù)，這在理論上是不可靠的。該研究通過實(shí)驗(yàn)發(fā)現(xiàn)，當(dāng)應(yīng)用于視覺時，這些方法都不是有效的（參見第 4.2 節(jié)）。該研究發(fā)現(xiàn)現(xiàn)有高效 Transformer 的局限性是由使用 softmax self-attention 引起的，并首次提出了一種無 softmax 的 Transformer。更具體地說，在所有現(xiàn)有的 Transformer（有或沒有線性化）中，在 token 特征向量之間的縮放點(diǎn)積之上需要一個 softmax 歸一化。保持這種 softmax 操作挑戰(zhàn)任何后續(xù)的線性化工作。

為了克服這個障礙，該研究提出了一種新的無 softmax 的自注意力機(jī)制，命名為 SOFT，在空間和時間上具有線性復(fù)雜度 O(n)。具體來說，SOFT 使用 Gaussian kernel 來定義相似度（self-attention）函數(shù)，不需要后續(xù)的 softmax 歸一化。有了這個 softmax-free 注意力矩陣，該研究進(jìn)一步引入了一種新的低秩矩陣分解算法來逼近。通過采用 Newton-Raphson 方法可靠地計算矩陣的 Moore-Penrose 逆，理論上可以保證近似的穩(wěn)健性。

該研究的主要貢獻(xiàn)包括：

· 提出了一種具有線性空間和時間復(fù)雜度的新型 softmax-free Transformer；

· 該研究的注意力矩陣近似是通過一種具有理論保證的新型矩陣分解算法來實(shí)現(xiàn)的；

· 為了評估該方法在視覺識別任務(wù)上的性能，該研究使用 SOFT 作為核心自注意力組件設(shè)計了一系列具有不同能力的通用骨干架構(gòu)。大量實(shí)驗(yàn)表明，具有線性復(fù)雜性（圖 1b），SOFT 模型可以將更長的圖像 token 序列作為輸入。因此，在模型大小相同的情況下，SOFT 在準(zhǔn)確度 / 復(fù)雜度權(quán)衡方面優(yōu)于 ImageNet [9] 分類上最先進(jìn)的 CNN 和 ViT 變體（圖 1a）。

下圖 2 給出了該模型的示意圖。

圖 2：所提出的無 softmax 自注意力 (SOFT) 方法的示意圖。P.E.：位置嵌入。虛線：線性投影。dh：每個注意力頭的隱藏暗淡。? 表示矩陣點(diǎn)積。

作者采用了兩個實(shí)驗(yàn)設(shè)置。在第一個設(shè)置下，對于所有方法，該研究使用相同的 Tiny（表 2）架構(gòu)進(jìn)行公平比較。也就是說，用每個基線自己的注意力塊替換 SOFT 中的核心自注意力塊，而架構(gòu)的其余部分保持不變。請注意，[35] 的空間縮減模塊是 Linformer [34] 的特例。研究者將減速比設(shè)置為與該方法相同。使用相同的統(tǒng)一采樣思想，該研究將 Nystr?mformer（用于 NLP 任務(wù)）的 1D 窗口平均替換為 2D 平均池化（用于圖像）。下采樣率與該研究的方法的保持一致。還值得一提的是，Reformer [19] 沒有官方代碼發(fā)布，本地敏感哈希（LSH）模塊對輸入 token 的長度有嚴(yán)格的要求，因此該研究的比較中不包括這種方法。

從下表 1 可以觀察到：

· 與 Tiny 架構(gòu)上的 Transformer 相比，Linear Transformer 方法大大減少了內(nèi)存和 FLOP，同時保持了相似的參數(shù)大小；

· SOFT 方法在所有線性化方法中實(shí)現(xiàn)了最好的分類精度；

· 該方法的推理速度與其他線性 Transformer 相當(dāng)，訓(xùn)練速度比 Nystromformer 稍慢，并且都比 Performer 和 Linformer 慢。

研究者指出：該模型的訓(xùn)練速度緩慢主要是由于 Newton-Raphson 迭代，它只能按順序應(yīng)用以確保 Moore-Penrose 逆的準(zhǔn)確性?？傊捎谕鹊耐评硭俣?，研究者認(rèn)為訓(xùn)練成本的增加是值得為卓越的準(zhǔn)確性付出的代價。

該研究與最先進(jìn)的替代方案進(jìn)行比較，并報告 ImageNet-1K 驗(yàn)證集上的 top-1 準(zhǔn)確率。FLOP 的計算批大小為 1024。從圖 1a 和表 3 中得出以下觀察結(jié)果：(i) 總體而言，ViT 及其變體比 CNN 產(chǎn)生更好的分類準(zhǔn)確度。(ii) 該研究在最近基于純視覺 Transformer 的方法中取得了最佳性能，包括 ViT [11] 和 DeiT [31]，以及最先進(jìn)的 CNN RegNet [26]。(iii)SOFT 在所有變體中都優(yōu)于最相似的（在架構(gòu)配置中）Transformer 對應(yīng)物 PVT [35]。由于注意力模塊是主要區(qū)別，這直接驗(yàn)證了該模型的有效性。(iv) 該方法還擊敗了旨在解決 ViT 效率限制的最新 ViT 變體 Twins，并且所需的參數(shù)和浮點(diǎn)計算都更少。

為了深入了解如何使用 SOFT 及替代方法學(xué)習(xí)注意力，圖 3 顯示了各種比較模型的注意力掩碼。對于每個模型，論文中給出了前兩個注意力頭的輸出。很明顯，SOFT 在捕捉像素之間的局部和長距離關(guān)系方面表現(xiàn)出魯棒性和多功能性。有趣的是，盡管 SOFT 在 ImageNet [9] 中的對象分類數(shù)據(jù)集上進(jìn)行了訓(xùn)練，但它似乎能夠?qū)W習(xí)同一類別中的實(shí)例之間共享的語義概念和實(shí)例特定的特征。