Hinton 新作「在線蒸餾」，提升深度學習分散式訓練表現的利器

雷鋒網 AI 科技評論按：Hinton 團隊最近發表了一篇關於「在線蒸餾」的新論文。論文里講了什麼呢？

我們為什麼需要在線蒸餾？

近年來，隨著深度學習技術的蓬勃發展，我們日常生活的每個角落都充滿了「人工智慧」的影子，並由此催生了許多新的業態、以人工智慧應用為核心產品的初創公司也如雨後春筍般在科技行業中嶄露頭角。是的，這也許是人工智慧（特別是深度學習）最好的時代。

然而，深度學習作為當今人工智慧領域皇冠上最璀璨的明珠，如果要想將其推向工業級的應用，則往往需要可望而不可及的強大算力！而這種算力則是由價格高昂的硬體、龐大而複雜的分散式計算環境、精妙高效的演算法作為支撐的。可想而知，除了 Google、Amazon、阿里巴巴、百度等掌握強大計算資源的科技巨頭外，想獲得這樣的計算資源對於小型企業或者個人用戶可謂是天方夜譚。實際上，在很多普通的高等院校和科研單位中，擁有像 Titan X 或 GTX 1080 Ti 這樣的民用顯卡已經是很奢侈的事情。更重要的是，由於基礎架構的限制（分散式計算節點的通信、同步、任務調度等問題）、對模型進行優化求解的障礙、集成學習環境下各模型的貢獻缺乏好的決策，使得盲目的堆砌計算資源也可能觸碰到模型性能的天花板（例如：分散式隨機梯度下降（distrtibuted SGD））。

為此，「深度學習泰斗」Geoffrey E.Hinton 近年來在知識蒸餾（或者「暗知識提取」）方面做了一系列工作，試圖通過這種從高計算能力要求、同時也具有高表現的模型中提取出一些隱含的知識，並且將其作為先驗，通過設計新的網路結構和目標函數將這種知識「傳授」給規模較小的深度學習網路，實現對網路模型的壓縮，以顯著縮小的網路規模和計算需求展現出盡量高的模型表現。最近，為了將這種思想部署在分散式環境中，用以突破如今常常被使用的分散式 SGD（同步和非同步形式）的瓶頸，Hinton 團隊又發布了名為「LARGE SCALE DISTRIBUTED NEURAL NETWORK TRAINING THROUGH ONLINE DISTILLATION」（通過在線蒸餾的神經網路大規模分散式訓練）的論文。在筆者看來，Hinton 的這一系列工作進一步降低了深度學習模型應用的門檻，之前的研究對深度學習模型在移動終端上的部署、這篇論文對提高大規模計算集群上的模型表現和計算效率都有深遠意義。

論文地址：https://arxiv.org/abs/1804.03235

要想理解「在線蒸餾」的概念，我們有必要回顧一下 Hinton 從 2014 年開始對 dark knowledge extraction（暗知識提取） 和 knowledge distillation（知識蒸餾）的相關工作，甚至更早的 Caruana et.al 所做的模型壓縮的工作。

為了提升神經網路模型的性能，Caruana 等人早在 2006 年（那時深度學習還沒有大火）就提出了一種壓縮大規模複雜網路的方法。由於集成學習在傳統的機器學習領域大獲成功，許多深度學習研究人員很自然地想要利用集成學習的思想，將大量的模型聚合到一個神經網路中，通過暴力的訓練，為不同的模型賦予不同的權值，對這些模型的輸出進行加權平均得到最終的結果，以便充分利用它們各自對於不同的任務所具有的優勢。然而，這種暴力的訓練過程和臃腫的網路結構需要消耗巨大的計算資源、造成額外的能源消耗。Caruana 等人提出了 MUNGE 的數據增強演算法，將大規模模型學習到的函數壓縮進規模更小、訓練更便捷的模型中。

受此啟發，時隔 8 年後，當深度學習迎來春天、人們沉醉於大規模深度學習網路帶來的人工智慧在計算機視覺、自然語言處理等方面的成功應用的時候，宗師 Hinton 則認為，是時候回過頭來看看 Caruana 的文章，思考如何將這些「醜陋的」大規模的模型壓縮到小而快的模型中去。

由此，Hinton 提出了「dark knowledge」的概念。在他看來，這種被稱為「暗知識」的東西才是深度學習本質上學到的知識（或許這也是尋求深度學習可解釋性的一種途徑）。Dark knowledge，顧名思義，就是隱藏在深度學習表面上所展現出來的網路結構、節點之間的連接權重、網路的輸出這些看得到的數據之下的知識。如果能夠找到一種途徑，使得我們能夠獲取這種知識，並且將其包裝成一種先驗概率，遷移到更小的模型中去，是否能夠提升小模型的性能呢？事實上，筆者認為，這也可以看作是在遷移學習的框架下，將大規模網路視作信息充分的 source domain，將小規模網路視作需要大規模網路補充信息的 target domain，而我們提取到的 dark knowledge 則是兩個任務之間的 common knowledge。

然而，該從哪裡下手，獲取這種知識呢？Hinton 敏銳地觀察到：我們在絕大多數的預測任務的深度學習網路中，都會使用 softmax layer 為大量的標籤分配概率分布。然而這種處理方式存在一個負作用：與正確標籤相比，模型為所有的誤標籤都分配了很小的概率；然而實際上對於不同的錯誤標籤，其被分配的概率仍然可能存在數個量級的懸殊差距。例如：在圖片分類問題中，我們要將圖片分成貓、狗、老虎三類。在一次訓練中，我們給三類分配的概率分別為 [0.0010, 0.0001, 0.9989]，從而最終得到 [0，0，1] 的 one-hot 編碼作為分類結果（即 hard-target），我們認為圖片所代表的是一隻老虎。然而，softmax 函數輸出的概率往往包含著類別之間潛在的相關性。在這個例子中，我們可以看到，圖片可能是貓的概率比圖片是狗的概率更接近圖片是老虎的概率，這說明貓和老虎之間存在的內在聯繫更加強大。類似地，Hinton 也舉例說：在識別一輛寶馬汽車的圖片時，分類器將該圖片識別為清潔車的概率是很小的，然而這種概率比起將其識別為胡蘿蔔的可能是會大出很多。由於在宏觀上由於這些概率都很小，這一部分的知識很容易在訓練過程中淹沒，這無疑是浪費了重要的可以用於將大規模網路的知識遷移到小規模網路中去的寶貴先驗概率。

為了充分利用這種類類別之間的相關性，我們需要通過某種方式去改變概率分布，使其更加平緩。而 Hinton 僅僅對我們經常使用的 softmax 函數進行了一點點修改，就達到了這一目標，他究竟是怎麼做的呢？

事實上，如下面的公式所示，Hinton 向 softmax 函數添加了一點「佐料」——參數「T, 溫度」（如今 T 已經成為了許多深度學習模型的標配，例如在生成文本的 RNN 中提高 T 可以增加生成文本的多樣性）：

其中，z 為每一個類別輸入的 logit。式中 T=1 時，退化成傳統的 softmax；T無窮大時，結果趨近於 1/C，即所有類別上的概率趨近於相等。T>1 時，我們就能獲得 soft target label。通過提高 T，softmax層的映射曲線更加平緩，因而實例的概率映射將更為集中，便使得目標更加地「soft」。

有了這個 distillation 的內核，Hinton 按照以下的步驟對大規模網路進行「蒸餾」：

1. 訓練大模型：先用 hard target（類似於 [0，0，1] 的 one-hot 編碼）的樣本訓練。

2. 計算 soft target：利用訓練好的大模型來計算 soft target 。也就是大模型「軟化後」再經過 softmax 的輸出。

3. 重新創建一個小的網路，該網路最後有兩個 loss，一個是 hard loss，即傳統的 softmax loss，使用 one-hot label；另外一個是 soft loss，即 T>1 的 softmax loss，使用我們第二步保存下來的 soft target label。

整體的 loss 如下式：

其中

為第二步輸出的 soft label。

用「軟化」訓練集訓練小模型。訓練小模型時 T 不變仍然較大，訓練完之後 T 改為1。

4. 預測時，將訓練好的小模型按常規方式使用。

現在我們可以把 Hinton 的方法和下圖所示的最初 knowledge distillation 的由來作個對比。

可見，本質上這相當於對數據進行了增強（augmentation），加入了類別之間的關聯性的先驗信息。將大規模網路學習到的這種關係包裝在數據中，用這種更強的數據來訓練小規模的模型，充分考慮到了類間的距離和類內的方差信息。從而提升了小規模模型的性能，達到了「蒸餾」的效果。與直接使用預訓練模型的結構和權重相比，這是一種相對更「高級」的知識遷移的方式。

時間的車輪一眨眼就駛向了 2018 年，隨著移動終端上的深度學習等應用的出現，網路模型壓縮成為了一個廣受關注的領域，大量的研究者在 Hinton 的啟發下，退出了 distillation 的諸多變形，進行了深入的優化。Hinton 則一如既往地嘗試開拓更多新的深度學習範式；當然也可能是谷歌的計算資源太過充足所以遇到了別人沒有機會遇到的問題（給跪），Hinton 開始嘗試在大規模分散式計算環境下使用「在線蒸餾（online distillation）」方法。這是由於目前的分散式 SGD 方法遇到了瓶頸。而本地的「蒸餾」演算法也因其對數據管道的計算操作過於複雜而暴露出越來越多的問題。

具體而言，在分散式 SGD 中，由於邊際效益遞減的規律，增加參與訓練的機器數量而獲得的計算效率的提升漸漸變小，直到毫無效果。另一方面，他們也想使用集成的深度學習模型提高預測的準確率。為了在不增加測試時間成本的情況下獲得與集成學習同等的收益，他們對一個 n 路集成模型進行蒸餾（distill），得到一個單一模型，這包括兩個階段：

• 使用 M 機器來訓練分散式 SGD 的 n 路集成模型，然後使用 M 機器（T 不變）來訓練 student 網路（小規模網路），這個小規模網路會模擬這個 n 路集成模型。

• 通過在訓練過程中使用更多機器，蒸餾會增加訓練時間和計算複雜度，以換取接近更大的 teacher 集成模型的質量改進。

Hinton 他們將這種在線的蒸餾方式稱為「codistillation」：即分散式環境中的每個節點之間都可以互為 teacher 和 student，並且互相提取內在的知識，用以提升其它節點的模型性能，具體的演算法如下：

如演算法 1 中所示，為了突破分散式 SGD 的瓶頸，Hinton 他們利用蒸餾演算法提高了模型的訓練效率。使用蒸餾方法更新一個網路的參數只需要對其他網路的預測結果，這些網路可以利用其他網路權重的副本進行本地計算。

值得注意的是，即使教師模型和學生模型是同一神經網路的兩個實例，蒸餾也有好處，只要它們有足夠的區別（比如，不同的初始化、以不同的順序接收輸入樣本；可以參見論文第 3 節的經驗證據）。這也說明這種「在線蒸餾」的方法是具有很強的普適性的。

Hinton 他們在這個工作中將蒸餾技術和分散式 SGD 相結合，從而使一個分散式 SGD 的工作組內的各個節點能夠交換檢查點（checkpoint）保存的網路信息，利用這種信息作為蒸餾出來的知識，通過「老師-學生」的訓練，加速學生網路的訓練。在這個過程中，學生節點和老師結點的角色是互換的，因此，各個網路互相促，進從而實現共同的蒸餾。

其實，就算「在線蒸餾」是新瓶裝舊酒，那也是一個非常恰當的，閃閃發光的新瓶子。它利用蒸餾技術，降低了分散式 SGD 的通信開銷，成功提高了預測的準確率，提升模型的計算性能！

筆者通過這篇文章和大家一起回顧了知識蒸餾的相關知識，並且粗淺地了解了 Hinton 在這個領域所做的最新工作。我們可以看到，Hinton 作為一代宗師，每次都能夠以超過常人的眼光發現研究的新方向，並且提出簡單、美、效果卓越的解決方案，這與他對自然、對生物的神經系統、對生活、對所面臨的問題的犀利的觀察是密不可分的。這激勵著所有的人工智慧研究者開拓視野、放飛心靈，充滿創造力地去探索新的未知領域。雷鋒網 AI 科技評論這樣的學術媒體也會不斷地把最新的學術研究進展介紹給大家。

論文地址：https://arxiv.org/abs/1804.03235，雷鋒網 AI 科技評論報道

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！