特征蒸餾：為對比學(xué)習(xí)正名！效果堪比圖像掩碼方法！

基于圖像掩碼（MIM，Masked image modeling）的自監(jiān)督方法讓ViT實(shí)現(xiàn)了更好的微調(diào)性能，比如基于MAE預(yù)訓(xùn)練的ViT-B可以在ImageNet-1K數(shù)據(jù)集達(dá)到83.6%準(zhǔn)確度，這要超過之前基于對比學(xué)習(xí)的模型，如基于DINO的ViT-B只能達(dá)到82.8%。近日微軟AI（SwinTransformer原團(tuán)隊(duì)）在論文Contrastive Learning Rivals Masked Image Modeling in Fine-tuning via Feature Distillation提出了一種簡單的后處理方法來優(yōu)化對比學(xué)習(xí)得到的預(yù)訓(xùn)練模型：通過簡單的特征蒸餾（FD，feature distillation）來將原來的模型轉(zhuǎn)換成和MIM有相似性質(zhì)的新模型。對對比學(xué)習(xí)的模型進(jìn)行特征蒸餾其在ImageNet-1K數(shù)據(jù)集上的微調(diào)性能可以和基于MIM的模型不相上下。這種特征蒸餾優(yōu)化方法具有普適性，不僅可以優(yōu)化對比學(xué)習(xí)模型，也可以用于有監(jiān)督模型，MIM模型以及基于圖像-文本對比學(xué)習(xí)的CLIP模型。對CLIP ViT-L模型進(jìn)行特征蒸餾，其可以在ImageNet-1K數(shù)據(jù)集上達(dá)到89.0%，這是目前ViT-L所能實(shí)現(xiàn)的SOTA，而對30億參數(shù)的SwinV2-G模型進(jìn)行特征蒸餾，可以在ADE20K語義分割數(shù)據(jù)集上達(dá)到新的SOTA：61.4 mIoU。

特征蒸餾

論文提出的特征蒸餾方法非常簡單，其整體架構(gòu)如下所示，這里預(yù)訓(xùn)練的模型作為teacher模型，而要轉(zhuǎn)換的新模型為student模型。這里的特征蒸餾主要有以下4個(gè)要點(diǎn)：（1）蒸餾的目標(biāo)采用特征圖

常規(guī)的知識蒸餾常常采用logits作為蒸餾目標(biāo)，而這里的特征蒸餾采用模型輸出的特征圖作為目標(biāo)，這也適用于那些沒有l(wèi)ogits輸出的模型如CLIP。為了讓teacher和student輸出的特征圖一致，teacher和student模型的輸入是原始圖像的同一數(shù)據(jù)增強(qiáng)視野（這里只采用RandomResizeAndCrop 0.08-1數(shù)據(jù)增強(qiáng)）。同時(shí)在student模型上加了一個(gè)projector，采用一個(gè)1x1卷積層，這其實(shí)也讓特征蒸餾變得具有普適性，比如teacher模型和student模型的輸出特征圖的維度不一致。采用特征圖作為蒸餾目標(biāo)，比采用logits或者GAP特征（全局池化后特征）在效果上也更好，ViT-B的對比結(jié)果如下所示：（2）對teacher的特征進(jìn)行白化（whiten）

不同的預(yù)訓(xùn)練模型輸出的特征的數(shù)量級不同，這樣不同的模型其蒸餾訓(xùn)練的超參數(shù)都需要調(diào)整。這里對teacher輸出的特征進(jìn)行白化處理來將特征歸一化同一量級，在實(shí)現(xiàn)上采用非訓(xùn)練的LayerNorm（去掉訓(xùn)練的weight和bias）。這里的損失函數(shù)采用smooth L1 loss，如下所示：

論文共選擇了5種不同的預(yù)訓(xùn)練策略：DINO，EsViT，CLIP， DeiT和MAE，其中DINO，EsViT和DeiT采用ViT-B和Swin-B進(jìn)行實(shí)驗(yàn)，而CLIP選擇ViT-B和ViT-L進(jìn)行實(shí)驗(yàn)。特征蒸餾在ImageNet-1K訓(xùn)練集上進(jìn)行，訓(xùn)練epochs為300。下表對比了不同的預(yù)訓(xùn)練模型在特征蒸餾后在ImageNet-1K數(shù)據(jù)集上微調(diào)和linear probe，以及在ADE20K數(shù)據(jù)集上遷移的效果?？梢钥吹紻INO，EsViT，CLIP， DeiT方法預(yù)訓(xùn)練的模型經(jīng)過特征蒸餾后均可以在ImageNet-1K數(shù)據(jù)集上微調(diào)效果上提升1～2%，也可以在ADE20K數(shù)據(jù)集上提升1～3.3 mIoU。這也使得對比學(xué)習(xí)的預(yù)訓(xùn)練模型效果和基于MIM的預(yù)訓(xùn)練模型（如MAE和BEiT）效果相當(dāng)了。

特征可視化

為了進(jìn)一步研究為什么特征蒸餾會起作用，論文對特征蒸餾前后的模型做了特征可視化分析，主要包括以下幾個(gè)部分：

（1）average attention distance per head

首先可視化了各個(gè)不同layer的各個(gè)head的average attention distance，average attention distance反應(yīng)的是attention head的感受野大小。如下左圖所示，可以看到對于特征蒸餾前的模型，其在深層的attention heads表現(xiàn)出了相似的感受野，這說明了不同的attention heads學(xué)習(xí)到了相似的模式，這是可能對模型容量的浪費(fèi)；而經(jīng)過特征蒸餾后，深層的attention heads的感受野相對發(fā)散一些，這也意味著模型的表達(dá)力可能更好。我們同樣也可以從右圖中的各個(gè)heads的average cosine similarity看出來差異（特征蒸餾后，深層的余弦相似度差異變得更大）。

（2）attention 模式的改變

第二個(gè)分析是對attention map可視化，這里不同layer的average attention maps（不同heads平均）如下所示，這里attention map主要有兩種模式：diagonal模式和column模式，其中diagonal模式說明模型學(xué)習(xí)到了不同patchs間的一些固定的相對位置，而column模式說明模型學(xué)習(xí)到的是patchs的固定位置。當(dāng)然diagonal模式要比column模式更好，我們也可以看到特征蒸餾后，模型的diagonal模式要更加突出，其中共享的RPB也是其中的一部分原因。（3）loss / accuracy landscapes

最后論文也分析了loss / accuracy landscapes，如下圖所示，我們可以看到特征蒸餾后模型的loss / accuracy landscapes變得更平坦，網(wǎng)絡(luò)更容易進(jìn)行優(yōu)化。

論文也對基于MAE的預(yù)訓(xùn)練模型進(jìn)行了特征蒸餾，不過效果提升的不是特別明顯（ViT-B從83.6% -> 83.8%）。如果同樣進(jìn)行特征可視化的話，可以看到MAE本身得到的模型其average attention distance就比較發(fā)散，同時(shí)loss / accuracy landscapes也比較平坦，特征蒸餾后變化不是特別大。這其實(shí)也側(cè)面反映了通過對對比學(xué)習(xí)得到的模型進(jìn)行特征蒸餾可以轉(zhuǎn)換為和MIM預(yù)訓(xùn)練模型具有相似性質(zhì)的新模型，從而提升其微調(diào)效果。在論文的最后，作者認(rèn)為一個(gè)好的預(yù)訓(xùn)練模型應(yīng)該具有好的優(yōu)化友好度（optimization friendliness），簡單來說，就是一個(gè)好的預(yù)訓(xùn)練模型來進(jìn)行初始化，它應(yīng)該比較容易進(jìn)行微調(diào)（比如有平坦的loss landscape）。而特征蒸餾則是增加了預(yù)訓(xùn)練模型的優(yōu)化友好度。

小結(jié)

同樣架構(gòu)的模型，其最終的性能也要取決于優(yōu)化策略。無論是有監(jiān)督，還是基于MIM的自監(jiān)督，甚或是基于對比學(xué)習(xí)的自監(jiān)督，我們?nèi)绻苷业胶玫膬?yōu)化技巧，它們應(yīng)該可以得到相似的性能：DeiT III證明了有監(jiān)督，而這個(gè)工作證明了對比學(xué)習(xí)自監(jiān)督。

參考

• Contrastive Learning Rivals Masked Image Modeling in Fine-tuning via Feature Distillation
• https://github.com/SwinTransformer/Feature-Distillation