尧图建网站 尧图建网站 YAOTU WEB BUILD 免费咨询
ARTICLE DETAIL

资讯详情

深耕网站建设与建站编程的一线实战洞察。

FFN、SwiGLU 、MoE

FFN、SwiGLU 、MoE FFN、SwiGLU 与 MoE1. FFNFFNFeed-Forward Network是 Transformer Block 中除 Attention 外的另一核心模块对每个 token 独立进行非线性特征变换。经典 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。不过这并不是强制要求在一些现代模型尤其是 MoE Expert 中也可能有dffdmodel d_{\text{ff}}d_{\text{model}}dff​dmodel​PyTorch 实现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 变体广泛应用于 LLaMA、Qwen 等现代大模型。核心公式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动态控制 (W_{\text{up}}x) 中哪些特征被保留。因此可以简单记为SwiGLU 带 SiLU 门控机制的 FFN。PyTorch 实现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。假设共有 (N) 个 ExpertE1,E2,…,EN E_1,E_2,\ldots,E_NE1​,E2​,…,EN​Router 首先根据 token hidden state (x) 计算各个 Expert 的路由分数rWrx rW_rxrWr​x经过 Softmaxpsoftmax⁡(r) p\operatorname{softmax}(r)psoftmax(r)其中pi p_ipi​表示当前 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 只经过其中两个 Expert而不是计算全部 8 个 Expert。MoE 的核心优势在于大量总参数少量激活参数 \boxed{ \text{大量总参数} \text{少量激活参数} }大量总参数少量激活参数​即模型可以拥有大量不同的 Expert 参数但一个 token 只需要计算其中少数几个 Expert。MoE 中的 Expert 本质上通常就是 FFN现代模型中也经常直接使用SwiGLU FFN作为 Expert。简单 PyTorch 实现下面以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 权重重新归一化topk_probstopk_probs/topk_probs.sum(dim-1,keepdimTrue)outputtorch.zeros_like(x)# 将 token 分配给对应 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 数8每个 token 激活2 个 Expert
返回列表