此博客为Datawhale 组队学习打卡笔记

手搓一个大模型

Llama2 架构总览

• Decoder-Only 架构。完全由 Transformer 解码器层堆叠而成
• 预归一化（Pre-Normalization）：与经典 Transformer 的后归一化不同，输入在进入注意力层和前馈网络之前，都会先经过一次 RMS Norm。这被认为是提升大模型训练稳定性的关键
• 组件升级：支持 Grouped-Query Attention（GQA），前馈网络采用 SwiGLU，归一化使用 RMSNorm。
• 旋转位置编码（RoPE）：位置信息并非在输入端与词嵌入相加，而是在注意力层内部，通过 RoPE 操作动态地施加于QK向量之上
• 残差连接

MoE 架构解析

• 如果我们尝试让一个网络同时学习多个截然不同的子任务（例如既学做菜又学修车），往往会出现强干扰效应
• 基于“分治（Divide and Conquer）”策略的系统架构
  • 专家网络：系统包含多个独立的神经网络（可以是简单的前馈网络）。每个专家不再需要处理全局任务，只需专注于输入空间中的一个局部区域或一类特定的子任务。
  • 门控网络：充当协调者的角色。它接收与专家相同的输入 x，并输出一组混合比例（Mixing Proportions）pi，即选择每个专家的概率。它就像一个软性的随机开关，决定当前的输入案例应该由哪位专家来主导处理。

损失函数设计

设计了特定的机制来鼓励专家之间的竞争（Competition），而非合作。如果仅仅将所有专家的输出简单线性相加来逼近目标 $\mathbf{y}$，即预测值 $\hat{\mathbf{y}}=\sum p_i{\mathbf{E}}_i$，并在最终输出上计算误差，专家们会倾向于“合作”。
$$Los{s}_{coop}=||\mathbf{y}-\sum _i{p}_i{\mathbf{E}}_i|{|}^2$$

在这种合作模式下，为了减小总误差，每个专家都会试图去弥补其他专家的残差。继而导致每个案例都牵动所有专家的权重，依然无法解决“干扰效应”。为了实现解耦，作者提出将系统视为一个随机生成器，并采用负对数似然作为目标函数：

$$Los{s}_{comp}=-\log \sum _i{p}_i{e}^{-\frac{1}{2}||\mathbf{y}-{\mathbf{E}}_i|{|}^2}$$

其中：

• $\mathbf{y}$ 是我们希望模型输出的真实目标。
• ${\mathbf{E}}_i$ 是第 $i$ 个专家的输出。
• $p_i$ 是门控网络分配给第 $i$ 个专家的权重（概率）。

在这个目标函数中，系统倾向于**“赢家通吃”**。当某个专家 ${\mathbf{E}}_i$ 的输出非常接近目标 $\mathbf{y}$ 时（即误差项 $||\mathbf{y}-{\mathbf{E}}_i|{|}^2$ 很小），指数项 $e^{-\dots }$ 会变得很大。为了最小化总 Loss，门控网络会倾向于显著增加这个“表现好”的专家的权重 $p_i$，而忽略其他专家。

深度神经网络中的 MoE

• MoE 不再是一个孤立的架构，而成为了一种可插拔的层。我们可以在一个深层网络的不同位置（例如第 1 层和第 2 层）分别插入 MoE 模块，每一层都有自己独立的门控网络和专家集合。
• 指数级增长的组合路径：通过这种堆叠，网络能够表达的有效“专家组合”数量呈指数级增长。如果第一层有 $N$ 个专家，第二层有 $M$ 个专家，那么网络潜在的组合路径就高达 $N\times M$ 种。每个输入样本都会根据其特性，动态地选择一条最适合的处理路径。

稀疏门控 MoE

• 条件计算：能否在不增加计算量的前提下，大幅增加模型的参数量？
• 稀疏门控混合专家层（Sparsely-Gated MoE Layer），可以在每个样本处理过程中只激活网络的一小部分。
• 稀疏性与负载均衡
  • 引入噪声的稀疏门控：在门控输入中加入可训练的高斯噪声，计算后仅保留权重最大的 k 个专家
  • 避免“赢家通吃”的负载均衡：在总损失函数中加入了额外的辅助损失（Auxiliary Loss），包含 Importance Loss 和 Load Loss。这些损失函数并不直接服务于预测准确率，而是专门用来惩罚“分配不均”的现象，

大模型时代的 MoE

MoE 层的构建与门控机制

在 GShard 中，MoE 层的应用变得更加标准化:

• 保留 Attention 不变
• 替换 FFN：将 Transformer Block 中的前馈神经网络替换为 MoE 层。
• 替换 FFN：将 Transformer Block 中的前馈神经网络替换为 MoE 层。
• Top-2 Gating（Top-2 门控） 策略
  • 第一专家：选择权重最高的专家，保证主要任务的处理。
  • 第二专家：根据权重概率随机选择或直接选择第二高权重的专家，引入辅助处理和一定的随机性，有助于负载均衡。

Switch Transformer 提出了 Switch Layer

• 用稀疏的 Switch FFN 层替换了标准 Transformer 中的稠密 FFN 层
• Top-1 Routing（单专家路由）
  • 减少路由计算
  • 降低通信成本
  • 减小专家批量：每个专家需要处理的 Token 数量（Expert Capacity）至少减半

参考

[1] https://datawhalechina.github.io/base-llm