大型語言模型（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
