arXiv:2503.07137

作者：Siyuan Mu (四川农业大学), Sen Lin (休斯顿大学)

研究背景

随着AI基础大模型的快速发展，现代数据集变得越来越多样化和复杂，包含多模态数据（文本、图像、音频）和复杂结构（图、层次关系）。这给大模型发展带来两大挑战：

1. 计算资源消耗巨大：训练和部署大模型的计算成本呈指数增长
2. 异构数据拟��困难：在单一模型中整合冲突或异构知识变得困难，导致训练不稳定和性能次优

混合专家模型（Mixture of Experts, MoE）通过动态选择和激活最相关的子模型来处理输入数据，成为解决这些挑战的有效方案。

研究目标

• 填补现有MoE综述的空白（过时或缺乏关键领域讨论）
• 全面总结MoE的基础设计、算法、理论和应用四大关键组件
• 为研究者提供系统性参考，激发进一步研究

核心概念

MoE基本原理

MoE采用”分而治之”（divide and conquer）策略，与传统密集模型不同：

• 传统模型：对每个输入激活所有参数
• MoE模型：根据输入特征动态选择和激活最相关的参数子集

MoE层数学表示

$$
\text{MoE}(x) = \sum_{i \in \mathcal{I}_D} w_i M_i(x)
$$
其中 $\mathcal{I}_D$ 是被选中专家的索引集，$w_i$ 是第 $i$ 个专家的权重，$M_i(x)$ 是专家网络输出。

研究方法

1. 门控函数（Gating Function）

线性门控（Softmax Gating）

$$
G(x)i = \text{softmax}(\text{TopK}(g(x) + R{noise}, k))i
$$
其中 $g(x)$ 是线性函数计算的门控值，$R{noise}$ 是鼓励专家探索的噪声。

非线性门控

• 余弦门控（GMoE）：
  $$
  G(x) = \text{TopK}\left(\text{softmax}\left(\frac{E^T W_{linear} x}{\tau |W_{linear} x| |E|}\right)\right)
  $$
• 指数族分布门控
• Soft MoE：使用加权平均而非离散分配

2. 专家网络（Expert Network）

3. 路由策略（Routing Strategy）

• Token级路由：基于token表示进行路由决策（最经典）
• 模态级路由：根据数据模态进行路由（多模态任务）
• 任务级路由：根据任务ID确定路由（多任务学习）

4. 训练策略

负载均衡损失（Switch Transformer）

$$
\mathcal{L}{aux} = \alpha \cdot N \cdot \sum{i=1}^{N} f_i \cdot Q_i
$$
其中 $f_i$ 是分配给专家 $i$ 的token比例，$Q_i$ 是路由概率比例。

MoA（Mixture-of-Attention）详解

基本架构

MoA将MoE机制引入多头注意力模块，每个注意力头视为一个”专家”。

工作流程

1. 输入token进入MoA层
2. 门控网络计算每个注意力头的重要性分数
3. 选择TopK个最相关的注意力头
4. 仅计算被选中头的输出并加权求和

代码实现核心

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

class MixtureOfAttention(nn.Module):
    def __init__(self, d_model, num_heads=8, head_dim=64, top_k=2):
        self.attention_experts = nn.ModuleList([
            AttentionExpert(d_model, head_dim) for _ in range(num_heads)
        ])
        self.router = AttentionRouter(d_model, num_heads, top_k)

    def forward(self, x):
        weights, indices, logits = self.router(x)
        output = torch.zeros_like(x)
        for k in range(self.top_k):
            for head_idx in range(self.num_heads):
                mask_k = (indices[:, :, k] == head_idx)
                if mask_k.any():
                    head_output = self.attention_experts[head_idx](x)
                    output[mask_k] += weights[:, :, k][mask_k].unsqueeze(-1) * head_output[mask_k]
        return output

MoA vs 标准多头注意力

主要发现

算法应用领域

应用领域

代表性大模型

讨论

优势

• 计算效率：通过稀疏激活显著降低计算成本
• 模型容量：可扩展至万亿参数而不成比例增加计算
• 专业化学习：不同专家专注于不同知识领域
• 可解释性：通过分析专家分配机制理解模型行为

局限性

• 训练稳定性：动态专家选择可能导致负载不均衡和模型崩溃
• 系统复杂性：All-to-All通信模式增加系统设计难度
• 内存需求：多专家参数存储可能超出单设备容量

未来方向

1. 训练稳定性与负载均衡：开发更鲁棒的训练策略
2. 训练与系统效率：优化硬件-软件协同设计
3. 架构设计：使用元学习或强化学习动态调整专家数量
4. 理论发展：深入理解专家路由决策和聚类机制
5. 定制算法设计：探索MoE与对比学习、自监督学习的结合
6. 新应用领域：医疗、机器人、自动驾驶、教育、金融

相关工作

• Switch Transformer：首个万亿参数MoE模型
• V-MoE：视觉领域MoE应用
• Mixtral 8×7B：高效MoE语言模型
• Soft MoE：软分配MoE新范式
• [[BAGEL-Unified-Multimodal-Pretraining]]

参考文献

• Fedus et al. (2022). Switch Transformers: Scaling to Trillion Parameter Models. JMLR.
• Shazeer et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. arXiv.
• Riquelme et al. (2021). Scaling Vision with Sparse Mixture of Experts. NeurIPS.
• Jiang et al. (2024). Mixtral of Experts. arXiv.

框架

论文使用**RSSM(Recurrent State Space Model)**：使用encoder来编码环境和动作生成latent state, 预测未来latent state，最后基于latent state预测奖励。

优势：

• 网络可以在 latent 中快速 roll-out 数千条 imagined trajectories
• 不用预测 pixel → 速度极快
• 潜在空间的 Markov 性保证了规划时的可微分性

重参数化 Reparameterization Trick

Dreamer 最关键的地方：

动作必须是可微的随机变量，这样梯度才能从 value 反传到 actor。

如果我们直接写 $a \sim \mathcal{N}(\mu, \sigma)$, 那么采样是不可微的 → 梯度断掉 → Actor 无法学习。
重参数化技巧的做法：$a=\mu+\sigma\cdot\epsilon$, $\epsilon\sim\mathcal{N}(0,1)$
现在：

• ε 是随机的
• μ 和 σ 是可微的网络输出
  所以动作对 actor 参数有梯度, 这就是可微规划（differentiable planning）的基础。

Introduction

IL是区别于传统手动编程来赋予机器人自主能力的方法。
IL 允许机器通过演示（人类演示专家行为）来学习所需的行为，从而消除了对显式编程或特定于任务的奖励函数的需要。
IL主要有两个类别：

• 行为克隆(BC)
• 反向强化学习(IRL)
  

Behavior Cloning

BC 是一种 IL 技术，它将学习行为的问题视为监督学习任务 。 BC 涉及通过建立环境状态与相应专家操作之间的映射来训练模型来复制专家的行为。专家的行为被记录为一组state-action pair，也称为演示。在训练过程中，模型学习一个函数，利用这些演示作为输入，将当前状态转换为相应的专家操作。经过训练，模型可以利用这个学习函数来生成遇到新状态的动作。

不需要了解环境的潜在动态，计算效率很高，相对简单的方法。

The covariate shift problem: 测试期间观察到的状态分布可能与训练期间观察到的状态分布有所不同，使得代理在遇到未见过的状态时容易出错，而对于如何进行操作缺乏明确的指导。BC监督方法的问题是，当智能体漂移并遇到分布外状态时，它不知道如何返回到演示的状态。

为了解决这个问题：

Inverse Reinforcement Learning

IRL 涉及一个学徒代理，其任务是推断观察到的演示背后的奖励函数，这些演示被认为源自表现最佳的专家 。然后使用推断的奖励函数通过 RL 训练学习代理的策略。

为了解决“政策->奖励函数“的模糊性，有以下三种IRL

• maximum-margin methods（奖励函数比任何其他策略在一定程度上更全面地解释最优策略。这本质上意味着找到一个最大化指定利润的解决方案，确保派生的奖励函数捕捉专家行为的本质。）
• maximum entropy（处理专家次优性和随机性的有前景的能力）
• guided cost learning（旨在优化策略优化内循环内的非线性奖励函数的方法。这种方法通过直接利用系统的原始状态来构建奖励函数，从而改变了传统的 IRL 范式，从而消除了广泛的特征工程的需要。）

Adversarial Imitation Learning

The agent strives to deceive the discriminator by generating trajectories closely resembling those of the expert.

Imitation From Observation

仅通过图像序列来学习，不需要具体的关节动作操作数据。

Unlike the traditional methods, IfO presents a more organic approach to learning from experts, mirroring how humans and animals approach imitation. Humans often learn new behaviors by observing others without detailed knowledge of their actions (e.g., the muscle commands). People learn a diverse range of tasks, from weaving to swimming to playing games, by watching online videos. Despite differences in body shapes, sensory inputs, and timing, humans exhibit an impressive ability to apply knowledge gained from the online demonstrations

将可学习的资源扩大到了线上的视频资源。

Latent Action Policies (LAPOs)

过分析观察到的动态，LAPO 推断出行动空间的底层结构，促进潜在行动策略的训练。然后，这些策略可以进行高效的微调，以达到专家级的性能，从而提供离线和在线场景的适应性。使用包含标记动作的小数据集进行离线微调是可行的，而在线微调可以使用奖励来完成。与依赖标记数据来训练逆动力学模型不同，LAPO直接从观察到的环境动态中导出潜在动作信息，而不需要任何标签。

Challenges And Limitations

。。。