MMoE論文筆記

Modeling Task Relationships in Multi-task Learning with Multi-gate Mixture-of-Experts

論文鏈接

論文筆記

摘要

為什么提出？

• 多任務學習旨在構建單個模型來同時學習多個目標和任務。

• 但是，通常任務之間的關系會極大地影響多任務模型的預測質量。

• 因此，學習task-specific objectives和inter-task relationships之間的權衡也非常重要。

  It is therefore important to study the modeling tradeoffs between task-specific objectives and inter-task relationships.

怎么做？

Multi-gate Mixture-of-Experts(MMoE) 能夠清晰地從數據中學習任務之間的關系：

• 將Mixture-of-Experts(MoE)的結構應用於多任務學習中
• 在所有任務共享expert submodels的同時，通過訓練gating network來優化每一個任務。

結果如何？

• 實驗表明在任務相關性比較弱的時候，MMoE會表現的更好。
• 同時能夠提升模型的可訓練性：an additional trainability benefit

一、引言

實際上，多任務學習模型的表現並不總是一定會由於單任務模型：

• 許多DNN-based multi-task learning models會對數據分布的不同、任務之間的關系等因素特別敏感。
• 源於不同任務之間的固有的沖突事實上會對於至少一部分任務造成不利的影響，尤其當模型的參數在各個任務中被廣泛共享的時候。

現有的解決方法：

• 基於一個特定的數據生成假設前提去探究多任務學習中任務的區別，並且依據任務區別的程度去提出建議 -> 真實的數據具有極為復雜的模式，很難去明確地衡量任務之間的區別並依此去應用這些建議
• 也有不需要給出明確的任務區別程度的方法，但是為了適應不同的任務，往往帶來更多的模型參數 -> 在資源有限的情況下，計算的成本是難以接受的

MMoE怎么樣：

• MMoE對於任務之間的關系進行明確地建模，並且學習任務特定的函數去平衡共享的表達
• MMoE能夠在不添加大量新的參數的情況下讓參數進行自動分配以去捕捉shared task information以及task-specific information。

MMoE explicitly models the task relationshipds and learns task-specific functionalities to leverage shared representations.

It allows parameters to be automatically allocated to capture either shared task information or task-specific inofrmation, avoiding the need of adding many new parameters per task.

MMoE結構概述：

• 基於Shared-Bottom multi-task DNN structure（如上圖(a)）: Shared-Bottom: input -> bottom layers (shared) -> tower network (每個任務各自的)
• 如上圖(c)，MMoE有一組bottom networks, 每一個叫做一個expert, 在本文中， expert network是一個feed-forward network
• 然后為每個任務引入一個gating network。Gating networks 的輸入是input features， 輸出是softmax gates，即各個expert的權重
• 加權之后的expert結果被輸入到task-specific 的tower networks
• 這樣的話，不同任務的gating networks能夠學到不同的專家混合方式，以此捕捉到任務之間的關系
• MMoE更容易訓練並且能夠收斂到一個更好的loss，因為近來有研究發現modulation和gating機制能夠提升訓練非凸深度神經網絡的可訓練性。

核心:

1. 清晰地描述任務之間的關系：通過modulation和gating networks，模型能夠自動調整學習shared information以及task-specific information之間的參數化過程。
2. 同時提升模型的表達力以及可訓練性。

MMoE explicitly models task relationships. Through modulation and gating networks, our model autoatically ajusts parameterization between modeling shared information and modeling task-specific information.

MMoE improves both model expressiveness and trainability.

二、文獻綜述

2.1 Multi-task Learning in DNNs

shared-bottom model struccture:

• 所有任務共享bottom hidden layers
• 減少過擬合的風險
• 但是會遭受由於任務的區分帶來的優化沖突，因為所有任務在shared-bottom layers使用相同的參數集

為task-specific 參數增加不同類型的約束：

• cross-stitch network, tensor factorization model
• 有更多task-specific的參數來解決由於task differences帶來的更新shared參數時的沖突，但是參數極大地增加，需要更多的訓練數據

2.2 Ensemble of Subnets & Mixture of Experts

Eigen 和 Shazeer將mixture-of-experts model轉變為基本的building blocks (MoE layer)，並且把它stack到一個DNN中：

• 在訓練和測試的時候，MOE layer會依據這一層的輸入來選擇subnets(experts)

PathNet:

• 是一個巨大的神經網絡，每一層都有multiple layers和multiple submodules

三、基本知識

3.1 Shared-bottom Multi-task Model

• 給定\(K\)個任務，模型由一個shared-bottom network (由 \(f\) 表達) 和 \(K\) 個 tower networks \(h^k\) 組成，其中 \(k=1,2,\cdots,K\)。

• shared-bottom network跟在輸入層后面，tower networks建立在shared-bottom 的輸出之上，最后由每個tower產生各自任務的輸出\(y_k\)

• 對於任務\(k\)，模型可以表達為：

\[y_k=h^k(f(x)) \]

四、Modeling Approaches

4.1 Mixture-of-Experts

最初的Mixture-of-Experts (MoE) Model 可以由下式表達：

\[y=\sum_{i=1}^{n}g(x)_if_i(x) \]

其中:

• \(\sum_{i=1}^{n}g(x)_i=1\), \(g(x)_i\) 表示expert \(f_i\) 的權重
• \(f_i, i=1,\cdots,n\) 是 \(n\) 個expert networs
• 表示一個用於綜合所有專家結果的gating network
• 換句話說，gating network \(g\) 基於input產生\(n\)個experts上的一個分布，而最后的結果是所有experts輸出的weighted sum。

MoE Layer:

• MoE layer和MOE model有着相同的結構，只不過是將前一層的輸出作為輸入，輸出則會繼續傳到下一層。整個模型以端到端的方式訓練
• MoE layer最開始提出的主要目的是進行conditional computation, 也就是說only parts of a network are active on a per-example basis
• 對於每個輸入的example，模型能夠通過gating network基於input選擇a subset of experts。

4.2 Multi-gate Mixture-of-Experts

• 本文提出一個新的MoE模型用於捕捉任務之間的區別，並且相對於shared-bottom多任務模型，本模型不會帶來參數的極大增加
• 新的模型叫做Multi-gate Mixture-of-Experts(MMoE) model, 其主要思想是用MoE layer來替代share bottom network \(f\)
• 此外，對於每個任務\(k\)，都加入一個單獨的gating network \(g^k\)
• 具體地，任務\(k\)的輸出是：

\[y_k = h^k(f^k(x)) \]

\[f^k(x)=\sum_{i=1}^{n}g^k(x)_if_i(x) \]

​ gating networks就是一個簡單加上softmax layer的線性變換：

\[g^k(x)=softmax(W_{gk}x) \]

​ 其中\(W_{gk}\in R^{n\times d}\)，\(n\)為experts的個數，\(d\) 為特征維度。

參數維度分析

github - Keras MMoE

• d: number of input features

• h: number of units per export

• n: number of experts

• k: number of tasks

  expert network: \(W_{n\times h \times d}\)

  gating network: \(W_{k\times n \times d}\)

expert networks:

n個expert networks的輸出\(y\in R^{n\times h}\)

\[y=f(x)=W_{n\times h \times d}x + b_{k\times h} \quad \quad(i=1,\dots,n) \]

k個gating networks的輸出 \(z \in R^{k\times n}\)

\[z= g(x)=softmax(W_{k\times n \times d}x + b_{k\times n}) \]

五、MMoE on Synthetic Data

5.1 Performance on Data with Different Task Correlations

• 所有的模型，任務相關性越高，表現越好
• 隨着任務相關性的差別造成的MMoE model表現的差異要小於OMoE model以及Shared-Bottom model，這種趨勢在比較MMoE model和OMoE model的時候尤為明顯：
  • 在極端的情況下，當兩個任務完全一樣的時候，MMoE和OMoE model的表現基本沒有差別
  • 但是隨着兩個任務相關性的減弱，OMoE model的表現出現明顯的下降，而對MMoE model僅有微弱的影響。可見，在低相關性的情況下，gate network要因任務而異是非常重要的
• OMoE 和 MMoE model在所有情境下都優於Shared-Bottom model

5.2 Trainability

近來，有一些研究發現gated RNN models(比如LSTM， GRU) 之所以比普通的RNN表現要好，是因為它們更容易去訓練，而不是有更好的模型能力。

我們希望對於MMoE的trainability進行深一步地探討。

通過使用不同的random seeds生成相同分布下的數據以及不同的模型初始化，下圖描述了MMoE, OMoE, Shared-Bottom最終的loss分布直方圖。

可以發現：

• Shared-Bottom model 有更大的variances, 也就是說Shared-Bottom in general have much more poor quality local minima
• task correlation =1時，OMoE和MMoE健壯性相當，但是當任務相關性下降，OMoE的健壯性出現顯著下滑，可見，multi-gate structure在resolve bad local minima caused by the confict from task difference是有效的

結論

• MMoE能夠對於任務之間的關系明確地建模
• gating networks是輕量的，而expert networks是共享的，因此並不會引入造成computation cost
• 更容易訓練