NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度

微軟研究院AI頭條發表於2020-11-26
從數億數量級的神經網路結構中搜尋出高效能的網路,是一個充滿挑戰但又令人嚮往的研究任務。正如深度學習變革了傳統的手工影像特徵設計,神經網路結構的設計也正在逐漸由手工設計轉變為演算法自動設計。

面對數以億級的網路結構,將每一個可能的結構都訓練收斂,並選擇其中最好的結構是不現實的。在近期的神經網路設計研究中,一個被廣泛使用的解決方法是先訓練一個包含了所有可能結構的超網(hypernetwork),當測試某一個網路結構的效能時,直接繼承超網訓練後的引數。這樣的做法省去了重新訓練的時間,大大加快了網路搜尋的速度。然而,雖然預訓練超網的方法能夠大幅度加速網路搜尋,但因為很難對所有的路徑(子模型)進行充分訓練,所以其給出的結果往往不能真正體現出眾多子網路正確的排序關係。

為了解決這一問題,微軟亞洲研究院的研究員們提出維護一個優先路徑組(prioritized path board)。也就是說,在訓練超網的某一條路徑時,使用元網路(meta-network)從組中選出一條效能較好的子網對其進行網路蒸餾(distillation),從而提升超網的收斂程度與效能。採用這種方式選出的網路在 ImageNet 上的分類準確率達到了80.0%, 超越了現有的 EfficientNet-B0/B1 和 MobileNetV3。該論文已被 NeurIPS 2020 接收

基於優先路徑的網路蒸餾


現有的超網訓練方式多種多樣,研究員們採用了簡單有效的均勻取樣單條路徑(single-path uniform sampling strategy)作為基礎,即每一次隨機從超網中採一條路徑進行訓練。與之不同的是,在訓練單條路徑時,會從優先路徑組中選出一條路徑對其進行蒸餾,從而提升訓練效果。

NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度

圖1:方法示意圖,左側為常規的蒸餾方法,即採用一個預訓練的模型進行蒸餾,右側為提出的基於優先路徑的蒸餾方法。

優先路徑組

優先路徑組是由少量效能優異的路徑構成的。超網訓練會對優先路徑組進行動態的維護,如果取樣出來的網路在效果和複雜度上均優於有限路徑組中的網路,那麼就會將其替換到優先路徑組中。不僅如此,維護這樣一個優先路徑組還使得超網訓練結束後可以直接從中選取最優路徑,從而節省以往方法在進行網路結構搜尋時運用強化學習方法或進化演算法(Evolution Algorithm)的時間。在選取優先路徑時,可根據公式

NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度

進行選擇,其中 NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度代表元網路,NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度代表網路最後一層輸出的特徵(logits),NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度代表元網路所預測的兩條路徑的匹配程度。

知識蒸餾

知識蒸餾是一種被廣泛應用的模型壓縮方法,透過讓小型網路來模仿預訓練大型網路的最後一層輸出特徵,可以使小型網路達到接近於大型網路的表現。研究員們透過優先路徑來進行蒸餾,從而無需提前訓練一個大型神經網路。對超網進行更新的具體公式如下:

NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度

其中 NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度 為正常的訓練損失,NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度 為蒸餾損失,NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度 仍為前面所提到的匹配程度。

元網路

由於不同的子網結構差異可能非常之大,因此研究員們希望能夠從優先路徑組中選出最有助於子網訓練的優先路徑對其進行知識蒸餾。經過嘗試,研究員們採用了元網路的技術,將取樣路徑和優先路徑最後一層輸出的特徵差輸入到元網路中,由元網路判斷優先路徑與取樣路徑的匹配程度。當訓練元網路時,研究員們會在驗證集上計算損失,透過匹配程度 NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度進行梯度回傳並更新元網路 NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度

NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度
實驗結果

對基於優先路徑蒸餾的網路結構搜尋演算法的測試是在 ImageNet 上進行的。實驗結果如圖2和表1所示。可以看出,在各種模型大小下,該方法的搜尋結果均超越了此前的 EfficientNet-B0/B1 和 MobileNetV3,實現了優越的效能。不僅如此,該方法搜尋所需要的時長也是各種網路結構搜尋演算法中最短的。

NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度

圖2:基於優先路徑蒸餾的網路結構搜尋得到的模型在 ImageNet 上的效能

NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度

表1:基於優先路徑蒸餾的網路結構搜尋得到的模型在 ImageNet 上的效能

除了影像分類任務外,研究員們還在物體檢測任務上對演算法進行了測試,結果如表2所示。可以看到,該模型同樣能夠泛化到物體檢測任務上,在各種指標下都比已有模型有近1%的提升。

NeurIPS 2020 | 百裡挑一:如何加速超網訓練的收斂和搜尋速度

表2:基於優先路徑蒸餾的網路結構搜尋模型在物體檢測任務上的表現

基於NNI工具介面的原始碼實現

NNI (Neural Network Intelligence) 是當下最熱門的開源自動機器學習(AutoML)工具之一,由微軟亞洲研究院與微軟(亞洲)網際網路工程院領銜開發。NNI 對機器學習生命週期的各個環節都做了較為全面的支援,包括特徵工程神經網路架構搜尋、超參調優和模型壓縮。

目前,微軟亞洲研究院的研究員們已將此次提出的基於優先路徑蒸餾的網路結構搜尋演算法透過 NNI 平臺的 One-Shot 演算法介面進行了實現,提供了完整的搜尋、重訓練和測試程式碼以及模型。由於 NNI 提供了統一的介面表達網路搜尋空間,所以有對比此演算法與其他神經網路架構搜尋結果需求的使用者可選擇這份程式碼實現做參考。程式碼以及更多技術細節,請參見:https://github.com/microsoft/nni。

結語

本篇 NeurIPS 2020 論文針對網路結構搜尋中超網訓練不充分的問題,提出了使用優先路徑進行蒸餾的方法,加快了超網訓練的收斂和搜尋速度,且搜尋得到的網路效能超越了大多數現有演算法。儘管如此,該方法目前仍處於初步的探索階段,未來可能還會將模型延遲考慮到優先路徑選擇中,並對優先路徑蒸餾的理論依據進行進一步的探索。

更多技術細節,詳見論文:Cream of the Crop: Distilling Prioritized Paths For One-Shot Neural Architecture Search

論文連結:https://arxiv.org/abs/2010.15821

程式碼連結:https://github.com/microsoft/Cream

NNI實現連結:https://github.com/microsoft/nni

參考文獻

[1] Hieu Pham, Melody Guan, Barret Zoph, Quoc Le, and Jeff Dean. Efficient neural architecture search via parameters sharing. In ICML, 2018

[2] Mingxing Tan and Quoc V. Le. Efficientnet: Rethinking model scaling for convolutional neural networks. In ICML, 2019

[3] Andrew Howard, Mark Sandler, Grace Chu, Liang-Chieh Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, et al. Searching for mobilenetv3. In ICCV, 2019.

[4] Zichao Guo, Xiangyu Zhang, Haoyuan Mu, Wen Heng, Zechun Liu, Yichen Wei, and Jian Sun. Single path one-shot neural architecture search with uniform sampling. In ECCV, 2020.

[5] Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han. Once for all: Train one network and specialize it for efficient deployment. In ICLR, 2020.

[6] Changlin Li, Jiefeng Peng, Liuchun Yuan, Guangrun Wang, Xiaodan Liang, Liang Lin, and Xiaojun Chang. Blockwisely supervised neural architecture search with knowledge distillation. In CVPR, 2020

[7] Jiahui Yu, Pengchong Jin, Hanxiao Liu, GabrielBender, Pieter-Jan Kindermans, Mingxing Tan, Thomas Huang, Xiaodan Song,Ruoming Pang, and Quoc Le. Bignas: Scaling up neural architecture search with big single-stage models. In ECCV, 2020.

[8] Hieu Pham, Qizhe Xie, Zihang Dai, and Quoc V Le. Meta pseudo labels. arXiv:2003.10580, 2020

[9] Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: Differentiable architecture search. In ICLR, 2019

相關文章