OVO：在線蒸餾一次視覺Transformer搜索

論文地址：https://arxiv.org/pdf/2212.13766.pdf

最近，Pure transformers在視覺任務(wù)方面顯示出巨大的潛力。然而，它們在中小數(shù)據(jù)集中的準確性并不令人滿意。盡管一些現(xiàn)有方法引入了CNN作為教師，通過蒸餾來指導(dǎo)訓(xùn)練過程，但教師和學(xué)生網(wǎng)絡(luò)之間的差距將導(dǎo)致次優(yōu)績效。在這項工作中，研究員提出了一種新的具有在線蒸餾的One-shot Vision transformer搜索框架，即OVO。OVO對教師和學(xué)生網(wǎng)絡(luò)的子網(wǎng)進行采樣，以獲得更好的蒸餾結(jié)果。得益于在線蒸餾，超網(wǎng)中的數(shù)千個子網(wǎng)訓(xùn)練有素，無需額外的微調(diào)或再訓(xùn)練。在實驗中，OVO Ti在ImageNet和CIFAR-100上分別達到73.32%和75.2%的top-1精度。

二、ViT超網(wǎng)訓(xùn)練的困境

在AutoFormer中的超網(wǎng)訓(xùn)練期間，在每個訓(xùn)練迭代中均勻地采樣子網(wǎng)α=（α（1）。。。α（i）。。。α（l））。采樣權(quán)重w=（w（1）。。。w（i）。。。w（l））被更新，而其余部分被凍結(jié)。然而，超網(wǎng)訓(xùn)練對于ViT來說并不穩(wěn)定，每個超網(wǎng)需要很長的訓(xùn)練周期（500個周期）才能獲得其組裝子網(wǎng)的滿意性能。其他方法使用三明治采樣策略，對多個子網(wǎng)絡(luò)（包括最大、最小和兩個隨機選擇）進行采樣，然后將每個小批次的梯度聚集在一起。當采用三明治采樣策略時，超網(wǎng)訓(xùn)練的計算成本更高。

三、超網(wǎng)在線蒸餾訓(xùn)練

當訓(xùn)練教師和學(xué)生網(wǎng)絡(luò)時，研究者在[Single path one-shot neural architecture search with uniform sampling]中提出的每次迭代中更新一個隨機采樣的子網(wǎng)絡(luò)。使用GT標簽訓(xùn)練來自教師超網(wǎng)絡(luò)的子網(wǎng)絡(luò)，同時使用KD訓(xùn)練采樣的學(xué)生網(wǎng)絡(luò)。配備在線蒸餾，一次性NAS能夠以高效和有效的方式搜索Transform架構(gòu)。與經(jīng)典的一次性NAS方法相比，新提出的方法有兩個優(yōu)點。1） 更快的收斂。來自CNN的知識提供了電感偏置，這有助于每個Transform塊比之前的獨立訓(xùn)練更快地收斂。2） 更好的子網(wǎng)性能。通過在線蒸餾訓(xùn)練的子網(wǎng)可以在中小數(shù)據(jù)集上獲得更好的性能。

四、Search Pipeline

在超網(wǎng)完成訓(xùn)練后，進行進化搜索，以最大化分類精度選擇子網(wǎng)絡(luò)。在進化搜索開始時，隨機抽取N個架構(gòu)作為種子。所有子網(wǎng)絡(luò)都基于從超網(wǎng)繼承的權(quán)重在驗證數(shù)據(jù)集上進行評估。選擇前k個架構(gòu)作為父代，通過交叉和變異生成下一代。對于交叉，兩個隨機選擇的父網(wǎng)絡(luò)交叉，以在每一代中產(chǎn)生一個新網(wǎng)絡(luò)。當進行突變時，父網(wǎng)絡(luò)首先以概率Pd對其深度進行突變，然后以概率Pm對每個塊進行突變，以產(chǎn)生新的架構(gòu)。

五、實驗

將搜索到的最優(yōu)模型的性能與ImageNet上最先進的CNN和ViT的性能進行了比較。在ImageNet-1K上訓(xùn)練OVO的超集合，并使用指定的參數(shù)大小搜索目標Transform模型。在超網(wǎng)完成訓(xùn)練后，子網(wǎng)直接繼承權(quán)重，而無需額外的再訓(xùn)練和其他后處理。下表中報告了性能。很明顯，OVO比其他最先進的模型具有更高的精度。