梯度手術-多任務學習優化方法[NeurIPS 2020]

  • 2021 年 1 月 15 日
  • AI

 分享一篇NeurIPS 2020上關於多任務學習的論文。


在之前的一篇文章《「一心多用」的方法——多任務學習》中,對多任務學習(Multi-task Learning,MTL)的背景、應用、挑戰做了一些簡單的介紹。在文章末尾,還提了我個人認為的待研究的方向:

針對這些挑戰,在多任務學習的領域中,可以從以下幾個研究方向考慮能做一些工作:

    1. 多任務網路結構設計

    2. 多任務損失函數設計

    3. 輔助學習(任務的篩選方法)

    4. 多任務網路的訓練方法

今天分享的這一篇論文,便是從第4點「多任務網路的訓練方法」著手,提出名為PCGrad的方法,在訓練多任務模型時,通過調整梯度,盡量降低任務之間的干擾。並且,這一方法與模型無關:它可以與其他多任務網路結構相結合,進一步提升多任務學習的性能

圖源[1]

首先需要明確的是,在多任務網路的訓練過程中,需要解決一些什麼問題。

假設有下式中的兩個優化任務,對應損失函數分別為L1和L2,θ對應模型的參數(在這個例子中做了簡化,只考慮兩個參數θ1和θ2)。對多任務學習而言,優化目標是學習模型的參數θ,以最小化L(L1與L2的和)。

式源[1]

定義了這一簡化的問題後,下圖(a)可視化了這一多任務學習問題的目標空間(即L和(θ1, θ2)的關係),圖(b)和圖(c)分別為任務1和任務2的目標空間,即為圖(a)的拆解版。其中,顏色越深的位置表示損失函數L越小。

圖源[1]

對於一個任務1+任務2的多任務問題,其目標空間為下圖(d):最深顏色的一條線為兩個任務共有的目標最優(損失最小)區域。

圖源[1] (d)中箭頭朝向表示梯度的方向,箭頭長度表示梯度的大小

為了解決這一多任務的優化問題,一個naive的做法是:使用經典的優化器,分別對兩個任務的損失函數進行梯度計算得到g1和g2,然後進行向量加法得到最終的梯度g (參見下面的Definition 4)。

這一做法存在顯著的問題,在上圖(d)中兩個箭頭的起點處,進行梯度計算,得到任務1梯度g1的紅色向量(在該點處任務1的目標空間比較陡峭(參見上圖b),故梯度大)和任務2的梯度g2的藍色向量(在該點處任務2的目標空間比較平緩(參見上圖c),故梯度小)。這兩個任務的梯度向量:

  • 數值相差巨大;

  • 方向南轅北轍。

因此,簡單的對梯度進行向量加法雖然看似能夠整體上優化損失函數,但並不能達到多任務學習的優勢,因為上述問題可能帶來:

  • 減速訓練過程;

  • 帶來性能下降,無法達到最優點。

具體來說,可以將問題分解為:

  1. 梯度衝突:不同任務梯度之間但更新方向不同,定義上來說,當cosθ<0時,兩個梯度存在衝突;

    圖源[1]

  2. 存在主導梯度:兩個任務的梯度,雖然方向並不衝突,但其中一個梯度數值上遠大於另一個梯度,則較大梯度會形成主導從而影響較小梯度。文中定義一個梯度幅值相似度(gradient magnitude similarity),用于衡量兩個梯度的大小差異;

  3. 高曲率:在高曲率的位置,主導任務的性能提升會被高估,其性能下降會被低估。


已有的一些改進方法,不能很好的解決上述的三個問題,例如:

  • 對不同任務的損失函數增加權重(通過經驗、或者通過Uncertainty的計算[2]),但這隻改變了梯度的數值差異。

這篇論文的作者提出的PCGrad (projecting conflicting gradients)方法,同時處理數值和方向的差異。下圖(e)中,通過PCGrad與Adam的結合,能夠沿著空間中正確的優化方向找到最小損失的點。

圖源[1]

簡單來說,就是要處理衝突的梯度(下圖a),將第i個任務的梯度投影到與第j個任務正交的方向(下圖b),或者相反(下圖c)。下圖d表示的是對不衝突的梯度,不做處理。

圖源[1]

演算法流程如下

  1. 對一個batch中對每個任務Ti,隨機取樣一個其他任務Tj,計算Ti和Tj對應梯度向量之間的餘弦相似度;

  2. 如果餘弦相似度為正,表示梯度無衝突,不做梯度修改。若為負,表示是相互衝突的梯度,根據下式計算投影;

  3. 重複上述步驟,直到這一batch中的任務遍歷完成;

  4. 最終的梯度方向為,投影后的梯度向量之和。

圖源[1] PCGrad的梯度更新流程

論文還提供了兩個定理及其詳細證明,論述了PCGrad方法的收斂性(定理一)以及其相較於標準梯度下降方法能達到更低誤差(定理二):

這裡不講定理和證明部分,分享一些我個人認為有意思的前提條件

  1. 任務的損失函數是凸的且可微的,凸(convexity)對於神經網路方法來說雖然不滿足,但是很多用於神經網路的優化方法也是基於凸的假設;

  2. 定理2的(a)表示,兩個梯度的方向衝突不能太小,且兩個梯度的幅值差異要足夠大,否則可能影響收斂效果;

  3. 定理2的(b)表示,多任務的曲率H有下界,且曲率需要足夠大;

  4. 定理2的(c)表示,學習的步長/學習率要足夠大,並且這個值與曲率有關,曲率越大,步長也要越大。因此,在使用PCGrad方法時,需要設置較大的學習率。

雖然並不是說不滿足上述前提條件就一定不work了,但在一定程度上對於演算法的應用效果能夠給予分析的參考。

最後,來看一看實際的實驗效果。作者在多個電腦視覺領域監督學習數據集和機器人領域強化學習數據集上做了實驗,均取得了好的效果。具體可以查閱原文,這裡摘錄一個在CityScapes和NYUv2數據集上的實驗結果:通過和其他多任務學習的方法(例如Cross-Stitch, MTAN,Uncertainty)進行結合,能夠取得更好的性能。

表源[1]

總結一下

  • 程式碼開源在[3]中;

  • 要說這一方法有什麼明顯的缺點或者代價呢,那就是會帶來額外的訓練時間和顯示記憶體佔用

  • 但這一方法總體思路清晰(提出了優化中的三個問題並針對性解決),實現簡單,並且模型無關,能夠與其他多任務方法共同使用;

  • 類似的思路,或許也可以用於其他領域,比如說持續學習、或者多模態學習中。

參考資料:

[1] Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgeryfor multi-task learning, 2020.

[2] Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scenegeometry and semantics, 2018.Al

[3] //github.com/tianheyu927/PCGrad


– END –

新朋友們可以看看我過往的相關文章

【相關推薦閱讀】


「一心多用」的方法——多任務學習

多模態:自動駕駛中的多模態學習

模式識別學科發展報告丨前言

深度學習在電腦視覺中應用的三大問題

從2018到2020的AI發展預測


長按二維碼關注我