訓練程式碼、中間 checkpoint、訓練日誌和訓練資料都已經開源。
儘管大語言模型 (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）在所有基準測試模型中具有最高平均值。