轉(zhuǎn)載自：新智元 | 編輯：小咸魚(yú) 好困
【導(dǎo)讀】谷歌改造Vision Transformer的新作被NeurIPS 2021收錄了。在這篇文章里，谷歌提出了TokenLearner方法，Vision Transformer用上它最多可以降低8倍計(jì)算量，而分類(lèi)性能反而更強(qiáng)！目前，Transformer模型在計(jì)算機(jī)視覺(jué)任務(wù)（包括目標(biāo)檢測(cè)和視頻分類(lèi)等任務(wù)）中獲得了最先進(jìn)的結(jié)果。不同于逐像素處理圖像的標(biāo)準(zhǔn)卷積方法，Vision Transformer（ViT）將圖像視為一系列patch token（即由多個(gè)像素組成的較小部分圖像）。這也就意味著在每一層神經(jīng)網(wǎng)絡(luò)中，ViT模型使用多頭自注意力（multi-head self-attention），基于每對(duì)token之間的關(guān)系來(lái)處理patch token。這樣，ViT模型就能夠構(gòu)建整個(gè)圖像的全局表示。在輸入端，將圖像均勻地分割成多個(gè)部分來(lái)形成token，例如，將512×512像素的圖像分割成16×16像素的patch token。在中間層，上一層的輸出成為下一層的token。這里插一句。如果處理的是視頻，則視頻「管道」如16x16x2視頻片段（2幀16x16圖像）就成為了token。視覺(jué)token的質(zhì)量和數(shù)量決定了Vision Transformer的整體性能。許多Vision Transformer結(jié)構(gòu)面臨的主要挑戰(zhàn)是，它們通常需要太多的token才能獲得合理的結(jié)果。例如，即使使用16x16patch token化，單個(gè)512x512圖像也對(duì)應(yīng)于1024個(gè)token。對(duì)于具有多個(gè)幀的視頻，每層可能都需要處理數(shù)萬(wàn)個(gè)token。考慮到Transformer的計(jì)算量隨著token數(shù)量的增加而二次方增加，這通常會(huì)使Transformer難以處理更大的圖像和更長(zhǎng)的視頻。這就引出了一個(gè)問(wèn)題：真的有必要在每一層處理那么多token嗎？谷歌在「TokenLearner：What Can 8 Learned Tokens Do for Images and Videos?」中提到了「自適應(yīng)」這個(gè)概念。這篇文章將在NeurIPS 2021上進(jìn)行展示。論文地址：https:///pdf/2106.11297.pdf項(xiàng)目地址：https://github.com/google-research/scenic/tree/main/scenic/projects/token_learner實(shí)驗(yàn)表明，TokenLearner可以自適應(yīng)地生成更少數(shù)量的token，而不是總是依賴(lài)于由圖像均勻分配形成的token，這樣一來(lái)，可以使Vision Transformer運(yùn)行得更快，性能更好。TokenLearner是一個(gè)可學(xué)習(xí)的模塊，它會(huì)獲取圖像張量（即輸入）并生成一小組token。該模塊可以放置在Vision Transformer模型中的不同位置，顯著減少了所有后續(xù)層中要處理的token數(shù)量。實(shí)驗(yàn)表明，使用TokenLearner可以節(jié)省一半或更多的內(nèi)存和計(jì)算量，而分類(lèi)性能卻并不會(huì)下降，并且由于其適應(yīng)輸入的能力，它甚至可以提高準(zhǔn)確率。TokenLearner其實(shí)是一種簡(jiǎn)單的空間注意力方法。為了讓每個(gè)TokenLearner學(xué)習(xí)到有用的信息，先得計(jì)算一個(gè)突出的重要區(qū)域的空間注意力圖（使用卷積層或MLP）。接著，這樣的空間注意力圖會(huì)被用來(lái)對(duì)輸入的每個(gè)區(qū)域進(jìn)行加權(quán)（目的是丟棄不必要的區(qū)域），并且結(jié)果經(jīng)過(guò)空間池化后，就可以生成最終的學(xué)習(xí)好了的token。應(yīng)用于單個(gè)圖像的TokenLearner模塊的直觀圖示TokenLearner學(xué)習(xí)在張量像素的子集上進(jìn)行空間處理，并生成一組適應(yīng)輸入的token向量。這種操作被并行重復(fù)多次，就可以從原始的輸入中生成n個(gè)（10個(gè)左右）token。換句話說(shuō)，TokenLearner也可以被視為基于權(quán)重值來(lái)執(zhí)行像素的選擇，隨后進(jìn)行全局平均。值得一提的是，計(jì)算注意力圖的函數(shù)由不同的可學(xué)習(xí)參數(shù)控制，并以端到端的方式進(jìn)行訓(xùn)練。這樣也就使得注意力函數(shù)可以在捕捉不同輸入中的空間信息時(shí)進(jìn)行優(yōu)化。在實(shí)踐中，模型將學(xué)習(xí)多個(gè)空間注意力函數(shù)，并將其應(yīng)用于輸入，并平行地產(chǎn)生不同的token向量。TokenLearner模塊學(xué)習(xí)為每個(gè)輸出標(biāo)記生成一個(gè)空間注意力圖，并使用它來(lái)抽象化輸入的token因此，TokenLearner使模型能夠處理與特定識(shí)別任務(wù)相關(guān)的少量token，而不是處理固定的、統(tǒng)一的token化輸入。也就是說(shuō)，TokenLearner啟用了自適應(yīng)token，以便可以根據(jù)輸入動(dòng)態(tài)選擇token，這一做法有效地減少了token的總數(shù)，大大減少了Transformer網(wǎng)絡(luò)的計(jì)算。而這些動(dòng)態(tài)自適應(yīng)生成的token也可用于標(biāo)準(zhǔn)的Transformer架構(gòu)，如圖像領(lǐng)域的ViT和視頻領(lǐng)域的ViViT（Video Vision Transformer）。構(gòu)建TokenLearner模塊后，下一步就必須要確定將其放置在哪個(gè)位置。首先，研究人員嘗試將它放置在標(biāo)準(zhǔn)ViT架構(gòu)中的不同位置，輸入圖像使用224x224的大小。TokenLearner生成的token數(shù)量為8個(gè)和16個(gè)，遠(yuǎn)遠(yuǎn)少于標(biāo)準(zhǔn)ViT使用的196個(gè)或576個(gè)token。下圖顯示了在ViT B/16中的不同相對(duì)位置插入TokenLearner的模型的ImageNet 5-shot分類(lèi)精度和FLOPs，其中ViT B/16是一個(gè)基礎(chǔ)模型，有12個(gè)注意力層。其運(yùn)行時(shí)使用16x16大小的patch token。在JFT 300M的預(yù)訓(xùn)練下，ImageNet的5-shot精度與ViT B/16中TokenLearner的相對(duì)位置有關(guān)位置0意味著TokenLearner被置于任何Transformer層之前。其中，baseline是標(biāo)準(zhǔn)的ViT B/16的ImageNet 5-shot分類(lèi)精度和FLOPs。計(jì)算量以數(shù)十億次浮點(diǎn)運(yùn)算（GFLOPS）衡量我們發(fā)現(xiàn)，在網(wǎng)絡(luò)的最初四分之一處（1/4處）插入TokenLearner，實(shí)現(xiàn)了與基線幾乎相同的準(zhǔn)確性，同時(shí)將計(jì)算量減少到基線的三分之一以下。此外，將TokenLearner放在后面一層（網(wǎng)絡(luò)的3/4之后），與不使用TokenLearner相比，取得了更好的性能，同時(shí)由于其適應(yīng)性，性能更快。由于TokenLearner前后的token數(shù)量相差很大（例如，前196個(gè)，后8個(gè)），TokenLearner模塊后的相對(duì)計(jì)算量幾乎可以忽略不計(jì)。將帶有TokenLearner的ViT模型和普通的ViT模型進(jìn)行對(duì)比，同時(shí)在ImageNet的few-shot上采用相同的設(shè)置。TokenLearner會(huì)被放置在每個(gè)ViT模型中間的不同位置，如網(wǎng)絡(luò)的1/2和3/4處。其中，模型通過(guò)JFT 300M進(jìn)行預(yù)訓(xùn)練。從圖上觀察可以得知，TokenLearner模型在準(zhǔn)確率和計(jì)算量方面的表現(xiàn)都比ViT要好。不同版本的ViT模型在ImageNet分類(lèi)上的表現(xiàn)在更大的ViT模型中插入TokenLearner，如具有24個(gè)注意力層，并以10x10（或8x8）個(gè)patch作為初始token的L/10和L/8。之后，將這兩個(gè)模型與48層的ViT G/14模型進(jìn)行比較。可以看到，在表現(xiàn)和G/14模型相當(dāng)?shù)那闆r下，TokenLearner只需要非常少的參數(shù)和計(jì)算量。左：大規(guī)模TokenLearner模型與ViT G/14在ImageNet數(shù)據(jù)集上的分類(lèi)精度對(duì)比；右：參數(shù)量和FLOPS的對(duì)比視頻理解是計(jì)算機(jī)視覺(jué)的關(guān)鍵挑戰(zhàn)之一，TokenLearner在多個(gè)視頻分類(lèi)數(shù)據(jù)集基準(zhǔn)上取得了SOTA的性能。其中，在Kinetics-400和Kinetics-600上的性能超過(guò)了以前的Transformer模型，在Charades和AViD上也超過(guò)了之前的CNN模型。通過(guò)與視頻視覺(jué)Transformer（Video Vision Transformer，ViViT）結(jié)合，TokenLearner會(huì)在每個(gè)時(shí)間段學(xué)習(xí)8（或16）個(gè)token。左：視頻分類(lèi)任務(wù)；右圖：不同模型的對(duì)比隨著時(shí)間的推移，當(dāng)人物在場(chǎng)景中移動(dòng)時(shí)，TokenLearner會(huì)注意到不同的空間位置變化從而進(jìn)行token化。雖然Vision Transformer是計(jì)算機(jī)視覺(jué)領(lǐng)域的一個(gè)強(qiáng)大模型，但大量的token及龐大的計(jì)算量一直是將ViT應(yīng)用于更大圖像和更長(zhǎng)視頻的瓶頸。本文中作者表明，保留如此大量的token并在整個(gè)層集上完全處理它們是沒(méi)有必要的。此外，作者還證明了通過(guò)學(xué)習(xí)一個(gè)基于輸入圖像自適應(yīng)提取token的模塊，可以在節(jié)省計(jì)算的同時(shí)獲得更好的性能。最后，多個(gè)公共數(shù)據(jù)集上的驗(yàn)證也表明了TokenLearner在視頻表征學(xué)習(xí)任務(wù)中的表現(xiàn)十分優(yōu)異。參考資料： https://ai./2021/12/improving-vision-transformer-efficiency.html?m=1
|
