模塊化大模型來了！IBM公開WastonX核心架構(gòu)技術(shù)細節(jié)

允中 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

大型語言模型（LLMs）的性能非常強大，但是現(xiàn)有的模型訓(xùn)練和部署成本都很高。而且在不忘記先前知識的前提，擴展它們?nèi)W(xué)習(xí)新的知識也很困難。也很難針對特定的任務(wù)去提取出輕量化的模型。

最近，來自MIT-IBM Waston AI Lab、清華大學(xué)、Mila的研究人員聯(lián)合提出了一種新的神經(jīng)網(wǎng)絡(luò)架構(gòu)ModuleFormer，利用模塊化來大幅提高大型語言模型的效率和靈活性。

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

與之前基于SMoE的模塊化語言模型不同，ModuleFormer可以通過其新的負載平衡（load balance）和負載集中（load concentration）損失函數(shù)從未經(jīng)篩選的數(shù)據(jù)中誘導(dǎo)出模塊化。

在實驗中，團隊發(fā)現(xiàn)模塊化架構(gòu)使得大型預(yù)訓(xùn)練語言模型具備了三個重要的能力：

1）效率

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

2）可擴展性

實驗表明，由于微調(diào)階段ModuleFormer只需要更新一部分的模塊，因此比常規(guī)語言模型更不容易發(fā)生災(zāi)難性遺忘，并且可以輕松通過新的模塊擴展以學(xué)習(xí)訓(xùn)練數(shù)據(jù)中不包含的新知識。

3）模塊特異化和篩選

在微調(diào)階段ModuleFormer，新提出的負載集中損失函數(shù)可以自動篩選一部分模塊，讓它們專注于目標(biāo)任務(wù)，而與任務(wù)無關(guān)的模塊可以被直接拋棄掉以實現(xiàn)輕量化部署。

導(dǎo)言

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

大多數(shù)LLM（例如Llama，F(xiàn)alcon）在推理和訓(xùn)練過程中都使用了它們的全部參數(shù)，我們稱這些模型為密集模型。

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

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

隨著模型規(guī)模的增長，為了領(lǐng)域適應(yīng)或持續(xù)學(xué)習(xí)而對整個模型進行微調(diào)變得代價高昂且計算資源受限，這使得那些計算預(yù)算較小的用戶無法實施。同時，更新所有參數(shù)也使得模型容易遭受災(zāi)難性遺忘（catastrophic forgetting）。

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

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

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

• 模型可以在輸入或任務(wù)上激活一組模塊條件，從而比密集激活整個模型需要更少的計算量；
• 在給定領(lǐng)域或任務(wù)的情況下，可以組裝一組與領(lǐng)域/任務(wù)相關(guān)的模塊，形成一個新的輕量級模型；
• 模型可以輕松添加新的模塊進行領(lǐng)域適應(yīng)或持續(xù)學(xué)習(xí)；
• 模型可能更不容易發(fā)生災(zāi)難性遺忘，因為只有與輸入相關(guān)的模塊在模型微調(diào)期間進行更新。

這篇論文提出了一種新的模塊化架構(gòu)，ModuleFormer（圖1a），以及在其中進行模塊操作的方法。ModuleFormer每層包含一個組前饋神經(jīng)網(wǎng)絡(luò)（MLP）專家模塊和一組新提出的Stickbreaking注意力專家模塊。

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

為了實現(xiàn)模塊修建，文章引入了一種新的負載集中損失函數(shù)，用于在微調(diào)的過程中自動選擇最適合執(zhí)行給定任務(wù)的模塊，同時通過微調(diào)進一步增強了這些模塊執(zhí)行該任務(wù)的能力。

△圖1

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

• 由于稀疏的模塊激活機制（圖1a）它在更低的延遲（50％）和更小的內(nèi)存占用下實現(xiàn)了與密集LLM相同的性能，因此ModuleFormer可以實現(xiàn)密集模型兩倍的吞吐量。
• 在對新領(lǐng)域進行微調(diào)后，它收到災(zāi)難性遺忘的影響較低，并且也可以輕松地通過添加新模塊來學(xué)習(xí)新語言和知識。（圖1b）。
• 它可以在下游任務(wù)上進行微調(diào)，將一部分模塊特異化成為處理該任務(wù)專用的模塊，而未使用的模塊可以被扔掉而不會犧牲模型在該任務(wù)上的性能（圖1c）。

詳解ModuleFormer模型

Sparse Mixture of Experts（SMoE）的基本結(jié)構(gòu)

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

每個SMoE層包含一組用于處理輸入產(chǎn)生輸出的專家模塊m_1, m_2, …, m_n，以及一個用于挑選專家的路由函數(shù)g。在本篇文章中，路由函數(shù)由一個一層的神經(jīng)網(wǎng)絡(luò)來建模：

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

其中不再topk內(nèi)的專家模塊不會被計算，同時對應(yīng)的概率g也會被用0替代。

在ModuleFormer中，作者使用了兩種不同的專家模塊來分別構(gòu)建傳統(tǒng)transformer解碼器中的自注意力層（self attention）和前饋層（MLP）。其中前饋層的專家模塊和常見的SMoE專家模塊一致，是一個單一隱藏層的全連通神經(jīng)網(wǎng)絡(luò)。而自注意力層的專家模塊則是一個新提出的stick-breaking注意力模塊。

Stick-Breaking注意力模塊

Stick-Breaking自注意力是設(shè)計用于取代Transformer解碼器的自注意力層，使用狄利克雷過程中的Stick-Breaking（折棍子）過程對每個字符xt與之前的字符x<t之間的注意力分布進行建模，而不是標(biāo)準(zhǔn)自注意力層中的softmax函數(shù)。

Stick-Breaking自注意力的優(yōu)勢在于它會自動關(guān)注最近的相關(guān)字符，而不需要引入額外的位置信息，比如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 是查詢投影矩陣。對于所有之前的步驟和當(dāng)前步驟 i ≤ t，計算時間步 i 的鍵與時間步 t 的查詢匹配的概率:

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

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

ModuleFormer中的模塊控制

預(yù)訓(xùn)練中的負載均衡

為了避免SMoE反復(fù)使用相同的模塊并浪費其他模塊的額外容量，一般采用負載平衡損失函數(shù)來調(diào)節(jié)每個專家的使用頻率。與之前的SMoE模型 不同，團隊希望最大化輸入字符和模塊之間的互信息（MI）：

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

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

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

微調(diào)中的負載集中

盡管團隊希望在預(yù)訓(xùn)練期間最大限度地利用每個專家的能力，但在微調(diào)期間希望將少量的模塊專注于下游任務(wù)。這樣可以移除未使用的模塊并減少微調(diào)后模型的參數(shù)數(shù)量。為了將負載集中在較少的模塊上，團隊引入了一個新的負載集中損失函數(shù)來最小化模塊的邊際熵：

這樣可以鼓勵模型使用更少的模塊來處理下游任務(wù)。在微調(diào)后，可以計算在訓(xùn)練或驗證集上使用的模塊頻率f_m。f_m代表了模塊m對于這個任務(wù)的重要性，可以通過移除f_m小于某個特定閾值的專家來輕松實現(xiàn)模型剪枝。

用新的模塊來學(xué)習(xí)新的知識

對于模塊化模型來說，插入新模塊是一種直接且參數(shù)高效的方法，可以在不對整個模型進行微調(diào)的情況下學(xué)習(xí)新知識。當(dāng)向每一層插入N_new個隨機初始化的模塊時，還需要擴展路由器（方程2中的A）中的模塊嵌入層A，使其包含一個形狀為（N_new，D_rtr）的新矩陣A’。因此，

新的路由函數(shù)可以寫成：

由于在微調(diào)期間其他的模塊參數(shù)被凍結(jié)，因此使用新模塊進行持續(xù)學(xué)習(xí)可以在很大程度上避免災(zāi)難性遺忘問題。

然而，災(zāi)難性遺忘仍然可能影響路由函數(shù)。當(dāng)新模塊在一個新領(lǐng)域進行訓(xùn)練時，如果路由函數(shù)錯誤地將來自舊領(lǐng)域的輸入路由到新專家，模型可能會遭受災(zāi)難性遺忘。

為了避免這種情況，團隊對路由函數(shù)進行了正則化以避免災(zāi)難性遺忘，并提出了兩種訓(xùn)練策略：

1）全面微調(diào)路由，公式9中A和B使用預(yù)訓(xùn)練參數(shù)進行初始化，而A’則是隨機初始化的。這個策略是為了訓(xùn)練數(shù)據(jù)中同時包含新舊數(shù)據(jù)的情況設(shè)計。

2）只訓(xùn)練A’，這個策略是為了連續(xù)學(xué)習(xí)（lifelong learning）的情況而設(shè)計的，不使用以前訓(xùn)練過的數(shù)據(jù)。由于這種情況可能導(dǎo)致新的模塊使用頻率過高，從而帶來災(zāi)難性遺忘。團隊引入了正則項來限制A’的范數(shù)：

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

評估

基于ModuleFormer，研究者在Pile數(shù)據(jù)集上預(yù)訓(xùn)練了三個不同體積和計算量的ModuleFormer Language Model（MoLM）語言模型：

基礎(chǔ)性能評估

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

對于零樣本和少樣本任務(wù)，目標(biāo)是在給定上下文的基礎(chǔ)上從一組給定選項中選擇最合適的完成部分。最終選擇在給定上下文下具有最高可能性的完成部分。

對于語言建模，在Wikitext數(shù)據(jù)集上進行測試。目標(biāo)是最小化下一個標(biāo)記預(yù)測的困惑度。

對于代碼生成，在HumanEval數(shù)據(jù)集上評估模型。HumanEval包含164個手寫的Python編程問題。模型需要根據(jù)任務(wù)描述提示完成一個函數(shù)，以便能夠通過所有提供的測試案例。

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

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

由于其稀疏計算結(jié)構(gòu)，MoLM處理每個字符的激活參數(shù)僅（等同于計算量）相當(dāng)于同等性能稠密模型的約25%。因此，它減少了50%的延遲，同時具有較低的內(nèi)存使用峰值，并在GPU內(nèi)存完全占用時將吞吐量提高了2倍。

通過增加模塊學(xué)習(xí)新語言

在本節(jié)中，我們測試了模型學(xué)習(xí)新語言的能力。主要研究兩種實驗設(shè)置：連續(xù)聯(lián)合預(yù)訓(xùn)練（continual joint pre-training）和連續(xù)終身預(yù)訓(xùn)練（continual lifelong pre-training）。

它們的區(qū)別在于是否有英文文本的存在。對于這兩種設(shè)置，我們通過在CC-100語料庫上進行語言模型任務(wù)，不斷地對ModuleFormer和GPT-Neo進行預(yù)訓(xùn)練。為了評估質(zhì)量，我們采用了由XGLM和mGPT引入的0-shot方法的mLAMA基準(zhǔn)測試。

持續(xù)聯(lián)合預(yù)訓(xùn)練：在這部分中，我們對聯(lián)合訓(xùn)練的模型進行持續(xù)預(yù)訓(xùn)練。具體而言，我們混合了英語和一種新語言來構(gòu)建一個新的訓(xùn)練語料庫，并保持嵌入層可訓(xùn)練。聯(lián)合訓(xùn)練[Caruana, 1997]是一種眾所周知的多任務(wù)學(xué)習(xí)方法，展示了對舊任務(wù)和新任務(wù)的熟練掌握。然而，它經(jīng)常在不同任務(wù)之間產(chǎn)生負面干擾。

表4顯示了持續(xù)訓(xùn)練模型獲得的結(jié)果。表格揭示了以下發(fā)現(xiàn)：

1）我們觀察到稀疏模型在Fully Tuned的情況下經(jīng)歷較少干擾，最終得到了最好的的性能；
2）ModuleFormer通過增加模塊（Insert New Expert）的能力，比之前的LoRA方法展示出了更好的少量參數(shù)（Parameter Efficient）調(diào)優(yōu)的能力。這些結(jié)果表明，稀疏架構(gòu)帶來了更強的抗遺忘能力。

持續(xù)終身預(yù)訓(xùn)練：對于這個實驗設(shè)定，模型僅在新語言文本上進行訓(xùn)練。Abraham和Robins [2005] 提出了穩(wěn)定性-可塑性困境，這解釋了模型面臨的一個困難挑戰(zhàn)：
1）模型應(yīng)具有較高的可塑性以學(xué)習(xí)新語言，
2）模型必須具有出色的穩(wěn)定性，考慮到在眾多的訓(xùn)練迭代中不會接觸到任何英語標(biāo)記。

表5顯示了LoRA基準(zhǔn)和我們的方法在不同的路由正則化損失權(quán)重下的結(jié)果。我們的ModuleFormer借助路由正則化損失表現(xiàn)出了強大的平衡穩(wěn)定性和可塑性的能力。

當(dāng)我們通過增加損失權(quán)重來限制新專家的使用時，模型獲得了穩(wěn)定性，但可塑性下降。相比之下，使用LoRA對GPT-Neo進行微調(diào)在穩(wěn)定性和可塑性方面都落后。

相比于1.33億可訓(xùn)練參數(shù)的高秩LoRA，低秩LoRA（減少訓(xùn)練參數(shù)到2400萬）和基本正則化都無法改善穩(wěn)定性。

微調(diào)和壓縮模型

在本節(jié)中，我們展示了ModuleFormer中的模塊可以被快速移除，以創(chuàng)建一個在尺寸上更小但性能不受損的任務(wù)專用模型。

我們首先從GitHub-code-clean數(shù)據(jù)集中創(chuàng)建了一個包含150億個字符的子集，該子集只包含Python代碼。然后，我們使用負載集中損失函數(shù)（權(quán)重為0.001）對MoLM-4B-K2模型在該數(shù)據(jù)集上進行精調(diào)。

在精調(diào)之后，我們在從精調(diào)數(shù)據(jù)集中隨機抽樣的小型評估集上，計算每個專家的激活頻率，然后通過將每層除以層內(nèi)最大頻率來進行歸一化。之后，我們設(shè)定一個閾值τ，并修剪了所有歸一化頻率低于該閾值的模塊。

我們在HumanEval數(shù)據(jù)集上測試了我們修剪后的MoLM-4B-K2模型。

圖2a說明了pass@k指標(biāo)與剩余參數(shù)比例之間的相關(guān)性。圖2b展示了剩余參數(shù)比例與閾值之間的關(guān)聯(lián)。我們觀察到：

1）修剪不必要的模塊對結(jié)果影響不大。我們可以修剪40%至50%的參數(shù)而不犧牲性能。相反，適當(dāng)?shù)男藜簦?3%）使精調(diào)后的模型在任務(wù)上表現(xiàn)更好。

2）模塊分布存在顯著差異，大約有一半的模塊的激活頻率低于最常使用的專家的0.3%。這個結(jié)果顯示了負載集中損失函數(shù)的有效性。

總結(jié)

在這篇論文中，我們提出了一種新的模塊化架構(gòu)ModuleFormer，以及與之相關(guān)的模塊操作方法。

ModuleFormer包括幾個新組件：新的Stickbreaking注意力機制、新的互信息負載平衡損失函數(shù)用于預(yù)訓(xùn)練，以及新的負載集中損失函數(shù)用于微調(diào)。

基于ModuleFormer，我們預(yù)訓(xùn)練了一個新的語言模型MoLM。我們的實驗結(jié)果顯示了MoLM的相對于稠密LLM展現(xiàn)出了一些新的能力：

1）它在更低的延遲（50%）和更小的內(nèi)存占用下實現(xiàn)了與密集LLM相同的性能；從而提高了吞吐量超過2倍；

2）在對整個模型進行微調(diào)以適應(yīng)新領(lǐng)域后，它對災(zāi)難性遺忘的魯棒性較強，并且也可以輕松擴展以學(xué)習(xí)新的語言和知識；

3）它可以在下游任務(wù)上進行微調(diào)，以使一部分模塊專注于任務(wù)，并且未被任務(wù)使用的模塊可以被修剪而不影響性能。

論文地址：
https://arxiv.org/abs/2306.04640