【LLM】深入理解MoE

1、为什么会有MoEDense Transformer 越来越大↓计算成本太高↓引入多个 Expert 提升容量↓Router 决定每个 token 去哪里↓Top-k 控制计算量↓loss中加入 Load Balance 保证所有 Expert 都能学到东西2、MoE的本质MoE本质上是Transformer 稀疏FFN其中只有一小部分 Expert 会被激活因此可以拥有数百亿甚至上万亿参数而计算量接近普通 Dense 模型。为什么替换的是FFN而不是Attention原因一FFN 是对每个 Token 独立进行非线性变换天然支持不同 Token 路由到不同 Expert而不会影响其他 Token原因二Attention 负责 Token 间的信息交互需要所有 Token 在统一表示空间计算 Q、K、V如果不同 Token 使用不同 Attention Expert会显著增加通信和实现复杂度原因三Transformer 中 FFN 通常占据更大的参数量因此将 FFN 扩展为多个 Expert可以显著增加模型容量同时通过 Top-k 路由保持计算成本基本不变。先回忆transformer干了什么Input↓Multi-Head Attention↓FFNFeed Forward↓Output每一层其实只有两件事Attention 和 FNN它们负责完全不同的工作。Attention负责 Token和 Token之间的信息交互例如对“我喜欢喝咖啡”处理到“咖啡”时它需要知道谁是主语 前面有没有修饰词“咖啡”指什么Attention做的是咖啡↓看看整句话其它 Token↓融合上下文所以Attention做的是Token Communication通信。Attention 输出以后每个 Token 已经拿到了上下文。FFN 不再看别人。咖啡↓只处理 咖啡 自己↓非线性变换↓得到新的表示所以 FFN 是Token Transformation特征变换。它完全是一个 Token → MLP → 一个Token。没有 Token 之间的信息交换。为什么MoE更适合替换FFN因为 FFN 天然就是 TokenA → FFNTokenB → FFN两者互不影响。因此完全可以让 TokenA → Expert3而 TokenB → Expert18互不干扰。Attention做不到吗理论上可以。实际上非常困难。因为 Attention 是Query、Key、Value所有 Token 两两计算。例如一句话 100 个 Token。Attention 要算 100 × 100。如果不同 Token 又走不同 Expert。例如Token1 → Attention Expert1Token2 → Attention Expert5。那么Token1 怎么和 Token2 通信马上就变复杂了。因为 Attention 本质要求所有 Token 在同一个空间里计算 Q、K、V。否则不同 Expert 之间 Q、K、V 根本没法对应。FFN则没有这个问题因为每一个Token都是独立的各自输出即可。Transformer 一层参数主要来自哪里很多人以为是 Attention其实不是。因为Transformer的FFN需要先升维再降维因为FFN做的是信息加工所以需要更大的中间空间来学习复杂的非线性变换提高表达能力。举个例子。假设 Hidden Size 4096FFN Expansion 4那么FFN4096 → 16384 → 4096参数量约4096 × 16384 × 2 ≈ 1.34 亿Attention 的 QKV4096 × 4096 × 3 ≈ 5000 万Output4096 × 4096 ≈1700 万合起来约 6700 万。所以FFN 参数几乎是 Attention 的 2 倍左右具体比例随模型设计有所变化。也就是说Transformer 最大头的参数其实就在 FFN。如果把 FFN 换成 64 个 Expert。参数马上变1.34 亿 → 85 亿但是一次只激活2个。计算量几乎没怎么变。这就是MoE 最赚的地方。可以把Transformer 的每一层理解成一个公司Attention 是前台负责和所有部门沟通。它必须统一工作方式所以不容易拆成很多独立专家。FFN 是后端专家每个员工拿到自己的任务后独立处理不需要再互相交流因此最适合扩展成很多专家。所以如果公司要扩招扩招前台大家还得一起协调收益有限扩招后端专家每个任务分给不同专家处理就能大幅增加能力而不显著增加每次工作的成本。3、MoE模型细节① 整体架构┌─────────────────────────────┐│ Input Tokens │└──────────────┬──────────────┘│Token Embedding│▼┌────────────────────────┐│ Transformer Block 1 ││ ││ Multi-Head Attention ││ │ ││ ▼ ││ MoE Layer │└────────────────────────┘│▼┌────────────────────────┐│ Transformer Block 2 ││ ││ Multi-Head Attention ││ │ ││ ▼ ││ MoE Layer │└────────────────────────┘│…│▼Final Hidden State│▼LM Head / Output可以看到Attention 仍然是普通 TransformerFFN 被 MoE Layer 替代因此很多论文都会说Dense Transformer Attention FFN而MoE Transformer Attention MoE FFN② 一个 MoE Layer 内部真正的创新发生在这里。Hidden State││┌──────▼──────┐│ Router ││ (Gate Network)│└──────┬──────┘│Top-k Routing│┌─────────────────────────────────┐│ │▼ ▼Expert 3 Expert 7(FFN) (FFN)▼ ▼ Output3 Output7 \ / \ / \ / ▼ ▼ Weighted Sum │ ▼ Layer Output这里有三个重要组件① RouterGate作用决定这个 Token 应该交给哪些 Expert。通常就是一个 Linear[score W_{gate}x]输出Expert1 : 0.03Expert2 : 0.61Expert3 : 0.12Expert4 : 0.20…然后取 Top-k。② Experts每个 Expert 都是一个完整 FFN。例如Expert1Linear↓GELU↓LinearExpert2Linear↓SiLU↓Linear实际上结构通常一样只是参数不同。例如一个 64 Expert 的模型Expert1Expert2Expert3…Expert64每个 Expert 都拥有自己的参数。③ Aggregation如果 Top-2例如Router:Expert5 0.72Expert18 0.28最终Output0.72×Expert5(x)0.28×Expert18(x)所以不是平均。而是Router 给权重。③ Token Routing 示意图例如一句话IlovecoffeetodayRouter 会这样分配ExpertsE1 E2 E3 E4 E5 E6I ●love ●coffee ●today ●如果 Top-2E1 E2 E3 E4 E5 E6I ● ●love ● ●coffee ● ●today ● ●所以每个 Token 激活不同 Expert。而不是一句话只选一个 Expert。④ DeepSeek / Mixtral 的 MoE例如Hidden│▼Router│┌──────────┴──────────┐│ │Top-2 Experts Others│▼Expert 12Expert 31│▼Weighted Combine│▼Output例如DeepSeek-V3256 Experts↓每个 Token↓只激活 8 个也就是说256个专家都存在但是一次推理只算8个所以虽然总参数巨大但 FLOPs 并没有增加很多。⑤ 与 Dense Transformer 对比一张完整的 MoE 架构图Input Tokens│Token Embedding│▼┌────────────────────┐│ Multi-Head Attention│└─────────┬──────────┘│▼LayerNorm│▼┌────────────────┐│ Router ││ (Gate Linear) │└────────┬────────┘│Top-k Experts Selected┌────────────┬────────────┬────────────┐▼ ▼ ▼Expert 5 Expert 12 Expert 31FFN FFN FFN│ │ │└──────┬─────┴──────┬──────┘▼ ▼Weighted Combine│▼Residual Add│▼Next Transformer Block这也是目前大多数 MoE 大语言模型如 DeepSeek 的 DeepSeek-V3、Mistral AI 的 Mixtral、Google 的 Switch Transformer 和 GShard的核心架构Attention 保持 Dense仅将 FFN 替换为由 Router 动态选择的稀疏 Expert 集合。