Stepped MoE: 段级路由与可配置的推理复杂度

TL;DR Stepped MoE 结合弹性嵌套子网络与稀疏专家路由，使单一模型可按部署约束和输入动态配置 1–4B 参数，精度比稠密基线高 2–5%，同时保持相近延迟并大幅节省设备存储，实现推理复杂度细粒度可控。

问题

问题背景

大型语言模型（LLM）推理成本已成为与训练同等重要的瓶颈，业界关注如何在保持精度的同时实现高效、灵活的部署，尤其针对边缘设备严格的内存和计算限制。

现有方法局限

• 弹性架构（如嵌套子网络）通过共享参数支持多个容量点，但推理时仍以稠密方式激活全部参数，计算开销随容量线性增长，无法利用输入稀疏性。
• 稀疏激活（Mixture of Experts, MoE）通过门控路由仅激活部分专家，实现输入自适应计算，但模型结构固定，难以根据部署环境动态调整容量。
• 两者独立发展：弹性模型缺乏稀疏计算带来的效率，MoE 模型缺乏结构弹性以适配不同硬件资源。此外，现有 MoE 路由多为 token 级，难以在序列或请求级别精确控制总计算预算，导致边缘设备资源利用率不稳定。

为什么这个问题难/重要

• 技术挑战：需要设计统一的训练框架，让模型同时学习弹性嵌套结构和稀疏门控路由；不同容量下路由决策需保持一致性，避免子网络切换导致精度波动。
• 业界关注：一个模型覆盖多种部署场景（从云端到边缘）可大幅降低磁盘占用、按需分配 DRAM 和算力，这是实际 AI 产品落地的核心诉求。

行业类比

就像移动应用需要同一代码库适配不同硬件配置，通过动态加载模块和分配计算资源来实现最佳性能。

核心洞察

• Stepped MoE 的核心创新在于将弹性架构（elastically-nested sub-networks）与稀疏专家路由统一到单一模型中，实现推理时联合调整模型容量与计算开销。以往工作要么独立设计弹性宽度/深度，要么独立设计输入自适应稀疏激活，二者割裂；该工作通过条件化于上下文和目标效率规格，让同一份参数在不同子网络内激活不同专家，跨越 1B 到 4B 容量点，同时保持输入自适应路由，是对弹性模型与 MoE 范式的实质性融合。
• 论文提出的 segment-level routing（段级路由）是基于序列段而非逐 token 的路由机制，使推理复杂度可配置且可预测。传统 MoE 逐 token 路由导致计算量随输入长度线性增长、难以满足固定资源预算；段级路由让一整段 token 共享路由决策，模型在推理时可显式接收目标效率规格（如 DRAM 或算力上限），据此激活相应规模的嵌套子网络和专家，实现精度-效率权衡的细粒度控制，直接回应了稀疏模型在边缘部署中的可服务性问题。
• 该工作证明一个共享参数的主干模型可以同时达到静态 MoE 的精度、接近稠密模型的延迟，并显著节省存储空间。实验显示 Stepped MoE 在 1B/2B/3B/4B 容量点上比对应稠密模型在知识密集型基准上高 2-5%，与独立训练的静态 MoE 版本持平，但推理延迟接近稠密模型（推测段级路由降低了通信开销），且仅需一份模型参数即可覆盖多个容量点，避免了为不同硬件维护多个模型副本的负担，为实际边缘部署提供了高性价比的弹性方案。

方法

输入与目标

模型接收 token 序列 与 目标效率规格（如期望激活的参数量或计算预算）作为联合输入。效率规格可以是一个标量或向量，表示服务设备上可用的 DRAM / 计算资源。

关键模块

1. 弹性嵌套子网络：骨干网络由多个容量层级（例如 1B / 2B / 3B / 4B 参数）组成，且小网络是大网络的严格子集（嵌套结构），所有层级共享参数。
2. 稀疏门控路由：借鉴 Mixture of Experts (MoE)，每个 token 仅激活部分专家子网络，而非全部参数。与普通 MoE 不同，这里路由不是 token 级独立决策，而是 段级路由——将连续 token 分组为 segment，段内共享路由决策，减少门控开销与通信成本。
3. 条件路由机制：路由模块不仅基于上下文，还 条件于目标效率规格。通过将效率规格嵌入并注入门控网络，可以动态调整路由概率或激活阈值，使得同一个模型能在不同规格下自适应选择不同规模的子网络组合。
4. 训练策略：预训练时采用 联合损失，包含语言建模损失和效率对齐损失（鼓励模型在所有规格下保持准确率）。随后进行 监督微调，使模型在下游任务上保持条件路由能力。

输出

给定输入与效率规格，模型输出 下一个 token 的预测分布，同时实际激活的参数总量满足目标规格，实现推理复杂度的可配置控制。

与同类方法的差异

现有工作将 弹性架构（用于部署灵活性）与 稀疏激活（用于输入自适应计算）独立处理，而 Stepped MoE 首次将两者统一在一个框架中，通过段级路由和效率条件实现单一模型跨越多个容量点，并在推理时同时感知上下文和设备约束。

实验

实验设计

论文提出 Stepped MoE，将 弹性结构 与 稀疏门控架构 统一在一个框架中。模型骨干同时基于上下文与目标效率规格进行条件化，在预训练与微调阶段学习激活弹性嵌套子网络中的任务相关参数。实验旨在验证单一模型能否在不牺牲准确率的前提下，覆盖 1B、2B、3B、4B 参数容量，并对比 dense 模型与静态 MoE 版本。

关键发现

• 模型可灵活切换 1–4B 参数规模，在知识密集型基准上比 dense 对应模型准确率高 2–5%。
• 性能与静态 MoE 版本持平，同时推理延迟与 dense 模型相近。
• 通过参数共享，大幅节省设备磁盘空间，并可根据 DRAM 与算力灵活部署。

对比解读

与以往将弹性结构与稀疏激活独立处理的工作不同，Stepped MoE 同时实现部署约束与任务需求的双重自适应。相比纯 dense 模型，它在相同参数规模下取得更高精度；相比静态 MoE，它在保持相当准确率的同时提供更灵活的推理配置。工程上，这一设计让边缘设备无需为不同容量点存储多个模型，且可动态调整计算负载，为资源受限场景提供了更实用的部署路径。

行业影响

落地场景

Stepped MoE 适合需要单一模型覆盖多种算力设备的场景：例如电商 App 内的智能客服、内容平台的个性化推荐与摘要生成、移动端文档/邮件助手。模型可在 1–4B 参数间切换，推理时根据设备 DRAM 和可用算力动态选择容量，无需为低端设备单独裁剪模型。

商业价值

• 降本：单个模型替代多个 dense 模型，节省磁盘空间和部署版本管理成本；边缘设备可利用 DRAM 上限灵活服务。
• 增收/体验提升：比同规模 dense 模型在知识密集型基准上高 2–5%，延迟与 dense 相当，用户感知更快更准；输入自适应路由进一步降低 token 平均计算量。

与现有工作流集成

模型以 MoE 为基础，兼容标准 Transformer 推理栈，可导出为 ONNX 或 TensorRT 引擎。推理时通过 config 指定目标参数量（如 num_active_params=2B）和延迟预算，框架在 segment 级路由选择激活子网络。现有 vLLM/TGI 等 MoE 服务框架可扩展支持动态容量切换。

具体落地 use case：

1. 电商客服：在促销高峰期为每个请求动态分配 2B/3B 参数版本，低峰期切换 1B 以省算力，模型共享权重无需重新部署。
2. 医疗摘要：医院内的平板设备运行 1B 离线模型处理电子病历摘要，桌面工作站可调用 4B 高精度版本，同一套权重满足 HIPAA 合规与算力差异。

局限

• **训练复杂度与稳定性**：同时优化弹性子网络和段级路由，需要引入额外的损失项和多容量采样策略，训练成本显著高于单一 MoE 或单一弹性模型。不同容量配置下的梯度冲突可能导致训练不稳定，需要精心调节超参数（如路由温度、负载均衡系数），实际收敛难度较大。论文未提供训练效率对比，难以评估其相对静态模型的额外开销。
• **路由与负载均衡的潜在问题**：段级路由与弹性嵌套结构结合后，低容量配置下部分专家可能长期不被激活，导致专家退化或路由崩溃。论文未深入讨论动态负载均衡机制在不同目标效率规格下的表现，也未验证路由决策的跨配置一致性。此外，推理时需显式指定效率规格，用户可能难以选择最优配置，自动化程度不足。
• **评估范围与对比基线有限**：实验主要基于知识密集型基准，缺乏对指令跟随、代码生成、多轮对话等任务的系统评估，泛化性存疑。与同类工作的对比不足，未与 **Nested Mixture of Experts**、**MatFormer** 等弹性稀疏模型直接比较，难以判断相对优势。延迟指标仅与稠密模型对照，未展示与静态 MoE 或纯弹性模型的实际硬件加速，磁盘空间节省也可能因多专家参数存储而打折扣。