Transformer模型能夠捕捉長(zhǎng)距離依賴和全局信息，在引入計(jì)算機(jī)視覺(jué)任務(wù)后，大多都取得了顯著的性能提升。

但Transformer的缺陷始終還是繞不過(guò)：時(shí)間和空間復(fù)雜度太高，都是輸入序列長(zhǎng)度的二次方。

通常情況下，一個(gè)輸入圖像被劃分為patch，然后flatten這些patch為一個(gè)token序列送入Transformer，序列越長(zhǎng)，復(fù)雜度也就越高。

所以，很多視覺(jué)任務(wù)中為了利用上Transformer，選擇將其應(yīng)用于低分辨率或?qū)⒆⒁饬C(jī)制限制在圖像局部。

但在高分辨率上應(yīng)用Transformer能夠帶來(lái)更廣闊的應(yīng)用前景和性能提升，因此，許多工作都在研究設(shè)計(jì)有效的Transformer以降低計(jì)算的復(fù)雜性。

有學(xué)者提出線性近似Transformer，用線性方法近似于標(biāo)準(zhǔn)的Transformer。然而，實(shí)驗(yàn)結(jié)果顯示這些線性Transformer在視覺(jué)任務(wù)中的性能是比較差的。為了降低計(jì)算成本，PVT使用降采樣的key和value，使得模型捕捉像素級(jí)細(xì)節(jié)的能力有所下降。相比之下，Swin變換器則是限制了全局注意力的交互信息來(lái)減少計(jì)算量

與以往的工作方向不同，來(lái)自西蒙菲莎大學(xué)和阿里巴巴AI Lab的研究人員提出了一個(gè)全新的注意力機(jī)制QuadTree attention，由粗到細(xì)地建立注意力機(jī)制，能夠同時(shí)包含全局交互和細(xì)粒度的信息，將時(shí)間復(fù)雜度降低為線性，論文已被ICLR 2022接收。

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

代碼地址：https://github.com/Tangshitao/QuadTreeAttention

當(dāng)我們看一張圖片的時(shí)候，可以發(fā)現(xiàn)，大多數(shù)圖像區(qū)域都是不相關(guān)的，所以我們可以建立一個(gè)token金字塔，以從粗到細(xì)的方式計(jì)算注意力。通過(guò)這種方式，如果對(duì)應(yīng)的粗級(jí)區(qū)域不相關(guān)，那么我們也可以快速跳過(guò)精細(xì)級(jí)別的不相關(guān)區(qū)域。

例如，第一層計(jì)算了圖像A中的藍(lán)色區(qū)域的注意力，即計(jì)算圖像A中的藍(lán)色patch與圖像B中的所有patch的注意力，并選擇前K個(gè)patch，把這些patch也被標(biāo)記為藍(lán)色，代表他們是相關(guān)的區(qū)域。

在第二層，對(duì)于圖像A中的第一層中藍(lán)色patch的四個(gè)子patch，我們只計(jì)算它們與對(duì)應(yīng)第一層圖像B中top K個(gè)patch的子patch的注意力，其他所有其他陰影的子patch都被跳過(guò)以減少計(jì)算量。我們將圖像A中的兩個(gè)patch用黃色和綠色表示，它們?cè)趫D像B中對(duì)應(yīng)的前K個(gè)patch也用同樣的顏色突出顯示。

整個(gè)過(guò)程在第三層迭代進(jìn)行，通過(guò)這種方式，既能獲得精細(xì)的注意力，還能夠保留長(zhǎng)距離的依賴連接。

最重要的是，在整個(gè)過(guò)程中只需要計(jì)算少量的注意力。因此，這種方法具有更低的內(nèi)存消耗和計(jì)算成本。

在實(shí)現(xiàn)上，研究人員采用了Quadtree 這種數(shù)據(jù)結(jié)構(gòu)去構(gòu)建注意力機(jī)制。

與傳統(tǒng)注意力機(jī)制一樣，首先將embeddings映射Q，K，V。然后用kernel size為2x2的pooling層或者卷積層將他們降采樣若干次構(gòu)建token金字塔。

從最粗的那層開(kāi)始，每一層只選K個(gè)最高注意力分?jǐn)?shù)的patch參與下一層的計(jì)算。根據(jù)計(jì)算方式的不同設(shè)計(jì)了2種機(jī)制：QuadTree-A與QuadTree-B。

在QuadTree-B方法中，對(duì)于最粗的那層，只需根據(jù)注意力公式計(jì)算。

對(duì)于其余幾層，則是從上一層選k個(gè)注意力分?jǐn)?shù)最高的patch，然后計(jì)算message passing。

最后把每層的信息結(jié)合起來(lái)即可，其中w_i是第i層可學(xué)參數(shù)。

實(shí)驗(yàn)結(jié)果

尋找不同圖像之間的特征對(duì)應(yīng)關(guān)系（feature corresponding）是一個(gè)經(jīng)典三維計(jì)算機(jī)視覺(jué)任務(wù)，通常的評(píng)估方式就是對(duì)應(yīng)點(diǎn)的相機(jī)姿態(tài)估計(jì)準(zhǔn)確率。

研究人員使用最近提出的SOTA框架LoFTR，其中包括一個(gè)基于CNN的特征提取器和一個(gè)基于Transformer的匹配器。

為了驗(yàn)證QuadTree Transformer的效果，研究人員將LoFTR中的線性變換器替換為QuadTree。此外，文章內(nèi)還實(shí)現(xiàn)了一個(gè)新版本的LoFTR與spatial reduction（SR）注意力進(jìn)行對(duì)比。

研究人員在包含1513個(gè)場(chǎng)景的ScanNet上進(jìn)行了實(shí)驗(yàn)。

對(duì)于QuadTree Transformer的參數(shù)，使用三層金字塔，最粗的分辨率為15×20個(gè)像素。在最精細(xì)的級(jí)別的參數(shù)K設(shè)置為8，而在較粗的級(jí)別上則為兩倍。對(duì)于SR注意力，將value和key tokens平均池化到8×8的大小以保證和QuadTree Attention相似的內(nèi)存消耗和flops。

在（5?，10?，20?）下相機(jī)姿勢(shì)誤差的AUC實(shí)驗(yàn)結(jié)果中可以看到，SR注意力與線性Transformer取得了類似的結(jié)果。相比之下，QuadTreeA 和QuadTreeB在很大程度上超過(guò)了線性Transformer和SR注意力，并且Quadtree-B 總體上比Quadtree-A表現(xiàn)得更好。

為了進(jìn)一步提高結(jié)果，研究人員還訓(xùn)練了一個(gè)K=16的模型，可以看到模型的性能得到進(jìn)一步提升。

在雙目視差估計(jì)（stereo matching）任務(wù)中，目的是在兩幅圖像之間找到對(duì)應(yīng)的線上的像素。最近的工作STTR將Transformer應(yīng)用于epipolar line之間的特征點(diǎn)，并取得了SOTA的性能。

在將STTR中的標(biāo)準(zhǔn)Transformer替換為QuadTree Transformer后，在Scene Flow FlyingThings3D合成數(shù)據(jù)集上進(jìn)行實(shí)驗(yàn)，該數(shù)據(jù)集包含25,466張分辨率為960×540的圖像。

研究人員建立了四層的金字塔來(lái)評(píng)估QuadTree Attention，實(shí)驗(yàn)結(jié)果可以看到非遮擋區(qū)域的EPE（End-Point-Error）和遮擋區(qū)域的IOU（Intersection-over- Union），表中還包括計(jì)算復(fù)雜性和內(nèi)存使用量也被報(bào)告。

與基于標(biāo)準(zhǔn)Transformer的STTR相比，QuadTree Transformer實(shí)現(xiàn)了類似的EPE（0.45 px vs 0.46 px）和更高的閉塞估計(jì)IOU，但計(jì)算和內(nèi)存成本低得多，只有52%的FLOPs和63%的內(nèi)存消耗。

研究人員還在基本的self-attention任務(wù)中測(cè)試了QuadTree Transformer的性能。

首先在ImageNet上的圖像分類任務(wù)實(shí)驗(yàn)結(jié)果中可以看到，基于PVTv2的模型，將其中的spatial reduction attention替換成quadtree attention，就能夠在ImageNet上實(shí)現(xiàn)了84.0%的top 1準(zhǔn)確度，在不同大小的模型上比PVTv2高0.4-1.5個(gè)百分點(diǎn)。

在COCO目標(biāo)檢測(cè)數(shù)據(jù)集的實(shí)驗(yàn)結(jié)果中可以看到，對(duì)于QuadTree Attention來(lái)說(shuō)，一個(gè)小的K就足夠捕捉粗到細(xì)的信息了。因此，在使用高分辨率的圖像時(shí)，可以減少更多計(jì)算量。

并且QuadTree-B實(shí)現(xiàn)了更高的性能，同時(shí)比PVTv2的flops少得多，而且性能也同時(shí)超過(guò)了ResNet和ResNeXt。QuadTree-B-b2的性能比ResNet101和ResNeXt101-32x4d分別高出7.7AP和6.3AP，骨干flops減少約40%。

在ADE20K的語(yǔ)義分割實(shí)驗(yàn)中，在相似的參數(shù)量與flops下，比PVTv2提升了0.8-1.3。

作者介紹

一作唐詩(shī)濤，現(xiàn)在西蒙菲莎大學(xué)三年級(jí)在讀博士，導(dǎo)師譚平，研究方向?yàn)樯疃葘W(xué)習(xí)，三維視覺(jué)。在ECCV、ICCV、CVPR、ICML、ICLR等會(huì)議上發(fā)表多篇論文。

共同一作張家輝，現(xiàn)任阿里巴巴算法工程師。2020年于清華大學(xué)取得博士學(xué)位，研究方向?yàn)槿S重建、三維深度學(xué)習(xí)。博士期間在Intel中國(guó)研究院及港科大實(shí)習(xí)或交流。在ECCV、ICCV、CVPR、ICLR、TPAMI、TVCG等會(huì)議或期刊上發(fā)表多篇論文。

朱思語(yǔ)博士，阿里云人工智能實(shí)驗(yàn)室算法團(tuán)隊(duì)負(fù)責(zé)人。他于香港科技大學(xué)獲得博士學(xué)位。在攻讀博士學(xué)位期間，共同創(chuàng)辦了3D視覺(jué)公司Altizure。朱思語(yǔ)博士在ICCV、CVPR、ECCV、PAMI等計(jì)算機(jī)視覺(jué)國(guó)際學(xué)術(shù)會(huì)議和期刊上發(fā)表30多篇論文。

譚平，現(xiàn)就職于阿里巴巴XR實(shí)驗(yàn)室，曾任西蒙菲莎大學(xué)終身副教授、新加坡國(guó)立大學(xué)副教授。主要研究興趣包括計(jì)算機(jī)視覺(jué)、計(jì)算機(jī)圖形學(xué)、機(jī)器人技術(shù)、3D 重建、基于圖像的建模、圖像和視頻編輯、照明和反射建模。

? ?