路由Mamba：MoE与状态空间模型的融合

概述

路由Mamba（Routing Mamba, RoM）将**混合专家（Mixture of Experts, MoE）**思想引入状态空间模型，通过稀疏混合线性投影专家来增强SSM的表达能力，同时保持线性复杂度。¹

背景与动机

MoE的成功

混合专家模型通过稀疏激活机制，在不增加推理成本的情况下大幅增加模型参数：

模型	总参数量	激活参数量	稀疏度
Mixtral-8x7B	46.7B	12.9B	72%
DBRX	132B	36B	73%
SWITCH-Transformer	1.6T	8B	99.5%

SSM的局限

传统SSM（如Mamba）在每个时间步使用固定的参数集，这限制了模型捕捉多样化模式的能力。SSM的瓶颈在于：

1. 固定状态转换：${\mathbf{h}}_{t+1}=\mathbf{A}{\mathbf{h}}_t+\mathbf{B}{\mathbf{x}}_t$
2. 单一动态模式：无法同时建模多种动态系统
3. 表达能力受限：参数数量受限于状态大小 $N$

路由Mamba架构

核心思想

RoM的核心是将SSM的输入投影参数 $\mathbf{B}$ 和 输出投影参数 $\mathbf{C}$ 分解为多个专家的加权组合：

$${\mathbf{B}}_{\text{routed}}=\sum _{e=1}^E{g}_e({\mathbf{x}}_t){\mathbf{B}}_e$$ $${\mathbf{C}}_{\text{routed}}=\sum _{e=1}^E{g}_e({\mathbf{x}}_t){\mathbf{C}}_e$$

其中：

• $E$ 是专家数量
• $g_e(\cdot )$ 是路由函数
• ${\mathbf{B}}_e,{\mathbf{C}}_e$ 是第 $e$ 个专家的参数

路由机制

1. 输入依赖路由

路由函数基于当前输入 ${\mathbf{x}}_t$ 计算每个专家的权重：

$$\mathbf{g}=\text{Softmax}\left(\frac{{\mathbf{W}}_g{\mathbf{x}}_t}{\tau }\right)$$

其中 $\tau$ 是温度参数，控制路由的稀疏程度。

2. Top-K稀疏路由

为了保持效率，RoM采用 Top-K 路由：

$${\tilde{g}}_e=\{\begin{array}{ll}{g}_e& \text{if }e\in \text{TopK}(\mathbf{g},K)\\ 0& \text{otherwise}\end{array}$$ $$g_e^{\text{norm}}=\frac{{\tilde{g}}_e}{{\sum }_{j\in \text{TopK}}{\tilde{g}}_j}$$

3. 线性投影专家

每个专家是一个线性投影：

$${\mathbf{B}}_e\in {\mathbb{R}}^{N\times d},{\mathbf{C}}_e\in {\mathbb{R}}^{d\times N}$$

这种设计比传统的MLP专家更轻量，同时与SSM的线性结构兼容。

完整前向传播

class RoutingMambaSSM(nn.Module):
    """
    路由Mamba SSM层
    
    核心思想：用稀疏混合专家增强SSM的输入/输出投影
    """
    def __init__(self, d_model, d_state=16, n_experts=8, topk=2):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.n_experts = n_experts
        self.topk = topk
        
        # 状态矩阵（共享）
        self.A = nn.Parameter(torch.randn(d_state, d_state))
        self.D = nn.Parameter(torch.ones(d_state))
        
        # 专家参数
        self.B_experts = nn.ParameterList([
            nn.Parameter(torch.randn(d_state, d_model))
            for _ in range(n_experts)
        ])
        self.C_experts = nn.ParameterList([
            nn.Parameter(torch.randn(d_model, d_state))
            for _ in range(n_experts)
        ])
        
        # 路由网络
        self.router = nn.Linear(d_model, n_experts)
        
        # 初始化
        self._init_parameters()
    
    def _init_parameters(self):
        # 状态矩阵初始化为对角（稳定动力学）
        nn.init.eye_(self.A)
        nn.init.normal_(self.B_experts[0], std=0.02)
        nn.init.normal_(self.C_experts[0], std=0.02)
    
    def forward(self, x, return_routing=False):
        """
        x: (batch, seq_len, d_model)
        """
        batch, seq_len, d_model = x.shape
        
        # 计算路由权重
        router_logits = self.router(x)  # (batch, seq_len, n_experts)
        routing_weights = F.softmax(router_logits, dim=-1)
        
        # Top-K稀疏路由
        if self.topk < self.n_experts:
            topk_values, topk_indices = torch.topk(routing_weights, self.topk, dim=-1)
            # 掩码
            mask = torch.zeros_like(routing_weights).scatter_(-1, topk_indices, 1.0)
            routing_weights = routing_weights * mask
            routing_weights = routing_weights / (routing_weights.sum(-1, keepdim=True) + 1e-6)
        
        # 状态更新
        h = torch.zeros(batch, self.d_state, device=x.device)
        outputs = []
        routing_stats = []
        
        for t in range(seq_len):
            # 加权专家组合
            B_t = sum(routing_weights[:, t, e:e+1] * B_e 
                      for e, B_e in enumerate(self.B_experts))
            C_t = sum(routing_weights[:, t, e:e+1] * C_e.T 
                      for e, C_e in enumerate(self.C_experts))
            
            # 记录路由统计
            if return_routing:
                routing_stats.append(routing_weights[:, t, :].clone())
            
            # SSM状态更新
            h = F.silu(self.A @ h.T).T + B_t @ x[:, t, :] * self.D
            
            # 输出
            y_t = h @ C_t.T
            outputs.append(y_t)
        
        y = torch.stack(outputs, dim=1)
        
        if return_routing:
            routing_stats = torch.stack(routing_stats, dim=1)
            return y, routing_stats
        return y

理论分析

表达能力提升

定理：路由SSM的表达能力

设 $E$ 个专家，状态大小为 $N$，则路由SSM的有效表达能力等价于：

$${\text{Expr}}_{\text{RoM}}\sim O(E\cdot N^2)$$

相比固定参数SSM的 $O(N^2)$，路由机制提供了 $E$ 倍的表达能力提升。

路由动态

引理：路由收敛性

在温和条件下（梯度有界、路由网络 Lipschitz 连续），路由权重以指数速率收敛到稳定的 Top-K 配置。

梯度流

路由Mamba的反向传播需要处理稀疏路由的不可微性：

$$\frac{\partial \mathcal{L}}{\partial {\mathbf{W}}_g}=\frac{\partial \mathcal{L}}{\partial g}\cdot \frac{\partial g}{\partial {\mathbf{W}}_g}$$

由于 Top-K 操作存在不可微点，使用 硬直通估计器（Straight-Through Estimator, STE）：

class StraightThroughTopK(Function):
    @staticmethod
    def forward(ctx, x, k):
        # 前向：稀疏选择
        values, indices = torch.topk(x, k, dim=-1)
        ctx.save_for_backward(indices)
        output = torch.zeros_like(x).scatter_(-1, indices, 1.0)
        return output
    
    @staticmethod
    def backward(ctx, grad_output):
        # 反向：STE
        indices = ctx.saved_tensors[0]
        grad_input = torch.zeros_like(grad_output).scatter_(-1, indices, grad_output)
        return grad_input, None

实验结果

语言建模

在Pile数据集上的结果：

模型	困惑度 $\downarrow$	参数量	激活量
Mamba-1.3B	10.31	1.3B	1.3B
RoM-1.3B (E=8, K=2)	9.72	2.1B	1.3B
改进	-5.7%	+62%	0%

关键发现：RoM在相同激活参数下实现显著困惑度提升。

专家利用率分析

训练过程中的专家利用率：

Step 0:    [███░░░░░░░]  Expert 0: 45.2%, Expert 1: 32.1%, ...
Step 1000: [██████░░░░]  Expert 0: 28.3%, Expert 1: 25.7%, ...
Step 5000: [██████████]  Expert 0: 22.1%, Expert 1: 21.3%, ...

专家利用率逐渐均匀化，表明模型学会了利用所有专家的能力。

长程依赖任务

模型	LRA Avg	ListOps	Pathfinder
Mamba	67.4%	58.3%	71.2%
RoM	70.8%	61.2%	73.9%
+对比	+3.4%	+2.9%	+2.7%

与其他方法的对比

RoM vs 标准MoE

特性	标准MoE	RoM
专家类型	MLP	线性投影
应用位置	FFN层	SSM层
参数量增加	显著	中等
表达力提升	高	中高
实现复杂度	中	低

RoM vs 混合SSM

混合SSM（如Jamba）通过交替使用SSM和Attention层来混合架构，而RoM在单一SSM层内实现混合，更加细粒度。

实现细节

负载均衡

为避免路由崩溃（所有样本路由到同一专家），引入辅助负载均衡损失：

$${\mathcal{L}}_{\text{load}}=\alpha \cdot \text{Var}(\text{load})$$

其中 ${\text{load}}_e={\sum }_{t}1[\text{TopK}(g_t)\ni e]$。

专家选择多样性

class DiversityRouter(nn.Module):
    """多样化路由，增加专家选择的多样性"""
    def __init__(self, d_model, n_experts, topk, entropy_coef=0.01):
        super().__init__()
        self.router = nn.Linear(d_model, n_experts)
        self.entropy_coef = entropy_coef
    
    def forward(self, x):
        logits = self.router(x)
        probs = F.softmax(logits, dim=-1)
        
        # 熵正则化
        entropy = -(probs * torch.log(probs + 1e-8)).sum(-1).mean()
        
        # Top-K选择
        topk_probs, topk_indices = torch.topk(probs, self.topk, dim=-1)
        output = torch.zeros_like(probs).scatter_(-1, topk_indices, 1.0)
        
        return output, entropy * self.entropy_coef

实践指南

超参数选择

参数	建议值	说明
$E$ (专家数)	4-16	8为常用值
$K$ (激活专家)	2-4	2为效率最优
温度 $\tau$	0.1-1.0	较低值更稀疏

训练技巧

1. 热身：前1000步使用全专家激活，然后逐渐引入稀疏路由
2. 梯度裁剪：防止路由权重剧烈变化
3. 专家多样化：使用熵正则化避免路由崩溃

局限性

1. 路由开销：路由计算带来少量额外开销
2. 内存占用：多个专家参数增加内存需求
3. 调优复杂：需要同时优化路由和SSM参数

总结

路由Mamba通过将MoE思想引入SSM，在保持线性复杂度的同时显著增强了模型的表达能力。稀疏混合线性投影专家的设计既轻量又高效，为SSM的进一步发展提供了新方向。

Footnotes

1. Routing Mamba论文: https://neurips.cc/virtual/2025/poster/116256 ↩