SFU、阿里提出通用QuadTree Attention,復雜度變線性,性能更強
Transformer模型能夠捕捉長距離依賴和全局信息,在引入計算機視覺任務(wù)后,大多都取得了顯著的性能提升。
但Transformer的缺陷始終還是繞不過:時間和空間復雜度太高,都是輸入序列長度的二次方。
通常情況下,一個輸入圖像被劃分為patch,然后flatten這些patch為一個token序列送入Transformer,序列越長,復雜度也就越高。
所以,很多視覺任務(wù)中為了利用上Transformer,選擇將其應(yīng)用于低分辨率或?qū)⒆⒁饬C制限制在圖像局部。
但在高分辨率上應(yīng)用Transformer能夠帶來更廣闊的應(yīng)用前景和性能提升,因此,許多工作都在研究設(shè)計有效的Transformer以降低計算的復雜性。
有學者提出線性近似Transformer,用線性方法近似于標準的Transformer。然而,實驗結(jié)果顯示這些線性Transformer在視覺任務(wù)中的性能是比較差的。為了降低計算成本,PVT使用降采樣的key和value,使得模型捕捉像素級細節(jié)的能力有所下降。相比之下,Swin變換器則是限制了全局注意力的交互信息來減少計算量
與以往的工作方向不同,來自西蒙菲莎大學和阿里巴巴AI Lab的研究人員提出了一個全新的注意力機制QuadTree attention,由粗到細地建立注意力機制,能夠同時包含全局交互和細粒度的信息,將時間復雜度降低為線性,論文已被ICLR 2022接收。
論文地址:https://arxiv.org/abs/2201.02767
代碼地址:https://github.com/Tangshitao/QuadTreeAttention
當我們看一張圖片的時候,可以發(fā)現(xiàn),大多數(shù)圖像區(qū)域都是不相關(guān)的,所以我們可以建立一個token金字塔,以從粗到細的方式計算注意力。通過這種方式,如果對應(yīng)的粗級區(qū)域不相關(guān),那么我們也可以快速跳過精細級別的不相關(guān)區(qū)域。
例如,第一層計算了圖像A中的藍色區(qū)域的注意力,即計算圖像A中的藍色patch與圖像B中的所有patch的注意力,并選擇前K個patch,把這些patch也被標記為藍色,代表他們是相關(guān)的區(qū)域。
在第二層,對于圖像A中的第一層中藍色patch的四個子patch,我們只計算它們與對應(yīng)第一層圖像B中top K個patch的子patch的注意力,其他所有其他陰影的子patch都被跳過以減少計算量。我們將圖像A中的兩個patch用黃色和綠色表示,它們在圖像B中對應(yīng)的前K個patch也用同樣的顏色突出顯示。
整個過程在第三層迭代進行,通過這種方式,既能獲得精細的注意力,還能夠保留長距離的依賴連接。
最重要的是,在整個過程中只需要計算少量的注意力。因此,這種方法具有更低的內(nèi)存消耗和計算成本。
在實現(xiàn)上,研究人員采用了Quadtree 這種數(shù)據(jù)結(jié)構(gòu)去構(gòu)建注意力機制。
與傳統(tǒng)注意力機制一樣,首先將embeddings映射Q,K,V。然后用kernel size為2x2的pooling層或者卷積層將他們降采樣若干次構(gòu)建token金字塔。
從最粗的那層開始,每一層只選K個最高注意力分數(shù)的patch參與下一層的計算。根據(jù)計算方式的不同設(shè)計了2種機制:QuadTree-A與QuadTree-B。
在QuadTree-B方法中,對于最粗的那層,只需根據(jù)注意力公式計算。
對于其余幾層,則是從上一層選k個注意力分數(shù)最高的patch,然后計算message passing。
最后把每層的信息結(jié)合起來即可,其中w_i是第i層可學參數(shù)。
實驗結(jié)果
尋找不同圖像之間的特征對應(yīng)關(guān)系(feature corresponding)是一個經(jīng)典三維計算機視覺任務(wù),通常的評估方式就是對應(yīng)點的相機姿態(tài)估計準確率。
研究人員使用最近提出的SOTA框架LoFTR,其中包括一個基于CNN的特征提取器和一個基于Transformer的匹配器。
為了驗證QuadTree Transformer的效果,研究人員將LoFTR中的線性變換器替換為QuadTree。此外,文章內(nèi)還實現(xiàn)了一個新版本的LoFTR與spatial reduction(SR)注意力進行對比。
研究人員在包含1513個場景的ScanNet上進行了實驗。
對于QuadTree Transformer的參數(shù),使用三層金字塔,最粗的分辨率為15×20個像素。在最精細的級別的參數(shù)K設(shè)置為8,而在較粗的級別上則為兩倍。對于SR注意力,將value和key tokens平均池化到8×8的大小以保證和QuadTree Attention相似的內(nèi)存消耗和flops。
在(5?,10?,20?)下相機姿勢誤差的AUC實驗結(jié)果中可以看到,SR注意力與線性Transformer取得了類似的結(jié)果。相比之下,QuadTreeA 和QuadTreeB在很大程度上超過了線性Transformer和SR注意力,并且Quadtree-B 總體上比Quadtree-A表現(xiàn)得更好。
為了進一步提高結(jié)果,研究人員還訓練了一個K=16的模型,可以看到模型的性能得到進一步提升。
在雙目視差估計(stereo matching)任務(wù)中,目的是在兩幅圖像之間找到對應(yīng)的線上的像素。最近的工作STTR將Transformer應(yīng)用于epipolar line之間的特征點,并取得了SOTA的性能。
在將STTR中的標準Transformer替換為QuadTree Transformer后,在Scene Flow FlyingThings3D合成數(shù)據(jù)集上進行實驗,該數(shù)據(jù)集包含25,466張分辨率為960×540的圖像。
研究人員建立了四層的金字塔來評估QuadTree Attention,實驗結(jié)果可以看到非遮擋區(qū)域的EPE(End-Point-Error)和遮擋區(qū)域的IOU(Intersection-over- Union),表中還包括計算復雜性和內(nèi)存使用量也被報告。
與基于標準Transformer的STTR相比,QuadTree Transformer實現(xiàn)了類似的EPE(0.45 px vs 0.46 px)和更高的閉塞估計IOU,但計算和內(nèi)存成本低得多,只有52%的FLOPs和63%的內(nèi)存消耗。
研究人員還在基本的self-attention任務(wù)中測試了QuadTree Transformer的性能。
首先在ImageNet上的圖像分類任務(wù)實驗結(jié)果中可以看到,基于PVTv2的模型,將其中的spatial reduction attention替換成quadtree attention,就能夠在ImageNet上實現(xiàn)了84.0%的top 1準確度,在不同大小的模型上比PVTv2高0.4-1.5個百分點。
在COCO目標檢測數(shù)據(jù)集的實驗結(jié)果中可以看到,對于QuadTree Attention來說,一個小的K就足夠捕捉粗到細的信息了。因此,在使用高分辨率的圖像時,可以減少更多計算量。
并且QuadTree-B實現(xiàn)了更高的性能,同時比PVTv2的flops少得多,而且性能也同時超過了ResNet和ResNeXt。QuadTree-B-b2的性能比ResNet101和ResNeXt101-32x4d分別高出7.7AP和6.3AP,骨干flops減少約40%。
在ADE20K的語義分割實驗中,在相似的參數(shù)量與flops下,比PVTv2提升了0.8-1.3。
作者介紹
一作唐詩濤,現(xiàn)在西蒙菲莎大學三年級在讀博士,導師譚平,研究方向為深度學習,三維視覺。在ECCV、ICCV、CVPR、ICML、ICLR等會議上發(fā)表多篇論文。
共同一作張家輝,現(xiàn)任阿里巴巴算法工程師。2020年于清華大學取得博士學位,研究方向為三維重建、三維深度學習。博士期間在Intel中國研究院及港科大實習或交流。在ECCV、ICCV、CVPR、ICLR、TPAMI、TVCG等會議或期刊上發(fā)表多篇論文。
朱思語博士,阿里云人工智能實驗室算法團隊負責人。他于香港科技大學獲得博士學位。在攻讀博士學位期間,共同創(chuàng)辦了3D視覺公司Altizure。朱思語博士在ICCV、CVPR、ECCV、PAMI等計算機視覺國際學術(shù)會議和期刊上發(fā)表30多篇論文。
譚平,現(xiàn)就職于阿里巴巴XR實驗室,曾任西蒙菲莎大學終身副教授、新加坡國立大學副教授。主要研究興趣包括計算機視覺、計算機圖形學、機器人技術(shù)、3D 重建、基于圖像的建模、圖像和視頻編輯、照明和反射建模。
? ?