LLM | Yuan 2.0-M32：帶注意力路由的專家混合模型
圖片
一、結(jié)論寫在前面
Yuan 2.0-M32，其基礎(chǔ)架構(gòu)與Yuan-2.0 2B相似，采用了一個(gè)包含32位專家的專家混合架構(gòu)，其中2位專家處于活躍狀態(tài)。提出并采用了一種新的路由網(wǎng)絡(luò)——注意力路由器，以更高效地選擇專家，相比采用經(jīng)典路由網(wǎng)絡(luò)的模型，其準(zhǔn)確率提升了3.8%。Yuan 2.0-M32從零開始訓(xùn)練，使用了2000B的token，其訓(xùn)練計(jì)算消耗僅為同等參數(shù)規(guī)模密集模型的9.25%。
Yuan 2.0-M32在編碼、數(shù)學(xué)及多個(gè)專業(yè)領(lǐng)域展現(xiàn)了競(jìng)爭(zhēng)力的能力，僅使用了400億總參數(shù)中的37億活躍參數(shù)，以及每token7.4 GFlops的前向計(jì)算，這兩項(xiàng)指標(biāo)均僅為L(zhǎng)lama3-70B的1/19。Yuan 2.0-M32在MATH和ARC-Challenge基準(zhǔn)測(cè)試中超越了Llama3-70B，準(zhǔn)確率分別達(dá)到55.89%和95.8%。Yuan 2.0-M32的模型及源代碼已在GitHub：https://github.com/IEIT-Yuan/Yuan2.0-M32。
二、論文的簡(jiǎn)單介紹
2.1 論文的背景
在每個(gè)token固定計(jì)算量的情況下，采用專家混合（MoE）結(jié)構(gòu)的模型可以通過(guò)增加專家數(shù)量輕松構(gòu)建得比密集模型更大規(guī)模，從而實(shí)現(xiàn)更高的準(zhǔn)確性表現(xiàn)。實(shí)際上，通常在有限的計(jì)算資源下訓(xùn)練模型，MoE被視為減少與模型、數(shù)據(jù)集規(guī)模和有限計(jì)算能力相關(guān)的巨大成本的良好候選方案。
MoE的概念可追溯至1991年。總損失是每個(gè)專家加權(quán)損失的組合，這些專家具有獨(dú)立判斷的能力。稀疏門控MoE的概念最初由Shazeer等人（2017年）在翻譯模型中提出。采用這種路由策略，推理時(shí)只有極少數(shù)專家被激活，而非所有專家同時(shí)被調(diào)用。這種稀疏性使得模型能夠在計(jì)算效率損失極小的情況下，在堆疊的LSTM層之間擴(kuò)展至1000倍。噪聲可調(diào)的Top-K門控路由網(wǎng)絡(luò)向softmax函數(shù)引入可調(diào)噪聲并保持前K值，以平衡專家利用率。近年來(lái)，隨著模型規(guī)模的不斷擴(kuò)大，路由策略在高效分配計(jì)算資源方面的作用受到了更多關(guān)注。
專家路由網(wǎng)絡(luò)是MoE結(jié)構(gòu)的核心。該結(jié)構(gòu)通過(guò)計(jì)算token分配給每個(gè)專家的概率來(lái)選擇參與計(jì)算的候選專家。目前，在大多數(shù)流行的MoE結(jié)構(gòu)中，普遍采用經(jīng)典路由算法，該算法執(zhí)行token與每個(gè)專家特征向量之間的點(diǎn)積，并選擇點(diǎn)積值最大的專家。在這種變換中，專家的特征向量是獨(dú)立的，忽略了專家之間的相關(guān)性。然而，MoE結(jié)構(gòu)通常每次選擇不止一個(gè)專家，并且多個(gè)專家經(jīng)常協(xié)同參與計(jì)算，這意味著專家之間應(yīng)存在固有相關(guān)性。如果在選擇專家的過(guò)程中考慮專家間的關(guān)系，無(wú)疑會(huì)提高模型的準(zhǔn)確性。
2.2 論文的方法
2.2.1模型架構(gòu)
Yuan 2.0-M32 基于 Yuan 2.0-2B 的模型結(jié)構(gòu)。Yuan 2.0 引入了基于局部過(guò)濾的注意力（LFA）以考慮輸入token的局部依賴性，從而提高模型的準(zhǔn)確性。在 Yuan 2.0-M32 中，每一層的密集前饋網(wǎng)絡(luò)（FFN）被替換為 MoE 組件。
圖1展示了論文模型中應(yīng)用的MoE層的架構(gòu)。以四個(gè)FFN為例（實(shí)際上有32個(gè)專家），每個(gè)MoE層由一組獨(dú)立的FFN作為專家組成。位于專家之前的路由網(wǎng)絡(luò)將輸入token分派給相關(guān)的專家。經(jīng)典的路由網(wǎng)絡(luò)本質(zhì)上為每個(gè)專家建立了一個(gè)特征向量，并計(jì)算輸入token與每個(gè)專家特征向量之間的點(diǎn)積，以獲得token與專家之間的具體似然度。具有最強(qiáng)似然度的專家被選中激活，并參與后續(xù)的計(jì)算。
圖片
圖1：Yuan 2.0-M32的說(shuō)明。左側(cè)圖展示了Yuan 2.0架構(gòu)中MoE層的擴(kuò)展情況。MoE層替代了Yuan 2.0中的前饋層。右側(cè)圖展示了MoE層的結(jié)構(gòu)。在論文的模型中，每個(gè)輸入token將被分配給總共32個(gè)專家中的2個(gè)，而在圖中論文以4個(gè)專家為例進(jìn)行展示。MoE的輸出是所選專家的加權(quán)和。N表示層的數(shù)量 
圖2展示了注意力路由器結(jié)構(gòu)的概覽
圖2(a)展示了經(jīng)典路由網(wǎng)絡(luò)的結(jié)構(gòu)。每個(gè)專家的特征向量彼此獨(dú)立，計(jì)算概率時(shí)忽略了專家之間的相關(guān)性。實(shí)際上，在大多數(shù)MoE模型中，通常會(huì)選擇兩個(gè)或更多的專家參與后續(xù)計(jì)算，這自然帶來(lái)了專家間的強(qiáng)相關(guān)性。考慮專家間的相關(guān)性無(wú)疑有助于提高準(zhǔn)確性。
圖2(b)展示了本工作提出的注意力路由器的架構(gòu)，該新穎的路由網(wǎng)絡(luò)通過(guò)采用注意力機(jī)制來(lái)整合專家間的相關(guān)性。構(gòu)建了一個(gè)表示專家間相關(guān)性的系數(shù)矩陣，并應(yīng)用于最終概率值的計(jì)算中。
表1：不同路由結(jié)構(gòu)的比較
表1列出了不同路由器的準(zhǔn)確性結(jié)果。論文的模型在8個(gè)可訓(xùn)練專家上測(cè)試了注意力路由器。經(jīng)典路由器模型有8個(gè)可訓(xùn)練專家，以確保相似的參數(shù)規(guī)模，并且路由結(jié)構(gòu)與應(yīng)用于Mixtral 8*7B的結(jié)構(gòu)相同，即一個(gè)線性層上的Softmax。共享專家路由器采用共享專家隔離策略與經(jīng)典路由架構(gòu)。有兩個(gè)固定專家捕捉通用知識(shí)，以及14個(gè)可選專家中前兩名作為專業(yè)專家。
MoE的輸出是固定專家和路由器選出的專家的組合。所有三個(gè)模型都使用30Btoken進(jìn)行訓(xùn)練，并使用另外10Btoken進(jìn)行測(cè)試。考慮到經(jīng)典路由器和共享專家路由器之間的結(jié)果，論文發(fā)現(xiàn)后者在訓(xùn)練時(shí)間增加了7.35%的情況下獲得了完全相同的測(cè)試損失。共享專家的計(jì)算效率相對(duì)較低，并沒(méi)有帶來(lái)比經(jīng)典MOE策略更好的訓(xùn)練準(zhǔn)確性。因此，在論文的模型中，論文采用了不帶任何共享專家的經(jīng)典路由策略。與經(jīng)典路由網(wǎng)絡(luò)相比，注意力路由器的測(cè)試損失提高了3.8%。
論文通過(guò)增加專家數(shù)量并固定每個(gè)專家的參數(shù)大小來(lái)測(cè)試模型的可擴(kuò)展性。訓(xùn)練專家數(shù)量的增加僅改變模型容量，而不改變實(shí)際激活的模型參數(shù)。所有模型均使用500億個(gè)token進(jìn)行訓(xùn)練，并使用另外的100億個(gè)token進(jìn)行測(cè)試。論文將激活的專家設(shè)置為2，三個(gè)模型的訓(xùn)練超參數(shù)相同。專家擴(kuò)展效果通過(guò)訓(xùn)練500億個(gè)token后的測(cè)試損失來(lái)衡量（表2）。與具有8個(gè)可訓(xùn)練專家的模型相比，具有16個(gè)專家的模型顯示出2%的損失降低，而具有32個(gè)專家的模型顯示出3.6%的損失降低。考慮到其準(zhǔn)確性，論文為Yuan 2.0-M32選擇了32個(gè)專家。
表2：擴(kuò)展實(shí)驗(yàn)結(jié)果
2.2.2 模型訓(xùn)練
Yuan 2.0-M32通過(guò)數(shù)據(jù)并行和流水線并行的組合進(jìn)行訓(xùn)練，但不使用張量并行或優(yōu)化器并行。圖3展示了損失曲線，最終訓(xùn)練損失為1.22。
圖3：Yuan2.0-M32在2000Btoken上的預(yù)訓(xùn)練損失
2.2.3 微調(diào)
在微調(diào)過(guò)程中，論文將序列長(zhǎng)度擴(kuò)展至16384。遵循CodeLLama（Roziere et al.， 2023）的工作，論文重置旋轉(zhuǎn)位置嵌入（RoPE）的基頻值，以避免隨著序列長(zhǎng)度增加，注意力分?jǐn)?shù)的衰減。論文沒(méi)有簡(jiǎn)單地將基值從1000增加到一個(gè)非常大的值（例如1000000），而是使用NTK感知（bloc97， 2023）計(jì)算新的基值。
論文還比較了預(yù)訓(xùn)練的Yuan 2.0-M32模型與NTK感知風(fēng)格的新基值，以及與其他基值在序列長(zhǎng)度高達(dá)16K的針檢索任務(wù)中的性能。論文發(fā)現(xiàn)NTK感知風(fēng)格的新基值40890表現(xiàn)更好。因此，在微調(diào)過(guò)程中應(yīng)用了40890。
2.2.4 預(yù)訓(xùn)練數(shù)據(jù)集
Yuan 2.0-M32 從零開始使用包含 2000B token 的雙語(yǔ)數(shù)據(jù)集進(jìn)行預(yù)訓(xùn)練。預(yù)訓(xùn)練的原始數(shù)據(jù)包含超過(guò) 3400B token，并根據(jù)數(shù)據(jù)質(zhì)量和數(shù)量調(diào)整每個(gè)類別的權(quán)重。
綜合預(yù)訓(xùn)練語(yǔ)料庫(kù)由以下內(nèi)容組成：
44個(gè)子數(shù)據(jù)集，涵蓋了網(wǎng)絡(luò)爬取數(shù)據(jù)、維基百科、學(xué)術(shù)論文、書籍、代碼、數(shù)學(xué)和公式以及特定領(lǐng)域的專業(yè)知識(shí)。其中一些是開源數(shù)據(jù)集，其余由Yuan 2.0創(chuàng)建。
部分常見(jiàn)的網(wǎng)絡(luò)爬蟲數(shù)據(jù)、中文書籍、對(duì)話及中文新聞數(shù)據(jù)繼承自 Yuan 1.0（吳等人，2021年）。Yuan 2.0 中的大部分預(yù)訓(xùn)練數(shù)據(jù)也得到了重新利用。
關(guān)于每個(gè)數(shù)據(jù)集的構(gòu)建和來(lái)源的詳細(xì)信息如下：
網(wǎng)絡(luò)（25.2%）：網(wǎng)站爬蟲數(shù)據(jù)是從開源數(shù)據(jù)集和論文之前工作（Yuan 1.0）中處理過(guò)的公共爬蟲數(shù)據(jù)中收集的。關(guān)于從網(wǎng)絡(luò)上下文中提取高質(zhì)量?jī)?nèi)容的Massive Data Filtering System（MDFS）的更多詳情，請(qǐng)參考Yuan 1.0。
百科全書（1.2%）、論文（0.84%）、書籍（6.49%）和翻譯（1.1%）：數(shù)據(jù)繼承自Yuan 1.0和Yuan 2.0數(shù)據(jù)集。
代碼（47.5%）：與Yuan 2.0相比，代碼數(shù)據(jù)集得到了極大的擴(kuò)展。論文采用了Stack v2（Lozhkov等人，2024年）中的代碼。Stack v2中的注釋被翻譯成中文。通過(guò)與Yuan 2.0相似的方法生成了代碼合成數(shù)據(jù)。
數(shù)學(xué)（6.36%）：所有來(lái)自Yuan 2.0的數(shù)學(xué)數(shù)據(jù)都被重新使用。這些數(shù)據(jù)主要來(lái)自開源數(shù)據(jù)集，包括proof-pile vl（Azerbayev，2022年）和v2（Paster等人，2023年），AMPS（Hendrycks等人，2021年），MathPile（Wang，Xia和Liu，2023年）以及StackMathQA（Zhang，2024年）。使用Python創(chuàng)建了一個(gè)數(shù)值計(jì)算的合成數(shù)據(jù)集，以利于四則運(yùn)算。
特定領(lǐng)域（1.93%）：這是一個(gè)包含不同背景知識(shí)的數(shù)據(jù)集。
2.2.5 微調(diào)數(shù)據(jù)集
微調(diào)數(shù)據(jù)集基于Yuan 2.0中應(yīng)用的數(shù)據(jù)集進(jìn)行了擴(kuò)展。
代碼指令數(shù)據(jù)集。所有帶有中文指令及部分帶有英文注釋的編程數(shù)據(jù)均由大型語(yǔ)言模型（LLMs）生成。約30%的代碼指令數(shù)據(jù)為英文，其余為中文。合成數(shù)據(jù)在提示生成和數(shù)據(jù)清洗策略上模仿了帶有中文注釋的Python代碼。
帶有英文注釋的Python代碼收集自Magicoder-Evol-Instruct-110K和CodeFeedback-Filtered-Instruction。從數(shù)據(jù)集中提取帶有語(yǔ)言標(biāo)簽（如“python”）的指令數(shù)據(jù)。
其他如C/C++/Go/Java/SQL/Shell等語(yǔ)言的代碼，帶有英文注釋，源自開源數(shù)據(jù)集，處理方式與Python代碼類似。清洗策略與Yuan 2.0中的方法相似。設(shè)計(jì)了一個(gè)沙箱以提取生成的代碼中可編譯和可執(zhí)行的行，并保留至少通過(guò)一個(gè)單元測(cè)試的行。
數(shù)學(xué)指令數(shù)據(jù)集。數(shù)學(xué)指令數(shù)據(jù)集全部繼承自Yuan 2.0中的微調(diào)數(shù)據(jù)集。為提高模型通過(guò)編程方法解決數(shù)學(xué)問(wèn)題的能力，論文構(gòu)建了Thoughts（PoT）提示的數(shù)學(xué)數(shù)據(jù)。PoT將數(shù)學(xué)問(wèn)題轉(zhuǎn)換為使用Python進(jìn)行計(jì)算的代碼生成任務(wù)。
安全指令數(shù)據(jù)集。除了元2.0的聊天數(shù)據(jù)集外，論文還基于一個(gè)開源的安全對(duì)齊數(shù)據(jù)集構(gòu)建了一個(gè)雙語(yǔ)安全對(duì)齊數(shù)據(jù)集。論文僅從公共數(shù)據(jù)集中提取問(wèn)題，并增加問(wèn)題的多樣性，利用大型語(yǔ)言模型重新生成中文和英文答案。
2.2.6 Tokenizer
對(duì)于 Yuan 2.0-M32，英文和中文分詞器繼承自 Yuan 2.0 中應(yīng)用的分詞器。
2.3 論文的效果
論文在HumanEval上評(píng)估了Yuan 2.0-M32的代碼生成能力，在GSM8K和MATH上評(píng)估了數(shù)學(xué)問(wèn)題解決能力，在ARC上評(píng)估了科學(xué)知識(shí)和推理能力，并在MMLU上作為一個(gè)綜合基準(zhǔn)進(jìn)行評(píng)估。
2.3.1 代碼生成
代碼生成能力的評(píng)估使用HumanEval基準(zhǔn)。評(píng)估方法和提示與元2.0中提到的相似。
表3：Yuan 2.0-M32與其他模型在HumanEval pass @1上的比較
模型預(yù)期在后完成函數(shù)。生成的函數(shù)將通過(guò)單元測(cè)試進(jìn)行評(píng)估。表3展示了Yuan 2.0-M32在零樣本學(xué)習(xí)中的結(jié)果，并與其它模型進(jìn)行了比較。Yuan 2.0-M32的結(jié)果僅次于DeepseekV2和Llama3-70B，并且遠(yuǎn)超其他模型，即使其活躍參數(shù)和計(jì)算消耗遠(yuǎn)低于其他模型。
與DeepseekV2相比，論文的模型使用的活躍參數(shù)不到其四分之一，每token的計(jì)算量不到其五分之一，同時(shí)達(dá)到了其超過(guò)90%的準(zhǔn)確度水平。與Llama3-70B相比，模型參數(shù)和計(jì)算量的差距更大，論文仍能達(dá)到其91%的水平。Yuan 2.0-M32展示了可靠的編程能力，通過(guò)了四分之三的問(wèn)題。Yuan 2.0-M32擅長(zhǎng)小樣本學(xué)習(xí)，通過(guò)14次嘗試將HumanEval的準(zhǔn)確率提高到78.0。
2.3.2 數(shù)學(xué)
Yuan 2.0-M32的數(shù)學(xué)能力通過(guò)GSM8K和MATH基準(zhǔn)進(jìn)行評(píng)估。GSM8K的提示和測(cè)試策略與應(yīng)用于Yuan 2.0的相似，唯一不同的是論文使用8次嘗試（表4）。
表4：Yuan 2.0-M32與其他模型在GSM8K和MATH上的比較
MATH是一個(gè)包含12，500個(gè)挑戰(zhàn)性數(shù)學(xué)競(jìng)賽問(wèn)答問(wèn)題的數(shù)據(jù)集。該數(shù)據(jù)集中的每個(gè)問(wèn)題都有一個(gè)完整的逐步解決方案，引導(dǎo)模型生成答案推導(dǎo)和解釋。問(wèn)題的答案可以是數(shù)值，或數(shù)學(xué)表達(dá)式（如y=2x+5，x-+2x-1，2a+b等）。Yuan 2.0-M32使用鏈?zhǔn)剿季S（CoT）方法，通過(guò)4次嘗試生成最終答案。答案將從分析中提取并轉(zhuǎn)換為統(tǒng)一格式。
對(duì)于數(shù)值結(jié)果，所有格式的數(shù)學(xué)等價(jià)輸出均被接受。例如，分?jǐn)?shù)1/2，12，0.5，0.50都轉(zhuǎn)換為0.5并被視為相同結(jié)果。對(duì)于數(shù)學(xué)表達(dá)式，論文移除制表符和空格符號(hào)，并統(tǒng)一了節(jié)奏或音符的正則表達(dá)式。55 '5'均被接受為相同答案。處理后的最終結(jié)果與標(biāo)準(zhǔn)答案進(jìn)行比較，并使用EM（精確匹配）分?jǐn)?shù)進(jìn)行評(píng)估。
從表4所示的結(jié)果可以看出，Yuan 2.0-M32在MATH基準(zhǔn)上得分最高。與Mixtral-8x7B相比，后者活躍參數(shù)是Yuan 2.0-M32的3.48倍，但Yuan的得分幾乎是其兩倍。在GSM8K上，Yuan 2.0-M32的得分也非常接近Llama 3-70B，并優(yōu)于其他模型。
2.3.3MMLU
大規(guī)模多任務(wù)語(yǔ)言理解（MMLU）涵蓋了STEM、人文科學(xué)、社會(huì)科學(xué)等57個(gè)學(xué)科，從基礎(chǔ)語(yǔ)言任務(wù)到高級(jí)邏輯推理任務(wù)不等。MMLU中的所有問(wèn)題都是英語(yǔ)的多選QA問(wèn)題。模型預(yù)期生成正確的選項(xiàng)或相應(yīng)的分析。
Yuan 2.0-M32的輸入數(shù)據(jù)組織如附錄B所示。之前的文本被發(fā)送給模型，所有與正確答案或選項(xiàng)標(biāo)簽相關(guān)的答案被視為正確。
最終準(zhǔn)確性通過(guò)MC1（表5）進(jìn)行衡量。MMLU上的結(jié)果展示了論文模型在不同領(lǐng)域的能力。Yuan 2.0-M32在性能上超過(guò)了Mixtral-8x7B、Phi-3-mini和Llama 3-8B。
表5：Yuan 2.0-M32與其他模型在MMLU上的比較
2.3.4 ARC
AI2推理挑戰(zhàn)（ARC）基準(zhǔn)是一個(gè)多選QA數(shù)據(jù)集，包含從3年級(jí)到9年級(jí)科學(xué)考試的問(wèn)題。它分為簡(jiǎn)單和挑戰(zhàn)兩部分，后者包含需要進(jìn)一步推理的更復(fù)雜部分。論文在挑戰(zhàn)部分測(cè)試論文的模型。
表 6：Yuan 2.0-M32 與其他模型在 ARC-Challenge 上的比較
問(wèn)題和選項(xiàng)直接連接并用 分隔。 之前的文本發(fā)送給模型，模型預(yù)期生成一個(gè)標(biāo)簽或相應(yīng)的答案。生成的答案與真實(shí)答案進(jìn)行比較，結(jié)果使用 MC1 目標(biāo)計(jì)算。
表 6 顯示的結(jié)果 ARC-C 表明，Yuan 2.0-M32 在解決復(fù)雜科學(xué)問(wèn)題方面表現(xiàn)出色——它在這一基準(zhǔn)上超越了 Llama3-70B。
圖片
表 7：Yuan 2.0-M32 與其他模型在質(zhì)量與大小上的比較。平均準(zhǔn)確度是根據(jù) GSM-8K、Math、Humaneval、MMLU 和 ARC-C 的分?jǐn)?shù)平均得出的
論文將論文的性能與三種MoE模型（Mixtral家族、Deepseek）和六種密集模型（Qwen（Bai等，2023）、Llama家族和Phi-3家族（Abdin等，2024））進(jìn)行比較，以評(píng)估Yuan 2.0-M32在不同領(lǐng)域的性能。表7展示了Yuan 2.0-M32與其他模型在準(zhǔn)確度與計(jì)算量之間的比較。Yuan 2.0-M32僅使用3.7B活躍參數(shù)和每token 22.2 GFlops進(jìn)行微調(diào)，這是最經(jīng)濟(jì)的，以獲得與表中列出的其他模型相當(dāng)甚至超越的結(jié)果。表7暗示了論文模型在推理過(guò)程中的卓越計(jì)算效率和性能。Yuan 2.0-M32的平均準(zhǔn)確度為79.15，與Llama3-70B相當(dāng)。而平均準(zhǔn)確度/每token GFlops的值為10.69，是Llama3-70B的18.9倍。
論文標(biāo)題：Yuan 2.0-M32: Mixture of Experts with Attention Router
論文鏈接：https://arxiv.org/pdf/2405.17976
