Jurgen、曼寧等大佬新作:MoE重塑6年前的Universal Transformer,高效升級

机器之心發表於2024-10-19
7 年前,谷歌在論文《Attention is All You Need》中提出了 Transformer。就在 Transformer 提出的第二年,谷歌又釋出了 Universal Transformer(UT)。它的核心特徵是透過跨層共享引數來實現深度迴圈,從而重新引入了 RNN 具有的迴圈表達能力。層共享使得 UT 在邏輯推理任務等組合問題上的表現優於 Transformer,同時還在小規模語言建模和翻譯任務上得到改進。

UT 已被證明具有更好的組合泛化特性,能夠在無監督的情況下解構結構化問題並泛化到更長的序列。因此與 Transformer 相比,UT 是一種具有卓越泛化特性的通用性更強的架構。

但 UT 的計算效率遠低於標準 Transformer,不適合當前語言建模等以引數為王的任務。那麼,我們能不能開發出計算效率更高的 UT 模型,並這類任務上實現比標準 Transformer 更具競爭力的效能呢?

近日,包括 LSTM 之父 Jürgen Schmidhuber、史丹佛大學教授 Christopher Manning 等在內的研究者從全新視角出發,提出瞭解決 UT 基礎計算引數比問題的最新方案。具體來講,他們提出 Mixture-of-Experts Universal Transformers(簡稱 MoEUT),它是一種混合專家(MoE)架構,允許 UT 以計算和記憶體高效的方式擴充套件。
圖片
  • 論文標題:MoEUT: Mixture-of-Experts Universal Transformers
  • 論文地址:https://arxiv.org/pdf/2405.16039
  • 程式碼地址:https://github.com/robertcsordas/moeut

在文中,研究者利用了 MoE 在前饋和自注意力層方面的各種最新進展,並將這些進展與以下兩項創新工作相結合:1)layer grouping,迴圈堆疊 MoE 層組;2)peri-layernorm 方案(位於 pre-layernorm 和 post-layernorm 之間),並且僅在緊接 sigmoid 或 softmax 啟用之前應用層範數。這兩者都是專為共享層 MoE 架構設計,並且有強有力的實證證據支援。

從其作用來講,MoEUT 允許構建引數和資源高效的 UT 語言模型,不僅在我們可以負擔得起的所有規模(最高 10 億引數)上對算力和記憶體的需求更低,效能也超越了標準 Transformer。

研究者在 C4、SlimPajama 和 peS2o 語言建模資料集、以及 The Stack 程式碼資料集上測試 MoEUT 的能力,結果表明,迴圈對於模型實現具有競爭力的效能至關重要。同樣地,研究者在 BLiMP 和兒童圖書測試、Lambada、HellaSwag、PIQA 和 ARC-E 等下游任務上展現了良好的零樣本效能。

MoEUT 架構概覽

如前文所述,MoEUT 是一種具有層共享引數的 Transformer 架構,其中使用 MoE 來解決引數計算比問題。雖然最近出現了很多關於 Transformer 語言模型的 MoE 方法研究,但要讓它們在引數相同的情況下與密集方法競爭,仍然是一項艱鉅的任務。

因此,研究者利用了 MoE 前饋網路塊(FFN)、MoE 自注意力層,並引入了兩種考慮到共享層模型特定屬性的新方法 ——layer grouping 和訊號傳播。這些技術的結合對於實現有效的共享層 MoE Transformer 發揮了巨大作用。

MoE 前饋塊

為了透過 MoE 來引數化共享層 Transformer 的前饋塊,研究者使用了 σ-MoE 並做了一些修改。σ-MoE 將前饋塊劃分為 N_E 個切片,稱為專家(expert)。每個專家都有兩組權重,分別是 圖片圖片,其中 e ∈{1,...,N_E} 是專家索引。

在每個 token 位置 t,給定層輸入 x_t ∈ R^d_model,MoE 前饋層會為每個專家計算一個分數,從而得到一個向量 s ∈ R^N_E,其計算如下:
圖片
MoE 層僅選擇與 s_t ∈ R^N_E 中 top-K 元素相對應的 K 個專家(從 N_E 中),來產生層輸出 y_t ∈ R^d_model,如下所示:
圖片
初步實驗表明,σ-MoE 的原始正則化往往不穩定,有時會導致訓練期間損失激增。為了避免這種情況,研究者僅在序列內應用正則化(而不是批次中的所有 token)。對於輸入序列 x_t,t ∈ {1,...,T},計算平衡損失 L 如下所示:
圖片
MoE 自注意力

為了將 MoE 引入自注意力層,研究者應用了 SwitchHead,它是一種將 σ-MoE 擴充套件到注意力層的 MoE 方法。與標準多頭注意力層一樣,SwitchHead 層中的每個頭包含四個轉換:查詢、鍵、值和輸出投影。但是,SwitchHead 使用 MoE 來引數化值和輸出投影。

也即,每個頭都有一個與之關聯的查詢和鍵投影以及 N_A 值和輸出投影,它們針對每個輸入進行動態選擇。

鍵和查詢「照常」計算:給定位置 t 處的一個輸入,x_t ∈ R^d_model,並且 k^h_t = x_tW^h_K、q^h_t = x_tW^h_Q,h ∈{1,...,H} 是頭索引。專家對這些值的選擇計算如下:
圖片
值和輸出的選擇是獨立的。輸出的選擇使用不同的權重矩陣 W^h_SO ∈ R^d_model× N_A 來類似地計算,即 圖片圖片。輸出 y ∈ R^d_model 計算如下:
圖片
層分組:MoE 高效層共享和操作內的子操作

研究者觀察到,基於 MoE 的原始 UT 只有一個共享層,在更大規模上往往難以獲得良好的效能。假設原因有二:首先,隨著網路規模的擴大,層中專家的數量會迅速增加,但我們無法以相同的速度增加活躍專家 K 數量而不大幅增加所需計算量。這就迫使我們降低活躍專家的比例,而這通常是不利的。其次,注意力頭的總數保持在相對較低的水平,這對於一個大型模型來說可能是不夠的。增加註意力頭的數量也同樣代價高昂。

因此,在增加註意力頭總數的同時,可以呼叫一組層,減少每個 σ-MoE 中的專家數量。最終的網路是透過重複堆疊這些共享相同引數的小組而得到的(從某種意義上說,將組重新定義為 UT 中的共享層)。

下圖 1 提供了一個示例,標記為「層 A」(或層 B)的所有層在整個網路中共享相同的引數。組 G 的大小是非共享層的數量。在研究者的實驗中,組大小在 2 到 4 之間,典型的迴圈步驟數為 8 或 9。
圖片
在 UT 中改進訊號傳播的新型層歸一化方案

研究者提出另一種方法來避免上述問題:在「主資料路徑」中不使用分層歸一化。這意味著,對於本文的 UT,在注意力值投影之前不使用分層矩陣,在 σ-MoE 層之前也不使用分層矩陣。相反,只有在緊跟著 sigmoid 或 softmax 啟用函式的線性層(在這些非線性層之前產生關鍵的重歸一化啟用)之前,即注意力中的查詢和關鍵投影、注意力層和前饋層上的專家選擇以及最終分類層之前,才會使用 layernorm。如圖 3 所示。

由於前饋層內的主資料路徑上只使用了 ReLU 啟用函式,因此輸出更新將與輸入成正比,從而有效地解決了殘差增長問題,同時也提供了高效的梯度流路徑。這種方案稱為 「peri-layernorm」方案,它介於「pre-layernorm」和「post-layernorm」方案之間,將 layernorm 定位在殘差連線的「周圍」(但不在其上)。具體如下圖 3 所示。
圖片
實驗結果

在論文中,研究者展示了使用流行的 C4 資料集進行語言建模時 MoEUT 效能和效率的主要實驗結果。為了證明 MoEUT 的通用性,研究者還展示了在 SlimPajama 和 peS2o 語言建模資料集上的主要結果,以及在 「The Stack」上的程式碼生成。

與標準 Transformer 對比 Scaling

MoEUT 的 Scaling 結果如圖 4 所示。y 軸顯示的是 C4 held-out 子集上的複雜度。在引數數量相同的情況下,MoEUT 模型略微優於密集模型(圖 4a),而且隨著規模的擴大,差距呈擴大趨勢。研究者還與非共享 σ-MoE 模型進行了比較,該模型的表現明顯不如 MoEUT,這表明共享層具有明顯的優勢。此外如圖 4b 顯示,就訓練期間所有前向傳遞所花費的總 MAC 運算次數而言,MoEUT 遠遠優於基線密集模型。
圖片
程式碼生成效能

為了證實 MoEUT 在不同任務領域的有效性,研究者在「The Stack」資料集的一個子集上對其進行了訓練,該資料集是一個程式碼生成任務。由於無法進行完整的 epoch 訓練,因此這裡只使用了幾種語言並混合使用了這些語言:Python、HTML、C++、Rust、JavaScript、Haskell、Scala 和彙編。研究者在資料集的一個 held-out 子集上評估了 MoEUT。結果如圖 5 所示,與自然語言領域的結果一致:MoEUT 的表現優於基線。
圖片
圖 6 展示了 layer grouping 對 244M 引數 MoEUT 模型的影響。研究者發現 G = 2 是最佳值,而且層維度的迴圈確實是有益的。

下游任務上的零樣本表現

研究者評估了 MoEUT 在六個不同下游任務中的零樣本效能:LAMBADA、BLiMP、Children's Book Test (CBT) 、HellaSwag、PIQA 和 ARC-E。結果見表 1,MoEUT 的表現往往優於基線,但在所有情況下差異都很小。
圖片
研究者還將 MoEUT 與另一個基準模型 Sparse Universal Transformer(SUT)進行了比較,SUT 是最近提出的一個 UT 模型,也使用了 MoE 層,且以前未在標準語言建模任務中進行過評估。雖然 MoEUT 和 SUT 都在前饋層和注意力層使用了 MoE,但這兩種方法在不同層面上存在一些技術差異:SUT 使用競爭性專家選擇(softmax)、多重負載平衡損失和更大的專家規模,且採用 post-layernorm 模式,不使用 layer grouping。與 MoEUT 的方法不同,SUT 在層維度上使用了自適應計算時間(ACT)。

結果如圖 7 所示。與 MoEUT 和引數匹配的密集基線相比, SUT 在效能上有明顯的劣勢。研究者認為這種效能低下的主要原因是作者將 ACT 機制作為其模型的主要組成部分之一。移除 ACT 後,效能顯著提高。然而,即使在這種設定下,它的效能仍然低於 MoEUT 和標準 Transformer 基線。
圖片
研究者還對「peri - 層歸一化」進行了評估。圖 8 顯示了結果。本文的層歸一化方案始終表現最佳。小模型的差距更大,而大模型的差距則越來越小(對於 719M 引數模型,peri-norm 和 post-norm 之間的差距微乎其微)。同時,隨著訓練步數的增加,peri-norm 和 post-norm 之間的差距也在增大,因此如果模型的訓練時間更長,就有可能獲得更高的收益。

調整專家選擇機制

為了更好地理解 MoEUT 的專家選擇,研究者分析了在 C4 上訓練的 244M 引數 MoEUT 模型的 MLP 塊中的專家選擇。本節中的所有實驗都是透過計算 C4 驗證集上 G = 2(即模型組中有兩層)模型的統計資料進行的。這裡只展示了模型組第一層的行為,因為研究者發現第二層的結果在本質上是相似的。

結果表明,MoEUT 能夠根據不同情況動態調整其專家選擇機制。有時,專家會被分配給流行的 token,而在其他情況下,專家會在各層之間共享或專門化,這取決於哪種方式更適合任務。

專家的跨層使用。如圖 9 所示,右下角的黃點表示一些專家主要被分配到最後一層。然而,對於其他專家來說,專家被啟用的層範圍很廣。專家似乎是在連續的層序列中被啟用的,這可以從縱向排列的寬闊結構中看出。因此可以得出這樣的結論:如果有必要,MoEUT 能夠專注於特定層,並可在各層之間共享權重
圖片
每個 token 專家選擇的多樣性。研究者分析了 MLP 各層針對給定輸入 token 在不同層和上下文中的專家選擇多樣性。為此,他們測量了不同層在不同位置 / 上下文下為單個 token 啟用的專家總數。結果如圖 10 所示。

各欄/位置的專家選擇動態。結果如圖 11 ,在後續層中,所選專家之間存在不可忽略的重疊;但是,這種重疊還遠遠沒有達到完全重疊的程度。這表明,專家通常在單列中動態變化,在不同層中執行不同的功能。
圖片
更多研究細節,可參考原論文。

相關文章