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

允中 發自 凹非寺
量子位 | 公眾號 QbitAI

大型語言模型（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<t之間的注意力分布進行建模，而不是標準自注意力層中的softmax函數。

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