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

然而，雖然Transformer架構(gòu)及其眾多變體在實踐中取得了巨大成功，但其設(shè)計大多是基于經(jīng)驗的，并沒有嚴格的數(shù)學(xué)解釋，也在一定程度上限制了研究人員的思路，無法開發(fā)出更高效、更具可解釋性的Transformer新變體。

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

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

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

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

項目鏈接：https://rayjryang.github.io/CRATE-alpha/

具體來說，CRATE中的ISTA模塊是限制進一步擴展的因素，為了克服這一限制，CRATE-α主要做了三個修改：

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

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

3. 添加了殘差連接。

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

例如，CRATE-α-B在ImageNet分類任務(wù)上的性能顯著超過了之前最好的CRATE-B模型，準確率提高了3.7%，達到了83.2%；進一步對模型進行擴展時，CRATE-α-L在ImageNet分類任務(wù)上達到了85.1%的準確率。

值得注意的是，模型性能的提升是在保持甚至增強了CRATE模型可解釋性的同時實現(xiàn)的，因為更大尺寸的CRATE-α模型學(xué)到的token表征能夠生成更高質(zhì)量的無監(jiān)督圖像分割。

實驗結(jié)果

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

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

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

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

最后，在ImageNet-1K的驗證集上評估模型的性能。

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

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

從大（large）到巨大（huge）

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

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

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

最后，為了評估CRATE-α模型的性能，研究人員采用了零樣本學(xué)習的方法，在ImageNet-1K數(shù)據(jù)集上測試模型的準確率，該方法可以有效地評估模型在面對未見過類別數(shù)據(jù)時的泛化能力，提供了一個衡量模型可擴展性和實用性的重要指標。

從實驗結(jié)果中可以看到，

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

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

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

節(jié)省計算資源的擴展策略

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

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

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

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

CRATE-α的語義可解釋性得到提升

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

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

在COCO val2017上的定量評估結(jié)果顯示，當為CRATE-α擴展模型大小時，大型模型在目標檢測和分割方面比base模型有所提高。