ICLR 2021 | 显存不够？不妨抛弃端到端训练
?作者｜王語霖
學校｜清華大學自動化系博士生
研究方向｜機器學習和計算機視覺
本文主要介紹我們被 ICLR 2021 接收的一篇文章，代碼已經(jīng)在?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
（太長不看版）本文研究了一種比目前廣為使用的端到端訓練模式顯存開銷更小、更容易并行化的訓練方法：將網(wǎng)絡拆分成若干段、使用局部監(jiān)督信號進行訓練。我們指出了這一范式的一大缺陷在于損失網(wǎng)絡整體性能，并從信息的角度闡明了，其癥結(jié)在于局部監(jiān)督傾向于使網(wǎng)絡在淺層損失對深層網(wǎng)絡有很大價值的任務相關(guān)信息。
為有效解決這一問題，我們提出了一種局部監(jiān)督學習算法：InfoPro。在圖像識別和語義分割任務上的實驗結(jié)果表明，我們的算法可以在不顯著增大訓練時間的前提下，有效節(jié)省顯存開銷，并提升性能。
研究動機及簡介
一般而言，深度神經(jīng)網(wǎng)絡以端到端的形式訓練。以一個 13 層的簡單卷積神經(jīng)網(wǎng)絡為例，我們會將訓練數(shù)據(jù)輸入網(wǎng)絡中，逐層前傳至最后一層，輸出結(jié)果，計算損失值（End-to-End Loss），再從損失求得梯度，將之逐層反向傳播以更新網(wǎng)絡參數(shù)。
▲ 圖1 端到端訓練（End-to-End Training）
盡管端到端訓練在大量任務中都穩(wěn)定地表現(xiàn)出了良好的效果，但其效率至少在以下兩方面仍然有待提升。其一，端到端訓練需要在網(wǎng)絡前傳時將每一層的輸出進行存儲，并在逐層反傳梯度時使用這些值，這造成了極大的顯存開銷，如下圖所示。
▲ 圖2 端到端訓練具有較大的的顯存開銷
其二，對整個網(wǎng)絡進行前傳-->反傳的這一范式是一個固有的線性過程。前傳時深層網(wǎng)絡必須等待淺層網(wǎng)絡的計算完成后才能開始自身的前傳過程；同理，反傳時淺層網(wǎng)絡需要等待來自深層網(wǎng)絡的梯度信號才能進行自身的運算。這兩點線性的限制使得端到端訓練很難進行并行化以進一步的提升效率。
▲ 圖3 端到端訓練難以并行化
為了解決或緩解上述兩點低效的問題，一個可能的方案是使用局部監(jiān)督學習，即將網(wǎng)絡拆分為若干個局部模塊（local module），并在每個模塊的末端添加一個局部損失，利用這些局部損失產(chǎn)生監(jiān)督信號分別訓練各個局部模塊，注意不同模塊間沒有梯度上的聯(lián)通。下圖給出了一個將網(wǎng)絡拆分為兩段的例子。
▲ 圖4 局部監(jiān)督學習（Locally Supervised Learning）
相較于端到端訓練的兩點不足，局部監(jiān)督學習在效率上先天具有顯著優(yōu)勢。其一，我們一次只需保存一個局部模塊內(nèi)的中間層輸出值，待此模塊完成反向傳播后，即可釋放存儲空間，進而復用同樣的空間用以存儲下一個局部模塊的中間層輸出值，如下圖所示。簡言之，理論上顯存開銷隨局部模塊數(shù)呈指數(shù)級下降。
▲ 圖5 局部監(jiān)督學習可有效降低顯存開銷
其二，不同局部模塊的反向傳播過程并沒有必然的前后依賴關(guān)系，在工程實現(xiàn)上，不同模塊的訓練可以自然的并行完成，例如分別使用不同的 GPU，如下圖所示。
▲ 圖6 局部監(jiān)督學習易于并行化完成
問題分析與假設(shè)
相信大家看到這里，都會有一個問題：既然局部監(jiān)督學習的效率自然地高于端到端訓練，為什么它現(xiàn)在沒有被大規(guī)模應用呢？其問題在于，局部監(jiān)督學習往往會損害網(wǎng)絡的整體性能。
以圖片識別為例，考慮一種簡單自然的情況，我們使用標準的線性分類器 +SoftMax+ 交叉熵作為每個局部模塊的損失函數(shù)，在 CIFAR-10 數(shù)據(jù)集上使用局部監(jiān)督學習訓練 ResNet-32，結(jié)果如下所示，其中??代表局部模塊的數(shù)目?？梢钥闯鲭S著 值的增長，網(wǎng)絡的測試誤差急劇上升。
▲ 圖7 局部監(jiān)督學習傾向于損害網(wǎng)絡性能
若能解決性能下降的問題，局部監(jiān)督學習就有可能作為一種更為高效的訓練范式而取代端到端訓練。出于這一點，我們探究和分析了這一問題的原因。
上述局部監(jiān)督學習和端到端訓練的一個顯著的不同點在于，前者對網(wǎng)絡的中間層特征直接加入了與任務直接相關(guān)的監(jiān)督信號，從這一點出發(fā)，一個自然的疑問是，由此引發(fā)的中間層特征在任務相關(guān)行為上的區(qū)別是怎樣的呢？
因此，我們固定了圖 7 中得到的模型，使用網(wǎng)絡每層的特征訓練了一個線性分類器，其測試誤差如下圖右側(cè)所示。其中，橫軸代碼取用特征的網(wǎng)絡層數(shù)，縱軸代表測試誤差，不同的曲線對應于不同的? 取值， 表示端到端的情形。
▲ 圖8 中間層特征的線性可分性
從結(jié)果中可以觀察到一個明顯的現(xiàn)象：局部監(jiān)督學習所得到的中間層特征在淺層時就體現(xiàn)出了極好的線性可分性，但當特征進一步經(jīng)過更深的網(wǎng)絡層時，其線性可分性卻沒有得到進一步的增長；相比而言，盡管在淺層時幾乎線性不可分，端到端訓練得到的中間層特征隨著層數(shù)的加深可分性逐漸增強，最終取得了更低的測試誤差。
于是便產(chǎn)生了一個非常有趣的問題：局部監(jiān)督學習中，深層網(wǎng)絡使用了分辨性遠遠強于端到端訓練的特征，為何它得到的最終效果卻遜于端到端訓練？難道基于可分性已經(jīng)很強的特征，訓練網(wǎng)絡以進一步提升其線性可分性，不應該得到更好的最終結(jié)果嗎？這似乎與一些之前的觀察（例如 deeply supervised net）矛盾。
為了解答這個疑問，我們進一步從信息的角度探究網(wǎng)絡特征在可分性之外的區(qū)別。我們分別估計了中間層特征??與輸入數(shù)據(jù)??和任務標簽??之間的互信息??和?，并以此作為? 中包含的全部信息和任務相關(guān)信息的度量指標。
▲ 圖9 估算互信息
其結(jié)果如下圖所示，其中橫軸為取用信息的層數(shù)，縱軸表示估計值。從中不難看出，端到端訓練的網(wǎng)絡中，特征所包含的總信息量逐層減少，但任務相關(guān)信息維持不變，說明網(wǎng)絡逐層剔除了與任務無關(guān)的信息。與之形成鮮明對比的是，局部監(jiān)督學習得到的網(wǎng)絡在淺層就丟失了大量的任務相關(guān)信息，特征所包含的總信息量也急劇下降。
我們猜測，這一現(xiàn)象的原因在于，僅憑淺層網(wǎng)絡難以如全部網(wǎng)絡一般有效分離和利用所有任務相關(guān)信息，因此索性去丟棄部分無法利用的信息換取局部訓練損失的降低。而在這種情況下，網(wǎng)絡深層接收到的特征相較網(wǎng)絡原始輸入本就缺少關(guān)鍵信息，自然難以基于其建立更有效的表征，也就難以取得更好的最終性能。
▲ 圖10 中間層特征包含的信息
基于上述觀察，我們可以總結(jié)得到：局部監(jiān)督學習之所以會損害網(wǎng)絡的整體性能，是因為其傾向于使網(wǎng)絡在淺層丟失與任務相關(guān)的信息，從而使得深層網(wǎng)絡空有更多的參數(shù)和更大的容量，卻因輸入特征先天不足而無用武之地。
方法詳述
為了解決損失信息的問題，本文提出了一種專為局部監(jiān)督學習定制的損失函數(shù)：InfoPro。首先，我們引入一個基本模型。如下圖所示，我們假設(shè)訓練數(shù)據(jù)受到兩個隨機變量影響，其一是任務標簽?，決定我們所關(guān)心的主體內(nèi)容；其二是無關(guān)變量?，用于決定數(shù)據(jù)中與任務無關(guān)的部分，例如背景、視角、天氣等。
▲ 圖11 變量作用關(guān)系假設(shè)
基于上述變量設(shè)置，我們將 InfoPro 損失函數(shù)定義為下面的結(jié)合形式。它用于作為局部監(jiān)督信號訓練局部模塊，由兩項組成。第一項用于推動局部模塊向前傳遞所有信息；在第二項中，我們使用一個滿足特殊條件無關(guān)變量?來建模中間層特征中的全部任務無關(guān)信息（無用信息），在此基礎(chǔ)上迫使局部模塊剔除這些與任務無關(guān)的信息。
▲ 圖12 InfoPro損失函數(shù)
InfoPro 與端到端訓練和在 2. Analysis 中所述的簡單局部監(jiān)督學習（Greedy Supervised Learning）的對比如下圖所示。簡言之，InfoPro 的目標是使得局部模塊能夠在保證向前傳遞全部有價值信息的條件下，盡可能丟棄特征中的無用信息，以解決局部監(jiān)督學習在淺層丟失任務相關(guān)信息、影響網(wǎng)絡最終性能的問題。
事實上，這也是我們前面觀察到的、端到端訓練對網(wǎng)絡淺層的影響形式。InfoPro 與其它局部學習方法最大的區(qū)別在于它是非貪婪的，并不直接對局部的任務相關(guān)行為（如Greedy Supervised Learning 中基于局部特征的分類損失）做出直接約束。
▲ 圖13 3種訓練算法的對比
在具體實現(xiàn)上，由于 InfoPro 損失的第二項比較難以估算，我們推導出了其的一個易于計算的上界，如下圖所示：
▲ 圖14 InfoPro損失的一個易于計算的上界
關(guān)于這一上界的具體推導過程、一些數(shù)學性質(zhì)和其實際上的計算方式，由于流程比較復雜且不關(guān)鍵，不在此贅述，歡迎感興趣的讀者參閱我們的文章~
實驗結(jié)果
●?在不同局部模塊數(shù)目的條件下，穩(wěn)定勝過 baseline
●?大量節(jié)省顯存，且不引入顯著的額外計算/時間開銷，效果相較端到端訓練略有提升
●?ImageNet 大規(guī)模圖像識別任務上的結(jié)果，節(jié)省顯存的效果同樣顯著，效果略有提升
●?Cityscapes 語義分割實驗結(jié)果，除節(jié)省顯存方面的作用外，我們還證明了，在相同的顯存限制下，InfoPro 可以使用更大的 batch size 或更大分辨率的輸入圖片
結(jié)語
總結(jié)來說，這項工作的要點在于：（1）從效率的角度反思端到端訓練范式；（2）指出了局部監(jiān)督學習相較于端到端的缺陷在于損失網(wǎng)絡性能，并從信息的角度分析了其原因；（3）在理論上提出了初步解決方案，并探討了具體實現(xiàn)方法。
歡迎大家 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} }如有任何問題，歡迎留言或者給我發(fā)郵件，附上我的主頁鏈接：
http://www.rainforest-wang.cool
更多閱讀
#投 稿?通 道#
?讓你的論文被更多人看到?
如何才能讓更多的優(yōu)質(zhì)內(nèi)容以更短路徑到達讀者群體，縮短讀者尋找優(yōu)質(zhì)內(nèi)容的成本呢？答案就是：你不認識的人。
總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術(shù)靈感相互碰撞，迸發(fā)出更多的可能性。?
PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優(yōu)質(zhì)內(nèi)容，可以是最新論文解讀，也可以是學習心得或技術(shù)干貨。我們的目的只有一個，讓知識真正流動起來。
?????來稿標準：
? 稿件確系個人原創(chuàng)作品，來稿需注明作者個人信息（姓名+學校/工作單位+學歷/職位+研究方向）?
? 如果文章并非首發(fā)，請在投稿時提醒并附上所有已發(fā)布鏈接?
? PaperWeekly 默認每篇文章都是首發(fā)，均會添加“原創(chuàng)”標志
?????投稿郵箱：
? 投稿郵箱：hr@paperweekly.site?
? 所有文章配圖，請單獨在附件中發(fā)送?
? 請留下即時聯(lián)系方式（微信或手機），以便我們在編輯發(fā)布時和作者溝通
????
現(xiàn)在，在「知乎」也能找到我們了
進入知乎首頁搜索「PaperWeekly」
點擊「關(guān)注」訂閱我們的專欄吧
關(guān)于PaperWeekly
PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術(shù)平臺。如果你研究或從事 AI 領(lǐng)域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。
總結(jié)
以上是生活随笔為你收集整理的ICLR 2021 | 显存不够？不妨抛弃端到端训练的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 怎样消灭家里的蟑螂？
- 下一篇： Python 和 C/C++ 拓展程序的