機器之心專欄
作者：馬昊宇
在這項研究中，來自百度研究院和香港大學的研究者重新思考了局部自注意力機制，提出了特征空間局部注意力（feature-space local attention或簡稱FSLA）。
Vision Transformer 舍棄了 ConvNet 先驗資訊，通過引入自注意力機制對遠距離特征依賴進行模組化，提升了模型的表征能力。然而 Vision Transformer 的自注意力機制在圖像分辨率較高時，計算複雜度過高。為了克服這個問題，研究人員使用局部視窗計算自注意力，在此稱之為圖像空間局部注意力（image-space local attention 或簡稱 ISLA）。盡管基于視窗的圖像空間局部注意力顯著提升了效率，但仍面臨難以捕捉遠距離特征依賴的問題。
在這項研究中，來自百度研究院和香港大學的研究者重新思考了局部自注意力機制，提出了特征空間局部注意力（feature-space local attention 或簡稱 FSLA）。這種局部注意力從圖像内容出發，把特征相似的 token 聚成類，并且隻在每類特征的内部計算自注意力，相比全局自注意力顯著降低了計算量，同時基本保留了原始的全局自注意力機制對遠距離特征依賴的模組化能力。
為了将特征空間局部注意力與圖像空間局部注意力相結合，本文作者進一步提出了雙邊局部注意力 ViT （簡稱 BOAT)，把特征空間局部注意力子產品加入到現有的基于視窗的局部注意力視覺 Transformer 模型中，作為圖像空間局部注意力的補充，大大提升了針對遠距離特征依賴的模組化能力，在幾個基準資料集上的大量實驗表明結合了特征空間局部注意力的模型明顯優于現有的 ConvNet 和 ViT 模型。
論文位址：https://arxiv.org/pdf/2201.13027v1
代碼位址：https://github.com/mahaoyuHKU/pytorch-boat
創新動機
為了保持更高的特征圖分辨率，同時不會帶來過高的運算複雜度，現有的圖像空間局部注意力機制将一個圖像劃分為多個局部視窗，自注意力隻在同一視窗的 token 間運算。這是一個合理的設計，因為一個 token 很可能與空間上鄰近的 token 相關聯。是以，局限于局部視窗的自注意力很可能不會顯著降低性能，但是可以顯著降低計算量。
本文重新思考了局部自注意力，從特征角度而非空間角度，實作了對局部視窗的劃分。具體來說，圖像空間局部自注意力的依據是：在空間上鄰近的 token 很可能對彼此施加更大的影響（圖 1 左）；而本文提出的特征空間局部自注意力機制的依據是：即使在圖像空間距離較遠但在特征空間距離較近的 token 同樣會對彼此有很大的影響，是以它在特征空間對 token 進行聚類，并且隻在每類特征的内部如同空間局部視窗一樣計算自注意力（圖 1 右）。
本文提出的特征空間局部自注意力僅計算特征空間内最近鄰的特征向量間的注意力，将距離較遠的特征向量間的影響直接設為 0。這本質上定義了一個分段相似度函數，将相似度小的特征向量間的注意力近似為 0，降低了運算複雜度。與圖像空間局部自注意力相比，特征空間局部自注意力在 ViT 模型中運用的較少。特征空間局部自注意力關注的是相似度較高的特征向量間的注意力，而不考慮空間上兩者的鄰近程度。是以，它是圖像空間局部自注意力的很好補充，能對因跨越空間局部視窗而被遺漏的遠距離特征依賴進行模組化。
方法概述
本文中的 ViT 采用了和 Swin 和 CSWin 相同的階層化金字塔架構，由一個 patch embedding 子產品和若幹個雙邊局部注意力（bilateral local attention）子產品組成。
本文與 Swin 和 CSWin 的主要差別是其中的 local attention 子產品被替換成了下圖所示的 bilateral local attention。而 patch embedding，position encoding 等設計皆和 Swin/CSWin 保持一緻，是以接下來對 bilateral local attention 進行詳細介紹。
Bilateral Local Attention
本文提出的 bilateral local attention 在基于視窗的圖像空間局部注意力（ISLA）模型中添加了特征空間局部注意力（FSLA）子產品。FSLA 子產品根據 ISLA 子產品的輸出計算在特征空間彼此鄰近的 token 之間的注意力：
最後，将 FSLA 子產品的輸出送入另一個歸一化層和一個 MLP 子產品進行處理，再通過一個短路連接配接得到整個 bilateral local attention 子產品的輸出：
FSLA 的重點是如何對特征進行聚類操作，并且在各個類内部計算自注意力。最直覺的方法是使用 K-means 聚類，但 K-means 聚類不能確定分組結果大小相同，這使得在 GPU 平台上難以有效地實作并行加速，同時也可能對自注意力計算的有效性産生負面影響。
是以本文提出均衡層次聚類，它進行 k 層聚類。在每一層，它進行均衡二分聚類，将上一層的各個類組均衡地劃分為兩個更小的類組。如下圖所示，所有 token 分成了 token 數量相同的 8 個類組，然後在每組内部計算自注意力，具體的自注意力參數和圖像空間局部注意力保持一緻。
假如某個類組原先有 2m 個 token，均衡二分聚類後得到的每組的 token 數量為 m。與 K-means 類似，均衡二分聚類是一個疊代算法并且依賴于聚類中心。如以下算法所示，在每次疊代對所有 token 進行分組時，先計算每個 token 到兩個聚類中心的距離比值，然後把所有 token 按距離比值的遞減順序排序，最後将排序清單前半部分 m 個 token 賦給第一組，後半部分 m 個 token 賦給第二組。
需要注意的是，這樣進行無重疊的均衡二分聚類可能會導緻兩個處于排序清單中段位置的、特征比較相似的 token 被配置設定到兩個不同的類組中，進而無法計算它們之間的互相影響。是以在實際計算中，為了避免遺漏鄰近特征間的影響，會保留一定程度的類間重疊，也就是把排序清單的最前面 m+n 個 token 賦給第一組，最後 m+n 個 token 賦給第二組。這樣兩組之間就存在 2n 個重用的 token，這樣的類間重疊會導緻額外的運算，是以實際隻在層次聚類的最後一層進行有重疊的均衡二分聚類。完成聚類以後，在每組 token 内部按照正常操作進行自注意力機制的計算即可。
值得注意的是，本文中的所有聚類都是臨時計算的，不包含任何可學習的參數，是以不存在對聚類算法本身進行梯度回傳的問題。此外，所有聚類運算都用 GPU 進行了加速，對模型的整體計算量影響不大。
實驗結果
BOAT 遵循與其它 ViT 相同的訓練政策。本文使用 ImageNet-1K 的訓練集訓練模型，輸入圖像使用 224×224 分辨率，并且沒有外部資料。
具體來說，訓練 300 個 epochs，使用 AdamW 優化器、餘弦學習速率排程器和一個線性預熱過程。BOAT 在多個資料集上都取得了 SOTA 的效果。比如，在 ImageNet-1K 測試集上，BOAT-CSWin-T 取得了 83.7 的 Top-1 分類準确率；在 ADE20K 語義分割測試集上，BOAT-CSWin-T 的 mIoU 達到了 50.5。