解析SwitchHead的Mixture-of-Experts注意力机制

加速Transformer：深入解析SwitchHead的Mixture-of-Experts注意力机制

在Transformer模型的快速发展中，计算和内存资源的瓶颈一直是研究的重点。Róbert Csordás等人在NeurIPS 2024发表的论文《SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention》提出了一种创新的Mixture-of-Experts（MoE）注意力机制——SwitchHead，不仅显著降低了计算和内存需求，还保持了与标准Transformer相当的语言建模性能。本文将面向Transformer研究者，详细解析SwitchHead的核心思想、数学公式及其实现细节，帮助大家理解其背后的设计逻辑。

Paper: https://arxiv.org/pdf/2312.07987

---

背景：为何需要SwitchHead？

Transformer模型以其强大的语言建模能力闻名，但其计算复杂度随着序列长度和模型规模的增加呈二次甚至更高次增长，尤其是在自注意力层（self-attention）。传统MoE方法主要应用于前馈层（MLP），而对注意力层的优化尝试往往无法兼顾性能和效率。SwitchHead通过将MoE思想引入注意力层，减少了注意力矩阵的计算量，同时保持模型的表达能力。

核心亮点：

• 计算效率：SwitchHead最多可减少8倍的注意力矩阵计算量。
• 内存节省：在262M参数模型上，仅用27%的内存即可匹配标准Transformer的性能。
• 性能匹配：在C4数据集上，SwitchHead的困惑度（perplexity）与基线相当，并在下游任务（如BLiMP）上提升了3.5%以上。

---

SwitchHead的核心思想

SwitchHead的核心是将注意力层的投影（key、query、value和output）通过MoE机制进行条件计算，减少冗余计算，同时保留注意力机制的表达能力。以下是其设计的关键点：

1. 减少注意力头数：传统Transformer使用多个注意力头（heads），每个头计算一个注意力矩阵。SwitchHead通过减少头数（$n_{\text{heads}}$），并为每个头引入多个专家（experts），大幅降低计算量。
2. MoE投影：对value和output投影应用MoE，允许动态选择专家，减少无效计算。
3. 非竞争性选择：使用sigmoid激活函数进行专家选择，避免专家坍塌问题，无需额外正则化。

---

数学公式解析

SwitchHead的数学公式较为复杂，但理解它们是掌握其核心的关键。以下逐步拆解公式，并解释其意义。

1. 标准Transformer注意力机制

标准Transformer的自注意力机制为每个头 $h$ 计算注意力矩阵 ${\mathbf{A}}^h$，并通过投影生成输出。公式如下（参考论文Eq. 3）：

$$\mathbf{y}[t,c]=\sum _{h=1}^{n_{\text{heads}}}\left({\mathbf{A}}^h{\mathbf{V}}^h{\mathbf{W}}_O^h\right)[t,c]$$

• $\mathbf{y}\in {\mathbb{R}}^{T\times d_{\text{model}}}$：输出矩阵，$T$ 是序列长度，$d_{\text{model}}$ 是模型维度。
• ${\mathbf{A}}^h\in {\mathbb{R}}^{T\times T}$：第 $h$ 个头的注意力矩阵，通常由 ${\mathbf{Q}}^h{\mathbf{K}}^{hT}$ 经softmax归一化得到。
• ${\mathbf{V}}^h\in {\mathbb{R}}^{T\times d_{\text{head}}}$：value投影，${\mathbf{W}}_O^h\in {\mathbb{R}}^{d_{\text{head}}\times d_{\text{model}}}$ 是output投影矩阵。
• $n_{\text{heads}}$：注意力头数。

计算复杂度：

• 注意力矩阵计算：$O(n_{\text{heads}}{d}_{\text{head}}{T}^2)$ MACs（乘累加操作）。
• 内存需求：存储 ${\mathbf{A}}^h$ 需要 $O(n_{\text{heads}}{T}^2)$，投影中间结果需要 $O(n_{\text{heads}}T{d}_{\text{head}})$。

2. SwitchHead的MoE机制

SwitchHead通过引入MoE机制，动态选择专家来计算投影矩阵，从而减少计算量。以下是其核心公式（参考论文Eq. 7-10）。

2.1 专家选择

SwitchHead为每个头 $h$ 定义 $E$ 个专家，动态选择 $k$ 个活跃专家。选择过程分为源端（source，涉及key和value）和目标端（destination，涉及query和output）。

源端专家选择：

$${\mathbf{s}}_S^h=\sigma \left(\mathbf{x}{\mathbf{W}}_S^h\right)$$

E S h = arg ⁡ topk ⁡ ( s S h , k ) , E S h ⊂ { 1 , … , E } \mathcal{E}_S^h = \arg \operatorname{topk} \left( \boldsymbol{s}_S^h, k \right), \quad \mathcal{E}_S^h \subset \{1, \ldots, E\} ESh​=argtopk(sSh​,k),ESh​⊂{1,…,E}

• $\mathbf{x}\in {\mathbb{R}}^{T\times d_{\text{model}}}$：输入序列。
• ${\mathbf{W}}_S^h\in {\mathbb{R}}^{d_{\text{model}}\times E}$：源端选择矩阵。
• $\sigma$：sigmoid激活函数，生成非竞争性得分 ${\mathbf{s}}_S^h$。
• $\arg \mathrm{topk}$：选择得分最高的 $k$ 个专家，${\mathcal{E}}_S^h$ 是选中的专家索引集合。

目标端专家选择类似：

$${\mathbf{s}}_D^h=\sigma \left(\mathbf{x}{\mathbf{W}}_D^h\right),{\mathcal{E}}_D^h=\arg \mathrm{topk}\left({\mathbf{s}}_D^h,k\right)$$

2.2 投影计算

SwitchHead对value和output投影应用MoE，key和query投影保持单一（非MoE），以优化参数分配。value投影公式如下：

$${\mathbf{V}}^h=\sum _{e\in {\mathcal{E}}_S^h}{\mathbf{s}}_S^h[e]\mathbf{x}{\mathbf{W}}_V^{h,e}$$

• ${\mathbf{W}}_V^{h,e}\in {\mathbb{R}}^{d_{\text{model}}\times d_{\text{head}}}$：第 $h$ 个头、第 $e$ 个专家的value投影矩阵。
• ${\mathbf{s}}_S^h[e]$：第 $e$ 个专家的权重，来自sigmoid得分。
• ${\mathbf{V}}^h\in {\mathbb{R}}^{T\times d_{\text{head}}}$：加权后的value投影。

output投影类似：

$$\mathbf{y}=\sum _{h=0}^{n_{\text{heads}}-1}\sum _{e\in {\mathcal{E}}_D^h}{\mathbf{s}}_D^h[e]{\mathbf{A}}^h{\mathbf{V}}^h{\mathbf{W}}_O^{h,e}$$

key和query投影不使用MoE，直接计算：

$${\mathbf{K}}^h=\mathbf{x}{\mathbf{W}}_K^h,{\mathbf{Q}}^h=\mathbf{x}{\mathbf{W}}_Q^h$$

2.3 注意力矩阵计算

注意力矩阵 ${\mathbf{A}}^h$ 的计算与标准Transformer类似，但由于 $n_{\text{heads}}$ 显著减少，总体计算量降低：

$${\mathbf{A}}^h=\text{softmax}\left(\frac{{\mathbf{Q}}^h{\mathbf{K}}^{hT}}{\sqrt{d_{\text{head}}}}\right)$$

3. 计算复杂度分析

SwitchHead的计算复杂度公式（参考论文Eq. 13）为：

$$N_{\text{MAC}}^{\text{SwitchHead}}=n_{\text{heads}}\left(2T{d}_{\text{head}}{d}_{\text{model}}+2Tk{d}_{\text{head}}(d_{\text{model}}+1)+2C{T}^2{d}_{\text{head}}+2CT{d}_{\text{head}}{d}_{\text{model}}\right)$$

• 对比标准Transformer（Eq. 11）：

$$N_{\text{MAC}}^{\text{XL}}=n_{\text{heads}}\left(4T{d}_{\text{head}}{d}_{\text{model}}+2C{T}^2{d}_{\text{head}}+2CT{d}_{\text{head}}{d}_{\text{model}}\right)$$

SwitchHead通过减少 $n_{\text{heads}}$ 和引入 $k$ 个专家，显著降低了计算量。例如，$n_{\text{heads}}$ 从16减少到2，$k=2$ 时，注意力矩阵计算量减少约8倍。

• 内存需求：SwitchHead的内存需求为：

$$N_{\text{mem}}^{\text{SwitchHead}}=n_{\text{heads}}\left(4T{d}_{\text{head}}+2C{T}^2+2CT{d}_{\text{head}}\right)$$

相比标准Transformer的 $N_{\text{mem}}^{\text{XL}}$，由于 $n_{\text{heads}}$ 减少，内存占用大幅降低。

---

实验结果与分析

1. 性能对比

SwitchHead在多个数据集（C4、Wikitext 103、peS2o、Enwik8）上测试，模型规模包括47M和262M参数。以下是关键结果（参考Table 2）：

• C4数据集（262M参数）：

  • SwitchHead：困惑度16.23，MACs 2.4G，内存5.6M。
  • 标准Transformer（16头）：困惑度16.28，MACs 5.4G，内存21M。
  • SwitchHead仅用44%的计算量和27%的内存，性能几乎相同。

• 下游任务（Table 4）：

  • 在BLiMP上，SwitchHead（262M）达到79.6%准确率，比基线（76.1%）提升3.5%。

2. SwitchAll：全MoE Transformer

SwitchHead可与MoE前馈层结合，形成“SwitchAll”模型（Table 3）。SwitchAll在大多数数据集上优于基线，例如在Wikitext 103上，47M参数模型的困惑度为12.17，低于基线的12.32。

3. 为什么选择value和output投影？

论文通过消融实验（Table 6）发现，value和output投影的MoE化对性能贡献最大，而key和query投影的MoE化可能浪费参数预算。这是因为在参数匹配设置下，额外专家会挤占其他部分的资源。

---

与MoA的对比

Mixture of Attention Heads（MoA）是与SwitchHead最相似的方法，但有以下区别：

• 注意力矩阵数量：MoA为每个选中的专家计算一个注意力矩阵，计算量较高。SwitchHead在专家选择前后进行加权平均，减少了矩阵数量。
• 选择机制：MoA使用竞争性softmax，需要复杂正则化；SwitchHead使用非竞争性sigmoid，无需正则化。
• 资源效率：SwitchHead在相同计算预算下，困惑度优于MoA（Table 1）。

---

实现细节与局限性

实现细节

• Triton内核：SwitchHead采用Triton内核实现MoE投影，当前效率约为cuBLAS的60%，但仍有优化空间（可达80-90%）。
• 超参数：SwitchHead通过系统性搜索设置 $n_{\text{heads}}$ 和 $k$，确保性能匹配基线（参考Sec. 3）。
• 数据集：使用SentencePiece分词器，词汇表大小为8k，训练100k批次。

局限性

• 模型规模：SwitchHead测试的模型规模（47M和262M）较小，未能验证在超大模型上的效果。
• 硬件依赖：当前实现依赖GPU，边缘设备上的适用性待验证。
• 并行训练：模型并行训练需要动态负载均衡，增加实现复杂度。

---

结论与展望

SwitchHead通过将MoE机制引入注意力层，显著降低了Transformer的计算和内存需求，同时保持了性能。其非竞争性选择机制和对value/output投影的优化设计，为高效Transformer提供了新思路。未来，研究者可以：

• 在更大规模模型上验证SwitchHead的效果。
• 结合FlashAttention等硬件优化技术，进一步提升效率。
• 探索SwitchHead在边缘设备上的部署方案。

对于Transformer研究者，SwitchHead不仅是一个实用的优化工具，也为MoE在注意力机制中的应用提供了理论启发。

疑问解答

k,q,v o 四个矩阵，多头注意力，比如8个头，会有4 x 8个矩阵，这里的switchhead的原理是k和q数量不变，v和o矩阵数量变少吗？

是的，你的理解基本正确，但让我详细解释一下SwitchHead的原理，特别是关于key（K）、query（Q）、value（V）和output（O）矩阵在多头注意力机制中的处理方式，以及SwitchHead如何通过MoE机制减少计算量和矩阵数量。

---

1. 标准Transformer多头注意力的矩阵

在标准Transformer的多头注意力机制中，假设有 $n_{\text{heads}}$ 个头（例如8个头），每个头都会计算自己的K、Q、V和O矩阵。具体来说：

• 输入：输入序列 $\mathbf{x}\in {\mathbb{R}}^{T\times d_{\text{model}}}$，其中 $T$ 是序列长度，$d_{\text{model}}$ 是模型维度。
• 投影：对于每个头 $h$（$h=1,\dots ,n_{\text{heads}}$），通过线性投影生成：
  • ${\mathbf{Q}}^h=\mathbf{x}{\mathbf{W}}_Q^h$，${\mathbf{Q}}^h\in {\mathbb{R}}^{T\times d_{\text{head}}}$
  • ${\mathbf{K}}^h=\mathbf{x}{\mathbf{W}}_K^h$，${\mathbf{K}}^h\in {\mathbb{R}}^{T\times d_{\text{head}}}$
  • ${\mathbf{V}}^h=\mathbf{x}{\mathbf{W}}_V^h$，${\mathbf{V}}^h\in {\mathbb{R}}^{T\times d_{\text{head}}}$
  • ${\mathbf{O}}^h=({\mathbf{A}}^h{\mathbf{V}}^h){\mathbf{W}}_O^h$，其中 ${\mathbf{A}}^h=\text{softmax}(\frac{{\mathbf{Q}}^h{\mathbf{K}}^{hT}}{\sqrt{d_{\text{head}}}})$ 是注意力矩阵，${\mathbf{W}}_O^h\in {\mathbb{R}}^{d_{\text{head}}\times d_{\text{model}}}$。
• 矩阵数量：对于8个头，共有：
  • 8个 ${\mathbf{W}}_Q^h$，8个 ${\mathbf{W}}_K^h$，8个 ${\mathbf{W}}_V^h$，8个 ${\mathbf{W}}_O^h$，即 $4\times 8=32$ 个投影矩阵。
  • 每个头计算一个注意力矩阵 ${\mathbf{A}}^h$，共8个 ${\mathbf{A}}^h$。

计算复杂度：

• 每个头的注意力矩阵 ${\mathbf{A}}^h$ 计算需要 $O(T^2{d}_{\text{head}})$，共 $O(n_{\text{heads}}{T}^2{d}_{\text{head}})$。
• 投影计算（K、Q、V、O）需要 $O(n_{\text{heads}}T{d}_{\text{model}}{d}_{\text{head}})$。

---

2. SwitchHead的原理

SwitchHead通过引入Mixture-of-Experts（MoE）机制，减少注意力矩阵 ${\mathbf{A}}^h$ 的数量（即减少 $n_{\text{heads}}$），并对部分投影矩阵（特别是V和O）应用MoE，从而降低计算和内存需求。以下是其核心设计：

2.1 减少注意力头数

• 在SwitchHead中，$n_{\text{heads}}$ 显著减少（例如从16减少到2或4）。每个头仍计算一个独立的注意力矩阵 ${\mathbf{A}}^h$，但由于头数减少，注意力矩阵的数量大幅降低。
• 为了保持模型的表达能力，SwitchHead为每个头引入 $E$ 个专家（experts），并动态选择 $k$ 个活跃专家（$k\ll E$）。这意味着每个头的V和O投影不再是单一矩阵，而是由多个专家的加权组合生成。

2.2 哪些投影矩阵使用MoE？

SwitchHead通过实验（参考论文Table 6）发现，将MoE应用于value（V）和output（O）投影效果最佳，而key（K）和query（Q）投影保持单一（非MoE）。具体来说：

• K和Q投影：

  • 数量不变：每个头 $h$ 仍有一个 ${\mathbf{W}}_K^h$ 和 ${\mathbf{W}}_Q^h$，共 $n_{\text{heads}}$ 个K矩阵和 $n_{\text{heads}}$ 个Q矩阵。
  • 计算方式：${\mathbf{K}}^h=\mathbf{x}{\mathbf{W}}_K^h$，${\mathbf{Q}}^h=\mathbf{x}{\mathbf{W}}_Q^h$，与标准Transformer相同。
  • 原因：实验表明，K和Q的MoE化会占用参数预算，降低性能（见论文Sec. 3.1）。

• V和O投影：

  • 数量变化：每个头 $h$ 有 $E$ 个专家的V和O投影矩阵（${\mathbf{W}}_V^{h,e}$ 和 ${\mathbf{W}}_O^{h,e}$，$e=1,\dots ,E$），但在推理时只选择 $k$ 个活跃专家。
  • 计算方式：
    • V投影：${\mathbf{V}}^h={\sum }_{e\in {\mathcal{E}}_S^h}{\mathbf{s}}_S^h[e]\mathbf{x}{\mathbf{W}}_V^{h,e}$，其中 ${\mathcal{E}}_S^h$ 是源端选择的 $k$ 个专家索引，${\mathbf{s}}_S^h[e]$ 是专家权重。
    • O投影：$\mathbf{y}={\sum }_{h=0}^{n_{\text{heads}}-1}{\sum }_{e\in {\mathcal{E}}_D^h}{\mathbf{s}}_D^h[e]{\mathbf{A}}^h{\mathbf{V}}^h{\mathbf{W}}_O^{h,e}$。
  • 效果：通过动态选择 $k$ 个专家，V和O的实际计算量减少到 $O(kT{d}_{\text{model}}{d}_{\text{head}})$，远小于标准Transformer的 $O(n_{\text{heads}}T{d}_{\text{model}}{d}_{\text{head}})$。

2.3 专家选择机制

• 源端选择（V投影）：${\mathbf{s}}_S^h=\sigma (\mathbf{x}{\mathbf{W}}_S^h)$，${\mathcal{E}}_S^h=\arg \mathrm{topk}({\mathbf{s}}_S^h,k)$，使用sigmoid生成非竞争性得分。
• 目标端选择（O投影）：类似地，${\mathbf{s}}_D^h=\sigma (\mathbf{x}{\mathbf{W}}_D^h)$，${\mathcal{E}}_D^h=\arg \mathrm{topk}({\mathbf{s}}_D^h,k)$。
• 非竞争性选择（sigmoid）避免了专家坍塌问题，无需额外正则化。

---

3. SwitchHead对矩阵数量的影响

以例子（8个头）为例，比较标准Transformer和SwitchHead：

标准Transformer（8头）

• 投影矩阵：
  • K：8个 ${\mathbf{W}}_K^h$
  • Q：8个 ${\mathbf{W}}_Q^h$
  • V：8个 ${\mathbf{W}}_V^h$
  • O：8个 ${\mathbf{W}}_O^h$
  • 共 $4\times 8=32$ 个矩阵。
• 注意力矩阵：8个 ${\mathbf{A}}^h$。

SwitchHead（假设 $n_{\text{heads}}=2$，$E=5$，$k=2$）

• 投影矩阵：
  • K和Q：数量不变，每个头一个矩阵。
    • K：2个 ${\mathbf{W}}_K^h$
    • Q：2个 ${\mathbf{W}}_Q^h$
  • V和O：每个头有 $E=5$ 个专家矩阵，但每次只计算 $k=2$ 个活跃专家的投影。
    • V：2头 $\times$ 5专家 = 10个 ${\mathbf{W}}_V^{h,e}$（存储），但实际计算只涉及 2头 $\times$ 2专家 = 4个 ${\mathbf{V}}^h$。
    • O：类似，10个 ${\mathbf{W}}_O^{h,e}$（存储），计算4个。
  • 总计：
    • 存储：2（K）+ 2（Q）+ 10（V）+ 10（O）= 24个矩阵。
    • 计算：2（K）+ 2（Q）+ 4（V）+ 4（O）= 12个矩阵。
• 注意力矩阵：只有2个 ${\mathbf{A}}^h$（因为 $n_{\text{heads}}=2$），比标准Transformer的8个减少了4倍。

总结：

• K和Q矩阵数量不变：SwitchHead保持每个头的K和Q矩阵单一，数量等于 $n_{\text{heads}}$。
• V和O矩阵数量减少：通过MoE机制，V和O的实际计算量从 $n_{\text{heads}}$ 减少到 $n_{\text{heads}}\times k$，存储的矩阵数量为 $n_{\text{heads}}\times E$（但仍少于标准Transformer的 $n_{\text{heads}}$ 个矩阵在计算时的开销）。
• 注意力矩阵减少：最显著的优化是 $n_{\text{heads}}$ 减少，导致 ${\mathbf{A}}^h$ 数量减少（例如从8到2）。

---

4. 为什么V和O矩阵数量减少有效？

SwitchHead通过实验（Table 6）发现，V和O投影的MoE化对性能贡献最大，原因可能包括：

• V和O的多样性：V和O投影直接影响注意力输出的内容和最终表示，MoE允许动态选择适合当前上下文的投影，增加表达能力。
• K和Q的稳定性：K和Q主要决定注意力分布（${\mathbf{A}}^h$），单一投影足以捕捉序列关系，MoE化可能引入不必要的复杂性。
• 参数效率：在参数匹配设置下，V和O的MoE化更有效地利用参数预算，而K和Q的MoE化可能挤占其他部分的资源。

---

5. 计算和内存节省

• 计算量：SwitchHead的计算复杂度为（参考论文Eq. 13）：
  $$N_{\text{MAC}}^{\text{SwitchHead}}=n_{\text{heads}}\left(2T{d}_{\text{head}}{d}_{\text{model}}+2Tk{d}_{\text{head}}(d_{\text{model}}+1)+2C{T}^2{d}_{\text{head}}+2CT{d}_{\text{head}}{d}_{\text{model}}\right)$$
  相比标准Transformer（Eq. 11），由于 $n_{\text{heads}}$ 减少（例如从8到2）和 $k\ll E$，计算量显著降低。例如，262M参数模型上，SwitchHead仅用44%的MACs。

• 内存：注意力矩阵 ${\mathbf{A}}^h$ 的存储从 $O(n_{\text{heads}}{T}^2)$ 减少到 $O(n_{\text{heads}}^{\text{SwitchHead}}{T}^2)$，V和O投影的中间结果也因 $k$ 的限制而减少。

---

6. 你的问题的答案

SwitchHead的原理是K和Q数量不变，V和O矩阵数量变少吗？

是的，但更精确地说：

• K和Q矩阵数量不变：每个头仍有一个K和Q投影矩阵，数量等于 $n_{\text{heads}}$，但 $n_{\text{heads}}$ 本身减少（例如从8到2）。
• V和O矩阵数量减少：
  • 存储上：V和O矩阵数量为 $n_{\text{heads}}\times E$，可能比标准Transformer的 $n_{\text{heads}}$ 多，但实际计算时只使用 $n_{\text{heads}}\times k$ 个专家的投影。
  • 计算上：由于只计算 $k$ 个专家的V和O投影，实际计算量远低于标准Transformer的 $n_{\text{heads}}$ 个V和O矩阵。
• 核心优化：通过减少 $n_{\text{heads}}$（从而减少注意力矩阵 ${\mathbf{A}}^h$ 数量）和MoE化的V和O投影，SwitchHead大幅降低计算和内存需求。

---

7. 举例说明

假设标准Transformer有8头，SwitchHead有2头，$E=5$，$k=2$：

• 标准Transformer：8个K、Q、V、O矩阵，共32个投影矩阵，8个注意力矩阵。
• SwitchHead：
  • K和Q：2个K矩阵，2个Q矩阵。
  • V和O：2头 $\times$ 5专家 = 10个V矩阵和10个O矩阵（存储），但每次只计算2头 $\times$ 2专家 = 4个V和4个O。
  • 注意力矩阵：2个 ${\mathbf{A}}^h$。
• 结果：SwitchHead的注意力矩阵数量减少4倍，V和O的计算量减少约2倍（4 vs. 8）。

---

σ-MoE

σ-MoE method是什么，non-competitive selection function是什么，和传统的做法相比有什么好处

将详细解释 σ-MoE 方法、非竞争性选择函数（non-competitive selection function） 的定义，以及它们与传统 MoE（Mixture-of-Experts）方法的区别和优势。以下内容基于《SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention》论文及其引用的相关工作，特别是 σ-MoE 方法的原始论文 [17]（Csordás et al., 2023）。将尽量清晰地解释这些概念，并突出其好处，面向熟悉 Transformer 和 MoE 的研究者。

---

1. σ-MoE 方法是什么？

σ-MoE 是一种基于 非竞争性激活函数（specifically sigmoid）的 Mixture-of-Experts（MoE）方法，最初由 Csordás 等人提出（参考 [17]），用于高效地扩展 Transformer 的前馈层（MLP），并在 SwitchHead 中被扩展到注意力层。MoE 的核心思想是通过多个“专家”网络（experts）处理输入，并动态选择部分专家进行计算，从而在保持模型容量（参数量）的同时降低计算成本。

σ-MoE 的核心机制

在 σ-MoE 中：

• 专家网络：模型包含多个专家（例如 $E$ 个专家），每个专家是一个小型神经网络（在 Transformer 中通常是 MLP 或投影矩阵）。
• 选择机制：通过一个路由函数（routing function）动态选择 $k$ 个活跃专家（$k\ll E$），只计算这些专家的输出。
• 激活函数：σ-MoE 使用 sigmoid 函数（记为 $\sigma$）作为路由函数的激活函数，生成每个专家的得分（score），并通过 $\arg \mathrm{topk}$ 选择得分最高的 $k$ 个专家。

具体公式（以 SwitchHead 的源端选择为例，参考论文 Eq. 7-8）：
$${\mathbf{s}}_S^h=\sigma \left(\mathbf{x}{\mathbf{W}}_S^h\right)$$
E S h = arg ⁡ topk ⁡ ( s S h , k ) , E S h ⊂ { 1 , … , E } \mathcal{E}_S^h = \arg \operatorname{topk} \left( \boldsymbol{s}_S^h, k \right), \quad \mathcal{E}_S^h \subset \{1, \ldots, E\} ESh​=argtopk(sSh​,k),ESh​⊂{1,…,E}

• $\mathbf{x}\in {\mathbb{R}}^{T\times d_{\text{model}}}$：输入序列。
• ${\mathbf{W}}_S^h\in {\mathbb{R}}^{d_{\text{model}}\times E}$：路由矩阵。
• ${\mathbf{s}}_S^h\in {\mathbb{R}}^{T\times E}$：每个专家的得分，$\sigma (x)=\frac{1}{1+e^{-x}}$ 是 sigmoid 函数。
• ${\mathcal{E}}_S^h$：选中的 $k$ 个专家索引。

输出通过加权求和计算，例如 value 投影：
$${\mathbf{V}}^h=\sum _{e\in {\mathcal{E}}_S^h}{\mathbf{s}}_S^h[e]\mathbf{x}{\mathbf{W}}_V^{h,e}$$

与传统 MoE 的区别

传统 MoE（例如 Shazeer et al. [11] 的 sparsely-gated MoE）通常使用 softmax 激活函数 进行专家选择：
$$\mathbf{s}=\text{softmax}\left(\mathbf{x}{\mathbf{W}}_S\right)$$

• softmax 是一种竞争性激活函数，专家得分的和为 1，导致专家之间相互竞争。
• 传统 MoE 通常需要额外的正则化（例如负载均衡损失）来防止“专家坍塌”（expert collapse），即模型只使用少数专家，其他专家被忽略。

而 σ-MoE 的 sigmoid 激活 是非竞争性的，每个专家的得分独立于其他专家，取值范围为 $[0,1]$，不要求总和为 1。

---

2. 非竞争性选择函数（Non-competitive Selection Function）是什么？

非竞争性选择函数 指的是 σ-MoE 中使用的 sigmoid 函数，区别于传统 MoE 的竞争性 softmax 函数。它的特点是：

• 独立性：每个专家的得分 ${\mathbf{s}}_S^h[e]$ 只取决于输入 $\mathbf{x}$ 和对应的路由矩阵 ${\mathbf{W}}_S^h$，不受其他专家得分的影响。
• 连续性：sigmoid 输出 $[0,1]$ 范围的连续值，允许专家以不同权重贡献输出，而不是二元选择（全用或不用）。
• 选择过程：通过 $\arg \mathrm{topk}$ 从 sigmoid 得分中选取 $k$ 个最高得分的专家，但得分本身用于加权输出（见 ${\mathbf{V}}^h$ 公式）。

在 SwitchHead 中，非竞争性选择函数用于源端（value 投影）和目标端（output 投影）的专家选择，确保动态选择的同时避免专家间的过度竞争。

---

3. σ-MoE 与传统 MoE 相比的好处

σ-MoE（以及其非竞争性选择函数）相较于传统 MoE（基于 softmax 的方法）有以下显著优势：

3.1 避免专家坍塌，无需复杂正则化

• 传统 MoE 的问题：
  • Softmax 的竞争性导致某些专家可能始终获得高得分，其他专家被忽略，造成“专家坍塌”。
  • 为解决此问题，传统 MoE 通常引入多种正则化技术，例如：
    • 负载均衡损失（load balancing loss）：强制专家被均匀使用。
    • 专家多样性正则化：鼓励不同输入选择不同专家。
    • 噪声注入：在路由过程中添加噪声以增加多样性。
  • 这些正则化增加了训练复杂度和超参数调优的难度。
• σ-MoE 的优势：
  • Sigmoid 的非竞争性允许每个专家独立评分，避免了专家间的直接竞争，天然减少了专家坍塌的风险。
  • 论文明确指出，σ-MoE（包括 SwitchHead）无需额外的正则化即可稳定训练（参考 Sec. 3.2）。
  • 实验结果（Table 1）显示，σ-MoE 在 WikiText-103 上比需要正则化的 MoA（Mixture of Attention Heads）表现更好，且资源占用更低。

3.2 更高的资源效率

• 传统 MoE：
  • 由于 softmax 的竞争性，路由函数可能导致不均衡的专家选择，增加计算开销（例如需要计算更多专家的输出以保证性能）。
  • 在注意力层（如 MoA [18]），传统 MoE 为每个选中的专家计算一个注意力矩阵，导致计算量和内存占用较高（例如 MoA 需要 $n_{\text{heads}}$ 个注意力矩阵）。
• σ-MoE 的优势：
  • 通过减少 $n_{\text{heads}}$（例如从 16 到 2）并为每个头引入 $E$ 个专家，σ-MoE 大幅降低注意力矩阵的数量（SwitchHead 最多减少 8 倍）。
  • 只计算 $k$ 个活跃专家的投影（V 和 O），计算量从 $O(n_{\text{heads}}T{d}_{\text{model}}{d}_{\text{head}})$ 减少到 $O(n_{\text{heads}}kT{d}_{\text{model}}{d}_{\text{head}})$。
  • 内存占用也显著减少，例如 262M 参数模型仅用 27% 的内存（Table 2）。

3.3 更灵活的专家选择

• 传统 MoE：
  • Softmax 的输出是概率分布，专家选择倾向于“硬性”分配（一个输入主要由少数专家处理）。
  • 这可能限制模型对输入的动态适应能力，尤其在注意力层需要捕获复杂模式时。
• σ-MoE 的优势：
  • Sigmoid 允许专家以连续权重（$[0,1]$）贡献输出，增加选择灵活性。
  • 在 SwitchHead 中，源端（V）和目标端（O）的独立选择机制（Eq. 7-10）进一步增强了模型对输入上下文的适应性。
  • 实验表明，这种灵活性使 SwitchHead 在参数匹配设置下（parameter-matched setting）能匹配甚至超越标准 Transformer 的性能（Table 2）。

3.4 简化实现和训练

• 传统 MoE：
  • 实现复杂，需要负载均衡、专家同步等机制，尤其在分布式训练中需要额外的通信开销。
  • 正则化超参数的调优增加了工程负担。
• σ-MoE 的优势：
  • 非竞争性选择简化了路由逻辑，Triton 内核实现（Sec. A.5）高效且易于优化（当前效率为 cuBLAS 的 60%，未来可达 80-90%）。
  • 无需正则化，训练过程更稳定，超参数搜索更简单（Sec. 3 描述了系统性超参数选择流程）。

3.5 适用于注意力层的扩展

• 传统 MoE：
  • 主要应用于前馈层（MLP），在注意力层的应用（如 MoA）效率较低，因为需要为每个专家计算注意力矩阵（Sec. 3.2）。
  • MoA 需要更多注意力头（$n_{\text{heads}}$）来匹配基线性能，导致资源消耗高（Table 1）。
• σ-MoE 的优势：
  • SwitchHead 将 σ-MoE 扩展到注意力层，通过在 V 和 O 投影上应用 MoE，显著减少注意力矩阵数量（$n_{\text{heads}}$ 减少到 2 或 4）。
  • 实验表明，SwitchHead 在 C4 数据集上仅用 44% 的计算量和 27% 的内存即可匹配基线 Transformer 的困惑度（Sec. 3.3）。

---

4. 具体实验证据

以下是论文中 σ-MoE（SwitchHead）与传统方法（MoA 和标准 Transformer）的对比结果，凸显其优势：

• Table 1（WikiText-103）：
  • SwitchHead（47M 参数，$n_{\text{heads}}=2$）：困惑度 12.27，MACs 170.4M，内存 0.8M。
  • MoA（47M 参数，$n_{\text{heads}}=4$）：困惑度 12.60，MACs 223.5M，内存 1.3M。
  • 标准 Transformer（47M 参数，$n_{\text{heads}}=10$）：困惑度 12.31，MACs 453.4M，内存 3.5M。
  • 结论：SwitchHead 用更少的资源（约 37% 的 MACs 和 23% 的内存）实现更好的困惑度。
• Table 5（壁钟时间和内存）：
  • SwitchHead（262M 参数）：训练时间 442ms/迭代，内存 12.5G。
  • MoA（262M 参数）：训练时间 851ms/迭代，内存 16.4G。
  • 标准 Transformer（262M 参数）：训练时间 670ms/迭代，内存 20.5G。
  • 结论：SwitchHead 训练速度提升约 1.5 倍，内存占用减少约 39%。

---

5. 总结

σ-MoE 方法 通过使用 sigmoid 非竞争性选择函数，在 MoE 框架中实现了高效的专家选择，特别适合 Transformer 的前馈层和注意力层（SwitchHead）。与传统 MoE（基于 softmax）相比，σ-MoE 的主要好处包括：

1. 避免专家坍塌：无需复杂正则化，训练更稳定。
2. 资源效率：减少计算量（例如 44% MACs）和内存占用（例如 27%），尤其在注意力层效果显著。
3. 灵活选择：连续权重增加模型对输入的适应性。
4. 简化实现：路由逻辑简单，易于优化和扩展。
5. 注意力层优化：SwitchHead 通过减少 $n_{\text{heads}}$ 和 MoE 化的 V/O 投影，实现高效注意力计算。

对于 Transformer 研究者，σ-MoE 提供了一种简单而强大的方法来构建高效的大规模模型，尤其在资源受限场景下。SwitchHead 的成功表明，σ-MoE 的非竞争性设计不仅适用于前馈层，还能显著提升注意力层的效率，值得进一步探索和应用。

---

后记

2025年5月16日于上海，在grok 3大模型辅助下完成。