作者 | Felix J.S. Bragman, Ryutaro Tanno等
譯者 | 李傑
出品 | AI科技大本營（ID:rgznai100）
多任務學習（MTL）的性能表現，很大程度上取決於任務共享的方式，而任務共享方式通常是依靠網絡架構的設計。共享方式是由網絡深度和任務數量綜合決定的，因此如果完全依靠人工設計的組合方式，可能得到的結果是耗時的，而且並不能保證是最優方案。
本文提出了一種用於多任務學習的CNNs中任務特定（task-specific）和共享表示（shared representations）的概率學習方法。具體來說，本文提出了一個隨機濾波分組（stochastic filter groups，SFG）的方法，SFG是一種將每一層中的卷積核分配給專有（specialist）或通用（generalist）組的機制，這些組分別針對不同的任務或在不同的任務之間共享。SFG確定了層與網絡中特定於任務和共享表示的結構之間的連接性，使用變分推斷（variational inference）來學習可能的卷積分組和網絡參數分組的後驗分佈。實驗證明，本文所提出的方法適用於多個任務，並且比基準方法有更好的性能。
引言
多任務學習（MTL）
多任務學習旨在通過同時解決多個相關任務，提升學習效率和網絡泛化性能。將卷積神經網絡(CNNs)嵌入到MTL中已經在廣泛的計算機視覺應用中顯示出良好的前景，諸如目標識別、檢測、分割等等。設計一個成功的MTL模型的關鍵是學習共享特徵和特定任務特徵表示的能力。理解任務之間的共性和差異的機制允許模型在任務之間傳遞信息，同時裁剪預測模型來描述單個任務的不同特徵。這些特徵表示的質量主要取決於MTL網絡架構的設計，如網絡中的哪些層是在所有任務之間共享的，哪些層是分開，專門針對於某個子任務的。
但是，這種組合方式是多樣的，如果我們只依靠自己手工設計，效率是非常低下的，而且不能保證得到的組合是最優的。如下圖所示，左邊的圖展示了一個典型的多任務體系結構，它由一個共享的“主幹”特徵提取器和特定於任務的“分支”網絡組成，但實際上，所需的共享和特定於任務的表示，以及它們在體系結構中的交互取決於單個任務的難度和它們之間的關係，在大多數情況下，這兩者都不是先驗已知的。這說明了手工設計適當的體系結構十分困難。而右邊的圖展示了一個可以用我們的方法學習的示例體系結構。基於此，本文提出隨機濾波分組（SFG），一種學習卷積核分配到特定任務和共享組的原則方法。
隨機濾波分組（SFG）
本文提出了隨機濾波分組機制（Stochastic Filter Groups SFG），它是一種概率機制，用於學習在MTL結構的每一層中所需的特定於任務和共享表示的數量。如下圖所示，SFG學習將每一層中的卷積核分配給專有（specialist）或通用（generalist）組，這些組分別針對不同的任務或在不同的任務之間共享。SFG爲網絡配備了一種機制來學習層間連接，從而學習特定任務和共享表示的結構。
背景知識
變分推斷（variational inference）
我們的任務是根據已有的數據去推斷需要的分佈P，但是在很多情況下，P不是規則的，不容易表達和直接求解，在這種情況下，我們可以轉換思路，尋找一個容易表達和求解的分佈Q，我們的目標是不斷優化Q，使得Q和原始分佈P的差距不斷縮小。這種情況下，可以把Q看做P的近似分佈，將Q的輸出近似作爲P的輸出。
MTL中網絡架構設計
多任務中主要通過設計相應的網絡結構，確定哪些特徵表示是共享的，哪些特徵表示是特定於某個子任務的，主要有兩大類方法：
（1）第一類是優化權重分配結構以使任務性能最大化的方法。這些方法旨在學習一組向量，這些向量控制哪些特性在一個層中共享，以及這些特性如何分佈，代表性網絡有十字繡網絡，如下圖所示，通過一個參數矩陣，決定特徵在每個任務之間的共享程度。
（2）第二類側重於根據任務相似度進行權重聚類。常見的做法是使用迭代算法來構建一個類似於樹的深度結構，將類似的任務按層次進行分組，或者根據任務之間的統計相關性來確定權重共享的程度。
本文提出的方法屬於第一大類即通過對特定任務和共享特性進行“硬”劃分來區分。本文的方法通過對卷積核進行分組來學習層之間連接性的分佈。使得模型能夠學習有意義的特定於任務的分組和共享特性。
方法
本文引入了一個新的方法來確定在多任務CNN架構中從哪裏學習特定任務和共享表示。這就是隨機濾波分組（stochastic filter groups SFG），它是一種概率機制，將每個卷積層中的卷積核劃分爲“專有”組和“共享”組，它們分別針對不同的任務或在不同的任務之間共享。本文使用變分推斷來學習可能的內核分組和網絡參數的分佈，這些參數決定了層與層之間的連接性以及共享的和特定於任務的特徵。
隨機濾波分組（stochastic filter groups）
SFG將一個稀疏連接結構引入到CNN的多任務學習架構中，從而將特徵分離爲特定於任務的和共享的組件。濾波器組的結構如下圖所示，將每個卷積層中的卷積核劃分爲組，每個組僅作用於特徵上的一個子集。這種稀疏性在不影響精度的情況下降低了計算成本和參數數量。
在SFG中，借鑑了濾波器分組的概念，將其引入到MTL中，在此基礎上提出一種擴展，使用一種額外的機制來學習最優的卷積核分組，而不是預先指定它們。
爲方便表示，我們以兩個子任務的多任務學習爲例，在一個CNN結構的第 L層，有 K個卷積核，可以用下列集合表示：
SFG可以通過兩個步驟完成：
（1）卷積核分配（FilterAssignment）：
如下圖所示，給定卷積層中的每個卷積核 {wk}都根據相關的分類分佈樣本概率（p1,p2,ps）分配給濾波組G1、G2、Gs中的一個。
如下圖所示，卷積核分配完成後，進行各自分組的卷積，得到不同組的特徵表示{F1，Fs，F2}。（i）表示所有的卷積核被均勻分組的情況；（ii）和（iii）表示卷積核在更深的層次上變得更加針對特定的任務的情況；（iv）表示一個任務間具有更多異構拆分的情況。
（2）特徵路由（FeatureRouting）：
如下圖所示，第 L層上經卷積核分組卷積得到的特徵{F1，Fs，F2}被路由到第L+1層的卷積核分組{G1，Gs，G2}中，然後得到新一層的特徵。
具體路由公式爲：
其中，h是一個激活函數，進行線性變化。*代表卷積操作。|表示特徵圖的維度拼接操作。
基於分析這個路由公式可知，當 L等於0時，也就是剛開始輸入的時候，輸入圖片 x會和第一層的分組卷積核進行卷積運算，第一層特徵輸出的通用表示爲：
下圖是兩個任務的損失L1和L2進行梯度回傳的路線，可以看到，特定於任務的卷積核分組G1、G2僅根據自身任務的損失進行更新，而共享組Gs則根據兩個任務的損失進行更新。
網絡結構
下圖是整體架構的示意圖，其中每個SFG模塊在每個卷積層中隨機生成濾波器分組，所得到的特徵按上述方式稀疏路由。黑色圓點代表融合模塊，通常是對來自特徵任務和來自共享任務的特徵圖進行維度拼接（concat）或者對應元素相加（element-wise add），然後將新的特徵圖傳遞到下一層的卷積核分組中。白色圓點表示一些附加的操作，通常是在激活函數前加入額外的卷積或者全連接層。
這種稀疏連接是確保任務性能和結構化表示的必要條件。我們可能會感到疑惑的一點是，第L層的共享分特徵路由到第L+1層的特定任務的卷積核分組中，會讓子任務的loss不能很好的分離。作者在這裏的考量是，這種稀疏路由連接方式對於特定任務學到更豐富的特徵是有幫助的。舉例來說，如下圖所示，如果在第一層有很大比例的卷積核被分配到了共享卷積分組中，並且沒有路由，那麼這會大大減少後續層中應用到特定任務的卷積核數量，在極端情況下，一層中所有的卷積核都被分配到共享卷積分組中去，後續層中的特定於任務的卷積核分組實際上未被使用。
另外一個需要關注的點是在這個過程中，特徵圖的維度變化。具體來說，在不同分組中卷積核的數量是會隨着訓練的輪數進行變化的，在這種情況下，每一個卷積核分組的輸出特徵{F1,F2,Fs}的維度也是在變化。本文不是直接處理不同大小的特性圖，而是在SFG模塊中，對輸入特徵，首先和全部卷積核進行卷積，然後根據任務分組，將原始來源於不同卷積核分組的通道置零。在特徵融合方面，如果沒有額外的轉化，進行元素加操作；如果有諸如殘差塊等複雜轉化，進行維度拼接操作，然後對融合特徵進行1x1卷積進行升降維，確保融合後的特徵與共享特徵維度相同。
優化卷積網絡參數和分組概率的方法
給定一個有L個SFG模塊的網絡，每一個模塊中包含K個卷積核。SFG模塊中卷積核的後驗概率可以表示爲：
由於SFG模塊中的卷積核的後驗分佈難以直接進行準確的求解，本文采取變分推斷的思想，用一個簡單的近似分佈qφ(W)代替，其中W是所有層的所有卷積核的集合。假定將qφ(W)分佈按照卷積層和每一層的卷積核分組進行分解，那麼該分佈的變分推斷可以表示爲：
紅框部分表示第L層中第k個卷積被路由到不同卷積核分組後的結果，它的詳細表示爲：
其中，函數z的作用是將第k個卷積核屬於哪一個卷積核分組（task1/task2/share）的概率轉化爲一個one-hot向量，函數M爲卷積核的參數。
此時別忘了，我們真正的後驗概率是上文中的P，我們提出一個簡單的近似分佈q來擬合它，怎樣優化q讓它們的分佈更相似呢，我們想到了KL散度，KL散度可以表示兩個概率分佈之間的相似程度。另一方面，似然函數能夠根據輸出，推斷最有可能產生這種輸出的參數，因此，我們定義一個優化函數L，它是由聯合似然估計和KL散度綜合而來，如下式：
其中，M是batch_size，W表示從近似分佈q抽樣的一組模型參數。而作者借鑑《Uncertaintyin Deep Learning》一文中的思想，將KL散度進一步表示爲：
其中，H函數爲分組概率的交叉熵。將（4）式帶入（3）式，可以得到最終的損失L爲：
λ1和λ2是正則化係數
實驗
本文在兩個多任務學習任務上對SFG模塊進行實驗：（1）基於UTKFace數據集上的人臉圖像年齡預測和性別分類；（2）基於醫學影像數據集的語義圖像合成與分割。
數據集
UTKFace dataset
UTKFace datasets由23703張帶年齡和性別標籤的人臉圖像組成。在訓練中分別按照70%，15%，15%的比例劃分爲訓練集，驗證集和測試集。同時構建一個次級數據集，其中只包含初始集10%的圖像，以便模擬一個數據匱乏的場景。
Medical imaging dataset
在這個實驗中，創建15個3D醫療圖像掃描，並使用語義3D標籤對器官進行掃描。創建了一個包含10名患者的訓練集，其餘5名用於測試。將網絡訓練在從軸向切片中隨機採樣的尺寸爲128x128的二維圖像上，並在測試時重建尺寸爲288x288x62的三維體塊。
基準
baseline共有四種，分別是：1）單任務網絡，2）硬參數共享的多任務網絡，先進行統一的共享特徵，然後進行特徵分離，3）MT-constantmask，每個任務中分配相同數量的卷積核，4）MT-constantp ,每個卷積核別分配到任務1，任務2或者任務共享分組中的概率相等，都是1/3。
實驗結果
年齡預測和性別分類
如下圖所示，MAE是預測年齡與真實年齡的平均預測誤差，Accuracy是分類正確率。在充足數據集情況下，本文提出的MT-SFG方法在年齡預測和性別分類上都超出了baseline性能表現，這說明了學習分配卷積核的好處。
即使在數據量匱乏的數據集上，MT-SFG依然獲得了最佳性能。
醫學圖像迴歸和語義分割
結構分析
通過分析嵌入SFG模塊的網絡的分組概率，可以將網絡連接可視化，從而瞭解所學習的MTL體系結構，爲了分析每一層卷積核的分組方式，計算了每一層卷積核類別概率的和。結果如下圖所示，藍，綠，粉分別代表卷積核分配到task1，task-share和task2中的概率，可以看到，從總體來說，隨着網絡深度的增加，網絡越來越傾向於把卷積核分配到特定任務的分組中去（隨着網絡層數加深，藍色和粉色的比例不斷增加）。
在第一層，幾乎所有的卷積核都被分配到共享卷積核分組中，這是因爲，初始的卷積得到的特徵是低階特徵，對於所有任務都會適用，但是隨着網絡層數的加深，特徵越來越高階和抽象，這時候卷積核更容易被分配到特定任務的卷積核分組中去，因爲抽象的特徵只有加上特定的任務屬性纔會有意義。
結論
本文提出了隨機濾波器分組(SFGs)來區分特定任務和共享任務特徵。SFGs在概率上定義了卷積核的分組，從而定義了CNNs中特性的連接性。在訓練過程中，我們使用變分推斷來近似給定訓練數據和可能架構上的樣本的連通性分佈。本文的方法可以看作是多任務體系結構學習的一種概率形式，通過後驗概率，推斷出最優的MTL結構。
論文鏈接：
ICCV2019 Open Access Repositoryopenaccess.thecvf.com
（*本文首發於「AI科技大本營」微信公衆號，更多幹貨內容請在公衆號查看。轉載請聯繫微信1092722531）