一、研究動機(jī)

掩碼建模（MIM, MAE）被證明是非常有效的自監(jiān)督訓(xùn)練方法。然而，如圖 1 所示，MIM 對于更大的模型效果相對更好。當(dāng)模型很小的時候（比如 ViT-T 5M 參數(shù)，這樣的模型對于現(xiàn)實(shí)世界非常重要），MIM 甚至可能一定程度上降低模型的效果。比如用 MAE 訓(xùn)練的 ViT-L 比普通監(jiān)督訓(xùn)練的模型在 ImageNet 上的分類效果提升 3.3%，但是用 MAE 訓(xùn)練的 ViT-T 比普通監(jiān)督訓(xùn)練的模型在 ImageNet 上的分類效果降低了 0.6%。

在這篇工作中我們提出了 TinyMIM，其在保持 ViT 結(jié)構(gòu)不變并且不修改結(jié)構(gòu)引入其他歸納偏置（inductive bias）的基礎(chǔ)上、用蒸餾的方法遷移大模型上的知識到小模型。

• 論文地址：https://arxiv.org/pdf/2301.01296.pdf
• 代碼地址：https://github.com/OliverRensu/TinyMIM

我們系統(tǒng)性的研究了蒸餾目標(biāo)、數(shù)據(jù)增強(qiáng)、正則化、輔助損失函數(shù)等對于蒸餾的影響。在嚴(yán)格的只用 ImageNet-1K 作為訓(xùn)練數(shù)據(jù)的情況下（包括 Teacher model 也只用 ImageNet-1K 訓(xùn)練）和 ViT-B 作為模型，我們的方法實(shí)現(xiàn)了當(dāng)前最好的性能。如圖所示：

把我們的方法（TinyMIM）和基于掩碼重建的方法 MAE，以及監(jiān)督式學(xué)習(xí)的方法從頭開始訓(xùn)練的 DeiT 作比較。MAE 在模型比較大的時候有顯著的性能提升，但是在模型比較小的時候提升幅度有限甚至?xí)δＰ偷淖罱K效果。我們的方法 TinyMIM 在不同模型的大小上都有大幅提升。

我們的貢獻(xiàn)如下：

1. 蒸餾的目標(biāo)（Distillation targets）:1）蒸餾 token 之間的關(guān)系比單獨(dú)蒸餾 class token 或者特征圖（feature map）更有效；2）用中間層作為蒸餾的目標(biāo)更有效。

2. 數(shù)據(jù)增強(qiáng)和模型正則化（Data and network regularization）：1）用帶掩碼的圖片效果更差；2）學(xué)生模型需要一點(diǎn) drop path，但是 teacher 模型不需要。

3. 輔助損失函數(shù)（auxiliary losses）：MIM 作為輔助損失函數(shù)沒有意義。

4. 宏觀蒸餾策略（Macro distillation strategy）：我們發(fā)現(xiàn)序列化的蒸餾（ViT-B -> ViT-S -> ViT-T）效果最好。

二、方法

我們系統(tǒng)性的調(diào)研了蒸餾的目標(biāo)，輸入的圖片，蒸餾目標(biāo)模塊。

2.1 影響蒸餾效果的因素

1）特征：

a. 中間 block 特征和輸出特征

當(dāng) i=L 時，指的是 Transformer 輸出層的特征。當(dāng) i< L 時，指的是 Transformer 中間層的特征。

b. 注意力（Attention）特征和前饋層（FFN）層特征

Transformer 每一個 block 有 Attention 層和 FFN 層，蒸餾不同的層會帶來不同的影響。

c.QKV 特征

在 Attention 層內(nèi)會有 Q，K，V 特征，這些特征用于計(jì)算注意力機(jī)制，我們也調(diào)研了直接蒸餾這些特征。

2）關(guān)系

Q，K，V 用于計(jì)算注意力圖，這些特征之間的關(guān)系也可以作為知識蒸餾的目標(biāo)。

3）輸入：是否帶掩碼

傳統(tǒng)的知識蒸餾是直接輸入完整的圖片。我們的方法為了探索蒸餾掩碼建模模型，所以我們也探索了帶掩碼的圖片是否適合作為知識蒸餾時候的輸入。

2.2 知識蒸餾方法對比

1）Class Token 蒸餾：

最簡單的方法就是類似 DeiT 直接蒸餾 MAE 預(yù)訓(xùn)練模型的 class token:

其中指學(xué)生模型的 class token，而 指老師模型的 class token。

2）特征蒸餾：我們直接參考了 feature distillation [1] 作為對比

3）關(guān)系蒸餾：我們提出了也是本文默認(rèn)的蒸餾策略

三、實(shí)驗(yàn)

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

我們的方法在 ImageNet-1K 上預(yù)訓(xùn)練，而且教師模型也是在 ImageNet-1K 預(yù)訓(xùn)練。然后我們將我們預(yù)訓(xùn)練的模型在下游任務(wù)（分類、語義分割）上進(jìn)行了微調(diào)。模型表現(xiàn)如圖：

我們的方法顯著超過之前基于 MAE 的方法，尤其是小模型。具體來講，對于超小的模型 ViT-T，我們的方法實(shí)現(xiàn)了 75.8% 的分類準(zhǔn)確性，相比 MAE 基線模型實(shí)現(xiàn)了 4.2 的提升。對于小模型 ViT-S，我們實(shí)現(xiàn)了 83.0% 的分類準(zhǔn)確性，比之前最好的方法提升了 1.4。對于 Base 尺寸的模型，我們的方法分別超過 MAE 基線模型和以前最好的模型 CAE 4.1 和 2.0。

同時我們也測試了模型的魯棒性，如圖所示：

TinyMIM-B 對比 MAE-B，在 ImageNet-A 和 ImageNet-R 分別提升了 + 6.4 和 +4.6。

3.2 消融實(shí)驗(yàn)

1）蒸餾不同關(guān)系

同時蒸餾 QK,VV 關(guān)系而且在計(jì)算關(guān)系的時候有 Softmax 實(shí)現(xiàn)了最好的效果。

2）不同的蒸餾策略

TinyMIM 這種蒸餾關(guān)系的方法實(shí)現(xiàn)了比 MAE 基線模型，class token 蒸餾，特征圖蒸餾都更好的效果，在各種尺寸的模型上都是如此。

3）蒸餾中間層

我們發(fā)現(xiàn)蒸餾第十八層實(shí)現(xiàn)了最好的效果。

四、結(jié)論

在本文中，我們提出了 TinyMIM，它是第一個成功地使小模型受益于掩碼重建建模（MIM）預(yù)訓(xùn)練的模型。我們沒有采用掩碼重建作為任務(wù)，而是通過以知識蒸餾的方式訓(xùn)練小模型模擬大模型的關(guān)系來預(yù)訓(xùn)練小模型。TinyMIM 的成功可以歸功于對可能影響 TinyMIM 預(yù)訓(xùn)練的各種因素的全面研究，包括蒸餾目標(biāo)、蒸餾輸入和中間層。通過大量的實(shí)驗(yàn)，我們得出結(jié)論，關(guān)系蒸餾優(yōu)于特征蒸餾和類標(biāo)記蒸餾等。憑借其簡單性和強(qiáng)大的性能，我們希望我們的方法能夠?yàn)槲磥淼难芯刻峁﹫?jiān)實(shí)的基礎(chǔ)。