?作者｜王語霖

學校｜清華大學自動化系博士生

研究方向｜機器學習和計算機視覺

本文主要介紹我們被 ICLR 2021 接收的一篇文章，代碼已經在?Github 上面開源。

論文標題：

Revisiting Locally Supervised Learning: an Alternative to End-to-end Training

論文鏈接：

https://openreview.net/forum?id=fAbkE6ant2

代碼鏈接：

https://github.com/blackfeather-wang/InfoPro-Pytorch

（太長不看版）本文研究了一種比目前廣為使用的端到端訓練模式顯存開銷更小、更容易并行化的訓練方法：將網絡拆分成若干段、使用局部監督信號進行訓練。我們指出了這一范式的一大缺陷在于損失網絡整體性能，并從信息的角度闡明了，其癥結在于局部監督傾向于使網絡在淺層損失對深層網絡有很大價值的任務相關信息。

為有效解決這一問題，我們提出了一種局部監督學習算法：InfoPro。在圖像識別和語義分割任務上的實驗結果表明，我們的算法可以在不顯著增大訓練時間的前提下，有效節省顯存開銷，并提升性能。

研究動機及簡介

一般而言，深度神經網絡以端到端的形式訓練。以一個 13 層的簡單卷積神經網絡為例，我們會將訓練數據輸入網絡中，逐層前傳至最后一層，輸出結果，計算損失值（End-to-End Loss），再從損失求得梯度，將之逐層反向傳播以更新網絡參數。

▲ 圖1 端到端訓練（End-to-End Training）

盡管端到端訓練在大量任務中都穩定地表現出了良好的效果，但其效率至少在以下兩方面仍然有待提升。其一，端到端訓練需要在網絡前傳時將每一層的輸出進行存儲，并在逐層反傳梯度時使用這些值，這造成了極大的顯存開銷，如下圖所示。

▲ 圖2 端到端訓練具有較大的的顯存開銷

其二，對整個網絡進行前傳-->反傳的這一范式是一個固有的線性過程。前傳時深層網絡必須等待淺層網絡的計算完成后才能開始自身的前傳過程；同理，反傳時淺層網絡需要等待來自深層網絡的梯度信號才能進行自身的運算。這兩點線性的限制使得端到端訓練很難進行并行化以進一步的提升效率。

▲ 圖3 端到端訓練難以并行化

為了解決或緩解上述兩點低效的問題，一個可能的方案是使用局部監督學習，即將網絡拆分為若干個局部模塊（local module），并在每個模塊的末端添加一個局部損失，利用這些局部損失產生監督信號分別訓練各個局部模塊，注意不同模塊間沒有梯度上的聯通。下圖給出了一個將網絡拆分為兩段的例子。

▲ 圖4 局部監督學習（Locally Supervised Learning）

相較于端到端訓練的兩點不足，局部監督學習在效率上先天具有顯著優勢。其一，我們一次只需保存一個局部模塊內的中間層輸出值，待此模塊完成反向傳播后，即可釋放存儲空間，進而復用同樣的空間用以存儲下一個局部模塊的中間層輸出值，如下圖所示。簡言之，理論上顯存開銷隨局部模塊數呈指數級下降。

▲ 圖5 局部監督學習可有效降低顯存開銷

其二，不同局部模塊的反向傳播過程并沒有必然的前后依賴關系，在工程實現上，不同模塊的訓練可以自然的并行完成，例如分別使用不同的 GPU，如下圖所示。

▲ 圖6 局部監督學習易于并行化完成

問題分析與假設

相信大家看到這里，都會有一個問題：既然局部監督學習的效率自然地高于端到端訓練，為什么它現在沒有被大規模應用呢？其問題在于，局部監督學習往往會損害網絡的整體性能。

以圖片識別為例，考慮一種簡單自然的情況，我們使用標準的線性分類器 +SoftMax+ 交叉熵作為每個局部模塊的損失函數，在 CIFAR-10 數據集上使用局部監督學習訓練 ResNet-32，結果如下所示，其中??代表局部模塊的數目。可以看出隨著 值的增長，網絡的測試誤差急劇上升。

▲ 圖7 局部監督學習傾向于損害網絡性能

若能解決性能下降的問題，局部監督學習就有可能作為一種更為高效的訓練范式而取代端到端訓練。出于這一點，我們探究和分析了這一問題的原因。

上述局部監督學習和端到端訓練的一個顯著的不同點在于，前者對網絡的中間層特征直接加入了與任務直接相關的監督信號，從這一點出發，一個自然的疑問是，由此引發的中間層特征在任務相關行為上的區別是怎樣的呢？

因此，我們固定了圖 7 中得到的模型，使用網絡每層的特征訓練了一個線性分類器，其測試誤差如下圖右側所示。其中，橫軸代碼取用特征的網絡層數，縱軸代表測試誤差，不同的曲線對應于不同的? 取值， 表示端到端的情形。

▲ 圖8 中間層特征的線性可分性

從結果中可以觀察到一個明顯的現象：局部監督學習所得到的中間層特征在淺層時就體現出了極好的線性可分性，但當特征進一步經過更深的網絡層時，其線性可分性卻沒有得到進一步的增長；相比而言，盡管在淺層時幾乎線性不可分，端到端訓練得到的中間層特征隨著層數的加深可分性逐漸增強，最終取得了更低的測試誤差。

于是便產生了一個非常有趣的問題：局部監督學習中，深層網絡使用了分辨性遠遠強于端到端訓練的特征，為何它得到的最終效果卻遜于端到端訓練？難道基于可分性已經很強的特征，訓練網絡以進一步提升其線性可分性，不應該得到更好的最終結果嗎？這似乎與一些之前的觀察（例如 deeply supervised net）矛盾。

為了解答這個疑問，我們進一步從信息的角度探究網絡特征在可分性之外的區別。我們分別估計了中間層特征??與輸入數據??和任務標簽??之間的互信息??和?，并以此作為? 中包含的全部信息和任務相關信息的度量指標。

▲ 圖9 估算互信息

其結果如下圖所示，其中橫軸為取用信息的層數，縱軸表示估計值。從中不難看出，端到端訓練的網絡中，特征所包含的總信息量逐層減少，但任務相關信息維持不變，說明網絡逐層剔除了與任務無關的信息。與之形成鮮明對比的是，局部監督學習得到的網絡在淺層就丟失了大量的任務相關信息，特征所包含的總信息量也急劇下降。

我們猜測，這一現象的原因在于，僅憑淺層網絡難以如全部網絡一般有效分離和利用所有任務相關信息，因此索性去丟棄部分無法利用的信息換取局部訓練損失的降低。而在這種情況下，網絡深層接收到的特征相較網絡原始輸入本就缺少關鍵信息，自然難以基于其建立更有效的表征，也就難以取得更好的最終性能。

▲ 圖10 中間層特征包含的信息

基于上述觀察，我們可以總結得到：局部監督學習之所以會損害網絡的整體性能，是因為其傾向于使網絡在淺層丟失與任務相關的信息，從而使得深層網絡空有更多的參數和更大的容量，卻因輸入特征先天不足而無用武之地。

方法詳述

為了解決損失信息的問題，本文提出了一種專為局部監督學習定制的損失函數：InfoPro。首先，我們引入一個基本模型。如下圖所示，我們假設訓練數據受到兩個隨機變量影響，其一是任務標簽?，決定我們所關心的主體內容；其二是無關變量?，用于決定數據中與任務無關的部分，例如背景、視角、天氣等。

▲ 圖11 變量作用關系假設

基于上述變量設置，我們將 InfoPro 損失函數定義為下面的結合形式。它用于作為局部監督信號訓練局部模塊，由兩項組成。第一項用于推動局部模塊向前傳遞所有信息；在第二項中，我們使用一個滿足特殊條件無關變量?來建模中間層特征中的全部任務無關信息（無用信息），在此基礎上迫使局部模塊剔除這些與任務無關的信息。

▲ 圖12 InfoPro損失函數

InfoPro 與端到端訓練和在 2. Analysis 中所述的簡單局部監督學習（Greedy Supervised Learning）的對比如下圖所示。簡言之，InfoPro 的目標是使得局部模塊能夠在保證向前傳遞全部有價值信息的條件下，盡可能丟棄特征中的無用信息，以解決局部監督學習在淺層丟失任務相關信息、影響網絡最終性能的問題。

事實上，這也是我們前面觀察到的、端到端訓練對網絡淺層的影響形式。InfoPro 與其它局部學習方法最大的區別在于它是非貪婪的，并不直接對局部的任務相關行為（如Greedy Supervised Learning 中基于局部特征的分類損失）做出直接約束。

▲ 圖13 3種訓練算法的對比

在具體實現上，由于 InfoPro 損失的第二項比較難以估算，我們推導出了其的一個易于計算的上界，如下圖所示：

▲ 圖14 InfoPro損失的一個易于計算的上界

關于這一上界的具體推導過程、一些數學性質和其實際上的計算方式，由于流程比較復雜且不關鍵，不在此贅述，歡迎感興趣的讀者參閱我們的文章~

實驗結果

●?在不同局部模塊數目的條件下，穩定勝過 baseline

●?大量節省顯存，且不引入顯著的額外計算/時間開銷，效果相較端到端訓練略有提升

●?ImageNet 大規模圖像識別任務上的結果，節省顯存的效果同樣顯著，效果略有提升

●?Cityscapes 語義分割實驗結果，除節省顯存方面的作用外，我們還證明了，在相同的顯存限制下，InfoPro 可以使用更大的 batch size 或更大分辨率的輸入圖片

結語

總結來說，這項工作的要點在于：（1）從效率的角度反思端到端訓練范式；（2）指出了局部監督學習相較于端到端的缺陷在于損失網絡性能，并從信息的角度分析了其原因；（3）在理論上提出了初步解決方案，并探討了具體實現方法。

歡迎大家 follow 我們的工作~

@inproceedings{wang2021revisiting,title?=?{Revisiting?Locally?Supervised?Learning:?an?Alternative?to?End-to-end?Training},author?=?{Yulin?Wang?and?Zanlin?Ni?and?Shiji?Song?and?Le?Yang?and?Gao?Huang},booktitle?=?{International?Conference?on?Learning?Representations?(ICLR)},year?=?{2021},url?=?{https://openreview.net/forum?id=fAbkE6ant2} }

如有任何問題，歡迎留言或者給我發郵件，附上我的主頁鏈接：

http://www.rainforest-wang.cool

更多閱讀

#投 稿?通 道#

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

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

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

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學習心得或技術干貨。我們的目的只有一個，讓知識真正流動起來。

?????來稿標準：

? 稿件確系個人原創作品，來稿需注明作者個人信息（姓名+學校/工作單位+學歷/職位+研究方向）?

? 如果文章并非首發，請在投稿時提醒并附上所有已發布鏈接?

? PaperWeekly 默認每篇文章都是首發，均會添加“原創”標志

?????投稿郵箱：

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

? 所有文章配圖，請單獨在附件中發送?

? 請留下即時聯系方式（微信或手機），以便我們在編輯發布時和作者溝通

????

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

進入知乎首頁搜索「PaperWeekly」

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

關于PaperWeekly

PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術平臺。如果你研究或從事 AI 領域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。

總結

以上是生活随笔為你收集整理的ICLR 2021 | 显存不够？不妨抛弃端到端训练的全部內容，希望文章能夠幫你解決所遇到的問題。

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