馬毅團(tuán)隊(duì)新作：白盒ViT成功實(shí)現(xiàn)

【導(dǎo)讀】CRATE模型完全由理論指導(dǎo)設(shè)計(jì)，僅用自監(jiān)督學(xué)習(xí)即可實(shí)現(xiàn)分割語義涌現(xiàn)。

基于Transformer的視覺基礎(chǔ)模型在各種下游任務(wù)，如分割和檢測(cè)中都展現(xiàn)出了非常強(qiáng)大的性能，并且DINO等模型經(jīng)過自監(jiān)督訓(xùn)練后已經(jīng)涌現(xiàn)出了語義的分割屬性。

不過奇怪的是，類似的涌現(xiàn)能力并沒有出現(xiàn)在有監(jiān)督分類訓(xùn)練后的視覺Transformer模型中。

最近，馬毅教授團(tuán)隊(duì)探索了基于Transformer架構(gòu)的模型中涌現(xiàn)分割能力是否僅僅是復(fù)雜的自監(jiān)督學(xué)習(xí)機(jī)制的結(jié)果，或者是否可以通過模型架構(gòu)的適當(dāng)設(shè)計(jì)在更通用的條件下實(shí)現(xiàn)相同的涌現(xiàn)。

在CVer微信公眾號(hào)后臺(tái)回復(fù)：CRATE，可以下載本論文pdf和代碼

Emergence of Segmentation with Minimalistic White-Box Transformers

代碼：https://github.com/Ma-Lab-Berkeley/CRATE

論文：https://arxiv.org/abs/2308.16271

通過大量的實(shí)驗(yàn)，研究人員證明了當(dāng)采用白盒Transformer模型CRATE時(shí)，其設(shè)計(jì)明確地模擬并追求數(shù)據(jù)分布中的低維結(jié)構(gòu)，整體和part級(jí)別的分割屬性已經(jīng)以最小化的監(jiān)督訓(xùn)練配方出現(xiàn)。

分層的細(xì)粒度分析表明，涌現(xiàn)屬性有力地證實(shí)了白盒網(wǎng)絡(luò)的設(shè)計(jì)數(shù)學(xué)功能。我們的結(jié)果提出了一條設(shè)計(jì)白盒基礎(chǔ)模型的途徑，該模型同時(shí)具有高性能和數(shù)學(xué)上完全可解釋性。

馬毅教授也表示，深度學(xué)習(xí)的研究將會(huì)逐漸從經(jīng)驗(yàn)設(shè)計(jì)轉(zhuǎn)向理論指導(dǎo)。

白盒CRATE的涌現(xiàn)屬性

DINO的分割涌現(xiàn)能力

智能系統(tǒng)中的表征學(xué)習(xí)旨在將世界的高維、多模態(tài)感官數(shù)據(jù)（圖像、語言、語音）轉(zhuǎn)換為更緊湊的形式，同時(shí)保留其基本的低維結(jié)構(gòu)，實(shí)現(xiàn)高效的識(shí)別（比如分類）、分組（比如分割）和追蹤。

深度學(xué)習(xí)模型的訓(xùn)練通常采用數(shù)據(jù)驅(qū)動(dòng)的方式，輸入大規(guī)模數(shù)據(jù)，以自監(jiān)督的方式進(jìn)行學(xué)習(xí)。

在視覺基礎(chǔ)模型中，自監(jiān)督Transformer架構(gòu)的DINO模型展現(xiàn)出了令人驚訝的涌現(xiàn)能力，即使沒有經(jīng)過有監(jiān)督分割訓(xùn)練，ViTs也能識(shí)別出顯式的語義分割信息。

后續(xù)有工作研究了如何在DINO模型中利用這種分割信息，并在下游任務(wù)中，如分割、檢測(cè)等實(shí)現(xiàn)了最先進(jìn)的性能，也有工作證明了用DINO訓(xùn)練的ViTs中的倒數(shù)第二層特征與視覺輸入中的顯著性信息強(qiáng)烈相關(guān)，如區(qū)分前景、背景和物體邊界，從而提升圖像分割和其他任務(wù)的性能。

為了讓分割屬性涌現(xiàn)，DINO需要在訓(xùn)練期間將自監(jiān)督學(xué)習(xí)、知識(shí)蒸餾和權(quán)重平均巧妙地結(jié)合起來。

目前還不清楚DINO中引入的每個(gè)組件是否對(duì)于分割遮罩的涌現(xiàn)來說必不可缺，盡管DINO也采用ViT架構(gòu)作為其主干，但在分類任務(wù)上訓(xùn)練的普通有監(jiān)督ViT模型中，并沒有觀察到分割涌現(xiàn)行為。

CRATE的涌現(xiàn)

基于DINO的成功案例，研究人員想要探究，復(fù)雜的自監(jiān)督學(xué)習(xí)pipeline對(duì)于獲得類似Transformer的視覺模型中的涌現(xiàn)屬性是否是必要的。

研究人員認(rèn)為，在Transformer模型中促進(jìn)分割屬性的一種有前途的方法是，在考慮輸入數(shù)據(jù)結(jié)構(gòu)的情況下設(shè)計(jì)Transformer模型架構(gòu)，也代表了表征學(xué)習(xí)經(jīng)典方法與現(xiàn)代數(shù)據(jù)驅(qū)動(dòng)的深度學(xué)習(xí)框架的結(jié)合。

與目前主流的Transformer模型對(duì)比，這種設(shè)計(jì)方法也可以叫做白盒Transformer模型。

基于馬毅教授組之前的工作，研究人員對(duì)白盒架構(gòu)的CRATE模型進(jìn)行了廣泛的實(shí)驗(yàn)，證明了CRATE的白盒設(shè)計(jì)是自注意力圖中分割屬性涌現(xiàn)的原因。

定性評(píng)估

研究人員采用基于[CLS] token的注意力圖方法對(duì)模型進(jìn)行解釋和可視化，結(jié)果發(fā)現(xiàn)CRATE中的query-key-value矩陣都是相同的。

可以觀察到CRATE模型的自注意力圖（self-attention map）可以對(duì)應(yīng)到輸入圖像的語義上，模型的內(nèi)部網(wǎng)絡(luò)對(duì)每個(gè)圖像都進(jìn)行了清晰的語義分割，實(shí)現(xiàn)了類似DINO模型的效果。

相比之下，在有監(jiān)督分類任務(wù)上訓(xùn)練的普通ViT卻并沒有表現(xiàn)出類似的分割屬性。

遵循之前關(guān)于可視化圖像學(xué)習(xí)的逐塊深度特征的工作，研究人員對(duì)CRATE和ViT模型的深度token表征進(jìn)行主成分分析（PCA）研究。

可以發(fā)現(xiàn)，CRATE可以在沒有分割監(jiān)督訓(xùn)練的情況下，依然可以捕捉到圖像中物體的邊界。

并且，主成分（principal compoenents）也說明了token和物體中相似部分的特征對(duì)齊，例如紅色通道對(duì)應(yīng)馬腿。

而有監(jiān)督ViT模型的PCA可視化結(jié)構(gòu)化程度相當(dāng)?shù)汀?/p>

定量評(píng)估

研究人員使用現(xiàn)有的分割和目標(biāo)檢測(cè)技術(shù)對(duì)CRATE涌現(xiàn)的分割屬性進(jìn)行評(píng)估。

從自注意力圖可以看到，CRATE用清晰的邊界顯式地捕獲了對(duì)象級(jí)的語義，為了定量測(cè)量分割的質(zhì)量，研究人員利用自注意力圖生成分割遮罩（segmentation mask），對(duì)比其與真實(shí)掩碼之間的標(biāo)準(zhǔn)mIoU（平均交并比）。

從實(shí)驗(yàn)結(jié)果中可以看到，CRATE在視覺和mIOU評(píng)分上都顯著優(yōu)于ViT，表明CRATE的內(nèi)部表征對(duì)于分割掩碼任務(wù)生成來說要更有效。

對(duì)象檢測(cè)和細(xì)粒度分割

為了進(jìn)一步驗(yàn)證和評(píng)估CRATE捕獲的豐富語義信息，研究人員采用了一種高效的目標(biāo)檢測(cè)和分割方法MaskCut，無需人工標(biāo)注即可獲得自動(dòng)化評(píng)估模型，可以基于CRATE學(xué)到的token表征從圖像中提取更細(xì)粒度的分割。

在COCO val2017上的分割結(jié)果中可以看到，有CRATE的內(nèi)部表征在檢測(cè)和分割指標(biāo)上都要好于有監(jiān)督ViT，有監(jiān)督ViT特征的MaskCut在某些情況下甚至完全不能產(chǎn)生分割掩碼。

CRATE分割能力的白盒分析

深度在CRATE中的作用

CRATE的每一層設(shè)計(jì)都遵循相同的概念目的：優(yōu)化稀疏速率降低，并將token分布轉(zhuǎn)換為緊湊和結(jié)構(gòu)化的形式。

假設(shè)CRATE中語義分割能力的涌現(xiàn)類似于「表征Z中屬于相似語義類別token的聚類」，預(yù)期CRATE的分割性能可以隨著深度的增加而提高。

為了測(cè)試這一點(diǎn)，研究人員利用MaskCut來定量評(píng)估跨不同層的內(nèi)部表征的分割性能；同時(shí)應(yīng)用PCA可視化來理解分割是如何隨深度加深而涌現(xiàn)的。

從實(shí)驗(yàn)結(jié)果中可以觀察到，當(dāng)使用來自更深層的表征時(shí)，分割分?jǐn)?shù)提高了，與CRATE的增量優(yōu)化設(shè)計(jì)非常一致。

相比之下，即使ViT-B/8的性能在后面的層中略有提高，但其分割分?jǐn)?shù)明顯低于CRATE，PCA結(jié)果顯示，從CRATE深層提取的表征會(huì)逐漸更關(guān)注前景對(duì)象，并且能夠捕捉紋理級(jí)別的細(xì)節(jié)。

CRATE的消融實(shí)驗(yàn)

CRATE中的注意力塊（MSSA）和MLP塊（ISTA）都不同于ViT中的注意力塊。

為了了解每個(gè)組件對(duì)CRATE涌現(xiàn)分割屬性的影響，研究人員選取了三個(gè)CRATE變體：CRATE, CRATE-MHSA, CRATE-MLP，分別表示ViT中的注意塊（MHSA）和MLP塊。

研究人員在ImageNet-21k數(shù)據(jù)集上應(yīng)用相同的預(yù)訓(xùn)練設(shè)置，然后應(yīng)用粗分割評(píng)估和遮罩分割評(píng)估來定量對(duì)比不同模型的性能。

實(shí)驗(yàn)結(jié)果顯示，CRATE在所有任務(wù)中都明顯優(yōu)于其他模型架構(gòu)，可以發(fā)現(xiàn)，盡管MHSA和MSSA之間的架構(gòu)差異很小，但只需要簡單地用CRATE中的MSSA替換ViT中的MHSA，可以顯著改善ViT的粗分割性能（即VOC Seg），證明了白盒設(shè)計(jì)的有效性。

識(shí)別注意頭的語義屬性

[CLS] token和圖像塊token之間的自注意力圖可以看到清晰的分段掩碼，根據(jù)直覺，每個(gè)注意力頭應(yīng)該都可以捕捉到數(shù)據(jù)的部分特征。

研究人員首先將圖像輸入到CRATE模型，然后由人來檢查、選擇四個(gè)似乎具有語義含義的注意力頭；然后在其他輸入圖像上在這些注意力頭上進(jìn)行自注意力圖可視化。

可以發(fā)現(xiàn)，每個(gè)注意力頭都捕捉到了物體的不同部分，甚至不同的語義：例如第一列中顯示的注意力頭可以捕捉到不同動(dòng)物的腿，最后一列中顯示的注意力頭捕捉的是耳朵和頭部。

自從可形變部件模型（deformable part model）和膠囊網(wǎng)絡(luò)發(fā)布以來，這種將視覺輸入解析為部分-整體層次結(jié)構(gòu)的能力一直是識(shí)別架構(gòu)的目標(biāo)，白盒設(shè)計(jì)的CRATE模型也具有這種能力。