不是所有圖像都值16×16個(gè)詞，清華與華為提出動(dòng)態(tài)ViT

水木番 發(fā)自 凹非寺
量子位 報(bào)道 | 公眾號(hào) QbitAI

在NLP中，Transformer以自注意力模型機(jī)制為法寶，在圖像識(shí)別問題上的成功已經(jīng)很廣泛了。

尤其是，ViT在大規(guī)模圖像網(wǎng)絡(luò)上性能特別高，因此應(yīng)用特別廣。

但隨著數(shù)據(jù)集規(guī)模的增長，會(huì)導(dǎo)致計(jì)算成本急劇增加，以及自注意力中的tokens數(shù)量逐漸增長！

最近，清華自動(dòng)化系的助理教授黃高的研究團(tuán)隊(duì)和華為的研究人員另辟蹊徑，提出了一種Dynamic Vision Transformer (DVT)，可以自動(dòng)為每個(gè)輸入圖像配置適當(dāng)數(shù)量的tokens，從而減少冗余計(jì)算，顯著通過效率。

該文以《Not All Images are Worth 16×16 Words: Dynamic Vision Transformers with Adaptive Sequence Length》為標(biāo)題，已發(fā)表在arXiv上。

提出動(dòng)態(tài)ViT

很明顯，當(dāng)前的ViT面臨計(jì)算成本和tokens數(shù)量的難題。

為了在準(zhǔn)確性和速度之間實(shí)現(xiàn)最佳平衡，tokens的數(shù)量一般是 14×14/16×16的。

研究團(tuán)隊(duì)觀察到：

一般樣本中會(huì)有很多的“簡單”圖像，它們用數(shù)量為 4×4 標(biāo)記就可以準(zhǔn)確預(yù)測(cè)，現(xiàn)在的計(jì)算成本（ 14×14）相當(dāng)于增加了8.5倍，而其實(shí)只有一小部分“困難”的圖像需要更精細(xì)的表征。

通過動(dòng)態(tài)調(diào)整tokens數(shù)量，計(jì)算效率在“簡單”和 “困難”樣本中的分配并不均勻，這里有很大的空間可以提高效率。

基于此，研究團(tuán)隊(duì)提出了一種新型的動(dòng)態(tài)ViT（DVT ）框架，目標(biāo)是自動(dòng)配置在每個(gè)圖像上調(diào)節(jié)的tokens數(shù)量 ，從而實(shí)現(xiàn)高計(jì)算效率。

這種DVT被設(shè)計(jì)成一個(gè)通用框架。

在測(cè)試時(shí)間時(shí)，這些模型以較少的tokens 開始依次被激活。

一旦產(chǎn)生了充分置信度的預(yù)測(cè)，推理過程就會(huì)立即終止。

模型的主體架構(gòu)采用目前最先進(jìn)的圖像識(shí)別Transformer，如ViT、DeiT和T2T-ViT，可以提高效率。

這種方法同時(shí)也具有很強(qiáng)的靈活性。

因?yàn)镈VT的計(jì)算量可以通過簡單的提前終止準(zhǔn)則進(jìn)行調(diào)整。

這一特性使得DVT適合可用計(jì)算資源動(dòng)態(tài)變化，或通過最小功耗來實(shí)現(xiàn)給定性能的情況。

這兩種情況在現(xiàn)實(shí)世界的應(yīng)用程序中都是普遍存在的，像搜索引擎和移動(dòng)應(yīng)用程序中都經(jīng)常能夠看到。

根據(jù)上面的流程圖，仔細(xì)的讀者還會(huì)發(fā)現(xiàn)：

一旦從上游往下游計(jì)算無法成功的時(shí)候，就會(huì)采取向先前信息或者上游信息進(jìn)行重用的方法實(shí)現(xiàn)進(jìn)一步數(shù)據(jù)訓(xùn)練。

研究團(tuán)隊(duì)在此基礎(chǔ)上，還進(jìn)一步提出了特征重用機(jī)制和關(guān)系重用機(jī)制，它們都能夠通過最小化計(jì)算成本來顯著提高測(cè)試精度，以減少冗余計(jì)算。

前者允許在先前提取的深度特征的基礎(chǔ)上訓(xùn)練下游數(shù)據(jù)，而后者可以利用現(xiàn)有的上游的自注意力模型來學(xué)習(xí)更準(zhǔn)確的注意力。

這種對(duì)“簡單”“困難”動(dòng)態(tài)分配的方法，其現(xiàn)實(shí)效果可以由下圖的實(shí)例給出說明。

那么，接下來讓我們來看看這兩種機(jī)制具體是怎么做的？

特征重用機(jī)制

DVT中所有的Transformer都有一個(gè)共同的目標(biāo)：提取特征信號(hào)來實(shí)現(xiàn)準(zhǔn)確的識(shí)別。

因此，下游模型應(yīng)該在之前獲得的深度特征的基礎(chǔ)上進(jìn)行學(xué)習(xí)，而不是從零開始提取特征。

在上游模型中執(zhí)行的計(jì)算對(duì)其自身和后續(xù)模型都有貢獻(xiàn)，這樣會(huì)使模型效率更高。

為了實(shí)現(xiàn)這個(gè)想法，研究團(tuán)隊(duì)提出了一個(gè)特征重用機(jī)制。

簡單來說，就是利用上游Transformer最后一層輸出的圖像tokens，來學(xué)習(xí)逐層的上下文嵌入，并將其集成到每個(gè)下游Transformer的MLP塊中。

關(guān)系重用機(jī)制

Transformer的一個(gè)突出優(yōu)點(diǎn)是：

自注意力塊能夠整合整個(gè)圖像中的信息，從而有效地模型化了數(shù)據(jù)中的長期依賴關(guān)系。

通常，模型需要在每一層學(xué)習(xí)一組注意力圖來描述標(biāo)記之間的關(guān)系。

除了上面提到的深層特征外，下游模型還可以獲得之前模型中產(chǎn)生的自注意力圖。

研究團(tuán)隊(duì)認(rèn)為，這些學(xué)習(xí)到的關(guān)系也能夠被重用，以促進(jìn)下游Transformer學(xué)習(xí)，具體采用的是對(duì)數(shù)的加法運(yùn)算。

效果如何？

多說無益，讓我們看看實(shí)際效果如何？

在ImageNet上的 Top-1準(zhǔn)確率v.s.計(jì)算量如下圖。

可以看出，DVT比DeiT和T2T-VIT計(jì)算效率要顯著更好：

當(dāng)計(jì)算開銷在0.5-2 GFLOPs內(nèi)時(shí)，DVT的計(jì)算量比相同性能的T2T-ViT少了1.7-1.9倍。

此外，這種方法可以靈活地達(dá)到每條曲線上的所有點(diǎn)，只需調(diào)整一次DVT的置信閾值即可。

CIFAR的 Top-1 準(zhǔn)確率 v.s. GFLOP 如下圖。

在ImageNet上的 Top-1準(zhǔn)確率v.s.吞吐量如下表。

在DVT中，“簡單”和 “困難”的可視化樣本如下圖。

從上面ImageNet、CIFAR-10 和 CIFAR-100 上的大量實(shí)證結(jié)果表明：

DVT方法在理論計(jì)算效率和實(shí)際推理速度方面，都明顯優(yōu)于其他方法。

看到這樣漂亮的結(jié)果，難道你還不心動(dòng)嗎？

有興趣的小伙伴歡迎去看原文哦~

傳送門

論文地址：

https://arxiv.org/abs/2105.15075

研究團(tuán)隊(duì)

黃高

目前年僅33歲，已是清華大學(xué)自動(dòng)化系助理教授，博士生導(dǎo)師。

獲得過2020年阿里巴巴達(dá)摩院青橙獎(jiǎng)，研究領(lǐng)域包括機(jī)器學(xué)習(xí)、深度學(xué)習(xí)、計(jì)算機(jī)視覺、強(qiáng)化學(xué)習(xí)等。