本文提出了 TinyMIM，它是第一個成功地使小模型受益于 MIM 預訓練的模型,。
掩碼建模（MIM, MAE）被證明是非常有效的自監(jiān)督訓練方法。然而,，如圖 1 所示,，MIM 對于更大的模型效果相對更好。當模型很小的時候（比如 ViT-T 5M 參數(shù),，這樣的模型對于現(xiàn)實世界非常重要）,，MIM 甚至可能一定程度上降低模型的效果,。比如用 MAE 訓練的 ViT-L 比普通監(jiān)督訓練的模型在 ImageNet 上的分類效果提升 3.3%,，但是用 MAE 訓練的 ViT-T 比普通監(jiān)督訓練的模型在 ImageNet 上的分類效果降低了 0.6%。在這篇工作中我們提出了 TinyMIM,，其在保持 ViT 結(jié)構(gòu)不變并且不修改結(jié)構(gòu)引入其他歸納偏置（inductive bias）的基礎(chǔ)上,、用蒸餾的方法遷移大模型上的知識到小模型。
- 論文地址：https:///pdf/2301.01296.pdf
- 代碼地址：https://github.com/OliverRensu/TinyMIM
我們系統(tǒng)性的研究了蒸餾目標,、數(shù)據(jù)增強,、正則化,、輔助損失函數(shù)等對于蒸餾的影響。在嚴格的只用 ImageNet-1K 作為訓練數(shù)據(jù)的情況下（包括 Teacher model 也只用 ImageNet-1K 訓練）和 ViT-B 作為模型,，我們的方法實現(xiàn)了當前最好的性能,。如圖所示：
把我們的方法（TinyMIM）和基于掩碼重建的方法 MAE，以及監(jiān)督式學習的方法從頭開始訓練的 DeiT 作比較,。MAE 在模型比較大的時候有顯著的性能提升,，但是在模型比較小的時候提升幅度有限甚至會傷害模型的最終效果。我們的方法 TinyMIM 在不同模型的大小上都有大幅提升,。1. 蒸餾的目標（Distillation targets）:1）蒸餾 token 之間的關(guān)系比單獨蒸餾 class token 或者特征圖（feature map）更有效,；2）用中間層作為蒸餾的目標更有效。2. 數(shù)據(jù)增強和模型正則化（Data and network regularization）：1）用帶掩碼的圖片效果更差,；2）學生模型需要一點 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)研了蒸餾的目標,，輸入的圖片,，蒸餾目標模塊。
當 i=L 時,，指的是 Transformer 輸出層的特征,。當 i< L 時，指的是 Transformer 中間層的特征,。b. 注意力（Attention）特征和前饋層（FFN）層特征
Transformer 每一個 block 有 Attention 層和 FFN 層,，蒸餾不同的層會帶來不同的影響。
在 Attention 層內(nèi)會有 Q,，K,，V 特征，這些特征用于計算注意力機制,，我們也調(diào)研了直接蒸餾這些特征,。
Q，K,，V 用于計算注意力圖,，這些特征之間的關(guān)系也可以作為知識蒸餾的目標。傳統(tǒng)的知識蒸餾是直接輸入完整的圖片,。我們的方法為了探索蒸餾掩碼建模模型,，所以我們也探索了帶掩碼的圖片是否適合作為知識蒸餾時候的輸入。最簡單的方法就是類似 DeiT 直接蒸餾 MAE 預訓練模型的 class token:
其中 指學生模型的 class token,，而  指老師模型的 class token,。2）特征蒸餾：我們直接參考了 feature distillation [1] 作為對比
3）關(guān)系蒸餾：我們提出了也是本文默認的蒸餾策略
我們的方法在 ImageNet-1K 上預訓練，而且教師模型也是在 ImageNet-1K 預訓練,。然后我們將我們預訓練的模型在下游任務(wù)（分類,、語義分割）上進行了微調(diào)。模型表現(xiàn)如圖：
我們的方法顯著超過之前基于 MAE 的方法,，尤其是小模型,。具體來講，對于超小的模型 ViT-T,，我們的方法實現(xiàn)了 75.8% 的分類準確性,，相比 MAE 基線模型實現(xiàn)了 4.2 的提升。對于小模型 ViT-S,，我們實現(xiàn)了 83.0% 的分類準確性,，比之前最好的方法提升了 1.4。對于 Base 尺寸的模型,，我們的方法分別超過 MAE 基線模型和以前最好的模型 CAE 4.1 和 2.0,。
TinyMIM-B 對比 MAE-B,，在 ImageNet-A 和 ImageNet-R 分別提升了 + 6.4 和 +4.6,。
同時蒸餾 QK,VV 關(guān)系而且在計算關(guān)系的時候有 Softmax 實現(xiàn)了最好的效果。
TinyMIM 這種蒸餾關(guān)系的方法實現(xiàn)了比 MAE 基線模型,，class token 蒸餾,，特征圖蒸餾都更好的效果，在各種尺寸的模型上都是如此,。
我們發(fā)現(xiàn)蒸餾第十八層實現(xiàn)了最好的效果,。在本文中,，我們提出了 TinyMIM,，它是第一個成功地使小模型受益于掩碼重建建模（MIM）預訓練的模型。我們沒有采用掩碼重建作為任務(wù),，而是通過以知識蒸餾的方式訓練小模型模擬大模型的關(guān)系來預訓練小模型,。TinyMIM 的成功可以歸功于對可能影響 TinyMIM 預訓練的各種因素的全面研究，包括蒸餾目標,、蒸餾輸入和中間層,。通過大量的實驗，我們得出結(jié)論,，關(guān)系蒸餾優(yōu)于特征蒸餾和類標記蒸餾等,。憑借其簡單性和強大的性能，我們希望我們的方法能夠為未來的研究提供堅實的基礎(chǔ),。[1] Wei, Y., Hu, H., Xie, Z., Zhang, Z., Cao, Y., Bao, J., ... & Guo, B. (2022). Contrastive learning rivals masked image modeling in fine-tuning via feature distillation. arXiv preprint arXiv:2205.14141.AI offer最優(yōu)解—2024屆秋招線上宣講會 「TalentAI」將于9月-10月舉辦「AI offer最優(yōu)解—2024屆秋招線上宣講會」,，邀請企業(yè)介紹各自的秋招計劃與崗位詳情。未來兩周直播安排如下,，對螞蟻集團,、科大訊飛研究院、第四范式,、百圖生科,、同花順感興趣的同學別錯過。
|
