MAESTRO:马尔可夫链剪除不良专家
速览
稀疏专家模型(MoE)虽高效但专家库常驻内存造成部署瓶颈,现有剪枝方法忽视路由依赖。MAESTRO将专家激活轨迹建模为遍历马尔可夫链,得出全局重要性。在安全、偏见等五领域测试,50%压缩率下平均性能保持领先10.61%,跨任务方差更低。表明全局路由一致的剪枝能带来更稳定的泛化。
AI 深度解读
背景
近年来,稀疏激活的混合专家模型(Sparsely-activated Mixture-of-Experts, MoE)凭借其独特的计算范式,在大语言模型领域取得了显著进展。这类模型在推理时仅对每个 token 激活一小部分专家(expert),从而大幅降低计算开销;然而,全部专家的参数量始终驻留在内存中,导致部署时出现难以承受的内存瓶颈。现有的结构化剪枝方法主要针对稠密 Transformer 设计,它们通常基于局部启发式指标(如激活频率、梯度幅度)评估专家重要性,但这些指标忽略了 MoE 路由(routing)过程中跨层专家之间的相互依赖关系。这种“局部视角”使得剪枝决策与 MoE 的全局路由行为脱节,从而在压缩率较高时造成严重的性能下降。
针对这一短板,本文提出了一种名为 MAESTRO(Markov-chain Approximated Expert Sparsification via Transition-based Routing)的结构化剪枝框架,专门面向 MoE 架构。MAESTRO 将自回归语言模型中专家激活的轨迹建模为遍历马尔可夫链(Ergodic Markov chains),并利用其平稳分布编码跨层依赖关系,从而得到一种全局感知的专家重要性指标。
核心内容
MAESTRO 的核心思想是:不再孤立地评估每个专家的重要性,而是将专家序列视为一个随机过程。在 MoE 模型中,每一层都会根据 token 的隐藏状态选择 top-k 个专家,这种逐层选择形成了 token 在专家集合上的转移轨迹。MAESTRO 将这些轨迹视为一个马尔可夫链,其中状态是各层的专家索引,转移概率由路由器的 softmax 输出近似得到。由于路由决策依赖于前一层的信息(在自回归模型中逐层传递),该马尔可夫链是时间齐次的,且在足够长的序列下具有遍历性。
通过求解该马尔可夫链的平稳分布,MAESTRO 能够获得每个专家在长期路由过程中被访问的稳态概率。这个概率不仅反映了专家自身的“受欢迎程度”,还隐含了跨层之间的转移模式:如果某个专家频繁地与其他重要专家共现,或处于高转移概率路径的关键节点上,其稳态概率就会更高。MAESTRO 以此作为专家重要性的全局度量,然后剪枝掉那些稳态概率最低的专家(即“坏专家”),保留对全局路由连贯性贡献最大的专家。
MAESTRO 的剪枝过程分为两步:首先,在少量校准数据上收集各层路由器的激活记录,统计专家之间的转移计数,构造转移矩阵;其次,求解该矩阵的平稳分布(对应于最大特征值 1 的左特征向量),并依据该分布排序,移除低分专家。剪枝后无需微调或重量级训练,仅需冻结剩余专家进行前向推理。
实验在五个不同领域进行评估,包括安全性(Safety)、偏见(Bias)、伦理(Ethics)等多样任务。在严格的 50% 压缩率(即剪掉一半专家参数)下,MAESTRO 在平均性能保留率上比现有最佳基线高出最多 10.61%,同时跨任务方差显著更低,表明全局路由一致的剪枝策略能产生在不同任务上更稳定的泛化能力。
关键要点
- 问题本质:现有局部剪枝方法忽略 MoE 路由的跨层依赖,导致高压缩率下性能损失严重。
- 创新方法:MAESTRO 将专家激活轨迹建模为遍历马尔可夫链,利用平稳分布作为全局重要性指标。
- 核心机制:通过校准数据收集层间专家转移概率,构造转移矩阵,求解平稳分布,实现“路由一致”的结构化剪枝。
- 实验设置:覆盖五个多样化领域(Safety, Bias, Ethics 等),采用严格的 50% 参数量压缩。
- 性能优势:平均性能保留率最高提升 10.61%,且跨任务方差低于所有基线,说明泛化更稳定。
- 无需重训:剪枝后直接使用剩余专家进行推理,无需微调或额外训练。
- 技术路线:充分利用自回归 MoE 的逐层路由顺序性,将序列建模为时间齐次 Markov 链,这是一种轻量级、可解释的近似方法。
意义与影响
MAESTRO 的提出为 MoE 模型的结构化压缩提供了一个全新的理论视角。以往的工作往往将剪枝视为独立于路由机制的后处理步骤,而 MAESTRO 将路由本身作为剪枝决策的一部分,揭示了专家之间的动态协同关系。这种“路由感知”的全局剪枝范式有望成为未来 MoE 部署优化的标准方法。
从工程角度看,MAESTRO 的校准数据需求低(仅需少量样本即可估计转移矩阵),计算开销小(求解平稳分布是典型的线性代数操作),剪枝后模型无需重新训练,非常适合资源受限的落地场景,如边缘设备推理或私有化部署。此外,由于剪枝保留了与路由共生的关键专家,模型在不同领域(特别是安全、偏见等敏感维度)上的表现更加稳定,有助于提升 AI 系统的可靠性和公平性。
从学术角度看,该工作将概率图模型中的 Markov 链理论与神经网络结构搜索/剪枝结合起来,开辟了“基于路由动态剪枝”的新方向。未来可进一步探索非齐次 Markov 建模(如考虑序列长度依赖性)、结合在线路由自适应的剪枝,或将其推广到其他具有分层路由结构的模型(如视觉 MoE、多模态 MoE)中。总体而言,MAESTRO 为克服 MoE 部署的内存瓶颈提供了一个优雅且高效的解决方案,具有重要的理论创新与实用价值。
