自知識(shí)蒸餾方法在2014年被首次提出以來(lái)，其開(kāi)始廣泛被應(yīng)用于模型壓縮領(lǐng)域。在更強(qiáng)大教師模型輔助監(jiān)督信息的幫助下，學(xué)生模型往往能夠?qū)崿F(xiàn)比直接訓(xùn)練更高的精度。然而，現(xiàn)有的知識(shí)蒸餾相關(guān)研究只考慮了同架構(gòu)模型的蒸餾方法，而忽略了教師模型與學(xué)生模型異構(gòu)的情形。例如，最先進(jìn)的MLP模型在ImageNet上僅能達(dá)到83%的精度，無(wú)法獲取精度更高的同架構(gòu)教師模型以使用知識(shí)蒸餾方法進(jìn)一步提高M(jìn)LP模型的精度。因此，對(duì)異構(gòu)模型知識(shí)蒸餾的研究具有實(shí)際應(yīng)用意義。

本文的研究者們分析了針對(duì)異構(gòu)模型（CNN，ViT，MLP）特征的差異性，指出特征中模型架構(gòu)相關(guān)的信息會(huì)阻礙知識(shí)蒸餾的過(guò)程。基于此觀察，研究者們提出了名為OFAKD異構(gòu)模型知識(shí)蒸餾方法：該方法將特征映射到架構(gòu)無(wú)關(guān)的統(tǒng)一空間進(jìn)行異構(gòu)模型蒸餾，并使用一種能夠自適應(yīng)增強(qiáng)目標(biāo)類別信息的損失函數(shù)。在CIFAR-100和ImageNet數(shù)據(jù)集上，該方法實(shí)現(xiàn)了對(duì)現(xiàn)有同架構(gòu)知識(shí)蒸餾方法的超越。

異構(gòu)模型間的特征差異

圖1 異構(gòu)模型學(xué)習(xí)到的特征對(duì)比

相比于僅使用logits的蒸餾方法，同步使用模型中間層特征進(jìn)行蒸餾的方法通常能取得更好的性能。然而在異構(gòu)模型的情況下，由于不同架構(gòu)模型對(duì)特征的不同學(xué)習(xí)偏好，它們的中間層特征往往具有較大的差異，直接將針對(duì)同架構(gòu)模型涉及的蒸餾方法遷移到異構(gòu)模型會(huì)導(dǎo)致性能下降。

通用的異構(gòu)模型蒸餾方法

圖2 異構(gòu)模型的知識(shí)蒸餾方法

為了在異構(gòu)模型蒸餾過(guò)程中利用中間層特征，需要排除特征中模型架構(gòu)相關(guān)信息的干擾，僅保留任務(wù)相關(guān)信息。基于此，研究者們提出通過(guò)將學(xué)生模型的中間層特征映射到logits空間，實(shí)現(xiàn)對(duì)模型架構(gòu)相關(guān)信息的過(guò)濾。此外通過(guò)在原始基于KL散度的蒸餾損失函數(shù)中引入一項(xiàng)額外的調(diào)節(jié)系數(shù)，修正后的損失函數(shù)能夠?qū)崿F(xiàn)對(duì)目標(biāo)類別信息的自適應(yīng)增強(qiáng)，進(jìn)一步減緩異構(gòu)模型蒸餾時(shí)無(wú)關(guān)信息的干擾。

圖3 原始蒸餾損失與改進(jìn)后蒸餾損失的對(duì)比

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

圖4 在ImageNet上的異構(gòu)模型蒸餾結(jié)果

上表展示了在ImageNet上的異構(gòu)蒸餾結(jié)果。在所有架構(gòu)的六種可能異構(gòu)組合中，本文OFAKD方法都得到了超越現(xiàn)有方法結(jié)果。

圖5 不同值對(duì)結(jié)果的影響

上表在ImageNet上比較了不同的值設(shè)置對(duì)結(jié)果的影響。可以看出，通過(guò)選取合適的值設(shè)置，改進(jìn)后的蒸餾損失函數(shù)能得到超越原始蒸餾損失函數(shù)的結(jié)果。

圖6 在ImageNet上的同構(gòu)模型蒸餾結(jié)果

本文在ImageNet上與傳統(tǒng)同構(gòu)模型蒸餾方法進(jìn)行了對(duì)比。在常見(jiàn)的ResNet34和ResNet18同構(gòu)教師學(xué)生模型組合上，OFAKD也具有與現(xiàn)有SOTA方法相當(dāng)?shù)谋憩F(xiàn)。

圖7 MLP學(xué)生模型的蒸餾結(jié)果

最后，本文比較了文首提到的MLP作為學(xué)生模型時(shí)的蒸餾結(jié)果。通過(guò)選用ViT架構(gòu)的BEiT v2-base作為教師模型，僅中等尺寸的CycleMLP-B3就刷新了MLP模型在ImageNet上的最佳結(jié)果。

結(jié)論

本文研究了異構(gòu)模型之間的知識(shí)蒸餾方法，通過(guò)將學(xué)生模型中間層特征映射到logits空間來(lái)擬合教師模型最終輸出，并使用在原始知識(shí)蒸餾損失函數(shù)基礎(chǔ)上改進(jìn)而來(lái)的自適應(yīng)目標(biāo)信息增強(qiáng)損失，提出的OFAKD方法在多種數(shù)據(jù)集和教師學(xué)生模型組合上實(shí)現(xiàn)了對(duì)現(xiàn)有方法的超越，擴(kuò)展了知識(shí)蒸餾的應(yīng)用范圍。