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

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

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

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

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

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

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

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

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

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

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

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

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

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

預訓練與自適應

預訓練架構

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

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

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

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

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

預訓練數據

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

下表中展示了預訓練數據的組成：

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

自適應

論文作者從指令調優和偏好調優兩方面，基于之前的開放模型，構造了 OLMoE-1B-7B-INSTRUCT。在指令調優集中，他們增加了更多的代碼和數學數據，以提高模型在這些領域的性能。

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

下表展示了 OLMoE-1B-7B-INSTRUCT 所使用的數據：

實驗

該研究的評估程序由三部分組成：預訓練期間、預訓練之后和自適應之后。

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

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

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

研究發現 OLMoE-1B-7B 在使用少于 2B 個激活參數的模型中表現最好，使其成為許多 LM 用例中最經濟的選擇。

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

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

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

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