本文介紹了一種名為“LLaMA-MoE”的方法，通過(guò)將現(xiàn)有的大型語(yǔ)言模型（LLMs）轉(zhuǎn)化為混合專家網(wǎng)絡(luò)（MoE），從而解決了訓(xùn)練MoE時(shí)遇到的數(shù)據(jù)饑餓和不穩(wěn)定性問(wèn)題。該方法基于著名的LLaMA-2 7B模型，并將其參數(shù)分為多個(gè)專家，然后對(duì)轉(zhuǎn)換后的MoE模型進(jìn)行持續(xù)預(yù)訓(xùn)練以進(jìn)一步提高性能。實(shí)驗(yàn)結(jié)果表明，在使用200B個(gè)標(biāo)記進(jìn)行訓(xùn)練后，LLaMA-MoE-3.5B模型在激活參數(shù)相似的情況下顯著優(yōu)于密集模型。

下載地址和項(xiàng)目代碼見文末

圖1：構(gòu)建LLaMA-MoE模型的主要框架。(a)LLaMA中的原始ffn被分為不同的專家。(b)在轉(zhuǎn)換后的LLaMA-MoE中，隱藏狀態(tài)是由部分選擇的專家而不是所有的專家進(jìn)行處理的。

專家構(gòu)建

1. 概述

專家構(gòu)建是將原始Feed-Forward Networks (FFNs)的參數(shù)劃分為多個(gè)專家。傳統(tǒng)的FFN層包含大量參數(shù)，這些參數(shù)可以被分解為多個(gè)子集，每個(gè)子集稱為一個(gè)專家。通過(guò)這種方法，可以減少每次計(jì)算所需的激活參數(shù)數(shù)量，從而在保證模型性能的同時(shí)顯著降低計(jì)算成本。

2. 參數(shù)劃分方法

在構(gòu)建專家時(shí)，常用的兩種方法是獨(dú)立神經(jīng)元?jiǎng)澐趾凸蚕砩窠?jīng)元?jiǎng)澐帧?/p>

• 獨(dú)立神經(jīng)元?jiǎng)澐郑哼@種方法將FFN中的神經(jīng)元均勻地劃分為多個(gè)子集，每個(gè)子集獨(dú)立組成一個(gè)專家。例如，可以通過(guò)隨機(jī)劃分或基于聚類的方法來(lái)實(shí)現(xiàn)這種劃分。隨機(jī)劃分是將所有神經(jīng)元隨機(jī)分配給不同的專家，而聚類方法則基于神經(jīng)元的特征將其分配給不同的專家。
• 共享神經(jīng)元?jiǎng)澐郑号c獨(dú)立神經(jīng)元?jiǎng)澐植煌蚕砩窠?jīng)元?jiǎng)澐衷试S多個(gè)專家共享部分神經(jīng)元。這種方法可以通過(guò)評(píng)估神經(jīng)元的重要性來(lái)決定哪些神經(jīng)元需要共享。共享神經(jīng)元的目標(biāo)是保留模型的表示能力，同時(shí)減少計(jì)算資源的消耗。

3. 實(shí)踐案例

在LLaMA-MoE模型的構(gòu)建過(guò)程中，研究人員嘗試了多種參數(shù)劃分方法，最終發(fā)現(xiàn)隨機(jī)劃分方法（IndependentRandom）在保持模型性能方面效果最佳。這種方法簡(jiǎn)單而高效，有助于平衡不同專家之間的負(fù)載，避免某些專家過(guò)度頻繁使用而其他專家很少被激活的問(wèn)題。

持續(xù)預(yù)訓(xùn)練

1. 必要性

由于轉(zhuǎn)換后的MoE模型在結(jié)構(gòu)上與原始的密集模型（dense model）有所不同，直接使用轉(zhuǎn)換后的模型可能會(huì)導(dǎo)致性能下降。為了恢復(fù)和提升模型的語(yǔ)言建模能力，必須對(duì)轉(zhuǎn)換后的MoE模型進(jìn)行持續(xù)預(yù)訓(xùn)練。

2. 預(yù)訓(xùn)練策略

在持續(xù)預(yù)訓(xùn)練階段，研究人員使用了兩種主要的數(shù)據(jù)采樣策略：靜態(tài)數(shù)據(jù)采樣和動(dòng)態(tài)數(shù)據(jù)采樣。

• 靜態(tài)數(shù)據(jù)采樣：使用固定的采樣權(quán)重，從預(yù)定義的數(shù)據(jù)集中提取訓(xùn)練數(shù)據(jù)。這種方法簡(jiǎn)單直接，但可能無(wú)法充分適應(yīng)不同訓(xùn)練階段的需求。
• 動(dòng)態(tài)數(shù)據(jù)采樣：根據(jù)模型在訓(xùn)練過(guò)程中的表現(xiàn)動(dòng)態(tài)調(diào)整采樣權(quán)重，以更好地優(yōu)化模型性能。盡管這種方法可以提高訓(xùn)練效率，但也增加了計(jì)算復(fù)雜度。

3. 數(shù)據(jù)過(guò)濾

為了提高訓(xùn)練數(shù)據(jù)的質(zhì)量，研究人員在預(yù)訓(xùn)練前對(duì)數(shù)據(jù)進(jìn)行了過(guò)濾，去除了約50%的廣告內(nèi)容和約15%的非流利文本。這一策略有助于加快模型的收斂速度，提高訓(xùn)練效果。

增強(qiáng)推理能力

1. 激活部分模型參數(shù)

在處理具體任務(wù)時(shí)，MoE模型通過(guò)激活部分模型參數(shù)來(lái)提高推理能力。每個(gè)輸入token僅激活與其最相關(guān)的幾個(gè)專家，從而減少了不必要的計(jì)算。這種稀疏激活方式不僅提高了計(jì)算效率，還能在保持高性能的同時(shí)降低推理成本。

2. 性能提升

實(shí)驗(yàn)證明，經(jīng)過(guò)200B tokens的預(yù)訓(xùn)練后，LLaMA-MoE-3.5B模型在多個(gè)下游任務(wù)上顯著優(yōu)于具有相同激活參數(shù)的密集模型。這一結(jié)果表明，通過(guò)適當(dāng)?shù)膶＜覄澐趾统掷m(xù)預(yù)訓(xùn)練，MoE模型能夠在保持語(yǔ)言能力的同時(shí)顯著提升推理性能。

提高可解釋性

1. 部分參數(shù)激活

由于每次僅有部分參數(shù)被激活，MoE模型在決策過(guò)程中的激活路徑更加清晰。這使得研究人員能夠更容易地追蹤和解釋模型的行為，了解模型是如何處理和響應(yīng)不同輸入的。

2. 實(shí)例分析

在實(shí)驗(yàn)中，研究人員觀察到深層網(wǎng)絡(luò)層比淺層網(wǎng)絡(luò)層有更強(qiáng)的路由偏好，這意味著深層網(wǎng)絡(luò)層捕捉更多任務(wù)特定的特征，而淺層網(wǎng)絡(luò)層則更關(guān)注通用特征。這一發(fā)現(xiàn)有助于進(jìn)一步優(yōu)化專家劃分策略，提高模型的整體性能。

降低計(jì)算成本

MoE（Mixture-of-Experts）模型相比傳統(tǒng)的密集模型，通過(guò)只激活部分參數(shù)來(lái)處理輸入，可以顯著降低計(jì)算成本。傳統(tǒng)的密集模型在處理每一個(gè)輸入時(shí)都需要使用所有參數(shù)，這樣隨著模型容量的增加，計(jì)算成本也會(huì)急劇上升。而MoE模型則通過(guò)引入專家網(wǎng)絡(luò)和門控網(wǎng)絡(luò)，只激活一部分專家，從而降低了計(jì)算成本。

例如，LLaMA-MoE模型在構(gòu)建過(guò)程中，將原始LLaMA模型的FFN（Feed-Forward Network）分割成多個(gè)專家網(wǎng)絡(luò)。這種分割方法在維持模型性能的前提下，顯著減少了需要激活的參數(shù)量，進(jìn)而減少了計(jì)算開銷。通過(guò)訓(xùn)練200B tokens，LLaMA-MoE-3.5B模型在激活參數(shù)量相當(dāng)?shù)那闆r下，顯著優(yōu)于類似的密集模型。

工程應(yīng)用

MoE模型在需要高效推理的實(shí)際應(yīng)用中展現(xiàn)出巨大的潛力。例如，在實(shí)時(shí)翻譯和智能助手等場(chǎng)景中，計(jì)算成本的降低和推理效率的提高尤為重要。MoE模型可以根據(jù)輸入動(dòng)態(tài)選擇合適的專家，從而實(shí)現(xiàn)快速而準(zhǔn)確的推理。

以實(shí)時(shí)翻譯為例，傳統(tǒng)模型可能需要大量計(jì)算資源來(lái)處理復(fù)雜的語(yǔ)言轉(zhuǎn)換，而MoE模型則能夠通過(guò)激活少量專家，快速處理翻譯任務(wù)，降低延遲并提高響應(yīng)速度。同樣地，在智能助手中，MoE模型可以根據(jù)用戶的不同需求，動(dòng)態(tài)分配計(jì)算資源，從而提供更加個(gè)性化和高效的服務(wù)。

理論研究

在模型架構(gòu)設(shè)計(jì)和優(yōu)化方面，MoE模型提供了新的思路和方法。傳統(tǒng)模型在擴(kuò)展過(guò)程中面臨著計(jì)算成本急劇上升的挑戰(zhàn)，而MoE模型通過(guò)稀疏激活部分參數(shù)，為解決這一問(wèn)題提供了有效的途徑。

研究表明，通過(guò)將密集模型的FFN參數(shù)隨機(jī)分割成多個(gè)專家，并在每一層引入MoE模塊，可以在保持模型性能的同時(shí)，減少計(jì)算開銷。例如，獨(dú)立隨機(jī)拆分方法在實(shí)驗(yàn)中取得了最佳性能。與其他方法相比，該方法在專家和門控網(wǎng)絡(luò)同時(shí)訓(xùn)練時(shí)，可以減少偏差，快速恢復(fù)模型的語(yǔ)言能力。

非重疊隨機(jī)拆分法

非重疊隨機(jī)拆分法通過(guò)隨機(jī)拆分原始FFN的參數(shù)來(lái)構(gòu)建專家，這種方法在實(shí)踐中取得了顯著效果。具體而言，給定一個(gè)包含所有中間神經(jīng)元索引的集合U，通過(guò)將U隨機(jī)分割成等大小的子集，從而構(gòu)建出多個(gè)專家網(wǎng)絡(luò)。這種方法能夠在保持模型原有表示能力的基礎(chǔ)上，減少計(jì)算復(fù)雜度。

在LLaMA-MoE模型的構(gòu)建中，采用了非重疊隨機(jī)拆分法，將FFN層中的中間神經(jīng)元均勻分割成多個(gè)子集，每個(gè)子集對(duì)應(yīng)一個(gè)專家網(wǎng)絡(luò)。實(shí)驗(yàn)結(jié)果表明，該方法不僅在減少計(jì)算成本方面表現(xiàn)出色，還能夠在持續(xù)預(yù)訓(xùn)練階段快速恢復(fù)模型的語(yǔ)言能力。研究還發(fā)現(xiàn)，對(duì)專家輸出進(jìn)行重新縮放操作，可以顯著提升MoE模型的性能。

通過(guò)這些研究和實(shí)踐，MoE模型不僅在理論上提供了新的研究方向，還在實(shí)際應(yīng)用中展現(xiàn)出了顯著優(yōu)勢(shì)，為未來(lái)的大規(guī)模語(yǔ)言模型發(fā)展提供了重要參考。

共享神經(jīng)元方法

共享神經(jīng)元方法通過(guò)結(jié)構(gòu)化裁剪來(lái)保留模型的部分表示能力，這種方法主要分為內(nèi)部共享和外部共享兩種策略。

1. 內(nèi)部共享

內(nèi)部共享策略主要是通過(guò)對(duì)神經(jīng)元的重要性進(jìn)行排序，并根據(jù)排序結(jié)果選擇部分神經(jīng)元進(jìn)行共享。具體來(lái)說(shuō)，首先對(duì)每個(gè)神經(jīng)元進(jìn)行重要性評(píng)估，可以使用一階泰勒展開來(lái)度量每個(gè)神經(jīng)元對(duì)損失變化的影響。然后，根據(jù)這些重要性分?jǐn)?shù)，將最重要的神經(jīng)元在不同的專家間共享，而其余神經(jīng)元?jiǎng)t分配給特定的專家。這種方法可以在不顯著降低模型表示能力的情況下，實(shí)現(xiàn)有效的參數(shù)裁剪和共享。

2. 外部共享

外部共享策略則是在不同專家間直接共享部分神經(jīng)元，而不進(jìn)行重要性排序。這種方法通過(guò)預(yù)先定義的規(guī)則，將一些神經(jīng)元設(shè)定為共享神經(jīng)元，并將其余神經(jīng)元分配給特定專家。這種方法的優(yōu)點(diǎn)在于實(shí)現(xiàn)簡(jiǎn)單且計(jì)算開銷較小，但可能在某些情況下無(wú)法達(dá)到內(nèi)部共享策略所帶來(lái)的性能提升。

數(shù)據(jù)采樣權(quán)重

在訓(xùn)練過(guò)程中，數(shù)據(jù)采樣權(quán)重的選擇對(duì)模型的收斂速度和最終性能有重要影響。我們研究了靜態(tài)和動(dòng)態(tài)兩種數(shù)據(jù)采樣策略，以期獲得最快的收斂速度和最佳的性能提升。

1. 靜態(tài)采樣

靜態(tài)采樣策略是指在訓(xùn)練過(guò)程中使用固定的采樣權(quán)重，不隨時(shí)間變化。這種方法的優(yōu)點(diǎn)在于實(shí)現(xiàn)簡(jiǎn)單且計(jì)算開銷較小，但可能無(wú)法適應(yīng)數(shù)據(jù)分布的動(dòng)態(tài)變化。

2. 動(dòng)態(tài)采樣

動(dòng)態(tài)采樣策略則會(huì)在訓(xùn)練過(guò)程中不斷調(diào)整采樣權(quán)重，以適應(yīng)當(dāng)前模型的訓(xùn)練需求和數(shù)據(jù)分布變化。具體來(lái)說(shuō)，可以每隔一段時(shí)間（例如每2.5B tokens）調(diào)整一次采樣權(quán)重，根據(jù)當(dāng)前模型在不同數(shù)據(jù)域上的表現(xiàn)進(jìn)行調(diào)整。這種方法雖然計(jì)算開銷較大，但可以顯著提升模型的收斂速度和性能。

數(shù)據(jù)過(guò)濾

為了加快模型的收斂速度，我們對(duì)訓(xùn)練數(shù)據(jù)進(jìn)行了嚴(yán)格的質(zhì)量過(guò)濾。具體來(lái)說(shuō)，我們過(guò)濾掉了低質(zhì)量的文本數(shù)據(jù)，如廣告和不流暢的文本。

1. 廣告過(guò)濾

廣告通常包含大量冗余和無(wú)關(guān)信息，對(duì)模型的訓(xùn)練效果影響較大。我們通過(guò)特定的規(guī)則和算法，過(guò)濾掉了大約50%的廣告數(shù)據(jù)，從而提升了數(shù)據(jù)集的整體質(zhì)量。

2. 不流暢文本過(guò)濾

不流暢的文本通常表現(xiàn)為語(yǔ)法錯(cuò)誤、拼寫錯(cuò)誤或邏輯不連貫。我們使用自然語(yǔ)言處理技術(shù)，過(guò)濾掉了大約15%的不流暢文本數(shù)據(jù)，從而進(jìn)一步提升了模型的訓(xùn)練效率和效果。

實(shí)驗(yàn)設(shè)置

我們的實(shí)驗(yàn)在112個(gè)A100 (80G) GPU上進(jìn)行訓(xùn)練，最大學(xué)習(xí)率為2e-4。訓(xùn)練數(shù)據(jù)集采用了SlimPajama，該數(shù)據(jù)集經(jīng)過(guò)清洗和去重處理，包含627B tokens的數(shù)據(jù)。訓(xùn)練過(guò)程中，我們?cè)O(shè)置了全局批次大小為15M tokens，最大上下文長(zhǎng)度為4096。在經(jīng)過(guò)100步的熱身訓(xùn)練后，學(xué)習(xí)率逐步下降到2e-5，采用余弦調(diào)度策略。整個(gè)訓(xùn)練過(guò)程中，我們對(duì)每個(gè)模型進(jìn)行了13.6k步（約200B tokens）的訓(xùn)練。

通過(guò)以上方法和設(shè)置，我們成功構(gòu)建并訓(xùn)練了LLaMA-MoE模型，并在多項(xiàng)任務(wù)中顯著超越了同類模型。

實(shí)驗(yàn)結(jié)果

LLaMA-MoE-3.5B在多個(gè)下游任務(wù)上的表現(xiàn)顯著優(yōu)于其他具有相似激活參數(shù)的開源模型，如Sheared-LLaMA和Open-LLaMA-3B-v2。具體來(lái)說(shuō)，LLaMA-MoE-3.5B（4/16）在各種任務(wù)中的平均分?jǐn)?shù)超過(guò)了最具競(jìng)爭(zhēng)力的模型Sheared-LLaMA 1.3分。此外，LLaMA-MoE-3.0B與Open-LLaMA-3B-v2表現(xiàn)相當(dāng)。

在ARC-c和HellaSwag數(shù)據(jù)集上的表現(xiàn)顯示，隨著訓(xùn)練過(guò)程的推進(jìn)，模型的性能穩(wěn)步提升。盡管ARC-c的結(jié)果波動(dòng)較大，但HellaSwag提供了較為平滑的結(jié)果。訓(xùn)練損失方面，LLaMA-MoE-3.0B和LLaMA-MoE-3.5B分別收斂到1.95和1.90，這兩個(gè)模型激活的參數(shù)較少，因此損失較LLaMA-2 7B略高。

專家構(gòu)建方法對(duì)比

在實(shí)驗(yàn)中，我們比較了四種不同的專家構(gòu)建方法。結(jié)果顯示，非重疊隨機(jī)拆分法（IndependentRandom）表現(xiàn)最佳。這種方法在訓(xùn)練200B tokens后，表現(xiàn)出最佳的平均分?jǐn)?shù)。相比之下，共享神經(jīng)元構(gòu)建方法（SharingInter和SharingInner）在初始階段表現(xiàn)良好，但隨著訓(xùn)練的進(jìn)行，其性能顯著下降。

我們還進(jìn)行了專家輸出重新縮放的消融研究，結(jié)果表明，重新縮放操作顯著提高了MoE模型的性能。這表明，專家構(gòu)建方法對(duì)模型最終性能有著重要影響，而重新縮放操作則進(jìn)一步提升了專家的表現(xiàn)能力。

數(shù)據(jù)采樣策略

在數(shù)據(jù)采樣策略的比較中，靜態(tài)采樣權(quán)重策略（StaticSheared）在性能上優(yōu)于動(dòng)態(tài)采樣策略。盡管StaticSheared在訓(xùn)練損失上并不是最低的，但其在下游任務(wù)上的表現(xiàn)最佳。動(dòng)態(tài)采樣權(quán)重策略（DynamicUniform和DynamicLLaMA）在訓(xùn)練損失上波動(dòng)較大，顯示出不穩(wěn)定性。

在數(shù)據(jù)采樣權(quán)重的變化中，我們發(fā)現(xiàn)不同策略對(duì)不同領(lǐng)域的數(shù)據(jù)有不同的權(quán)重分配。靜態(tài)采樣策略的權(quán)重在整個(gè)訓(xùn)練過(guò)程中保持不變，而動(dòng)態(tài)采樣策略的權(quán)重則隨著訓(xùn)練的進(jìn)行逐漸變化。這表明，數(shù)據(jù)采樣策略的選擇對(duì)模型的訓(xùn)練效率和最終性能有著重要影響。

數(shù)據(jù)過(guò)濾策略

數(shù)據(jù)過(guò)濾策略在提高模型性能方面也起到了關(guān)鍵作用。通過(guò)過(guò)濾掉廣告和不流暢文本，訓(xùn)練損失顯著降低。具體而言，過(guò)濾掉廣告數(shù)據(jù)的方法在下游任務(wù)上的表現(xiàn)不如過(guò)濾不流暢文本的方法。這可能是由于廣告數(shù)據(jù)中的知識(shí)和信息較多，被過(guò)濾掉的數(shù)量較大，從而影響了模型的性能。

基于這些結(jié)果，我們最終選擇使用過(guò)濾掉不流暢文本的數(shù)據(jù)集進(jìn)行訓(xùn)練。盡管沒(méi)有引入新的數(shù)據(jù)集，但通過(guò)過(guò)濾部分低質(zhì)量數(shù)據(jù)，我們加快了模型的收斂速度，并提高了模型的整體表現(xiàn)。

地址：https://arxiv.org/pdf/2406.16554    

代碼：https://github.com/pjlab-sys4nlp/llama-moe

本文轉(zhuǎn)載自 AI論文解讀，作者：柏企