Mixtral 8x7B 的推出, 使我們開始更多地關注 基于MoE 的大模型架構, 那么,什么是MoE呢?
1. MoE溯源
MoE的概念起源于 1991 年的論文 Adaptive Mixture of Local Experts(https://www.cs.toronto.edu/~hinton/absps/jjnh91.pdf)。這個概念與集成學習方法相似,旨在為由多個單獨網絡組成的系統并建立一個監管機制。在這種系統中,每個網絡處理訓練樣本的不同子集,專注于輸入空間的特定區域。
后來,論文(https://arxiv.org/abs/1312.4314)探索了將 MoE 作為更深層網絡的一個組件。這種方法允許將 MoE 嵌入到多層網絡中的某一層,使得模型既大又高效。而且, 由研究開始探索基于輸入令牌動態激活或停用網絡組件的方法。2017年,論文(https://arxiv.org/abs/1701.06538)將這一概念應用于 137B 的 LSTM ,通過引入稀疏性,這項工作在保持高規模的同時實現了快速的推理速度。
總之,MoE 的引入使得訓練具有數千億甚至萬億參數的模型成為可能。
2.什么是MoE?
從溯源中可以看到,MoE 的基本思想在于“人多力量大”,這意味著利用多個專家網絡,每個專家網絡專門處理數據的不同方面,以獲得更好的性能。其目的是創建一個動態系統,在這個系統中,可以根據輸入數據利用各種專家的優勢,從而比單一模型能夠實現的預測更加靈活和準確。
在技術上,MoE是一種集成學習方法,把一個大任務分解成更小的部分,讓不同的專家處理每個部分。然后,有一個聰明的決策者,會根據情況決定采納哪位專家的建議,所有這些建議都被混合在一起。
圖片
這種方法有四個主要步驟:
1. 把大問題分成小塊;
2. 為每一塊訓練一個超級聰明的專家。
3. 引入一個決策者,也就是眾所周知的門控模型,來選擇應該由哪位專家帶頭。
4. 收集專家的意見和決策者的選擇,提出最終的預測。
典型地,MoE 是一類transformer模型, 使用“稀疏”方法,其中每個輸入只使用模型組件的一個子集。這種設置允許更有效的預訓練和更快的推理,同時管理一個更大的模型大小。每個專家模型一般都是一個神經網絡,通常是一個前饋網絡(FFN) ,專門處理輸入數據的不同方面,使模型能夠更有效地處理范圍更廣的任務。
圖片
MoE的優勢在于更少的計算工作使得模型的預訓練更快,與為更多步驟訓練更小的模型相比,為更少步驟訓練更大的模型可以產生更好的結果。在推理過程中,與其他大模型相比,參數數目相同的 MoE 模型表現出更快的推理速度。盡管有許多參數,但只使用了一個子集,從而導致更快的預測。
3. MoE的組成
在MoE系統中,傳統 Transformer 模型中的每個前饋網絡 (FFN) 層替換為 MoE 層,其中 MoE 層由兩個核心部分組成: 一個門控網絡和若干數量的專家。
同時,token在封裝輸入數據中的特征和信息方面發揮著重要作用。token表示為高維向量,作為一種綜合的結構化格式,用于編碼輸入的相關特征。每個token由多個維度組成,每個維度表示數據的特定特征或方面。通過以高維向量格式組織信息,系統可以有效地捕捉輸入數據中的復雜性和細微差別,為進一步分析提供豐富的數據來源。
token的高維表示使系統內的專家能夠從數據中提取詳細的見解和模式。通過將token作為高維向量處理,專家模型可以利用先進的算法和技術來發現隱藏的相關性和關系,提高MoE系統的整體分析能力。
3.1 專家模型
專家模型是針對較小問題創建的模型,不僅僅局限于神經網絡,也可以是任何類型的模型。對每個專家模型進行訓練,使其能夠在輸入空間的指定區域內專門預測數據點。當然,可以神經網絡來同時扮演決策者和專家的角色,這時一般會稱為“混合密度網絡”。所有這些專家都得到了相同的數據集 ,都會嘗試在此基礎上做出預測。
“專家”通常指的是模型的一個組成部分,專注于數據中特定類型的任務或模式(如代碼生成、推理、匯總) 。其專業化程度取決于他們接受的訓練數據和模型本身的結構,更多的是關于計算任務的本質(例如,識別特定的模式,處理特定類型的輸入) ,而不是特定領域的知識。
MoE 模型中,每個專家模型一般是一個經過訓練的神經網絡,能夠很好地處理整個數據或任務的子集。這些專家之前的共同骨干網絡充當特征提取器,將原始輸入轉化為專家可以更有效地工作的高級表示。這種設置允許每個專家關注它最適合處理的模式或數據特征,從而實現更加細致和準確的預測。
3.2 門控模型
在給定的情況下,我們需要一個決策者來決定遵循哪位專家的建議,這個決策者被稱為門控模型。門控模型決定了對于給定的輸入應該使用哪個專家模型的預測。門控模型學習將輸入路由到最合適的專家當門控模型是一個神經網絡的時候, 有時也稱為門控網絡。門控網絡基本上起到決策者的作用,根據不同專家的熟練程度和與輸入數據的相關性為他們分配權重。也就是說,門控網絡根據給出的信息動態調整專家模型的權重。
圖片
門控網絡是 MoE 區別于一個簡單的模型集合的關鍵。它是一個可訓練的組件,學會根據輸入數據為每個專家的輸出分配適當的注意力(或權重)。這種分配不是任意的; 它是由門控網絡對哪位專家最有可能為給定輸入產生正確輸出的評估決定的。門控機制有效地引導集合,將輸入導向最相關的專家,并將他們的輸出合成為一個內聚預測。這種動態路由功能允許 MoE 模型具有高度的適應性和高效性,并為每個任務利用最適當的資源。
當使用神經網絡模型時,門控網絡和專家模型一起訓練,這種訓練就像是教他們如何像一個團隊一樣工作。過去,用一種叫做“期望最大化”的方法來解決這個問題。門控網絡可能會給每個專家一個看起來像概率的分數,表明它對每個專家的預測有多么信任。
隨著Softmax和噪聲top-k 門機制的引入,MoE模型有了長足的進展。Softmax是一種流行的選擇,以確定每個專家模型的貢獻模型的預測。它根據每個專家與輸入數據的相關性為其分配權重,較高的權重給予更相關的專家。噪聲 Top-k 門控機制通過選擇一個Top-k 專家子集進行預測來引入隨機性。這種隨機元素有助于在專家選擇過程中引入多樣性,提高模型的魯棒性和泛化能力。
4. MoE 的實現
實現MoE系統的關鍵在于體系結構的設計,其中一個比較前沿的方法涉及到利用Switch Transformer結構來集成混合的專家模型層。該體系結構使系統能夠利用多個專家模型在作出決定和預測方面的綜合專門知識。
圖片
在構建Switch Transformer體系結構時,必須保證專家混合層的無縫集成。這些層負責從不同的角度處理和分析數據,使系統能夠根據各種不同的見解作出明智的決定。典型的實現步驟如下:
- 結合Switch Transformer框架,設計MoE系統的總體架構;
- 將專家模型混合層集成到系統中,以支持并行的數據處理;
- 實現根據上下文和信息類型將數據路由到適當專家的機制;
- 優化訓練過程,微調專家權重和切換機制。
實現MoE系統的另一個關鍵點是動態路由和負載均。此策略確保token在整個系統中有效分布,優化性能和資源利用率。動態路由涉及到根據工作量和信息的復雜性實時分配任務和數據給不同的專家模型。通過動態路由token,系統能夠適應不斷變化的需求,在運行過程中保持最佳效率。負載均衡對于在專家之間均勻分配計算任務和數據處理起著至關重要的作用。這可以防止瓶頸和資源浪費,最終提高系統的整體性能。一般的策略實現方式如下:
- 監控每個專家模型的工作負載,并動態地重新分發token以保持平衡。
- 實現基于緊迫性和復雜性對任務進行優先排序的算法,以優化處理效率。
- 利用反饋機制不斷調整負載均衡參數。
圖片
下面這些開源項目可以用于訓練MoE:
- Megablocks: https://github.com/stanford-futuredata/megablocks
- Fairseq: https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm
- OpenMoE: https://github.com/XueFuzhao/OpenMoE
5. MoE 的訓練
MoE模型的訓練是一個微妙的過程,其中每個專家在特定類型的推理中變得更好。它不像直接在特定任務或領域中訓練一個專家那樣簡單。
5.1 MoE 系統訓練的特點
由于模型的體系結構和它所暴露的數據組合,MoE 模型中專家的專業化通常在訓練過程中自然出現,特點如下:
- 多樣化數據: 該模型是在包含廣泛任務或領域的多樣化數據集上進行訓練的。這種多樣性至關重要,因為它使每個專家面臨不同類型的數據和問題。
- 門控機制: MoE 模型有一個門控機制,它決定哪個專家處理輸入數據的哪個部分。在訓練期間,這個門控網絡學會根據不同專家模型的專長向他們發送不同類型的數據。
- 專家模型: 隨著訓練的進行,每個專家模型逐漸變得更加擅長處理特定類型的數據或任務。發生這種專門化是因為專家接收并學習他們在處理數據時最有效的數據類型。
- 反饋回路: 有一個反饋回路在起作用的時候,如果一個專家模型在某種類型的數據方面做得更好,門控機制將更有可能向該專家發送類似的數據。這加強了每個專家的專業化。
- 正規化和損失函數: 訓練過程通常包括正規化技術和專門的損失函數,以鼓勵有效的學習,避免出現一個專家成為“萬事通”的情況,從而確保專業化的分布。
- 能力限制: 通過對模型施加能力限制,該模型確保沒有任何一位專家任務負擔過重,促進所有專家之間學習的平衡分配。
- 微調和調整: 模型可能經歷微調階段,其中強調某些類型的任務,進一步完善每個組件的專門知識。指令調優已經成為超越傳統稠密模型性能的關鍵策略。通過優化模型中每個專家執行的指令,以在不增加模型復雜性的情況下獲得更高的精度和效率。
專家模型通常不會只專注于一個狹窄的任務類型,“專業知識”應該被理解為某些類型的數據或任務的效率或有效性的相對提高,而不是嚴格的領域限制。
5.2 聯合訓練
優化MoE模型的一個重要策略是對門控網絡與其他組件進行聯合訓練。在聯合訓練過程中,門控網絡的參數通過整個模型的反向傳播進行更新。這個過程允許門控網絡根據從專家模型收到的反饋來調整其路由決策。該模型通過對所有組件進行集成優化,可以更好地平衡不同專家的貢獻,優化路由機制,達到最優結果。
5.3 指令調優
指令調優方法同樣可用于訓練MoE系統。這些方法著重于調整路由機制和專家模型,以提高模型在不同任務和輸入數據分布中的性能。常見的微調策略包括:
- 軟路由: 軟路由技術使用連續概率來分配每個專家模型對最終輸出的貢獻。通過將路由決策建模為概率,特別是在多個專家模型可能提供互補信息的情況下,該模型可以執行更平滑和更靈活的路由。
- 硬路由: 與軟路由不同,硬路由涉及為給定輸入選擇單個專家模型的離散路由決策。硬路由策略更易于實現和解釋,因此適用于需要顯式專家選擇的任務。
- 正則化技術: 正則化方法,如 L1或 L2正則化可以應用到路由參數,以防止過擬合并提高泛化能力。通過懲罰過于復雜的路由決策,正則化可以鼓勵模型關注更健壯和可解釋的路由策略。
- 自適應路由: 自適應路由機制根據輸入數據和模型的當前狀態動態調整路由概率或決策。自適應路由策略使模型能夠根據數據中不斷變化的模式調整自己的路由行為,從而實現更加自適應和高效的專家選擇。
6. MoE 的典型應用
MoE 架構已經被用來為語言模型建立大型的、計算效率高的神經網絡。它們可以處理計算模式或任務,如代碼生成、推理和匯總。MoE已經應用于強化學習問題,每個專家都可以專門處理環境的不同方面。在遷移學習中,MoE可以用來將在一個領域學到的知識應用到另一個相關領域。
另外,MoE正在個性化推薦系統方面取得長足進展,它們可以通過讓不同的專家參與不同類型的內容或用戶簡介來迎合用戶的不同偏好。這種基于輸入數據進行適應和專門化的能力使得 MoE 對于未來的應用程序特別有吸引力。
舉例而言,谷歌的 Gemini 1.5采用了MoE架構,這種架構使Gemini 1.5更有效的訓練和服務。它將您的請求路由到一組較小的“專家”神經網絡,這樣響應速度更快,質量更高。Gemini 1.5 Pro,這是一個中等大小的多模態模型,針對不同任務的縮放進行了優化。它的性能與1.0 Ultra 相當,后者可能是迄今為止谷歌最重要的模型。Gemini 1.5 Pro 還在長語境理解方面引入了一個突破性的實驗特性。它提供了一個標準的128,000 token 上下文窗口,但是有限的開發人員和企業客戶可以使用最多100萬token的上下文窗口來進行嘗試。
另一個應用示例是Mixtral 8x7B ,它使用稀疏混合專家架構(SMoE)。它的架構與Mixtral 7B 相同,不同之處在于每一層包含八個前饋網絡。對于每一個token,在每一層,門控網絡選擇兩個專家來處理當前狀態并合并他們的輸出。因此,每個令牌都可以訪問47B 參數,但是在推理期間只使用13B 活動參數。Mistral 8x7B 在所有評估的基準測試中都優于或匹配 Llama 270B 和 GPT-3.5。特別爹,Mixtral 在數學、代碼生成和多語言基準測試方面遠遠超過了Llama 270B 。
圖片
7. MoE 的簡單示例
MoE 模型通常用于處理復雜的數據分布,并根據不同的專家子模型進行預測。下面使用 TensorFlow/Kera 實現一個基本的 MoE 模型步驟。
(1)定義專家模型, 可以是任何典型的神經網絡結構
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import Dense, Input, Lambda, Layer, Softmax
from tensorflow.keras.models import Model
def create_expert_model(input_dim, output_dim):
inputs = Input(shape=(input_dim,))
x = Dense(64, activatinotallow='relu')(inputs)
x = Dense(32, activatinotallow='relu')(x)
outputs = Dense(output_dim, activatinotallow='softmax')(x)
model = Model(inputs=inputs, outputs=outputs)
return model
(2)定義門控網絡,該網絡接受輸入并輸出每個專家的門控系數(權重),這些系數決定了每個專家對最終預測的貢獻。
def create_gating_network(input_dim, num_experts):
inputs = Input(shape=(input_dim,))
x = Dense(32, activatinotallow='relu')(inputs)
x = Dense(num_experts, activatinotallow='softmax')(x)
outputs = x
model = Model(inputs=inputs, outputs=outputs)
return model
(3)創建 MoE 模型,使用門控機制來計算專家模型預測的加權和
def create_moe_model(input_dim, output_dim, num_experts):
input_layer = Input(shape=(input_dim,))
expert_models = [create_expert_model(input_dim, output_dim) for _ in range(num_experts)]
gating_network = create_gating_network(input_dim, num_experts)
expert_outputs = [expert(input_layer) for expert in expert_models]
gating_coefficients = gating_network(input_layer)
def moe_function(args):
expert_outputs, gating_coefficients = args
return tf.reduce_sum(expert_outputs * tf.expand_dims(gating_coefficients, axis=-1), axis=1)
moe_output = Lambda(moe_function)([expert_outputs, gating_coefficients])
model = Model(inputs=input_layer, outputs=moe_output)
return model
(4)選擇適當的損失函數,編譯模型并在數據集上對其進行訓練
input_dim = X_train.shape[1]
output_dim = len(np.unique(y_train))
num_experts = 5
moe_model = create_moe_model(input_dim, output_dim, num_experts)
moe_model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
moe_model.fit(X_train, y_train, epochs=10, batch_size=32)
一句話小結
在復雜的門控機制指導下,智能地結合多個專家網絡,基于MoE 系統的大模型提供了強大的靈活性、效率和準確性。
【參考文獻】
- Mixtral of Experts,https://arxiv.org/abs/2401.04088
- Improving Expert Specialization in Mixture of Experts,https://arxiv.org/abs/2302.14703
- Merging Experts into One: Improving Computational Efficiency of Mixture of Experts,https://arxiv.org/abs/2310.09832
- https://huggingface.co/blog/zh/moe
- https://blog.google/technology/ai/google-gemini-next-generation-model-february-2024/