模組化大模型來了！IBM公開WastonX核心架構技術細節

允中 發自 凹非寺

大型語言模型（LLMs）的性能非常強大，但是現有的模型訓練和部署成本都很高。而且在不忘記先前知識的前提，擴展它們去學習新的知識也很困難。也很難針對特定的任務去提取出輕量化的模型。

最近，來自MIT-IBM Waston AI Lab、清華大學、Mila的研究人員聯合提出了一種新的神經網路架構ModuleFormer，利用模組化來大幅提高大型語言模型的效率和靈活性。

ModuleFormer是一種基於稀疏專家混合(SMoE)的模組化架構，包括兩種不同類型的模組，即新的stick-breaking注意力專家模組和傳統的MLP專家模組。在訓練和推理過程中，根據輸入的向量，不同的模組會被稀疏地激活。

與之前基於SMoE的模組化語言模型不同，ModuleFormer可以通過其新的負載平衡（load balance）和負載集中（load concentration）損失函數從未經篩選的資料中誘匯出模組化。

在實驗中，團隊發現模組化架構使得大型預訓練語言模型具備了三個重要的能力：

1）效率‍

因為ModuleFormer只對每個輸入只激活一小部分模組，因此可以以兩倍以上的吞吐量達到與常規語言模型相同的性能。

2）可擴展性

實驗表明，由於微調階段ModuleFormer只需要更新一部分的模組，因此比常規語言模型更不容易發生災難性遺忘，並且可以輕鬆通過新的模組擴展以學習訓練資料中不包含的新知識。

3）模組特異化和篩選‍

在微調階段ModuleFormer，新提出的負載集中損失函數可以自動篩選一部分模組，讓它們專注於目標任務，而與任務無關的模組可以被直接拋棄掉以實現輕量化部署。

導言

儘管現代大型語言模型（LLM）在某些任務上取得了顯著的成果，甚至超過了人類的表現，但其效率和靈活性仍然不高。

大多數LLM（例如Llama，Falcon）在推理和訓練過程中都使用了它們的全部參數，我們稱這些模型為密集模型。

然而，先前的研究已經表明，在執行任何特定任務時，神經模型中的大部分參數都可以在不影響模型性能的前提下被剪枝掉。

此外，一旦訓練完成，LLM就會「定格在某個時間點」，但許多實際應用情況要求LLM具有最新的知識。因此模型進行continue leanring的擴展能力也十分重要。

隨著模型規模的增長，為了領域適應或持續學習而對整個模型進行微調變得代價高昂且計算資源受限，這使得那些計算預算較小的使用者無法實施。同時，更新所有參數也使得模型容易遭受災難性遺忘（catastrophic forgetting）。

為此，像LoRA這樣僅更新一小部分原始參數修正量的輕量級適應方法正在變得流行。

然而實驗表明，這種方法仍然可能遭受災難性遺忘影響，並且LoRA並不擅長需要模型學習大量新知識的場景，比如讓模型學習一種新的語言。

文章作者認為模組化是解決前述問題的一個好方法。模組化模型具有以下幾個優點：

• 模型可以在輸入或任務上激活一組模組條件，從而比密集激活整個模型需要更少的計算量；

• 在給定領域或任務的情況下，可以組裝一組與領域/任務相關的模組，形成一個新的輕量級模型；

• 模型可以輕鬆添加新的模組進行領域適應或持續學習；

• 模型可能更不容易發生災難性遺忘，因為只有與輸入相關的模組在模型微調期間進行更新。

這篇論文提出了一種新的模組化架構，ModuleFormer（圖1a），以及在其中進行模組操作的方法。ModuleFormer每層包含一個組前饋神經網路（MLP）專家模組和一組新提出的Stickbreaking注意力專家模組。

為了平衡在訓練過程中不同模組的負載，作者提出了一種新的互資訊損失函數。此外，文章還展示瞭如何在ModuleFormer中插入新模組（圖1b）和進行模組修剪（圖1c）。

為了實現模組修建，文章引入了一種新的負載集中損失函數，用於在微調的過程中自動選擇最適合執行給定任務的模組，同時通過微調進一步增強了這些模組執行該任務的能力。

△圖1

實驗結果顯示ModuleFormer相對於密集模型在以下方面有顯著的提升：

• 由於稀疏的模組激活機制（圖1a）它在更低的延遲（50％）和更小的記憶體佔用下實現了與密集LLM相同的性能，因此ModuleFormer可以實現密集模型兩倍的吞吐量。

• 在對新領域進行微調後，它收到災難性遺忘的影響較低，並且也可以輕鬆地通過添加新模組來學習新語言和知識。（圖1b）。

• 它可以在下游任務上進行微調，將一部分模組特異化成為處理該任務專用的模組，而未使用的模組可以被扔掉而不會犧牲模型在該任務上的性能（圖1c）。

詳解ModuleFormer模型

Sparse Mixture of Experts（SMoE）的基本結構

SMoE最早由Shazeer在《Outrageously large neural networks: The sparsely-gated mixture-of-experts layer》一文中提出。

每個SMoE層包含一組用於處理輸入產生輸出的專家模組m_1, m_2, …, m_n，以及一個用於挑選專家的路由函數g。在本篇文章中，路由函數由一個一層的神經網路來建模：

在給定一個輸入向量x之後，路由函數g會計算一個專家模組的概率分佈g(m|x)，然後模型會自動選擇分佈中前top k的專家模組來處理輸入x。SMoE層的輸出就是專家模組的輸出通過對應的路由概率家和得到：

其中不再topk內的專家模組不會被計算，同時對應的概率g也會被用0替代。

在ModuleFormer中，作者使用了兩種不同的專家模組來分別構建傳統transformer解碼器中的自注意力層（self attention）和前饋層（MLP）。其中前饋層的專家模組和常見的SMoE專家模組一致，是一個單一隱藏層的全連通神經網路。而自注意力層的專家模組則是一個新提出的stick-breaking注意力模組。

Stick-Breaking注意力模組

Stick-Breaking自注意力是設計用於取代Transformer解碼器的自注意力層，使用狄利克雷過程中的Stick-Breaking（折棍子）過程對每個字符xt與之前的字符x

Stick-Breaking自注意力的優勢在於它會自動關注最近的相關字符，而不需要引入額外的位置資訊，比如position embedding和relative position bias。

給定一個包含 t 個時間步的輸入向量序列 x1, x2, …, xt，每個輸入被投影到一系列key向量 k1, k2, …, kt 和一系列value向量 v1, v2, …, vt。為了計算時間 t 的注意力，輸入 x_t 被投影到一個查詢向量 q_t = W_q x_t，其中 W_q 是查詢投影矩陣。對於所有之前的步驟和當前步驟 i ≤ t，計算時間步 i 的鍵與時間步 t 的查詢匹配的概率:

需要注意的是，這個查詢概率使用了sigmoid激活函數，所以沒有歸一化。接下來通過stick-breaking過程來對查詢概率進去歸一化：

這樣，注意力就會自動分配給離t時刻最近，且具有較大查詢概率的時刻。使得自注意力機制在沒有額外的位置資訊的情況下，也能對於相對位置進行有效的建模。最終，自注意力模組的輸出是由注意力權重對歷史的value向量進行加和並且投影得到：

ModuleFormer中的模組控制

預訓練中的負載均衡

為了避免SMoE反覆使用相同的模組並浪費其他模組的額外容量，一般採用負載平衡損失函數來調節每個專家的使用頻率。與之前的SMoE模型 不同，團隊希望最大化輸入字符和模組之間的互資訊（MI）：

為了簡化起見，假設在批次X中的令牌分佈是均勻的，因此p(x) = 1/X。在去除所有常數成分後，可以簡化互資訊損失（公式6）為p(m)的熵與p(m | x)的條件熵之間的差異。

在上述內容中，p(m) = sum_x(g(m|x)p(x))，其中p(x)是批處理中每個字符的概率，H(m)是模組分佈的邊際熵，H(m | x)是模組在給定輸入字符x的條件下的熵，|X |是輸入字符的數量。對於長度為T的batch大小為B的小批量，字符的數量是|X | = BT，字符的概率是p(x) = 1/|X |。

直觀地說，互資訊損失最大化了模組的概率分佈的邊際熵，並最小化了給定輸入x的模組條件分佈的商。它平衡了整個batch中每個專家的負載（最大化H(m)），同時也鼓勵每個輸入x將其路由概率集中在較少的模組上（最小化H(m | x)）。

微調中的負載集中

儘管團隊希望在預訓練期間最大限度地利用每個專家的能力，但在微調期間希望將少量的模組專注於下游任務。這樣可以移除未使用的模組並減少微調後模型的參數數量。為了將負載集中在較少的模組上，團隊引入了一個新的負載集中損失函數來最小化模組的邊際熵：

這樣可以鼓勵模型使用更少的模組來處理下游任務。在微調後，可以計算在訓練或驗證集上使用的模組頻率f_m。f_m代表了模組m對於這個任務的重要性，可以通過移除f_m小於某個特定閾值的專家來輕鬆實現模型剪枝。

用新的模組來學習新的知識

對於模組化模型來說，插入新模組是一種直接且參數高效的方法，可以在不對整個模型進行微調的情況下學習新知識。當向每一層插入N_new個隨機初始化的模組時，還需要擴展路由器（方程2中的A）中的模組嵌入層A，使其包含一個形狀為（N_new，D_rtr）的新矩陣A’。因此，新的路由函數可以寫成：

由於在微調期間其他的模組參數被凍結，因此使用新模組進行持續學習可以在很大程度上避免災難性遺忘問題。

然而，災難性遺忘仍然可能影響路由函數。當新模組在一個新領域進行訓練時，如果路由函數錯誤地將來自舊領域的輸入路由到新專家，模型可能會遭受災難性遺忘。

為了避免這種情況，團隊對路由函數進行了正則化以避免災難性遺忘，並提出了兩種訓練策略：

1）全面微調路由，公式9中A和B使用預訓練參數進行初始化，而A’則是隨機初始化的。這個策略是為了訓練資料中同時包含新舊資料的情況設計。

2）只訓練A’，這個策略是為了連續學習（lifelong learning）的情況而設計的，不使用以前訓練過的資料。由於這種情況可能導致新的模組使用頻率過高，從而帶來災難性遺忘。團隊引入了正則項來限制A’的範數：

與被指出存在缺陷的傳統連續學習正則化方法（如衰減或L2損失）不同，路由正則化不限制專家的能力，而只限制對新專家的使用趨勢。

評估

基於ModuleFormer，研究者在Pile資料集上預訓練了三個不同體積和計算量的ModuleFormer Language Model（MoLM）語言模型：

基礎性能評估

團隊使用Language Model Evaluation Harness來評估零樣本、少樣本和語言建模任務中的語言模型。

對於零樣本和少樣本任務，目標是在給定上下文的基礎上從一組給定選項中選擇最合適的完成部分。最終選擇在給定上下文下具有最高可能性的完成部分。

對於語言建模，在Wikitext資料集上進行測試。目標是最小化下一個標記預測的困惑度。

對於程式碼生成，在HumanEval資料集上評估模型。HumanEval包含164個手寫的Python程式設計問題。模型需要根據任務描述提示完成一個函數，以便能夠通過所有提供的測試案例。

表2和表3顯示了MoLM和基準語言模型在常識推理、閉卷問答和程式碼生成基準上的性能。

總體而言，MoLM-4B-K2模型的性能與大約13億參數的稠密模型相當，MoLM-4B-K4和MoLM-8B-K2模型的性能與大約27億參數的稠密模型相當。

由於其稀疏計算結構，MoLM處理每個字符的激活參數僅（等同於計算量）相當於同等性能稠密模型的約25%。因此，它減少了50%的延遲，同時具有較低的記憶體使用峰值，並在GPU記憶體完全佔用時將吞吐量提高了2倍。

通過增加模組學習新語言

在本節中，我們測試了模型學習新語言的能力。主要研究兩種實驗設置：連續聯合預訓練（continual joint pre-training）和連續終身預訓練（continual lifelong pre-training）。

它們的區別在於是否有英文文字的存在。對於這兩種設置，我們通過在CC-100語料庫上進行語言模型任務，不斷地對ModuleFormer和GPT-Neo進行預訓練。為了評估質量，我們採用了由XGLM和mGPT引入的0-shot方法的mLAMA基準測試。

持續聯合預訓練：在這部分中，我們對聯合訓練的模型進行持續預訓練。具體而言，我們混合了英語和一種新語言來構建一個新的訓練語料庫，並保持嵌入層可訓練。聯合訓練[Caruana, 1997]是一種眾所周知的多工學習方法，展示了對舊任務和新任務的熟練掌握。然而，它經常在不同任務之間產生負面干擾。

表4顯示了持續訓練模型獲得的結果。表格揭示了以下發現：

1）我們觀察到稀疏模型在Fully Tuned的情況下經歷較少干擾，最終得到了最好的的性能；

2）ModuleFormer通過增加模組（Insert New Expert）的能力，比之前的LoRA方法展示出了更好的少量參數（Parameter Efficient）調優的能力。這些結果表明，稀疏架構帶來了更強的抗遺忘能力。

持續終身預訓練：對於這個實驗設定，模型僅在新語言文字上進行訓練。Abraham和Robins [2005] 提出了穩定性-可塑性困境，這解釋了模型面臨的一個困難挑戰：

1）模型應具有較高的可塑性以學習新語言，2）模型必須具有出色的穩定性，考慮到在眾多的訓練迭代中不會接觸到任何英語標記。

表5顯示了LoRA基準和我們的方法在不同的路由正則化損失權重下的結果。我們的ModuleFormer藉助路由正則化損失表現出了強大的平衡穩定性和可塑性的能力。

當我們通過增加損失權重來限制新專家的使用時，模型獲得了穩定性，但可塑性下降。相比之下，使用LoRA對GPT-Neo進行微調在穩定性和可塑性方面都落後。

相比於1.33億可訓練參數的高秩LoRA，低秩LoRA（減少訓練參數到2400萬）和基本正則化都無法改善穩定性。

微調和壓縮模型

在本節中，我們展示了ModuleFormer中的模組可以被快速移除，以創建一個在尺寸上更小但性能不受損的任務專用模型。

我們首先從GitHub-code-clean資料集中創建了一個包含150億個字符的子集，該子集只包含Python程式碼。然後，我們使用負載集中損失函數（權重為0.001）對MoLM-4B-K2模型在該資料集上進行精調。

在精調之後，我們在從精調資料集中隨機抽樣的小型評估集上，計算每個專家的激活頻率，然後通過將每層除以層內最大頻率來進行歸一化。之後，我們設定一個閾值τ，並修剪了所有歸一化頻率低於該閾值的模組。

我們在HumanEval資料集上測試了我們修剪後的MoLM-4B-K2模型。

圖2a說明了pass@k指標與剩餘參數比例之間的相關性。圖2b展示了剩餘參數比例與閾值之間的關聯。我們觀察到：

1）修剪不必要的模組對結果影響不大。我們可以修剪40%至50%的參數而不犧牲性能。相反，適當的修剪（33%）使精調後的模型在任務上表現更好。

2）模組分佈存在顯著差異，大約有一半的模組的激活頻率低於最常使用的專家的0.3%。這個結果顯示了負載集中損失函數的有效性。

總結

在這篇論文中，我們提出了一種新的模組化架構ModuleFormer，以及與之相關的模組操作方法。

ModuleFormer包括幾個新元件：新的Stickbreaking注意力機制、新的互資訊負載平衡損失函數用於預訓練，以及新的負載集中損失函數用於微調。

基於ModuleFormer，我們預訓練了一個新的語言模型MoLM。我們的實驗結果顯示了MoLM的相對於稠密LLM展現出了一些新的能力：

1）它在更低的延遲（50%）和更小的記憶體佔用下實現了與密集LLM相同的性能；從而提高了吞吐量超過2倍；

2）在對整個模型進行微調以適應新領域後，它對災難性遺忘的魯棒性較強，並且也可以輕鬆擴展以學習新的語言和知識；

3）它可以在下游任務上進行微調，以使一部分模組專注於任務，並且未被任務使用的模組可以被修剪而不影響性能。

論文地址：

https://arxiv.org/abs/2306.04640