FFN、SwiGLU 與 MoE1. FFNFFNFeed-Forward Network是 Transformer Block 中除 Attention 外的另一核心模塊對每個 token 獨立進行非線性特征變換。經(jīng)典 FFNFFN(x)W2σ(W1x) \mathrm{FFN}(x)W_2\sigma(W_1x)FFN(x)W2?σ(W1?x)通常維度變化為dmodel→dff→dmodel d_{\text{model}} \rightarrow d_{\text{ff}} \rightarrow d_{\text{model}}dmodel?→dff?→dmodel?一般有dffdmodel d_{\text{ff}}d_{\text{model}}dff?dmodel?例如4096→16384→4096 4096\rightarrow16384\rightarrow40964096→16384→4096即先升維擴展特征再降回 hidden size。不過這并不是強制要求在一些現(xiàn)代模型尤其是 MoE Expert 中也可能有dffdmodel d_{\text{ff}}d_{\text{model}}dff?dmodel?PyTorch 實現(xiàn)importtorchimporttorch.nnasnnclassFFN(nn.Module):def__init__(self,d_model,d_ff):super().__init__()self.upnn.Linear(d_model,d_ff)self.activationnn.GELU()self.downnn.Linear(d_ff,d_model)defforward(self,x):xself.up(x)xself.activation(x)xself.down(x)returnx2. SwiGLUSwiGLU 可以看作一種帶門控機制的 FFN 變體廣泛應(yīng)用于 LLaMA、Qwen 等現(xiàn)代大模型。核心公式SwiGLU(x)Wdown[SiLU(Wgatex)⊙(Wupx)] SwiGLU(x) W_{\text{down}} \left[ \mathrm{SiLU}(W_{\text{gate}}x) \odot (W_{\text{up}}x) \right]SwiGLU(x)Wdown?[SiLU(Wgate?x)⊙(Wup?x)]相比普通 FFNSwiGLU 增加了一條 Gate 分支gSiLU(Wgatex) g\mathrm{SiLU}(W_{\text{gate}}x)gSiLU(Wgate?x)uWupx uW_{\text{up}}xuWup?x然后進行逐元素相乘hg⊙u hg\odot uhg⊙u最后再映射回模型維度yWdownh yW_{\text{down}}hyWdown?h其中包含三個線性投影Wgate,Wup,Wdown W_{\text{gate}}, \quad W_{\text{up}}, \quad W_{\text{down}}Wgate?,Wup?,Wdown?其中SiLU(Wgatex) \mathrm{SiLU}(W_{\text{gate}}x)SiLU(Wgate?x)可以理解為一個Gate動態(tài)控制 (W_{\text{up}}x) 中哪些特征被保留。因此可以簡單記為SwiGLU 帶 SiLU 門控機制的 FFN。PyTorch 實現(xiàn)importtorchimporttorch.nnasnnimporttorch.nn.functionalasFclassSwiGLU(nn.Module):def__init__(self,d_model,d_ff):super().__init__()self.gate_projnn.Linear(d_model,d_ff)self.up_projnn.Linear(d_model,d_ff)self.down_projnn.Linear(d_ff,d_model)defforward(self,x):gateF.silu(self.gate_proj(x))upself.up_proj(x)xgate*up xself.down_proj(x)returnx3. MoEMoEMixture of Experts混合專家可以理解為將 Transformer 中的一個 FFN 替換成多個 FFN Expert再通過 Router 為每個 token 選擇少量 Expert。假設(shè)共有 (N) 個 ExpertE1,E2,…,EN E_1,E_2,\ldots,E_NE1?,E2?,…,EN?Router 首先根據(jù) token hidden state (x) 計算各個 Expert 的路由分?jǐn)?shù)rWrx rW_rxrWr?x經(jīng)過 Softmaxpsoftmax?(r) p\operatorname{softmax}(r)psoftmax(r)其中pi p_ipi?表示當(dāng)前 token 被分配到第 (i) 個 Expert 的概率。隨后只選擇概率最高的 Top-(K) 個 ExpertE(x)TopK?(p) \mathcal E(x)\operatorname{TopK}(p)E(x)TopK(p)最終輸出y∑i∈E(x)piEi(x) y \sum_{i\in\mathcal E(x)} p_iE_i(x)yi∈E(x)∑?pi?Ei?(x)例如共有 8 個 Expert使用 Top-2 Routing則每個 token 只經(jīng)過其中兩個 Expert而不是計算全部 8 個 Expert。MoE 的核心優(yōu)勢在于大量總參數(shù)少量激活參數(shù) \boxed{ \text{大量總參數(shù)} \text{少量激活參數(shù)} }大量總參數(shù)少量激活參數(shù)?即模型可以擁有大量不同的 Expert 參數(shù)但一個 token 只需要計算其中少數(shù)幾個 Expert。MoE 中的 Expert 本質(zhì)上通常就是 FFN現(xiàn)代模型中也經(jīng)常直接使用SwiGLU FFN作為 Expert。簡單 PyTorch 實現(xiàn)下面以Top-2 MoE為例importtorchimporttorch.nnasnnimporttorch.nn.functionalasFclassExpert(nn.Module):def__init__(self,d_model,d_ff):super().__init__()self.gate_projnn.Linear(d_model,d_ff)self.up_projnn.Linear(d_model,d_ff)self.down_projnn.Linear(d_ff,d_model)defforward(self,x):gateF.silu(self.gate_proj(x))upself.up_proj(x)returnself.down_proj(gate*up)classMoE(nn.Module):def__init__(self,d_model,d_ff,num_experts8,top_k2):super().__init__()self.num_expertsnum_experts self.top_ktop_k# Routerself.routernn.Linear(d_model,num_experts)# Expertsself.expertsnn.ModuleList([Expert(d_model,d_ff)for_inrange(num_experts)])defforward(self,x):# x: [batch, seq_len, d_model]batch_size,seq_len,d_modelx.shape# 將 token 展平xx.reshape(-1,d_model)# [num_tokens, d_model]# Router 計算每個 Expert 的概率router_logitsself.router(x)router_probsF.softmax(router_logits,dim-1)# 選擇 Top-K Experttopk_probs,topk_indicestorch.topk(router_probs,kself.top_k,dim-1)# Top-K 權(quán)重重新歸一化topk_probstopk_probs/topk_probs.sum(dim-1,keepdimTrue)outputtorch.zeros_like(x)# 將 token 分配給對應(yīng) Expertforexpert_id,expertinenumerate(self.experts):forkinrange(self.top_k):mask(topk_indices[:,k]expert_id)ifmask.any():expert_inputx[mask]expert_outputexpert(expert_input)weighttopk_probs[mask,k].unsqueeze(-1)output[mask](weight*expert_output)returnoutput.reshape(batch_size,seq_len,d_model)例如moeMoE(d_model4096,d_ff1024,num_experts8,top_k2)Hidden Size4096每個 Expert 的中間維度1024總 Expert 數(shù)8每個 token 激活2 個 Expert FFN、SwiGLU 、MoE
FFN、SwiGLU 與 MoE1. FFNFFNFeed-Forward Network是 Transformer Block 中除 Attention 外的另一核心模塊對每個 token 獨立進行非線性特征變換。經(jīng)典 FFNFFN(x)W2σ(W1x) \mathrm{FFN}(x)W_2\sigma(W_1x)FFN(x)W2?σ(W1?x)通常維度變化為dmodel→dff→dmodel d_{\text{model}} \rightarrow d_{\text{ff}} \rightarrow d_{\text{model}}dmodel?→dff?→dmodel?一般有dffdmodel d_{\text{ff}}d_{\text{model}}dff?dmodel?例如4096→16384→4096 4096\rightarrow16384\rightarrow40964096→16384→4096即先升維擴展特征再降回 hidden size。不過這并不是強制要求在一些現(xiàn)代模型尤其是 MoE Expert 中也可能有dffdmodel d_{\text{ff}}d_{\text{model}}dff?dmodel?PyTorch 實現(xiàn)importtorchimporttorch.nnasnnclassFFN(nn.Module):def__init__(self,d_model,d_ff):super().__init__()self.upnn.Linear(d_model,d_ff)self.activationnn.GELU()self.downnn.Linear(d_ff,d_model)defforward(self,x):xself.up(x)xself.activation(x)xself.down(x)returnx2. SwiGLUSwiGLU 可以看作一種帶門控機制的 FFN 變體廣泛應(yīng)用于 LLaMA、Qwen 等現(xiàn)代大模型。核心公式SwiGLU(x)Wdown[SiLU(Wgatex)⊙(Wupx)] SwiGLU(x) W_{\text{down}} \left[ \mathrm{SiLU}(W_{\text{gate}}x) \odot (W_{\text{up}}x) \right]SwiGLU(x)Wdown?[SiLU(Wgate?x)⊙(Wup?x)]相比普通 FFNSwiGLU 增加了一條 Gate 分支gSiLU(Wgatex) g\mathrm{SiLU}(W_{\text{gate}}x)gSiLU(Wgate?x)uWupx uW_{\text{up}}xuWup?x然后進行逐元素相乘hg⊙u hg\odot uhg⊙u最后再映射回模型維度yWdownh yW_{\text{down}}hyWdown?h其中包含三個線性投影Wgate,Wup,Wdown W_{\text{gate}}, \quad W_{\text{up}}, \quad W_{\text{down}}Wgate?,Wup?,Wdown?其中SiLU(Wgatex) \mathrm{SiLU}(W_{\text{gate}}x)SiLU(Wgate?x)可以理解為一個Gate動態(tài)控制 (W_{\text{up}}x) 中哪些特征被保留。因此可以簡單記為SwiGLU 帶 SiLU 門控機制的 FFN。PyTorch 實現(xiàn)importtorchimporttorch.nnasnnimporttorch.nn.functionalasFclassSwiGLU(nn.Module):def__init__(self,d_model,d_ff):super().__init__()self.gate_projnn.Linear(d_model,d_ff)self.up_projnn.Linear(d_model,d_ff)self.down_projnn.Linear(d_ff,d_model)defforward(self,x):gateF.silu(self.gate_proj(x))upself.up_proj(x)xgate*up xself.down_proj(x)returnx3. MoEMoEMixture of Experts混合專家可以理解為將 Transformer 中的一個 FFN 替換成多個 FFN Expert再通過 Router 為每個 token 選擇少量 Expert。假設(shè)共有 (N) 個 ExpertE1,E2,…,EN E_1,E_2,\ldots,E_NE1?,E2?,…,EN?Router 首先根據(jù) token hidden state (x) 計算各個 Expert 的路由分?jǐn)?shù)rWrx rW_rxrWr?x經(jīng)過 Softmaxpsoftmax?(r) p\operatorname{softmax}(r)psoftmax(r)其中pi p_ipi?表示當(dāng)前 token 被分配到第 (i) 個 Expert 的概率。隨后只選擇概率最高的 Top-(K) 個 ExpertE(x)TopK?(p) \mathcal E(x)\operatorname{TopK}(p)E(x)TopK(p)最終輸出y∑i∈E(x)piEi(x) y \sum_{i\in\mathcal E(x)} p_iE_i(x)yi∈E(x)∑?pi?Ei?(x)例如共有 8 個 Expert使用 Top-2 Routing則每個 token 只經(jīng)過其中兩個 Expert而不是計算全部 8 個 Expert。MoE 的核心優(yōu)勢在于大量總參數(shù)少量激活參數(shù) \boxed{ \text{大量總參數(shù)} \text{少量激活參數(shù)} }大量總參數(shù)少量激活參數(shù)?即模型可以擁有大量不同的 Expert 參數(shù)但一個 token 只需要計算其中少數(shù)幾個 Expert。MoE 中的 Expert 本質(zhì)上通常就是 FFN現(xiàn)代模型中也經(jīng)常直接使用SwiGLU FFN作為 Expert。簡單 PyTorch 實現(xiàn)下面以Top-2 MoE為例importtorchimporttorch.nnasnnimporttorch.nn.functionalasFclassExpert(nn.Module):def__init__(self,d_model,d_ff):super().__init__()self.gate_projnn.Linear(d_model,d_ff)self.up_projnn.Linear(d_model,d_ff)self.down_projnn.Linear(d_ff,d_model)defforward(self,x):gateF.silu(self.gate_proj(x))upself.up_proj(x)returnself.down_proj(gate*up)classMoE(nn.Module):def__init__(self,d_model,d_ff,num_experts8,top_k2):super().__init__()self.num_expertsnum_experts self.top_ktop_k# Routerself.routernn.Linear(d_model,num_experts)# Expertsself.expertsnn.ModuleList([Expert(d_model,d_ff)for_inrange(num_experts)])defforward(self,x):# x: [batch, seq_len, d_model]batch_size,seq_len,d_modelx.shape# 將 token 展平xx.reshape(-1,d_model)# [num_tokens, d_model]# Router 計算每個 Expert 的概率router_logitsself.router(x)router_probsF.softmax(router_logits,dim-1)# 選擇 Top-K Experttopk_probs,topk_indicestorch.topk(router_probs,kself.top_k,dim-1)# Top-K 權(quán)重重新歸一化topk_probstopk_probs/topk_probs.sum(dim-1,keepdimTrue)outputtorch.zeros_like(x)# 將 token 分配給對應(yīng) Expertforexpert_id,expertinenumerate(self.experts):forkinrange(self.top_k):mask(topk_indices[:,k]expert_id)ifmask.any():expert_inputx[mask]expert_outputexpert(expert_input)weighttopk_probs[mask,k].unsqueeze(-1)output[mask](weight*expert_output)returnoutput.reshape(batch_size,seq_len,d_model)例如moeMoE(d_model4096,d_ff1024,num_experts8,top_k2)Hidden Size4096每個 Expert 的中間維度1024總 Expert 數(shù)8每個 token 激活2 個 Expert