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

對于許多學(xué)者和開發(fā)人員來說，高性能的 LM 是無法訪問的，因為它們的構(gòu)建和部署成本過高。改善成本 - 性能的一種方法是使用稀疏激活混合專家 (MoE)。MoE 在每一層都有幾個專家，每次只激活其中的一個子集（參見圖 2）。這使得 MoE 比具有相似參數(shù)量的密集模型更有效，因為密集模型為每個輸入激活所有參數(shù)。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

對于數(shù)據(jù)的處理，論文作者使用了過濾器去除了包含太多重復(fù) token 的內(nèi)容、GitHub 上星標(biāo)少于 2 的項目以及某些詞出現(xiàn)頻率過高的文檔。他們將在每輪訓(xùn)練開始前隨機混洗數(shù)據(jù)，總計超過 5 萬億個 token。在「退火」階段（最后 100B 個 token），他們首先重新混洗整個數(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)集中，他們增加了更多的代碼和數(shù)學(xué)數(shù)據(jù)，以提高模型在這些領(lǐng)域的性能。

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

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

實驗

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

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

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

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

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

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

研究發(fā)現(xiàn)，盡管每條前向傳播所需的計算量減少了約 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)最先進的。

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

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