Visual Transformer（ViT）作為計(jì)算機(jī)視覺領(lǐng)域的新興霸主，已經(jīng)在各個(gè)研究任務(wù)中逐漸替換掉了卷積神經(jīng)網(wǎng)絡(luò)CNN。

?

ViT與CNN之間存在著許多不同點(diǎn)，例如ViT的輸入是image patch，而非像素；分類任務(wù)中，ViT是通過對(duì)類標(biāo)記（class token）進(jìn)行決策等等。

?

class token實(shí)際上是ViT論文原作者提出，用于整合模型輸入信息的token。class token與每個(gè)patch進(jìn)行信息交互后，模型就能了解到具體的分類信息。

?

并且在自注意力機(jī)制中，最后一層中的softmax可以作為注意力圖，根據(jù)class token和不同patch之間的交互程度，就能夠了解哪些patch對(duì)最終分類結(jié)果有影響及具體程度，也增加了模型可解釋性。

?

但這種可解釋性目前仍然是很弱的，因?yàn)閜atch和最后一層的softmax之間還隔著很多層和很多個(gè)header，信息之間的不斷融合后，很難搞清楚最后一層softmax是否真的可以解釋分類。

?

所以如果ViT和CNN一樣有視覺屬性就好了！

?

最近Meta AI就提出了一個(gè)新模型，用attention map來增強(qiáng)卷積神經(jīng)網(wǎng)絡(luò)，說簡(jiǎn)單點(diǎn)，其實(shí)就是用了一個(gè)基于注意力的層來取代常用的平均池化層。

?

仔細(xì)一想，池化層和attention好像確實(shí)很配啊，都是對(duì)輸入信息的加權(quán)平均進(jìn)行整合。加入了注意力機(jī)制以后的池化層，可以明確地顯示出不同patch所占的權(quán)重。

?

并且與經(jīng)典ViT相比，每個(gè)patch都會(huì)獲得一個(gè)單一的權(quán)重，無需考慮多層和多頭的影響，這樣就可以用一個(gè)簡(jiǎn)單的方法達(dá)到對(duì)注意力可視化的目的了。

?

?

在分類任務(wù)中更神奇，如果對(duì)每個(gè)類別使用不同顏色進(jìn)行單獨(dú)標(biāo)記的話，就會(huì)發(fā)現(xiàn)分類任務(wù)也能識(shí)別出圖片中的不同物體。

?

?

基于Attention的池化層

文章中新提出的模型叫做PatchConvNet，核心組件就是可學(xué)習(xí)的、基于attention的池化層。

?

?

模型架構(gòu)的主干是一個(gè)卷積網(wǎng)絡(luò)，相當(dāng)于是一個(gè)輕量級(jí)的預(yù)處理操作，它的作用就是把圖像像素進(jìn)行分割，并映射為一組向量，和ViT中patch extraction操作對(duì)應(yīng)。

?

最近也有研究表明，采用卷積的預(yù)處理能讓模型的性能更加穩(wěn)定。

?

模型的第二部分column（主干trunk），包含了整個(gè)模型中的大部分層、參數(shù)和計(jì)算量，它由N個(gè)堆疊的殘差卷積塊組成。每個(gè)塊由一個(gè)歸一化、1*1卷積，3*3卷積用來做空間處理，一個(gè)squeeze-and-excitation層用于混合通道特征，最后在殘差連接前加入一個(gè)1*1的卷積。

?

研究人員對(duì)模型塊的選擇也提出了一些建議，例如在batch size夠大的情況下，BatchNorm往往效果比LayerNorm更好。但訓(xùn)練大模型或者高分辨率的圖像輸入時(shí)，由于batch size更小，所以BatchNorm在這種情況下就不太實(shí)用了。

?

下一個(gè)模塊就是基于注意力的池化層了。

?

在主干模型的輸出端，預(yù)處理后的向量通過類似Transformer的交叉注意力層（cross attention layer）的方式進(jìn)行融合。

?

注意力層中的每個(gè)權(quán)重值取決于預(yù)測(cè)patch與可訓(xùn)練向量（CLS）之間的相似度，結(jié)果和經(jīng)典ViT中的class token類似。

?

然后將產(chǎn)生的d維向量添加到CLS向量中，并經(jīng)過一個(gè)前饋網(wǎng)絡(luò)處理。

?

與之前提出的class-attention decoder不同之處在于，研究人員僅僅只用一個(gè)block和一個(gè)head，大幅度簡(jiǎn)化了計(jì)算量，也能夠避免多個(gè)block和head之間互相影響，從而導(dǎo)致注意力權(quán)重失真。

?

因此，class token和預(yù)處理patch之間的通信只發(fā)生在一個(gè)softmax中，直接反映了池化操作者如何對(duì)每個(gè)patch進(jìn)行加權(quán)。

?

也可以通過將CLS向量替換為k×d矩陣來對(duì)每個(gè)類別的attention map進(jìn)行歸一化處理，這樣就可以看出每個(gè)塊和每個(gè)類別之間的關(guān)聯(lián)程度。

?

但這種設(shè)計(jì)也會(huì)增加內(nèi)存的峰值使用量，并且會(huì)使網(wǎng)絡(luò)的優(yōu)化更加復(fù)雜。通常只在微調(diào)優(yōu)化的階段以一個(gè)小的學(xué)習(xí)率和小batch size來規(guī)避這類問題。

?

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

在圖像分類任務(wù)上，研究人員首先將模型與ImageNet1k和ImageNet-v2上的其他模型從參數(shù)量，F(xiàn)LOPS，峰值內(nèi)存用量和256張圖像batch size下的模型推理吞吐量上進(jìn)行對(duì)比。

?

?

實(shí)驗(yàn)結(jié)果肯定是好的，可以看到PatchConvNet的簡(jiǎn)單柱狀結(jié)構(gòu)（column architecture）相比其他模型更加簡(jiǎn)便和易于擴(kuò)展。對(duì)于高分辨率圖像來說，不同模型可能會(huì)針對(duì)FLOPs和準(zhǔn)確率進(jìn)行不同的平衡，更大的模型肯定會(huì)取得更高的準(zhǔn)確率，相應(yīng)的吞吐量就會(huì)低一些。

?

在語義分割任務(wù)上，研究人員通過ADE20k數(shù)據(jù)集上的語義分割實(shí)驗(yàn)來評(píng)估模型，數(shù)據(jù)集中包括2萬張訓(xùn)練圖像和5千張驗(yàn)證圖像，標(biāo)簽超過150個(gè)類別。由于PatchConvNet模型不是金字塔式的，所以模型只是用模型的最后一層輸出和UpperNet的多層次網(wǎng)絡(luò)輸出，能夠簡(jiǎn)化模型參數(shù)。研究結(jié)果顯示，雖然PatchConvNet的結(jié)構(gòu)更簡(jiǎn)單，但與最先進(jìn)的Swin架構(gòu)性能仍處于同一水平，并且在FLOPs-MIoU權(quán)衡方面優(yōu)于XCiT。

?

?

?

在檢測(cè)和實(shí)例分割上，研究人員在COCO數(shù)據(jù)集上對(duì)模型進(jìn)行評(píng)估，實(shí)驗(yàn)結(jié)果顯示PatchConvNet相比其他sota架構(gòu)來說，能夠在FLOPs和AP之間進(jìn)行很好的權(quán)衡。

?

?

在消融實(shí)驗(yàn)中，為了驗(yàn)證架構(gòu)問題，研究人員使用不同的架構(gòu)對(duì)比了Transformer中的class attention和卷積神經(jīng)網(wǎng)絡(luò)的平均池化操作，還對(duì)比了卷積主干和線性投影之間的性能差別等等。實(shí)驗(yàn)結(jié)果可以看到卷積主干是模型取得最佳性能的關(guān)鍵，class-attention幾乎沒有帶來額外的性能提升。

?

?

另一個(gè)重要的消融實(shí)驗(yàn)時(shí)attention-based pooling和ConvNets之間的對(duì)比，研究人員驚奇地發(fā)現(xiàn)可學(xué)習(xí)的聚合函數(shù)甚至可以提高一個(gè)ResNet魔改后模型的性能。

?

通過把a(bǔ)ttention添加到ResNet50中，直接在Imagenet1k上獲得了80.1%的最高準(zhǔn)確率，比使用平均池化層的baseline模型提高了+0.3%的性能，并且attention-based只稍微增加了模型的FLOPs數(shù)量，從4.1B提升到4.6B。