OLMoE: 開源的MoE語言模型(預訓練&效果)
一、結論寫在前面
論文標題：OLMoE: Open Mixture-of-Experts Language Models
論文鏈接：??https://arxiv.org/pdf/2409.02060??
Weights：??https://hf.co/allenai/OLMoE-1B-7B-0924??
Data：??https://hf.co/datasets/allenai/OLMoE-mix-0924??
Code：??https://github.com/allenai/OLMoE??
Logs：??https://wandb.ai/ai2-llm/olmoe/reports/OLMoE-1B-7B-0924--Vmlldzo4OTcyMjU3??
論文開源了OLMoE-1B-7B和OLMoE-1B-7B-INSTRUCT，包括模型、數(shù)據(jù)、代碼和日志。OLMOE-1B-7B擁有7B參數(shù)，但每個輸入token僅使用1B參數(shù)。論文在5T token上對其進行預訓練，并進一步Adaptation以創(chuàng)建OLMoE-1B-7B-INSTRUCT。
論文的模型在所有具有相似活躍參數(shù)的可用模型中表現(xiàn)最佳，甚至超越了Llama2-13B-Chat和DeepSeekMoE-16B等更大模型。
論文分享了各種訓練實驗，并定義和分析了路由器飽和度、專家協(xié)同激活、領域和詞匯專業(yè)化。通過論文的完全開源發(fā)布，論文希望幫助該領域構建更好的專家混合模型。論文對OLMoE的新迭代感到興奮，以縮小前沿模型與完全開源模型之間的差距。
圖1：開放式MoE和密集LMs的性能、成本和開放程度。模型名稱包含四舍五入的參數(shù)計數(shù)：MoE為model-active-total，密集LMs為model-total。#ckpts是可用的中間檢查點數(shù)量。論文突出顯示MMLU作為整體性能的總結；更多結果見S3。OLMOE-1B-7B在具有相似活躍參數(shù)計數(shù)的模型中表現(xiàn)最佳，是最開放的MoE。
二、論文的簡單介紹
2.1 論文的背景
高性能的語言模型對許多學者和開源開發(fā)者來說是很難接觸，因為它們的構建和部署成本過高。改善成本性能權衡的一種方法是使用稀疏激活的專家混合(Mixture-of-Experts，MoEs)模型。MoEs在每一層都有多個專家，每次只激活其中的一個子集(見圖2)。這使得MoEs比具有相似總參數(shù)數(shù)量的密集模型更加高效，因為密集模型需要為每個輸入激活所有參數(shù)。出于這個原因，行業(yè)前沿模型使用MoEs，包括Gemini-1.5和據(jù)報道的GPT-4。
然而，大多數(shù)MoE模型都是閉源的:雖然有些模型公開發(fā)布了模型權重，但它們提供的關于訓練數(shù)據(jù)、代碼或配方的信息有限或沒有(見圖1)。雖然之前有一些努力使語言建模研究完全可訪問，但它們主要局限于密集LMs。這種情況出現(xiàn)的原因是，MoEs需要更多的開放性，因為它們?yōu)長Ms增加了復雜的新設計問題，例如使用多少總參數(shù)與活躍參數(shù)，是使用多個小專家還是少數(shù)大專家，專家是否應該共享，以及使用什么路由算法。缺乏關于這些細節(jié)的開放資源和發(fā)現(xiàn)，阻礙了該領域構建接近閉源前沿模型能力的高效開放MoEs。
為了解決這些問題，論文引入了OLMOE，一個完全開放的專家混合語言模型，在同等規(guī)模的模型中具有最先進的性能。具體來說，論文預訓練了OLMOE-1B-7B，總共訓練了5.1T token，總參數(shù)為6.9B，其中只有1.3B參數(shù)為每個輸入token激活。這導致推理成本與使用約1B參數(shù)的密集模型(如OLMo 1B或TinyLlama 1B)相似，但需要更多GPU內存來存儲其7B總參數(shù)。
論文的實驗表明，MoEs的訓練速度比具有相同活躍參數(shù)數(shù)量的密集LMs快約2倍。在圖1中，論文展示了OLMOE-1B-7B顯著優(yōu)于所有開放的10億參數(shù)模型，并展現(xiàn)出與推理成本和內存存儲顯著更高的密集模型相當?shù)男阅?例如，在MMLU評分上與Llama2-13B相似，后者的成本約為前者的10倍)。通過指令和偏好微調，論文創(chuàng)建了OLMOE-1B-7B-INSTRUCT，論文發(fā)現(xiàn)它在常見基準測試(MMLU、GSM8k、HumanEval等)上超過了各種更大的指令模型，包括Llama2-13B-Chat、OLMo-7B-Instruct (0724)和DeepSeekMoE-16B。
論文全面的對照實驗突出了MoEs(見表1)和LMs的關鍵設計選擇。使MoEs性能出色的一個關鍵設計決策是使用細粒度路由和細粒度專家:論文在每層使用64個小專家，其中8個被激活。路由算法的選擇也很重要:論文發(fā)現(xiàn)無丟棄的基于token的路由優(yōu)于基于專家的路由。論文的發(fā)現(xiàn)還包括一些挑戰(zhàn)先前工作的結果，例如共享專家的無效性，以及將預訓練的密集LM稀疏地升級為MoE的有限益處，除非在小規(guī)模計算預算下。最后，論文分析了OLMOE-1B-7B的路由行為，發(fā)現(xiàn)路由在預訓練早期就飽和了，專家很少被共同激活，并且專家表現(xiàn)出領域和詞匯專門化。
圖2：密集型語言模型（LMs）與OLMoE等MoE模型的架構對比。圖中省略了一些細節(jié)，例如，OLMOE-1B-7B還使用了QK-Norm。
表1: MoE的關鍵設計選擇及基于論文實驗的OLMOE-1B-7B設置
表2：OLMOE-1B-7B預訓練數(shù)據(jù)的組成。StarCoder、peS2o和Wikipedia部分來自Dolma 1.7 。
表3：OLMOE-13.7的Adaptation訓練數(shù)據(jù)。
2.2 預訓練和Adaptation(Pretraining and Adaptation)
預訓練架構 OLMOE是一個僅有解碼器的語言模型，由NL個transformer層組成。在像OLMo這樣的密集模型中的前饋網絡(FFN)被替換為一個MoE模塊，該模塊由NE個較小的FFN模塊組成，這些模塊被稱為專家。對于每個處理的輸入token x，會激活其中的k個專家子集(另見圖2):
其中 r，稱為路由器，是一個從輸入logits映射到所選 k 個專家的學習線性層。softmax應用于路由器輸出，以計算所有 N_E個專家的路由概率。每個選定的專家 E_i 處理輸入 x，其輸出隨后與其相應的路由概率相乘。然后將所有選定的Top- k 專家的結果相加以構成MoE模塊的輸出，該模塊是模型 N_L層中的一層。設計MoE模型的關鍵決策包括確定激活和總參數(shù)的數(shù)量、專家的設計（例如，粒度，是否包含共享專家）以及路由算法的選擇。此外，訓練MoE模型可能涉及從密集模型（稀疏升級）初始化并更改訓練目標，例如包括輔助負載均衡和路由z-losses；表1顯示了論文的最終決策。
總結來說，論文使用了總共6.9B參數(shù)中的1.3B活躍參數(shù)，每層激活了8個專家，總共64個專家。論文采用了無丟棄的token選擇路由[58]：對于每個輸入token，學習到的路由網絡決定由8個專家來處理。論文從零開始訓練OLMOE-1B-7B，使用了兩個輔助損失：負載均衡損失( L_LB) 和路由z-損失( L_RZ) 。論文將它們分別乘以各自的損失權重alpha和beta，并與交叉熵損失(L_CE )線性相加，得到最終的訓練損失：
OLMOE-1B-7B的完整預訓練配置在附錄B中。
預訓練數(shù)據(jù) 論文使用了來自DCLM [89]和Dolma 1.7 [161]的混合數(shù)據(jù)，包括以下內容：
(1) 經過質量過濾的Common Crawl子集，稱為DCLM-Baseline;
(2) 在DCLM和Dolma 1. 7中都使用的StarCoder、Algebraic Stack和arXiv;
(3) 來自Dolma 1.7的peS2o和Wikipedia。論文將預訓練數(shù)據(jù)集稱為OLMOE-MIX。
對于上述所有來源，論文應用了一個過濾器，移除所有包含32個或更多重復n-gram的文檔，其中n-gram是任意1到13個token的跨度。對于StarCoder子集，論文還移除了任何來自GitHub上少于2顆星的倉庫的文檔，或者其最頻繁詞占據(jù)文檔超過30%的內容，或者其前兩個最頻繁詞占據(jù)文檔超過50%的內容。
論文在每個epoch開始時隨機打亂所有樣本，并總共訓練5.133T個token（根據(jù)Muennighoff等人[120]的方法，訓練1.3個epoch）。在論文的退火階段（最后100B個token），論文首先重新打亂整個數(shù)據(jù)集，然后按照先前的工作64， 89，將學習率線性衰減到0。論文的預訓練數(shù)據(jù)統(tǒng)計如表2所示。
Adaptation（Adaptation） 論文通過遵循標準的Adaptation配方創(chuàng)建了OLMoE-1B-7B-INSTRUCT，該配方分為指令微調和偏好微調，基于先前的開源模型。在論文的指令微調數(shù)據(jù)集中，論文增加了更多的代碼和數(shù)學數(shù)據(jù)，以提升下游編碼和數(shù)學應用的性能。其他模型，如GPT-4和Llama 3，在預訓練期間包含了來自GSM8k或MATH等數(shù)學數(shù)據(jù)集的樣本。論文還包含了No Robots和Daring Anteater的一個子集，因為它們質量高且增加了多樣性，這是成功Adaptation的兩個關鍵因素。論文在表3中描述了論文的Adaptation數(shù)據(jù)集。
2.3 結果
論文的評估程序包括三個部分：預訓練期間、預訓練后和Adaptation后(After adaptation)。
圖3: OLMOE-1B-7B和當前最佳OLMo模型在預訓練期間的評估。OLMOE-1B-7B與OLMo模型在MoE架構、幾個訓練超參數(shù)和訓練數(shù)據(jù)集方面有所不同。更多結果、日志和配置:??https://wandb.ai/ai2-llm/olmoe/reports/Plot-OLMoE-1B-7B-vs-OLMo-7B-vs-OLMo-1B--Vmlldzo4OTcyMjEz??
預訓練期間(During pretraining) 在圖3中，論文對比了OLMOE-1B-7B在預訓練期間與當前最佳OLMo模型在常用下游任務上的性能表現(xiàn)。論文發(fā)現(xiàn)，在所有任務中，OLMOE-1B-7B都以更少的計算量(FLOPs)達到了比密集OLMo模型更好的性能。
盡管OLMOE-1B-7B在訓練中使用的FLOPs不到一半，且只使用了1B的活躍參數(shù)，但在訓練結束時，它的表現(xiàn)與OLMo-7B相當或更優(yōu)。這可能是由于論文對OLMo設置進行的數(shù)據(jù)集和建模變更，包括MoE相關的改變、穩(wěn)定性和性能改進，這些在附錄B中有詳細說明。
表4: 預訓練后的OLMOE-1B-7B與更大的MoEs和密集LMs的比較。論文將其與在活躍參數(shù)(1B，近似速度和成本)或總參數(shù)(7B，近似內存需求)上與OLMOE-1B-7B相似的密集LMs進行比較。模型名稱包含四舍五入的參數(shù)計數(shù):MoEs的格式為model-active-total，密集LMs的格式為model-total(這導致與官方名稱有些差異，例如，雖然被稱為"Gemma2-2B"，但它實際上有26億活躍和總參數(shù)[175])。Chall. = Challenge(挑戰(zhàn))。論文自己使用5個少樣本示例進行所有評估
預訓練后(After pretraining) 在表4中，論文對OLMoE-1B-7B在常見的下游任務上進行了基準測試。論文發(fā)現(xiàn)，在使用的活躍參數(shù)少于2B的模型中，OLMoE-1B-7B表現(xiàn)最佳，使其成為許多語言模型使用場景中最經濟的選擇。對預算更大的情況，Qwenl.5 3B-14B的性能更強，但其活躍參數(shù)和總參數(shù)數(shù)量是OLMoE-1B-7B的兩倍以上。論文發(fā)現(xiàn)，盡管每次前向傳遞所需的計算量減少了6-7倍，OLMoE-1B-7B仍優(yōu)于一些具有7B參數(shù)的密集模型，如Llama2-7B ，但在其他模型如Llama3.1-8B面前稍顯不足。圖1比較了OLMoE-1B-7B與其他語言模型在MMLU性能與活躍參數(shù)（作為模型成本的代理）上的表現(xiàn)。OLMoE-1B-7B在其成本范圍內處于領先地位。
Adaptation后(After adaptation) 在表5中，論文基準測試了OLMoE-1B-7B的指令（SFT）和偏好（DPO）調優(yōu)。SFT在所有測量的任務上提升了論文的模型性能。論文在GSM8k上觀察到>10×的增益，這可能歸因于論文在預訓練期間加入了額外的數(shù)學數(shù)據(jù)，以彌補相對較少的數(shù)學數(shù)據(jù)（S2）。DPO在大多數(shù)任務上都有幫助，尤其是在AlpacaEval上，這與先前的工作[186，75，121]中的發(fā)現(xiàn)一致。
論文的DPO模型，論文稱之為OLMoE-1B-7B-INSTRUCT在所有基準模型中具有最高的平均分。論文發(fā)現(xiàn)它優(yōu)于Qwen1.5-3B-14B的聊天版本，盡管Qwen的參數(shù)量超過2倍，并且在表4中其預訓練模型優(yōu)于OLMoE-1B-7B。在AlpacaEval上的84%得分也優(yōu)于排行榜上更大的密集模型，例如Llama2-13B-Chat。
表5：OLMOE-1B-7BAdaptation后與其他模型的比較。論文發(fā)現(xiàn)JetMoE聊天模型（???https://hf.co/jetmoe/jetmoe-8b-chat??）存在問題，導致隨機評分，因此論文將其排除。名稱包含四舍五入的參數(shù)計數(shù)：模型-活躍-總參數(shù)（對于MoEs）和模型-總參數(shù)（對于密集LMs）。論文自行運行所有評估。模型使用不同的Adaptation混合，例如，OLMoE是在用于OLMo模型的改進管道上訓練的。
