一、結(jié)論寫在前面

論文標(biāo)題：OLMoE: Open Mixture-of-Experts Language Models

論文鏈接：??https://arxiv.org/pdf/2409.02060??

Weights：??https://hf.co/allenai/OLMoE-1B-7B-0924??

Data：??https://hf.co/datasets/allenai/OLMoE-mix-0924??

Code：??https://github.com/allenai/OLMoE??

Logs：??https://wandb.ai/ai2-llm/olmoe/reports/OLMoE-1B-7B-0924--Vmlldzo4OTcyMjU3??

論文開源了OLMoE-1B-7B和OLMoE-1B-7B-INSTRUCT，包括模型、數(shù)據(jù)、代碼和日志。OLMOE-1B-7B擁有7B參數(shù)，但每個(gè)輸入token僅使用1B參數(shù)。論文在5T token上對其進(jìn)行預(yù)訓(xùn)練，并進(jìn)一步Adaptation以創(chuàng)建OLMoE-1B-7B-INSTRUCT。

論文的模型在所有具有相似活躍參數(shù)的可用模型中表現(xiàn)最佳，甚至超越了Llama2-13B-Chat和DeepSeekMoE-16B等更大模型。

論文分享了各種訓(xùn)練實(shí)驗(yàn)，并定義和分析了路由器飽和度、專家協(xié)同激活、領(lǐng)域和詞匯專業(yè)化。通過論文的完全開源發(fā)布，論文希望幫助該領(lǐng)域構(gòu)建更好的專家混合模型。論文對OLMoE的新迭代感到興奮，以縮小前沿模型與完全開源模型之間的差距。

圖1：開放式MoE和密集LMs的性能、成本和開放程度。模型名稱包含四舍五入的參數(shù)計(jì)數(shù)：MoE為model-active-total，密集LMs為model-total。#ckpts是可用的中間檢查點(diǎn)數(shù)量。論文突出顯示MMLU作為整體性能的總結(jié)；更多結(jié)果見S3。OLMOE-1B-7B在具有相似活躍參數(shù)計(jì)數(shù)的模型中表現(xiàn)最佳，是最開放的MoE。    

二、論文的簡單介紹

2.1 論文的背景

高性能的語言模型對許多學(xué)者和開源開發(fā)者來說是很難接觸，因?yàn)樗鼈兊臉?gòu)建和部署成本過高。改善成本性能權(quán)衡的一種方法是使用稀疏激活的專家混合(Mixture-of-Experts，MoEs)模型。MoEs在每一層都有多個(gè)專家，每次只激活其中的一個(gè)子集(見圖2)。這使得MoEs比具有相似總參數(shù)數(shù)量的密集模型更加高效，因?yàn)槊芗Ｐ托枰獮槊總€(gè)輸入激活所有參數(shù)。出于這個(gè)原因，行業(yè)前沿模型使用MoEs，包括Gemini-1.5和據(jù)報(bào)道的GPT-4。

然而，大多數(shù)MoE模型都是閉源的:雖然有些模型公開發(fā)布了模型權(quán)重，但它們提供的關(guān)于訓(xùn)練數(shù)據(jù)、代碼或配方的信息有限或沒有(見圖1)。雖然之前有一些努力使語言建模研究完全可訪問，但它們主要局限于密集LMs。這種情況出現(xiàn)的原因是，MoEs需要更多的開放性，因?yàn)樗鼈優(yōu)長Ms增加了復(fù)雜的新設(shè)計(jì)問題，例如使用多少總參數(shù)與活躍參數(shù)，是使用多個(gè)小專家還是少數(shù)大專家，專家是否應(yīng)該共享，以及使用什么路由算法。缺乏關(guān)于這些細(xì)節(jié)的開放資源和發(fā)現(xiàn)，阻礙了該領(lǐng)域構(gòu)建接近閉源前沿模型能力的高效開放MoEs。

為了解決這些問題，論文引入了OLMOE，一個(gè)完全開放的專家混合語言模型，在同等規(guī)模的模型中具有最先進(jìn)的性能。具體來說，論文預(yù)訓(xùn)練了OLMOE-1B-7B，總共訓(xùn)練了5.1T token，總參數(shù)為6.9B，其中只有1.3B參數(shù)為每個(gè)輸入token激活。這導(dǎo)致推理成本與使用約1B參數(shù)的密集模型(如OLMo 1B或TinyLlama 1B)相似，但需要更多GPU內(nèi)存來存儲其7B總參數(shù)。

論文的實(shí)驗(yàn)表明，MoEs的訓(xùn)練速度比具有相同活躍參數(shù)數(shù)量的密集LMs快約2倍。在圖1中，論文展示了OLMOE-1B-7B顯著優(yōu)于所有開放的10億參數(shù)模型，并展現(xiàn)出與推理成本和內(nèi)存存儲顯著更高的密集模型相當(dāng)?shù)男阅?例如，在MMLU評分上與Llama2-13B相似，后者的成本約為前者的10倍)。通過指令和偏好微調(diào)，論文創(chuàng)建了OLMOE-1B-7B-INSTRUCT，論文發(fā)現(xiàn)它在常見基準(zhǔn)測試(MMLU、GSM8k、HumanEval等)上超過了各種更大的指令模型，包括Llama2-13B-Chat、OLMo-7B-Instruct (0724)和DeepSeekMoE-16B。

論文全面的對照實(shí)驗(yàn)突出了MoEs(見表1)和LMs的關(guān)鍵設(shè)計(jì)選擇。使MoEs性能出色的一個(gè)關(guān)鍵設(shè)計(jì)決策是使用細(xì)粒度路由和細(xì)粒度專家:論文在每層使用64個(gè)小專家，其中8個(gè)被激活。路由算法的選擇也很重要:論文發(fā)現(xiàn)無丟棄的基于token的路由優(yōu)于基于專家的路由。論文的發(fā)現(xiàn)還包括一些挑戰(zhàn)先前工作的結(jié)果，例如共享專家的無效性，以及將預(yù)訓(xùn)練的密集LM稀疏地升級為MoE的有限益處，除非在小規(guī)模計(jì)算預(yù)算下。最后，論文分析了OLMOE-1B-7B的路由行為，發(fā)現(xiàn)路由在預(yù)訓(xùn)練早期就飽和了，專家很少被共同激活，并且專家表現(xiàn)出領(lǐng)域和詞匯專門化。    

圖2：密集型語言模型（LMs）與OLMoE等MoE模型的架構(gòu)對比。圖中省略了一些細(xì)節(jié)，例如，OLMOE-1B-7B還使用了QK-Norm。

表1: MoE的關(guān)鍵設(shè)計(jì)選擇及基于論文實(shí)驗(yàn)的OLMOE-1B-7B設(shè)置    

表2：OLMOE-1B-7B預(yù)訓(xùn)練數(shù)據(jù)的組成。StarCoder、peS2o和Wikipedia部分來自Dolma 1.7 。

表3：OLMOE-13.7的Adaptation訓(xùn)練數(shù)據(jù)。

2.2 預(yù)訓(xùn)練和Adaptation(Pretraining and Adaptation)

預(yù)訓(xùn)練架構(gòu) OLMOE是一個(gè)僅有解碼器的語言模型，由NL個(gè)transformer層組成。在像OLMo這樣的密集模型中的前饋網(wǎng)絡(luò)(FFN)被替換為一個(gè)MoE模塊，該模塊由NE個(gè)較小的FFN模塊組成，這些模塊被稱為專家。對于每個(gè)處理的輸入token x，會激活其中的k個(gè)專家子集(另見圖2):

其中 r，稱為路由器，是一個(gè)從輸入logits映射到所選 k 個(gè)專家的學(xué)習(xí)線性層。softmax應(yīng)用于路由器輸出，以計(jì)算所有 N_E個(gè)專家的路由概率。每個(gè)選定的專家 E_i 處理輸入 x，其輸出隨后與其相應(yīng)的路由概率相乘。然后將所有選定的Top- k 專家的結(jié)果相加以構(gòu)成MoE模塊的輸出，該模塊是模型 N_L層中的一層。設(shè)計(jì)MoE模型的關(guān)鍵決策包括確定激活和總參數(shù)的數(shù)量、專家的設(shè)計(jì)（例如，粒度，是否包含共享專家）以及路由算法的選擇。此外，訓(xùn)練MoE模型可能涉及從密集模型（稀疏升級）初始化并更改訓(xùn)練目標(biāo)，例如包括輔助負(fù)載均衡和路由z-losses；表1顯示了論文的最終決策。    

總結(jié)來說，論文使用了總共6.9B參數(shù)中的1.3B活躍參數(shù)，每層激活了8個(gè)專家，總共64個(gè)專家。論文采用了無丟棄的token選擇路由[58]：對于每個(gè)輸入token，學(xué)習(xí)到的路由網(wǎng)絡(luò)決定由8個(gè)專家來處理。論文從零開始訓(xùn)練OLMOE-1B-7B，使用了兩個(gè)輔助損失：負(fù)載均衡損失( L_LB) 和路由z-損失( L_RZ) 。論文將它們分別乘以各自的損失權(quán)重alpha和beta，并與交叉熵?fù)p失(L_CE )線性相加，得到最終的訓(xùn)練損失：

OLMOE-1B-7B的完整預(yù)訓(xùn)練配置在附錄B中。

預(yù)訓(xùn)練數(shù)據(jù) 論文使用了來自DCLM [89]和Dolma 1.7 [161]的混合數(shù)據(jù)，包括以下內(nèi)容：

(1) 經(jīng)過質(zhì)量過濾的Common Crawl子集，稱為DCLM-Baseline;

(2) 在DCLM和Dolma 1. 7中都使用的StarCoder、Algebraic Stack和arXiv;

(3) 來自Dolma 1.7的peS2o和Wikipedia。論文將預(yù)訓(xùn)練數(shù)據(jù)集稱為OLMOE-MIX。

對于上述所有來源，論文應(yīng)用了一個(gè)過濾器，移除所有包含32個(gè)或更多重復(fù)n-gram的文檔，其中n-gram是任意1到13個(gè)token的跨度。對于StarCoder子集，論文還移除了任何來自GitHub上少于2顆星的倉庫的文檔，或者其最頻繁詞占據(jù)文檔超過30%的內(nèi)容，或者其前兩個(gè)最頻繁詞占據(jù)文檔超過50%的內(nèi)容。

論文在每個(gè)epoch開始時(shí)隨機(jī)打亂所有樣本，并總共訓(xùn)練5.133T個(gè)token（根據(jù)Muennighoff等人[120]的方法，訓(xùn)練1.3個(gè)epoch）。在論文的退火階段（最后100B個(gè)token），論文首先重新打亂整個(gè)數(shù)據(jù)集，然后按照先前的工作64， 89，將學(xué)習(xí)率線性衰減到0。論文的預(yù)訓(xùn)練數(shù)據(jù)統(tǒng)計(jì)如表2所示。

Adaptation（Adaptation） 論文通過遵循標(biāo)準(zhǔn)的Adaptation配方創(chuàng)建了OLMoE-1B-7B-INSTRUCT，該配方分為指令微調(diào)和偏好微調(diào)，基于先前的開源模型。在論文的指令微調(diào)數(shù)據(jù)集中，論文增加了更多的代碼和數(shù)學(xué)數(shù)據(jù)，以提升下游編碼和數(shù)學(xué)應(yīng)用的性能。其他模型，如GPT-4和Llama 3，在預(yù)訓(xùn)練期間包含了來自GSM8k或MATH等數(shù)學(xué)數(shù)據(jù)集的樣本。論文還包含了No Robots和Daring Anteater的一個(gè)子集，因?yàn)樗鼈冑|(zhì)量高且增加了多樣性，這是成功Adaptation的兩個(gè)關(guān)鍵因素。論文在表3中描述了論文的Adaptation數(shù)據(jù)集。

2.3 結(jié)果

論文的評估程序包括三個(gè)部分：預(yù)訓(xùn)練期間、預(yù)訓(xùn)練后和Adaptation后(After adaptation)。    

圖3: OLMOE-1B-7B和當(dāng)前最佳OLMo模型在預(yù)訓(xùn)練期間的評估。OLMOE-1B-7B與OLMo模型在MoE架構(gòu)、幾個(gè)訓(xùn)練超參數(shù)和訓(xùn)練數(shù)據(jù)集方面有所不同。更多結(jié)果、日志和配置:??https://wandb.ai/ai2-llm/olmoe/reports/Plot-OLMoE-1B-7B-vs-OLMo-7B-vs-OLMo-1B--Vmlldzo4OTcyMjEz??

預(yù)訓(xùn)練期間(During pretraining) 在圖3中，論文對比了OLMOE-1B-7B在預(yù)訓(xùn)練期間與當(dāng)前最佳OLMo模型在常用下游任務(wù)上的性能表現(xiàn)。論文發(fā)現(xiàn)，在所有任務(wù)中，OLMOE-1B-7B都以更少的計(jì)算量(FLOPs)達(dá)到了比密集OLMo模型更好的性能。

盡管OLMOE-1B-7B在訓(xùn)練中使用的FLOPs不到一半，且只使用了1B的活躍參數(shù)，但在訓(xùn)練結(jié)束時(shí)，它的表現(xiàn)與OLMo-7B相當(dāng)或更優(yōu)。這可能是由于論文對OLMo設(shè)置進(jìn)行的數(shù)據(jù)集和建模變更，包括MoE相關(guān)的改變、穩(wěn)定性和性能改進(jìn)，這些在附錄B中有詳細(xì)說明。    

表4: 預(yù)訓(xùn)練后的OLMOE-1B-7B與更大的MoEs和密集LMs的比較。論文將其與在活躍參數(shù)(1B，近似速度和成本)或總參數(shù)(7B，近似內(nèi)存需求)上與OLMOE-1B-7B相似的密集LMs進(jìn)行比較。模型名稱包含四舍五入的參數(shù)計(jì)數(shù):MoEs的格式為model-active-total，密集LMs的格式為model-total(這導(dǎo)致與官方名稱有些差異，例如，雖然被稱為"Gemma2-2B"，但它實(shí)際上有26億活躍和總參數(shù)[175])。Chall. = Challenge(挑戰(zhàn))。論文自己使用5個(gè)少樣本示例進(jìn)行所有評估

預(yù)訓(xùn)練后(After pretraining) 在表4中，論文對OLMoE-1B-7B在常見的下游任務(wù)上進(jìn)行了基準(zhǔn)測試。論文發(fā)現(xiàn)，在使用的活躍參數(shù)少于2B的模型中，OLMoE-1B-7B表現(xiàn)最佳，使其成為許多語言模型使用場景中最經(jīng)濟(jì)的選擇。對預(yù)算更大的情況，Qwenl.5 3B-14B的性能更強(qiáng)，但其活躍參數(shù)和總參數(shù)數(shù)量是OLMoE-1B-7B的兩倍以上。論文發(fā)現(xiàn)，盡管每次前向傳遞所需的計(jì)算量減少了6-7倍，OLMoE-1B-7B仍優(yōu)于一些具有7B參數(shù)的密集模型，如Llama2-7B ，但在其他模型如Llama3.1-8B面前稍顯不足。圖1比較了OLMoE-1B-7B與其他語言模型在MMLU性能與活躍參數(shù)（作為模型成本的代理）上的表現(xiàn)。OLMoE-1B-7B在其成本范圍內(nèi)處于領(lǐng)先地位。

Adaptation后(After adaptation) 在表5中，論文基準(zhǔn)測試了OLMoE-1B-7B的指令（SFT）和偏好（DPO）調(diào)優(yōu)。SFT在所有測量的任務(wù)上提升了論文的模型性能。論文在GSM8k上觀察到>10×的增益，這可能歸因于論文在預(yù)訓(xùn)練期間加入了額外的數(shù)學(xué)數(shù)據(jù)，以彌補(bǔ)相對較少的數(shù)學(xué)數(shù)據(jù)（S2）。DPO在大多數(shù)任務(wù)上都有幫助，尤其是在AlpacaEval上，這與先前的工作[186，75，121]中的發(fā)現(xiàn)一致。    

論文的DPO模型，論文稱之為OLMoE-1B-7B-INSTRUCT在所有基準(zhǔn)模型中具有最高的平均分。論文發(fā)現(xiàn)它優(yōu)于Qwen1.5-3B-14B的聊天版本，盡管Qwen的參數(shù)量超過2倍，并且在表4中其預(yù)訓(xùn)練模型優(yōu)于OLMoE-1B-7B。在AlpacaEval上的84%得分也優(yōu)于排行榜上更大的密集模型，例如Llama2-13B-Chat。

表5：OLMOE-1B-7BAdaptation后與其他模型的比較。論文發(fā)現(xiàn)JetMoE聊天模型（???https://hf.co/jetmoe/jetmoe-8b-chat??）存在問題，導(dǎo)致隨機(jī)評分，因此論文將其排除。名稱包含四舍五入的參數(shù)計(jì)數(shù)：模型-活躍-總參數(shù)（對于MoEs）和模型-總參數(shù)（對于密集LMs）。論文自行運(yùn)行所有評估。模型使用不同的Adaptation混合，例如，OLMoE是在用于OLMo模型的改進(jìn)管道上訓(xùn)練的。

本文轉(zhuǎn)載自??AI帝國??，作者： 無影寺 ????