微軟亞研提出TinyMIM，用知識蒸餾改進小型ViT

作者：機器之心 2023-09-01 14:49:09

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

一、研究動機

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

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

我們系統性的研究了蒸餾目標、數據增強、正則化、輔助損失函數等對于蒸餾的影響。在嚴格的只用 ImageNet-1K 作為訓練數據的情況下（包括 Teacher model 也只用 ImageNet-1K 訓練）和 ViT-B 作為模型，我們的方法實現了當前最好的性能。如圖所示：

把我們的方法（TinyMIM）和基于掩碼重建的方法 MAE，以及監督式學習的方法從頭開始訓練的 DeiT 作比較。MAE 在模型比較大的時候有顯著的性能提升，但是在模型比較小的時候提升幅度有限甚至會傷害模型的最終效果。我們的方法 TinyMIM 在不同模型的大小上都有大幅提升。

我們的貢獻如下：

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

2. 數據增強和模型正則化（Data and network regularization）：1）用帶掩碼的圖片效果更差；2）學生模型需要一點 drop path，但是 teacher 模型不需要。

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

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

二、方法

我們系統性的調研了蒸餾的目標，輸入的圖片，蒸餾目標模塊。

2.1 影響蒸餾效果的因素

1）特征：

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

當 i=L 時，指的是 Transformer 輸出層的特征。當 i< L 時，指的是 Transformer 中間層的特征。

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

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

c.QKV 特征

在 Attention 層內會有 Q，K，V 特征，這些特征用于計算注意力機制，我們也調研了直接蒸餾這些特征。

2）關系

Q，K，V 用于計算注意力圖，這些特征之間的關系也可以作為知識蒸餾的目標。

3）輸入：是否帶掩碼

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

2.2 知識蒸餾方法對比

1）Class Token 蒸餾：

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

其中指學生模型的 class token，而指老師模型的 class token。

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

3）關系蒸餾：我們提出了也是本文默認的蒸餾策略

三、實驗

3.1 主要實驗結果

我們的方法在 ImageNet-1K 上預訓練，而且教師模型也是在 ImageNet-1K 預訓練。然后我們將我們預訓練的模型在下游任務（分類、語義分割）上進行了微調。模型表現如圖：

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

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

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

3.2 消融實驗

1）蒸餾不同關系

同時蒸餾 QK,VV 關系而且在計算關系的時候有 Softmax 實現了最好的效果。

2）不同的蒸餾策略

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

3）蒸餾中間層

我們發現蒸餾第十八層實現了最好的效果。

四、結論

在本文中，我們提出了 TinyMIM，它是第一個成功地使小模型受益于掩碼重建建模（MIM）預訓練的模型。我們沒有采用掩碼重建作為任務，而是通過以知識蒸餾的方式訓練小模型模擬大模型的關系來預訓練小模型。TinyMIM 的成功可以歸功于對可能影響 TinyMIM 預訓練的各種因素的全面研究，包括蒸餾目標、蒸餾輸入和中間層。通過大量的實驗，我們得出結論，關系蒸餾優于特征蒸餾和類標記蒸餾等。憑借其簡單性和強大的性能，我們希望我們的方法能夠為未來的研究提供堅實的基礎。

責任編輯：張燕妮來源：機器之心

AI 微軟