研究者表示，他們提出的多頭高效注意力 Hydra Attention 保留了注意力的可解釋性等優點，能夠同時提升基線 DeiT-B 模型的準確率和速度。
得益於自身的泛化性以及從大規模資料中學習的能力，Transformers 成為過去幾年自然語言處理領域的主導技術。並且隨著 Vision Transformers（ViTs）的出現，視覺領域也出現了類似的趨勢。但我們應該看到，在 NLP 中使用 BERT 或在視覺中使用 ViT 得到的 Transformers 純實例化並不是計算機視覺任務的主導。相反，更多基於 vision-specialized 注意力的架構（如 Swin 或 MViT）或注意力 – 卷積混合架構（如 LeViT）正在被使用。
究其原因，specialized ViT 可以通過更少算力實現更好的性能，它們要麼添加摺積層、要麼使用 vision-specific 局部窗口注意力，要麼使用其他方式低成本地添加視覺歸納偏置。雖然純 ViTs 在大型資料集上表現良好（如 ImageNet 上取得 90.45% top-1），但是當我們在多個下游任務所需的大圖像上應用一個模型時，純 Transformers 的主要機制（即多頭自注意力）可能面臨很大的瓶頸。
NLP 領域已經引入很多技術來解決這一問題。一些工作或通過「核心技巧」重新安排計算順序或通過對映到與 token 無關的低秩空間或二者兼而有之，來引入線性注意力（就 token 而言）。但是，這些線性注意力方法中的大多數將跨 token 的計算轉換成跨特徵的計算，導致成本很高。最近，Flash Attention 已經展示了多頭自注意力的 IO 高效實現優於大多數線性注意力方法，即便 token 數量 dadao 數千也是如此。
此外也有工作在視覺空間中嘗試使用高效注意力，但還沒有人在傳統的 ViT shell 中單獨進行探索。舉例而言，PolyNL 將注意力視為高效的三階多項式，但尚未在 ViT 架構中探索。Attention Free Transformer 提出一個同樣高效的 AFTSimple 變體，但在純 ViT 中表現不佳且需要卷積和位置編碼的額外支持。因此我們亟需一個真正高效、準確率高且泛化性強的多頭自注意力。
近日， Meta AI 和佐治亞理工學院的研究者填補了這一空白，提出了 Hydra Attention。該方法源於線性注意力中有些自相矛盾的行為，即在使用標準多頭自注意力時，向模型中添加更多的 head 會保持計算量不變。但是，線上性注意力中改變操作順序後，添加更多的 head 實際上降低了層的計算成本。
因此，研究者將這一觀察結果發揮到了極致，將模型中 head 的數量設置為與特徵數相同，從而創建了一個在計算上對於 token 和特徵都呈線性的注意力模組。
論文地址：https://arxiv.org/pdf/2209.07484.pdf
下圖為標準注意力、線性注意力和 Hydra Attention 的結構比較。
Hydra Attention 不僅比先前高效注意力工作的泛化性更強，而且在使用正確的核心時能夠顯著提升準確率。事實上，當與標準多頭注意力混合時，Hydra Attention 可以提高基線 DeiT-B 模型的準確率同時速度也更快。由於源自多頭注意力，Hydra Attention 保留了注意力的幾個優點，比如可解釋性以及對不同任務的泛化性。
然而，雖然 Hydra Attention 對於大圖像高效且泛化強，但在本文中只關注使用 DeiT-B 的 ImageNet 分類任務，通常只使用較小的 224×224 和 384×384 圖像。雖然對於這些圖像的效率提升不是太大（根據圖像大小不同可以實現 10-27% 的提升），但其他高效注意力方法遭受較大的準確率下降，Hydra Attention 則沒有。
未來，研究者希望可以探索在其他更多 token 密集領域中的應用，比如檢測、分割或視訊。作為一種通用技術，Hydra attention 不對 token 之間的關係做任何假設，因而可以進一步用來提升 token 稀疏應用的速度，比如掩碼預訓練或 token 剪枝。
研究者還希望 Hydra Attention 可以為更強大、更高效、更通用的純 transformers 模型打下基礎。
Hydra Attention
在標準的多頭自注意力中，其計算量隨圖像中 token 的數量呈二次增長。更具體地說，如果 T 是 token 數量，D 是特徵維度數量，那麼形成的注意力矩陣是 O(T^2D)的。當 T 很大時，這種操作在計算上變得不可行。
因此開始有研究者引入線性注意力。給定查詢 Q、鍵 K、值 V 等，標準 softmax 自注意力計算為：
根據文獻 [16] 的研究，公式（1）還可以用 sim(·)函數寫成這樣：
如果選擇一個特徵表示為ϕ(·)的可分解核心，那麼 sim(x, y) = ϕ(x)ϕ(y)^T，又可以得到：
然後通過結合律，改變計算順序，最終得到：
多頭注意力
然而公式（1）到（4）為假設為單頭注意力。想要多頭注意力，則（1）應變為：
公式（4）變為：
到底增加多少頭合適呢？大多數的研究使用 6 到 16 個頭，但頭的數量超過這個範圍，會發生什麼?
為了找出答案，該研究在 ImageNet-1k 上訓練 DeiT-B ，並使用帶有 softmax 的標準多頭自注意力（公式 5，MSA）或帶有餘弦相似度的多頭線性注意力來改變頭數 H (公式 7, MLA) 進行實驗，結果繪製在圖 2 中。在記憶體使用方面，MSA 在 H > 96 時記憶體不足，而 MLA 在 H < 3 時記憶體不足。
在性能方面，對於 MSA，當 H > 12 時，Softmax 注意力的準確率似乎會崩潰，而具有餘弦相似度的 MLA 的準確率基本保持一致，直到 H=768。令人驚訝的是，處於這個頭數時，H 等於 D，這意味著每個頭只有一個標量特徵可以使用！
hydra
如圖 2 所示，只要相似度函數 sim(x, y)不是 softmax，H 任意放大都是可行的。為了利用這一點，該研究引入了 hydra ，即設置 H = D：
在這種情況下，每個 Q_h、K_h、V_h 都是 R^T ×1 中的列向量，然後得到：
Hydra Attention 通過一個全局瓶頸混合資訊，而不是像標準的自注意力那樣進行顯式的 token-to-token 混合。這導致計算複雜度為：
實驗結果
對於所有實驗，除非另有說明，該研究使用 DeiT-B，並在 ImageNet-1k 上進行默認設置訓練。
在表 1 中，該研究探索了其他可能的 kernel 選擇。
表 2 展示了使用 Hydra attention 與標準注意以及其他方法對比的結果：
較大的圖像：在表 3 中，我們看到 Hydra Attention 在 2 層時可以提高 0.59% 的準確率和 4.1% 的吞吐量，或者在 7 層時保持準確率基本不變，提高 15.4% 的吞吐量。
侷限性：當 T（token） = 197 時，Hydra attention 比標準注意力快 197 倍，那麼為什麼最大 FLOP 計數減少只有 4%？事實證明，對於 ViT-B/16 224 × 224 圖像（T = 197，D = 768），只有 4.10% 的模型 FLOP 用於創建和應用注意力矩陣。使用 Hydra attention，這一比例降低到 0.02%，基本上消除了模型中的注意力成本。
