?作者 | 韓翔宇

學校 |?南昌大學

研究方向 | 神經網絡架構搜索

概述

本文是 ICML?2021 收錄的 NAS（神經網絡結構搜索）領域的重磅論文，其作者是 AttentiveNas 的原作者，在 AttentiveNas 的基礎上，添加了 α- 散度損失函數，在 ImageNet（NAS 方向）排行榜中取得了 SOTA 結果。

論文標題：

AlphaNet: Improved Training of Supernets with Alpha-Divergence

論文鏈接：

https://arxiv.org/abs/2102.07954

代碼鏈接：

https://github.com/facebookresearch/AlphaNet

解決的問題

盡管權值共享的 NAS 有較好的效果，因為這種方法構建了一個超網，利用子網來訓練超網。然而，權重共享 NAS 的成功很大程度上依賴于將超網絡的知識提取到子網絡。如果廣泛使用蒸餾中的散度，例如 KL 散度，可能導致學生網絡高估或者低估教師網絡的不確定性，導致了子網的效果變差。論文提出了更廣義的 α- 散度來改進超網訓練。通過自適應選擇 α- 散度，避免了高估或者低估教師模型不確定性。改進之后的 AlphaNet 在 ImageNet top-1 的精度達到了 80%，且參數量只有 444M。

簡介

傳統的 NAS 方法代價非常大，需要數百個網絡結構從頭訓練、驗證效果。超網將所有候選體系結構組裝成一個權重共享網絡，每個網絡結構對應一個子網絡。通過同時訓練子網和超網，子網可以直接從超網中獲得的權重，用來重新訓練和驗證，而不需要單獨訓練或微調每個結構，因此成本大大降低。為了穩定超網訓練和提高子網絡的性能，大家廣泛使用知識蒸餾的方式。知識蒸餾用超網中最大的子網預測的軟標簽來監督所有其他子網，提取教師模型的知識來提高子網性能。

一般情況下，知識蒸餾使用 KL 散度衡量師生網絡之間的差異。但是如果學生網絡對教師網絡的某些部分覆蓋不完整，對學生網絡的懲罰大。因此，學生模型往往高估了教師模型的不確定性，不能準確近似教師模型的正確預測。

為了解決這個問題，論文提出用更廣義的 α- 散度代替 KL 散度，具體來講，通過自適應控制散度度量中的 α，可以同時懲罰對教師網絡高估或者低估教師網絡不確定性的行為，鼓勵學生網絡更好地近似教師網絡。直接優化 α 散度可能會受到梯度的高方差的影響，論文游提出了一種簡單的梯度裁剪技術，穩定訓練過程。通過提出的自適應 α- 散度，我們能夠訓練高質量的 alphanet，在 200 到 800 MFLOPs 范圍表現都達到 SOTA 效果。

關于知識蒸餾的相關知識

在權值共享的 NAS 中，知識蒸餾是重要的方法。假設超網有可訓練的參數 θ，訓練超網的目的是學習 θ，讓所有的子網都能同時得到優化，達到較高的準確率。

上圖描述了采用知識蒸餾的超網訓練過程。在每個訓練步驟中，給定一小批數據，對超網和幾個子網絡采樣。當超網使用真實標簽訓練時，所有抽樣的子網絡都使用超網預測的軟標簽進行監督訓練。然后對所有采樣的網絡的梯度進行融合，更新超網參數。在第 t 個 step 時，超參數 θ 被更新為如下（其中 ε 是 step 的數量）：

g 定義如下：其中 LD 是數據集 D 上超網的損失函數（交叉熵）， 為權重系數， 是超網和采樣出子網的 KL 散度， 和 表示：輸入 x 的超網和子網 s 的輸出概率。

知識蒸餾中KL散度的局限性（α-散度解決的問題）

KL 散度廣泛用于衡量教師網絡和學生網絡之間的差異性。但是 KL 散度的缺點在于，當學生網絡高估了教師網絡的不確定性時，不能充分懲罰學生網絡。我們首先列舉一下 KL 散度的公式：

▲ KL散度公式（公式2為f散度，即KL散度的一般形式）

我們用 P 代表教師網絡的某個輸出概率，用 Q 代表學生網絡的某個輸出概率。KL 散度有個避零性質，當 P=0 時，我們看等式右半部分 log(P/Q)，因為分子為 0，我們知道 log 函數趨向于 0 的時候，值趨向于負無窮。無論分母 Q 怎么變大（也就是是學生網絡某個輸出概率變大），值始終是負無窮，對應 KL 散度低，也就是對學生網絡的 Q 懲罰小，所以學生網絡就可以肆無忌憚地預測，即使預測錯誤很離譜，也不會懲罰。

被如下圖所示，針對某個輸出，橙色的教師網絡主要預測的類別是 3（概率最大的是 3，所以輸出為 3），其次是類別 4，預測其他類別的概率小的可以忽略；再看綠色的學生網絡，預測的最主要的類別是 4（概率最大的是 4，所以輸出為 4，輸出的類別就已經和教師網絡有偏差了），而且學生網絡預測其他類別的概率也比教師網絡高，所以學生網絡的不確定性比教師網絡大，換句話說，學生網絡高估了教師網絡的不確定性。

反過來想，我們用 Q 代表教師網絡的某個輸出概率，用 P 代表學生網絡的某個輸出概率（和上個假設掉換個位置）。這次讓 Q=0，這時就要再講 KL 散度的另一個性質：零強制性。用這個例子簡單來簡單解釋，就是分母 Q 為 0，P 也要為 0。這很好理解，當分母為 0 的時候，分子如果不為 0 就會報除 0 錯誤，但是當分子分母都為 0 時，P/Q=1， log(P/Q)=log1=0。

說完這個性質以后，我們想一想，如果 Q 為 0，KL 散度的值為 0，我們希望損失越小越好，我們希望損失比 0 更小，也就是負數。因此最小化 KL 散度會鼓勵 P 避免趨向 0（因為當 P 為 0 的時候，Q 可能為 0，Q一旦為 0，損失函數就一定為 0，也就不能繼續變小了，我們希望損失函數越小越好，當然 P 不愿意趨向 0）。

從論文的角度來說，也就是學生模型會避免低概率模式，也就是學生模型會傾向于較高的概率。如下圖所示，學生模型預測類別 2 的概率為 100%，幾乎沒有不確定性可言，而教師網絡除了預測類別最多的 2，其他類別也有概率分布，所以教師網絡的不確定性比學生網絡大，也就是學生網絡低估了教師網絡的不確定性。

使用α-散度訓練的超網

為了解決上述 KL 散度的不確定性，論文提出了一個靈活的 α- 散度（α∈R 且 α≠0 且 α≠1）：

和 代表每個類的離散分布，共有 個類別。當 時， 的極限就是 ，同樣， 是 時 的極限。 散度的關鍵點在于，可以通過選擇不同的 值來集中懲罰不同類型的差異（低估或高估）。

如圖所示，當阿爾法為負值的時候，藍色線代表了學生網絡 高估了教師網絡 的不確定性的情況，這時 很大。紫色線表示學生網絡 低估了教師網絡 的不確定性的情況，此時 很小。當 為正數的時候，情況正好相反。

為了同時緩解超網訓練時的高估和低估的問題，用一個正的 和一個負的 ，在知識蒸餾損失函數中使用 和 中最大那個，也就是：

總的 KL 散度為：

進一步改進的穩定的α-散度

人們傾向于將 和 設置為較大，以確保學生模型在低估或高估教師模型的不確定性時受到足夠的懲罰。但是直接通過增加 絕對值這種方法，會讓優化變得困難，我們首先來看一下 散度的梯度：

不難看出，如果 很大，則 也可能變得很大，從而影響了訓練的穩定性。為了讓訓練變得穩定，我們把 的最大值限定為 （如果小于 則不變，最大不能超過 ），重新定義梯度表達式為：

該梯度等價于 p 和 q 的 f 散度（f 散度是 KL 散度的一般形式，如果你認真看了這篇文章，應該會注意到 KL 散度的公式有兩個，沒錯，第二個公式就是 f 散度），進行梯度更新相當于最小化有效散度。通過裁剪重要性權重的值，我們優化的仍然是一個散度度量，但對基于梯度的優化更友好。

具體實現

自適應的 α- 散度的設置：α- 和 α+ 分別控制對過高估計和過低估計的懲罰幅度。并且，β 控制了教師模型和學生模型之間的密度比率范圍（也就是子網占了超網多大一部分）。通過實驗發現 AlphaNet 的方法在 α?，α+ 和 β 的選擇范圍內表現相對穩健。實驗中選擇 α?=?1，α+=1 和 β=5.0 作為默認值。作者在 Silimmable Network 和權值共享的 NAS 兩個層面分別做了實驗：

8.1 Slimmable Network部分

Slimmable 網絡是支持選擇多種通道寬度的超網，其搜索空間包含不同寬度的網絡和所有其他參數（深度、卷積類型、kernal 大小）都是相同的。Slimmable 允許不同的設備或應用程序根據設備上的資源限制，自適應地調整模型寬度，以實現最佳精度與能效的權衡。

使用 Slimmable MobileNet v1 和 v2 測試，其中 v1 的寬度范圍是 [0.25,1]，v2 的寬度是 [0.35,1]。在每次訓練迭代中，分別對 channel 寬度最大的最大子網絡、channel 寬度最小的最小子網絡和兩個隨機子網絡進行采樣，累積梯度（這種方式被稱為三明治法則）。使用 gt 標簽訓練超網，使用知識蒸餾訓練采樣到的子網，設置知識蒸餾中的系數 γ=3，作為抽樣獲得的子網數量。

為了驗證自適應 α- 散度的有效性，使用它替換 baseline 中的 KL散度。使用 SGD 優化器訓練 360epoch，動量為 0.9，重量衰減為 10?5,dropout 為 0.2。我們使用余弦學習速率衰減，初始學習速率為 0.8，batch_size 為 2048，使用 16 塊 GPU。

在 ImageNet 上進行訓練，上表是訓練得到的最佳精度，可以看到，無論在任何寬度上，自適應的α. 散度效果均優于 KL 散度（KL-KD）和不加 KL 散度（KD）。

8.2 權值共享的NAS部分

8.2.1 簡介

大多數基于權值共享的 NAS 由以下兩個階段組成（基于強化學習的NAS也一樣）：

• 階段 1：使用可微分權值共享或者看做一個黑盒的優化

• 階段 2：從頭開始訓練深度神經網絡，以獲得最佳的準確率和最終的效果

但是這類 NAS 有缺陷：如果需要不同的硬件約束條件，需要重新搜索。而且要求對所有的候選結果從頭訓練，達到理想的準確率。因此，顯著增加了 NAS 的搜索成本。

論文使用的權重共享的 NAS 是基于超網的權重共享 NAS，搜索過程如下：

階段1：聯合優化搜索空間中的超網和所有可能被采樣的子網，使所有可搜索網絡在訓練結束時都能獲得較好的性能。

階段2：然后所有的子網絡同時被優化。然后可以使用典型的搜索算法，比如進化算法，來搜索感興趣的最佳模型。每個子網絡的模型權值直接從預訓練的超網絡繼承而來，無需再進行再訓練或微調。

與基于 RL 的 NAS 算法和可微 NAS 算法相比，基于超網絡的權重共享的優勢主要有：

1. 只需要對超網進行一次訓練。搜索空間中定義的所有子網絡在第1階段完全優化后即可使用。不需要再訓練或微調；

2. 在階段 1 中，所有不同模型規模的子網絡進行聯合優化，找到一組帕累托最優模型，這一組模型天生支持各種需要考慮的資源。

注意，權重共享 NAS 的一個主要步驟是同時訓練搜索空間中指定的所有子網絡收斂。與訓練 Slimmable 神經網絡類似，也是用基于 KL 散度的知識蒸餾方法，強制所有的子網絡使從超網中學習來實現的。

8.2.2 具體訓練

為了簡單起見，我們使用均勻抽樣策略，使用和 Silimable 的三明治法則相同的方式，每次迭代訓練四個網絡。使用 SGD 和余弦淬火策略衰減學習率，使用 AutoAugment 進行數據增強，設置數據的標簽平滑為 0.1。

上表是論文采用的搜索空間，MBConv 是 mobileNet 的倒殘差模塊。使用 swish 激活函數，channel width 表示塊的輸出維度。輸入的分辨率表示候選的分辨率，為了簡化數據加載過程，預選選取固定大小 224*224，然后再使用雙三次插值將它們重新縮放到我們的目標分辨率。

8.2.3 驗證

我們比較通過知識蒸餾策略得到不同超網，評價采用準確率和 FLOPs 的帕累托最優，主要包括以下三點：

1. 首先從超網中隨機抽取 512 個子網，并估計它們在 ImageNet 驗證集上的準確性。

2. 對性能最好的 128 個子網應用交叉和隨機變異，將交叉規模和變異規模都固定為 128，共產生 256 個新的子網。然后我們評估了這些子網的性能。

3. 重復步驟 20 次步驟 2，得到 5376 個子網。

8.2.4 最終效果

上表可以看出， AlphaNet 相比于普通的 KL 散度和不使用 KL 散度的訓練方法，在各個參數量階段的準確率都得到了提升。

上圖分別是不用 KL 散度（棕色），使用 KL 散度（綠色），AlphaNet（紅色）的訓練收斂曲線，可以看出，使用 AlphaNet 訓練在前期收斂速度會慢于普通的 KL 散度，但是在后期收斂速度會超過普通的 KL 散度。

論文將 AlphaNet 和現有的 NAS 做對比，相比于 EfficientNet、MobileV3、MNasNet、BigNas 等模型都有較為明顯的提升，取得了 SOTA 的效果。

更多閱讀

#投 稿?通 道#

?讓你的文字被更多人看到?

如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。

總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學術熱點剖析、科研心得或競賽經驗講解等。我們的目的只有一個，讓知識真正流動起來。

📝?稿件基本要求：

? 文章確系個人原創作品，未曾在公開渠道發表，如為其他平臺已發表或待發表的文章，請明確標注?

? 稿件建議以?markdown?格式撰寫，文中配圖以附件形式發送，要求圖片清晰，無版權問題

? PaperWeekly 尊重原作者署名權，并將為每篇被采納的原創首發稿件，提供業內具有競爭力稿酬，具體依據文章閱讀量和文章質量階梯制結算

📬?投稿通道：

? 投稿郵箱：hr@paperweekly.site?

? 來稿請備注即時聯系方式（微信），以便我們在稿件選用的第一時間聯系作者

? 您也可以直接添加小編微信（pwbot02）快速投稿，備注：姓名-投稿

△長按添加PaperWeekly小編

🔍

現在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關注」訂閱我們的專欄吧

·

總結

以上是生活随笔為你收集整理的ICML 2021 | AlphaNet：基于α-散度的超网络训练方法的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。