在過(guò)去的幾年里，Transformer架構(gòu)在自然語(yǔ)言處理（NLP）、圖像處理和視覺(jué)計(jì)算領(lǐng)域的深度表征學(xué)習(xí)中取得了顯著的成就，幾乎成為了AI領(lǐng)域的主導(dǎo)技術(shù)。

然而，雖然Transformer架構(gòu)及其眾多變體在實(shí)踐中取得了巨大成功，但其設(shè)計(jì)大多是基于經(jīng)驗(yàn)的，并沒(méi)有嚴(yán)格的數(shù)學(xué)解釋?zhuān)苍谝欢ǔ潭壬舷拗屏搜芯咳藛T的思路，無(wú)法開(kāi)發(fā)出更高效、更具可解釋性的Transformer新變體。

為了填補(bǔ)這一空白，馬毅教授團(tuán)隊(duì)曾發(fā)布過(guò)白盒Transformer模型CRATE，其架構(gòu)的每一層都是通過(guò)數(shù)學(xué)推導(dǎo)得到的，可以完全解釋為展開(kāi)的梯度下降迭代；此外，CRATE學(xué)習(xí)到的模型和特征在語(yǔ)義上也比傳統(tǒng)的Transformer模型具有更好的可解釋性，例如，即使模型僅在分類(lèi)任務(wù)上進(jìn)行訓(xùn)練，可視化圖像的特征也能自然地形成該圖像的零樣本分割。

然而，到目前為止，CRATE的應(yīng)用規(guī)模仍然相對(duì)有限，CRATE-Large只包含77.6M參數(shù)，與標(biāo)準(zhǔn)Vision Transformer（ViTs）的22B參數(shù)量形成了鮮明對(duì)比。

最近，加利福尼亞大學(xué)圣克魯斯分校和伯克利分校的研究團(tuán)隊(duì)聯(lián)合提出了CRATE-α，首次探索了不同規(guī)模的CRATE用于視覺(jué)任務(wù)（從Tiny到Huge）時(shí)的模型性能，研究人員在CRATE架構(gòu)設(shè)計(jì)中對(duì)稀疏編碼塊進(jìn)行了策略性但最小化的（strategic yet minimal）修改，并設(shè)計(jì)了一種輕量級(jí)的訓(xùn)練方法，以提高CRATE的可擴(kuò)展性。

論文鏈接：https://arxiv.org/pdf/2405.20299

項(xiàng)目鏈接：https://rayjryang.github.io/CRATE-alpha/

具體來(lái)說(shuō)，CRATE中的ISTA模塊是限制進(jìn)一步擴(kuò)展的因素，為了克服這一限制，CRATE-α主要做了三個(gè)修改：

1. 大幅擴(kuò)展了通道，對(duì)稀疏編碼塊進(jìn)行過(guò)參數(shù)化（overparameterized），使用過(guò)完備字典（overcomplete dictionary）對(duì)token表征進(jìn)行稀疏化。

2. 解耦了關(guān)聯(lián)矩陣，在稀疏編碼塊的最后一部中引入一個(gè)解耦字典（decoupled dictionary）

3. 添加了殘差連接。

實(shí)驗(yàn)結(jié)果證明，CRATE-α能夠隨著模型尺寸和訓(xùn)練數(shù)據(jù)集的增大而擴(kuò)展，性能可以持續(xù)提升。

例如，CRATE-α-B在ImageNet分類(lèi)任務(wù)上的性能顯著超過(guò)了之前最好的CRATE-B模型，準(zhǔn)確率提高了3.7%，達(dá)到了83.2%；進(jìn)一步對(duì)模型進(jìn)行擴(kuò)展時(shí)，CRATE-α-L在ImageNet分類(lèi)任務(wù)上達(dá)到了85.1%的準(zhǔn)確率。

值得注意的是，模型性能的提升是在保持甚至增強(qiáng)了CRATE模型可解釋性的同時(shí)實(shí)現(xiàn)的，因?yàn)楦蟪叽绲腃RATE-α模型學(xué)到的token表征能夠生成更高質(zhì)量的無(wú)監(jiān)督圖像分割。

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

從基礎(chǔ)尺寸（base）到大尺寸（large）

ImageNet-21K是一個(gè)廣泛用于圖像識(shí)別和分類(lèi)任務(wù)的大型數(shù)據(jù)集，文中用于訓(xùn)練的數(shù)據(jù)集版本包含19,000個(gè)類(lèi)別和大約1300萬(wàn)張圖片，由于數(shù)據(jù)丟失，比標(biāo)準(zhǔn)數(shù)據(jù)集（包含21,000個(gè)類(lèi)別和大約1400萬(wàn)張圖片）的數(shù)據(jù)量要少一點(diǎn)。

在預(yù)訓(xùn)練時(shí)，從數(shù)據(jù)集中隨機(jī)選取1%作為驗(yàn)證集。

預(yù)訓(xùn)練完成后，在ImageNet-1K數(shù)據(jù)集上對(duì)模型進(jìn)行微調(diào)，其中ImageNet-1K是一個(gè)更小的子集，包含1000個(gè)類(lèi)別，通常用于模型的最終評(píng)估。在微調(diào)階段，模型會(huì)針對(duì)這1000個(gè)類(lèi)別進(jìn)行更精細(xì)的訓(xùn)練，以提高其在特定任務(wù)上的性能。

最后，在ImageNet-1K的驗(yàn)證集上評(píng)估模型的性能。

研究人員對(duì)比了在32、16和8像素塊大小下的CRATE-α-B和CRATE-α-L，從實(shí)驗(yàn)結(jié)果中可以看到，CRATE-α-L在所有像素塊大小上都取得了顯著的改進(jìn)，但從CRATE-B增加到CRATE-L只能帶來(lái)0.5%的性能提升，表明了收益遞減的情況，證明了CRATE-α模型的可擴(kuò)展性顯著優(yōu)于普通CRATE

同時(shí)，預(yù)訓(xùn)練階段的訓(xùn)練損失顯示，隨著模型容量的增加，訓(xùn)練損失的趨勢(shì)可預(yù)測(cè)地得到改善。

從大（large）到巨大（huge）

多模態(tài)數(shù)據(jù)集DataComp1B包含14億圖文對(duì)，可以提供足夠的數(shù)據(jù)來(lái)訓(xùn)練和擴(kuò)展模型。

研究人員采用對(duì)比學(xué)習(xí)的方法來(lái)訓(xùn)練CRATE-α，不僅能夠利用上龐大的圖文對(duì)數(shù)據(jù)集，還能在模型尺寸從大到巨大的提升過(guò)程中，觀察到顯著的性能提升。

然而，直接訓(xùn)練一個(gè)類(lèi)似CLIP的模型需要巨大的計(jì)算資源，研究人員采用了優(yōu)化后的CLIPA協(xié)議，可以在減少計(jì)算資源消耗的同時(shí)，可以保持與CLIP相當(dāng)?shù)男阅堋?/span>

最后，為了評(píng)估CRATE-α模型的性能，研究人員采用了零樣本學(xué)習(xí)的方法，在ImageNet-1K數(shù)據(jù)集上測(cè)試模型的準(zhǔn)確率，該方法可以有效地評(píng)估模型在面對(duì)未見(jiàn)過(guò)類(lèi)別數(shù)據(jù)時(shí)的泛化能力，提供了一個(gè)衡量模型可擴(kuò)展性和實(shí)用性的重要指標(biāo)。

從實(shí)驗(yàn)結(jié)果中可以看到，

1. 模型尺寸的影響：CRATE-α-CLIPA-L/14在預(yù)訓(xùn)練和微調(diào)階段的ImageNet-1K零樣本準(zhǔn)確率上，分別比CRATE-α-CLIPA-B/16高出11.3%和9.0%，表明學(xué)習(xí)到的表征質(zhì)量可能受到模型尺寸的限制，即增加模型尺寸可以利用上更多數(shù)據(jù)。

2. 擴(kuò)展模型尺寸的益處：當(dāng)繼續(xù)增加模型尺寸時(shí)，可以觀察到CRATE-α-CLIP-H/14從更大的訓(xùn)練數(shù)據(jù)集中繼續(xù)獲益，在預(yù)訓(xùn)練和微調(diào)階段的ImageNet-1K零樣本準(zhǔn)確率上，分別比CRATE-α-CLIP-L/14高出3.1%和2.5%，證明了CRATE-α模型的強(qiáng)大可擴(kuò)展性。

3. 性能上限的探索：為了探索性能的上限，研究人員從頭開(kāi)始訓(xùn)練了一個(gè)標(biāo)準(zhǔn)的ViT-CLIPA-H/14，并觀察到了性能的提升。

節(jié)省計(jì)算資源的擴(kuò)展策略

在追求模型擴(kuò)展的效率和計(jì)算資源的優(yōu)化方面，研究人員發(fā)現(xiàn)，通過(guò)調(diào)整預(yù)訓(xùn)練階段的圖像token序列長(zhǎng)度，可以在極大減少計(jì)算資源消耗的同時(shí)，保持模型性能。

具體來(lái)說(shuō)，研究人員嘗試了一種新的方法：在預(yù)訓(xùn)練時(shí)使用較長(zhǎng)序列長(zhǎng)度的CRATE-α-L/32，在微調(diào)時(shí)切換到較短序列長(zhǎng)度的CRATE-α-L/14或CRATE-α-L/8，不僅大幅度降低了預(yù)訓(xùn)練階段的計(jì)算成本，而且在微調(diào)后，模型在ImageNet-1K數(shù)據(jù)集上的準(zhǔn)確率仍然非常接近全尺寸模型的性能。

例如，使用CRATE-α-L/32進(jìn)行預(yù)訓(xùn)練，然后微調(diào)到CRATE-α-L/14，可以節(jié)省約70%的計(jì)算資源，而準(zhǔn)確率只是略有下降；更進(jìn)一步，當(dāng)從CRATE-α-L/32預(yù)訓(xùn)練后微調(diào)到CRATE-α-L/8時(shí)，僅使用了原模型所需訓(xùn)練時(shí)間的10%，準(zhǔn)確率依然達(dá)到了84.2%，與全尺寸模型的85.1%相差無(wú)幾。

上述結(jié)果表明，通過(guò)精心設(shè)計(jì)預(yù)訓(xùn)練和微調(diào)階段的策略，可以在資源有限的情況下，有效地?cái)U(kuò)展CRATE-α模型。

CRATE-α的語(yǔ)義可解釋性得到提升

除了可擴(kuò)展性，文中還研究了不同模型大小的CRATE-α的可解釋性，使用MaskCut來(lái)驗(yàn)證和評(píng)估模型捕獲的豐富語(yǔ)義信息，包括定性和定量結(jié)果。

為CRATE-α、CRATE和ViT在COCO val2017上提供了分割可視化后，可以發(fā)現(xiàn)，CRATE-α模型保持甚至提高了CRATE的（語(yǔ)義）可解釋性?xún)?yōu)勢(shì)。

在COCO val2017上的定量評(píng)估結(jié)果顯示，當(dāng)為CRATE-α擴(kuò)展模型大小時(shí)，大型模型在目標(biāo)檢測(cè)和分割方面比base模型有所提高。