第一個100%開源的MoE大模型，7B的參數，1B的推理成本

機器之心報道

機器之心編輯部

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

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

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

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

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

特別的，該研究使用 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 個專家會被選中並被激活，負責處理這個輸入。

預訓練數據

訓練數據方面，論文作者使用了來自兩個不同來源的數據集：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）在所有基準測試模型中具有最高平均值。