引入特征空間，顯著降低計(jì)算量：雙邊局部注意力ViT性能媲美全局注意力

在這項(xiàng)研究中，來(lái)自百度研究院和香港大學(xué)的研究者重新思考了局部自注意力機(jī)制，提出了特征空間局部注意力（feature-space local attention或簡(jiǎn)稱(chēng)FSLA）。

Vision Transformer 舍棄了 ConvNet 先驗(yàn)信息，通過(guò)引入自注意力機(jī)制對(duì)遠(yuǎn)距離特征依賴(lài)進(jìn)行建模，提升了模型的表征能力。然而 Vision Transformer 的自注意力機(jī)制在圖像分辨率較高時(shí)，計(jì)算復(fù)雜度過(guò)高。為了克服這個(gè)問(wèn)題，研究人員使用局部窗口計(jì)算自注意力，在此稱(chēng)之為圖像空間局部注意力（image-space local attention 或簡(jiǎn)稱(chēng) ISLA）。盡管基于窗口的圖像空間局部注意力顯著提升了效率，但仍面臨難以捕捉遠(yuǎn)距離特征依賴(lài)的問(wèn)題。
在這項(xiàng)研究中，來(lái)自百度研究院和香港大學(xué)的研究者重新思考了局部自注意力機(jī)制，提出了特征空間局部注意力（feature-space local attention 或簡(jiǎn)稱(chēng) FSLA）。這種局部注意力從圖像內(nèi)容出發(fā)，把特征相似的 token 聚成類(lèi)，并且只在每類(lèi)特征的內(nèi)部計(jì)算自注意力，相比全局自注意力顯著降低了計(jì)算量，同時(shí)基本保留了原始的全局自注意力機(jī)制對(duì)遠(yuǎn)距離特征依賴(lài)的建模能力。
為了將特征空間局部注意力與圖像空間局部注意力相結(jié)合，本文作者進(jìn)一步提出了雙邊局部注意力 ViT （簡(jiǎn)稱(chēng) BOAT)，把特征空間局部注意力模塊加入到現(xiàn)有的基于窗口的局部注意力視覺(jué) Transformer 模型中，作為圖像空間局部注意力的補(bǔ)充，大大提升了針對(duì)遠(yuǎn)距離特征依賴(lài)的建模能力，在幾個(gè)基準(zhǔn)數(shù)據(jù)集上的大量實(shí)驗(yàn)表明結(jié)合了特征空間局部注意力的模型明顯優(yōu)于現(xiàn)有的 ConvNet 和 ViT 模型。

• 論文地址：https://arxiv.org/pdf/2201.13027v1

• 代碼地址：https://github.com/mahaoyuHKU/pytorch-boat

創(chuàng)新動(dòng)機(jī)
為了保持更高的特征圖分辨率，同時(shí)不會(huì)帶來(lái)過(guò)高的運(yùn)算復(fù)雜度，現(xiàn)有的圖像空間局部注意力機(jī)制將一個(gè)圖像劃分為多個(gè)局部窗口，自注意力只在同一窗口的 token 間運(yùn)算。這是一個(gè)合理的設(shè)計(jì)，因?yàn)橐粋€(gè) token 很可能與空間上鄰近的 token 相關(guān)聯(lián)。因此，局限于局部窗口的自注意力很可能不會(huì)顯著降低性能，但是可以顯著降低計(jì)算量。
本文重新思考了局部自注意力，從特征角度而非空間角度，實(shí)現(xiàn)了對(duì)局部窗口的劃分。具體來(lái)說(shuō)，圖像空間局部自注意力的依據(jù)是：在空間上鄰近的 token 很可能對(duì)彼此施加更大的影響（圖 1 左）；而本文提出的特征空間局部自注意力機(jī)制的依據(jù)是：即使在圖像空間距離較遠(yuǎn)但在特征空間距離較近的 token 同樣會(huì)對(duì)彼此有很大的影響，因此它在特征空間對(duì) token 進(jìn)行聚類(lèi)，并且只在每類(lèi)特征的內(nèi)部如同空間局部窗口一樣計(jì)算自注意力（圖 1 右）。

本文提出的特征空間局部自注意力僅計(jì)算特征空間內(nèi)最近鄰的特征向量間的注意力，將距離較遠(yuǎn)的特征向量間的影響直接設(shè)為 0。這本質(zhì)上定義了一個(gè)分段相似度函數(shù)，將相似度小的特征向量間的注意力近似為 0，降低了運(yùn)算復(fù)雜度。與圖像空間局部自注意力相比，特征空間局部自注意力在 ViT 模型中運(yùn)用的較少。特征空間局部自注意力關(guān)注的是相似度較高的特征向量間的注意力，而不考慮空間上兩者的鄰近程度。因此，它是圖像空間局部自注意力的很好補(bǔ)充，能對(duì)因跨越空間局部窗口而被遺漏的遠(yuǎn)距離特征依賴(lài)進(jìn)行建模。
方法概述
本文中的 ViT 采用了和 Swin 和 CSWin 相同的層次化金字塔架構(gòu)，由一個(gè) patch embedding 模塊和若干個(gè)雙邊局部注意力（bilateral local attention）模塊組成。

本文與 Swin 和 CSWin 的主要區(qū)別是其中的 local attention 模塊被替換成了下圖所示的 bilateral local attention。而 patch embedding，position encoding 等設(shè)計(jì)皆和 Swin/CSWin 保持一致，所以接下來(lái)對(duì) bilateral local attention 進(jìn)行詳細(xì)介紹。

Bilateral Local Attention
本文提出的 bilateral local attention 在基于窗口的圖像空間局部注意力（ISLA）模型中添加了特征空間局部注意力（FSLA）模塊。FSLA 模塊根據(jù) ISLA 模塊的輸出計(jì)算在特征空間彼此鄰近的 token 之間的注意力：

最后，將 FSLA 模塊的輸出送入另一個(gè)歸一化層和一個(gè) MLP 模塊進(jìn)行處理，再通過(guò)一個(gè)短路連接得到整個(gè) bilateral local attention 模塊的輸出：

FSLA 的重點(diǎn)是如何對(duì)特征進(jìn)行聚類(lèi)操作，并且在各個(gè)類(lèi)內(nèi)部計(jì)算自注意力。最直覺(jué)的方法是使用 K-means 聚類(lèi)，但 K-means 聚類(lèi)不能確保分組結(jié)果大小相同，這使得在 GPU 平臺(tái)上難以有效地實(shí)現(xiàn)并行加速，同時(shí)也可能對(duì)自注意力計(jì)算的有效性產(chǎn)生負(fù)面影響。
因此本文提出均衡層次聚類(lèi)，它進(jìn)行 k 層聚類(lèi)。在每一層，它進(jìn)行均衡二分聚類(lèi)，將上一層的各個(gè)類(lèi)組均衡地劃分為兩個(gè)更小的類(lèi)組。如下圖所示，所有 token 分成了 token 數(shù)量相同的 8 個(gè)類(lèi)組，然后在每組內(nèi)部計(jì)算自注意力，具體的自注意力參數(shù)和圖像空間局部注意力保持一致。

假如某個(gè)類(lèi)組原先有 2m 個(gè) token，均衡二分聚類(lèi)后得到的每組的 token 數(shù)量為 m。與 K-means 類(lèi)似，均衡二分聚類(lèi)是一個(gè)迭代算法并且依賴(lài)于聚類(lèi)中心。如以下算法所示，在每次迭代對(duì)所有 token 進(jìn)行分組時(shí)，先計(jì)算每個(gè) token 到兩個(gè)聚類(lèi)中心的距離比值，然后把所有 token 按距離比值的遞減順序排序，最后將排序列表前半部分 m 個(gè) token 賦給第一組，后半部分 m 個(gè) token 賦給第二組。

需要注意的是，這樣進(jìn)行無(wú)重疊的均衡二分聚類(lèi)可能會(huì)導(dǎo)致兩個(gè)處于排序列表中段位置的、特征比較相似的 token 被分配到兩個(gè)不同的類(lèi)組中，從而無(wú)法計(jì)算它們之間的相互影響。因此在實(shí)際計(jì)算中，為了避免遺漏鄰近特征間的影響，會(huì)保留一定程度的類(lèi)間重疊，也就是把排序列表的最前面 m+n 個(gè) token 賦給第一組，最后 m+n 個(gè) token 賦給第二組。這樣兩組之間就存在 2n 個(gè)重用的 token，這樣的類(lèi)間重疊會(huì)導(dǎo)致額外的運(yùn)算，因此實(shí)際只在層次聚類(lèi)的最后一層進(jìn)行有重疊的均衡二分聚類(lèi)。完成聚類(lèi)以后，在每組 token 內(nèi)部按照常規(guī)操作進(jìn)行自注意力機(jī)制的計(jì)算即可。
值得注意的是，本文中的所有聚類(lèi)都是臨時(shí)計(jì)算的，不包含任何可學(xué)習(xí)的參數(shù)，因此不存在對(duì)聚類(lèi)算法本身進(jìn)行梯度回傳的問(wèn)題。此外，所有聚類(lèi)運(yùn)算都用 GPU 進(jìn)行了加速，對(duì)模型的整體計(jì)算量影響不大。
實(shí)驗(yàn)結(jié)果
BOAT 遵循與其它 ViT 相同的訓(xùn)練策略。本文使用 ImageNet-1K 的訓(xùn)練集訓(xùn)練模型，輸入圖像使用 224×224 分辨率，并且沒(méi)有外部數(shù)據(jù)。
具體來(lái)說(shuō)，訓(xùn)練 300 個(gè) epochs，使用 AdamW 優(yōu)化器、余弦學(xué)習(xí)速率調(diào)度器和一個(gè)線性預(yù)熱過(guò)程。BOAT 在多個(gè)數(shù)據(jù)集上都取得了 SOTA 的效果。比如，在 ImageNet-1K 測(cè)試集上，BOAT-CSWin-T 取得了 83.7 的 Top-1 分類(lèi)準(zhǔn)確率；在 ADE20K 語(yǔ)義分割測(cè)試集上，BOAT-CSWin-T 的 mIoU 達(dá)到了 50.5。

轉(zhuǎn)載來(lái)源：機(jī)器之心