盡管大語(yǔ)言模型 (LM) 在各種任務(wù)上取得了重大進(jìn)展，但在訓(xùn)練和推理方面，性能和成本之間仍然需要權(quán)衡。

對(duì)于許多學(xué)者和開發(fā)人員來(lái)說，高性能的 LM 是無(wú)法訪問的，因?yàn)樗鼈兊臉?gòu)建和部署成本過高。改善成本 - 性能的一種方法是使用稀疏激活混合專家 (MoE)。MoE 在每一層都有幾個(gè)專家，每次只激活其中的一個(gè)子集（參見圖 2）。這使得 MoE 比具有相似參數(shù)量的密集模型更有效，因?yàn)槊芗Ｐ蜑槊總€(gè)輸入激活所有參數(shù)。

出于這個(gè)原因，行業(yè)前沿模型包括 Gemini-1.5、 GPT-4 等在內(nèi)的模型都使用了 MoE。

然而，大多數(shù) MoE 模型都是閉源的，雖然有些模型公開發(fā)布了模型權(quán)重，但有關(guān)訓(xùn)練數(shù)據(jù)、代碼等的信息卻很有限，甚至有些研究沒有提供這些信息。由于缺乏開放資源和對(duì)研究細(xì)節(jié)的深入探索，在 MoE 領(lǐng)域無(wú)法構(gòu)建具有成本效益的開源模型，從而接近閉源前沿模型的能力。

為了解決這些問題，來(lái)自艾倫人工智能研究院、 Contextual AI 等機(jī)構(gòu)的研究者引入了  OLMoE ，這是一個(gè)完全開源的混合專家語(yǔ)言模型，在類似大小的模型中具有 SOTA 性能。

• 論文地址：https://arxiv.org/pdf/2409.02060
• 論文標(biāo)題：OLMoE: Open Mixture-of-Experts Language Models

特別的，該研究使用 5.1 萬(wàn)億個(gè) token 預(yù)訓(xùn)練了 OLMoE-1B-7B 模型，該模型總共擁有 69 億參數(shù)，其中每個(gè)輸入 token 只激活 13 億參數(shù)。 

結(jié)果是與使用具有約 1B 參數(shù)的密集模型（例如 OLMo 1B 或 TinyLlama 1B ）實(shí)現(xiàn)了類似的推理成本，只是需要更多的 GPU 內(nèi)存來(lái)存儲(chǔ)約 7B 的總參數(shù)。實(shí)驗(yàn)表明，MoE 的訓(xùn)練速度比具有等效激活參數(shù)的密集 LM 快 2 倍左右。

如圖 1 所示，OLMoE-1B-7B 顯著優(yōu)于所有開源 1B 模型，并且與推理成本和內(nèi)存存儲(chǔ)明顯更高的密集模型相比表現(xiàn)出了競(jìng)爭(zhēng)力。

通過指令和偏好調(diào)優(yōu)，該研究還創(chuàng)建了 OLMoE-1B-7B-INSTRUCT，它在常見基準(zhǔn) MMLU、GSM8k、HumanEval 等上超越了各種更大的指令模型，包括 Llama2-13B-Chat 、OLMo-7B-Instruct (0724) 和 DeepSeekMoE-16B。

受控實(shí)驗(yàn)強(qiáng)調(diào)了 MoE（見表 1）和一般 LM 的關(guān)鍵設(shè)計(jì)選擇。結(jié)果表明使 MoE 性能卓越的一個(gè)關(guān)鍵設(shè)計(jì)決策是使用細(xì)粒度路由和粒度專家（granular experts）：在每一層使用 64 個(gè)小專家，其中 8 個(gè)被激活。

此外，路由算法的選擇也很重要：該研究發(fā)現(xiàn)無(wú)丟棄（dropless）基于 token 的路由優(yōu)于基于專家的路由。最后，該研究分析了 OLMoE-1B-7B 中的路由行為，發(fā)現(xiàn)路由在預(yù)訓(xùn)練的早期就飽和了，專家很少被共同激活，并且專家表現(xiàn)出領(lǐng)域和詞匯的專業(yè)化。

最后，作者希望這個(gè)完全開源的 MoE 能夠促進(jìn)更多研究和分析，從而提高對(duì)這些模型的理解。訓(xùn)練代碼、中間檢查點(diǎn)（每 5000 step ）、訓(xùn)練日志和訓(xùn)練數(shù)據(jù)都已經(jīng)開源。

論文作者 Niklas Muennighoff 表示：OLMoE 是第一個(gè) 100% 開源的混合專家 LLM。

預(yù)訓(xùn)練與自適應(yīng)

預(yù)訓(xùn)練架構(gòu)

OLMoE 是由 N_L 個(gè) transformer 層組成的語(yǔ)言模型，僅包含解碼器。對(duì)于 OLMo 這樣的密集模型，原本模型中單一的前饋網(wǎng)絡(luò)被 N_E 個(gè)小型前饋網(wǎng)絡(luò)（專家）組成的混合專家網(wǎng)絡(luò)所替代，對(duì)于每個(gè)輸入 token x，只有 k 個(gè)專家會(huì)被選中并被激活，負(fù)責(zé)處理這個(gè)輸入。

其中，路由器（r）是一個(gè)經(jīng)過訓(xùn)練的線性層，將輸入的原始數(shù)據(jù)映射到被選中的 k 個(gè)專家上。對(duì)路由器的輸出應(yīng)用 softmax 函數(shù)，計(jì)算 N_E 個(gè)專家的路由概率。然后，每個(gè)被指定的專家 E_i 處理輸入 x，其輸出乘以其各自的路由概率。再將所有選定的專家的結(jié)果相加，構(gòu)成模型單個(gè)層的 MoE 模塊輸出。

MoE 模型的訓(xùn)練往往涉及對(duì)一個(gè)已經(jīng)存在的大型密集模型轉(zhuǎn)換成一個(gè)稀疏模型，也就是所謂的「稀疏升級(jí)」。這個(gè)過程中，需要改變模型的訓(xùn)練目標(biāo)，比如調(diào)整 auxiliary load balancing 以及路由器的損失函數(shù)。具體的方法如下表所示：

在這項(xiàng)研究中，論文作者使用了總計(jì) 69 億參數(shù)中的 13 億活躍參數(shù)，每層有 64 個(gè)專家，其中有 8 個(gè)被激活。他們使用了一種名為「無(wú)丟棄 token」的路由方法：對(duì)于每個(gè)輸入 token，路由器網(wǎng)絡(luò)將分配 8 個(gè)專家來(lái)處理它。

論文作者引入了兩個(gè)輔助損失函數(shù)：負(fù)載平衡損失（）和路由器 z 損失（），來(lái)訓(xùn)練 OLMoE-1B-7B。他們給這兩個(gè)損失函數(shù)分別設(shè)定了權(quán)重（α 和 β），然后把它們和模型的主要學(xué)習(xí)目標(biāo)（交叉熵?fù)p失結(jié)合起來(lái)，最終計(jì)算的損失函數(shù)為：

預(yù)訓(xùn)練數(shù)據(jù)

訓(xùn)練數(shù)據(jù)方面，論文作者使用了來(lái)自兩個(gè)不同來(lái)源的數(shù)據(jù)集：DCLM 和 Dolma 1.7。這些數(shù)據(jù)集包括了多種類型的數(shù)據(jù)，比如網(wǎng)絡(luò)爬取的數(shù)據(jù)、編程問題解答、數(shù)學(xué)問題解答和學(xué)術(shù)論文等。他們將這些數(shù)據(jù)混合起來(lái)，創(chuàng)建了一個(gè)名為 OLMOE-MIX 的新數(shù)據(jù)集。

下表中展示了預(yù)訓(xùn)練數(shù)據(jù)的組成：

對(duì)于數(shù)據(jù)的處理，論文作者使用了過濾器去除了包含太多重復(fù) token 的內(nèi)容、GitHub 上星標(biāo)少于 2 的項(xiàng)目以及某些詞出現(xiàn)頻率過高的文檔。他們將在每輪訓(xùn)練開始前隨機(jī)混洗數(shù)據(jù)，總計(jì)超過 5 萬(wàn)億個(gè) token。在「退火」階段（最后 100B 個(gè) token），他們首先重新混洗整個(gè)數(shù)據(jù)集，然后按照此前 OLMo 論文中的方法，將學(xué)習(xí)率線性衰減到 0。

自適應(yīng)

論文作者從指令調(diào)優(yōu)和偏好調(diào)優(yōu)兩方面，基于之前的開放模型，構(gòu)造了 OLMoE-1B-7B-INSTRUCT。在指令調(diào)優(yōu)集中，他們?cè)黾恿烁嗟拇a和數(shù)學(xué)數(shù)據(jù)，以提高模型在這些領(lǐng)域的性能。

GPT-4 和 Llama 3 在預(yù)訓(xùn)練階段使用了像 GSM8k 或 MATH 這樣的數(shù)學(xué)數(shù)據(jù)集的樣本。按照這個(gè)思路，論文作者還添加了「No Robots」和「Daring Anteater」的一個(gè)子集。這些數(shù)據(jù)集不僅質(zhì)量高還更多樣，這是拓展模型適應(yīng)性的兩個(gè)關(guān)鍵因素。

下表展示了 OLMoE-1B-7B-INSTRUCT 所使用的數(shù)據(jù)：

實(shí)驗(yàn)

該研究的評(píng)估程序由三部分組成：預(yù)訓(xùn)練期間、預(yù)訓(xùn)練之后和自適應(yīng)之后。

預(yù)訓(xùn)練期間：如圖 3 所示，該研究在預(yù)訓(xùn)練期間使用當(dāng)前最佳 OLMo 模型在常用下游任務(wù)上對(duì) OLMoE-1B-7B 的性能進(jìn)行了基準(zhǔn)測(cè)試。

研究團(tuán)隊(duì)發(fā)現(xiàn)，在所有任務(wù)中，OLMoE-1B-7B 比密集 OLMo 模型以更少的計(jì)算量 (FLOP) 獲得了更好的性能。盡管 OLMoE-1B-7B 使用了不到一半的 FLOP 進(jìn)行訓(xùn)練并且僅使用 1B 個(gè)激活參數(shù)，但 OLMoE-1B-7B 在訓(xùn)練結(jié)束時(shí)可與 OLMo-7B 媲美，甚至優(yōu)于 OLMo-7B。

預(yù)訓(xùn)練之后：在表 4 中，該研究在常見的下游任務(wù)上對(duì) OLMoE-1B-7B 進(jìn)行基準(zhǔn)測(cè)試。

研究發(fā)現(xiàn) OLMoE-1B-7B 在使用少于 2B 個(gè)激活參數(shù)的模型中表現(xiàn)最好，使其成為許多 LM 用例中最經(jīng)濟(jì)的選擇。

如果預(yù)算較大，Qwen1.5-3B-14B 具有更強(qiáng)的性能，但其激活參數(shù)和總參數(shù)比 OLMoE-1B-7B 多一倍以上。

研究發(fā)現(xiàn)，盡管每條前向傳播所需的計(jì)算量減少了約 6-7 倍，但 OLMoE-1B-7B 的性能優(yōu)于一些具有 7B 參數(shù)的密集 LM，例如 Llama2-7B ，但不如其他 LM，例如 Llama3.1-8B 。上圖 1 比較了 OLMoE-1B-7B 和其他 LM 的 MMLU 性能和激活參數(shù)，表明 OLMoE-1B-7B 是其成本范圍內(nèi)最先進(jìn)的。

自適應(yīng)之后：在表 5 中，該研究對(duì) OLMoE-1B-7B 的指令 (SFT) 和偏好 (DPO) 調(diào)優(yōu)進(jìn)行了基準(zhǔn)測(cè)試。SFT 在所有測(cè)量任務(wù)上都改進(jìn)了本文的模型。

DPO 在大多數(shù)任務(wù)上都有幫助，尤其是 AlpacaEval，這與先前研究的結(jié)果一致。DPO 模型（稱之為 OLMoE-1B-7B-INSTRUCT）在所有基準(zhǔn)測(cè)試模型中具有最高平均值。