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

作者：允中 2023-06-25 10:19:49

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

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

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

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

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

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

1）效率

因?yàn)镸oduleFormer只對(duì)每個(gè)輸入只激活一小部分模塊，因此可以以兩倍以上的吞吐量達(dá)到與常規(guī)語(yǔ)言模型相同的性能。

2）可擴(kuò)展性

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

3）模塊特異化和篩選

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

導(dǎo)言

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

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

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

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

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

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

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

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

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

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

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

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

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

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

詳解ModuleFormer模型

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

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

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

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

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

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

Stick-Breaking注意力模塊

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

Stick-Breaking自注意力的優(yōu)勢(shì)在于它會(huì)自動(dòng)關(guān)注最近的相關(guān)字符，而不需要引入額外的位置信息，比如position embedding和relative position bias。

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

需要注意的是，這個(gè)查詢概率使用了sigmoid激活函數(shù)，所以沒(méi)有歸一化。接下來(lái)通過(guò)stick-breaking過(guò)程來(lái)對(duì)查詢概率進(jìn)去歸一化：

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

ModuleFormer中的模塊控制

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

評(píng)估

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

基礎(chǔ)性能評(píng)估

團(tuán)隊(duì)使用Language Model Evaluation Harness來(lái)評(píng)估零樣本、少樣本和語(yǔ)言建模任務(wù)中的語(yǔ)言模型。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

在精調(diào)之后，我們?cè)趶木{(diào)數(shù)據(jù)集中隨機(jī)抽樣的小型評(píng)估集上，計(jì)算每個(gè)專家的激活頻率，然后通過(guò)將每層除以層內(nèi)最大頻率來(lái)進(jìn)行歸一化。之后，我們?cè)O(shè)定一個(gè)閾值τ，并修剪了所有歸一化頻率低于該閾值的模塊。