跳转至

MobileMoE: Scaling On-Device Mixture of Experts

会议: NeurIPS 2026
arXiv: 2605.27358
领域: LLM效率
关键词: 混合专家模型, 端侧部署, 缩放定律, 量化感知训练, 算子融合

一句话总结

针对端侧算力与 DRAM 严苛约束,提出首个亚十亿激活参数的端侧混合专家模型家族 MobileMoE,通过软硬件联合约束缩放定律确立“适度稀疏 + 细粒度专家 + 共享专家”最优架构,结合四阶段训练与 ExecuTorch 融合算子,在匹配 INT4 内存下实现比密集基线快 1.8-3.8 倍的端侧预填和 2.2-3.4 倍的解码加速。

研究背景与动机

混合专家(MoE)架构已成为百亿至千亿级前沿大语言模型(如 DeepSeek-V3、Qwen3-MoE、Grok)事实上的标准骨架,其通过条件计算将总参数量与单 token 推理激活算力解耦,兼顾模型高容量与低推理开销。然而在智能手机、可穿戴设备及具身智能体等边缘计算场景中,端侧模型长期被紧凑型密集结构(如 MobileLLM、MobileLLM-Pro、SmolLM)垄断。受限于亚十亿(sub-billion)激活参数量级与严苛的硬件约束,MoE 在端侧移动硬件上的有效缩放范式、架构超参选择与部署推理机制一直未被系统性探索。

在端侧落地 MoE 存在不可回避的核心工程矛盾:虽然手机端 DRAM 容量已逐步提升(如旗舰机达到 12GB 至 16GB),为容纳数十亿总参数的稀疏模型提供了静态驻留空间,但终端应用可用内存上限(通常受控在 3-5GB 预算内)与移动处理器极其受限的内存带宽、无原生稀疏加速算子的现状,使得云端通用的超大专家数或粗粒度路由策略完全失效。如果专家拆分过粗或激活路径过于离散,分散的小算子访存碎片化和频繁的内核发射延迟将彻底吞噬稀疏计算省下的理论 FLOPs 收益,造成严重的端侧推理减速。此外,现有 MoE 缩放定律(如 Clark et al.、Ludziejewski et al.)均以海量云端多卡并行与千亿参数为假设,未曾建模包含静态 INT4 权重和动态长上下文 KV Cache 的端侧物理 DRAM 联合约束。

为打破这一瓶颈,本文切入的角度是:系统化建立由移动端内存与计算物理边界驱动的端侧 MoE 缩放定律,指导亚十亿激活区间的专家拓扑设计,并辅以全流程四阶段训练优化与底层算子融合。核心 idea:推导联合移动端 DRAM 与计算开销的端侧 MoE 缩放定律,确定“适度稀疏(8专家)+ 细粒度划分(8细粒度子专家)+ 共享专家通路”的最优解,配合四阶段训练配方与端侧紧凑分组矩阵融合算子,打破端侧大模型长期依赖密集架构的固有范式。

方法详解

整体框架

MobileMoE 旨在设计、训练并在真实商用移动硬件上高效运行亚十亿激活规模的稀疏大语言模型。模型整体输入离散 Token 序列,经由嵌入层投影后逐层穿过基于深度瘦长(deep-and-thin)比例设计的解码器层,最终输出预测分布。每个 MoE 解码器层将传统的密集前馈网络(FFN)替换为由路由器门控的细粒度专家组与常驻共享专家组成的混合结构。整体流程涵盖:端侧缩放定律驱动的架构搜索、四阶段(预训练 → 中期训练 → 指令微调 → INT4量化感知训练)递进优化训练流、以及面向移动端硬件(CPU/GPU)的高效内存映射与融合内核部署。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入 Token 序列"] --> B["端侧缩放定律架构指导<br/>约束内存与计算解耦最优拓扑"]
    B --> C["路由门控与细粒度专家分发<br/>Sigmoid 独立评分与顶k选择"]
    B --> D["常驻共享专家通用计算<br/>无路由直连 dense 路径"]
    C --> E["端侧算子级融合与执行<br/>计数排序重排与分组矩阵乘"]
    D --> E
    E --> F["输出隐藏状态与下一 Token"]

关键设计

1. 联合内存与计算约束的端侧 MoE 缩放定律:确立 Pareto 最优模型拓扑 云端 MoE 架构追求无限制扩展总容量以换取泛化极限,而端侧推理必须在固定的物理内存预算(静态 INT4 权重与指定上下文长度的 KV Cache 总和)与固定单步计算开销下取得最低损失。本文提出了涵盖模型活跃参数量 \(N_{\text{act}}\)、训练数据量 \(D\)、专家数量单调变换 \(\hat{E}\) 以及结构超参 \(x\)(专家细粒度 \(g\) 与共享专家 \(s\))的端侧缩放定律: $$ \mathcal{L}(N_{\text{act}}, D, \hat{E}, x) = A_{x}\hat{E}^{\delta_{x}}N_{\text{act}}^{\alpha_{x}+\gamma_{x}\ln\hat{E}} + B_{x}\hat{E}^{\omega_{x}}D^{\beta_{x}+\zeta_{x}\ln\hat{E}} + c_{x} $$ 其优化目标是在满足训练算力 \(F_{\text{train}}=6N_{\text{act}}D\)、推理前向算力 \(F_{\text{inf}}=2N_{\text{act}}\) 以及端侧内存预算 \(\mathcal{M}(N_{\text{total}}, T) = \frac{b_{w}}{8}N_{\text{total}} + \frac{b_{\text{kv}}}{8}\cdot 2Tn_{l}n_{\text{kv}}d_{h} \leq M\)(其中 \(M \leq 5\,\text{GB}\))的多重约束下最小化预测损失。通过解耦扫描实验,论文证实了三个关键法则:其一,在控制总内存预算时,单纯增加总专家数 \(E\) 收益迅速饱和甚至在 \(E=32\) 时由于优化退化而出现负效应,确定 \(E=8\) 为兼顾内存与训练效率的甜点区;其二,将每个专家细分成 \(g=8\) 个子专家(共 64 个路由子专家,激活 Top-8)能在不增加活跃与总参数的前提下赋予组合路由更高的灵活性,显著降低困惑度;其三,将 4 个激活子专家槽位替换为容量相等的常驻共享专家(即 60 个细粒度路由专家选 4 个,外加 1 个相当于 4 倍子专家尺寸的共享专家),不仅提升跨域通用表征能力,还大幅提升多卡专家并行训练的通信吞吐量。基于该准则,衍生出 MobileMoE-S(0.3B 激活 / 1.3B 总量)、MobileMoE-M(0.5B 激活 / 2.8B 总量)与 MobileMoE-L(0.9B 激活 / 5.3B 总量)三款主力规格。

2. 兼顾稳定性与吞吐的四阶段端侧 MoE 训练体系:化解稀疏优化脆弱性 在亚十亿极小激活规模下,MoE 路由门控极易陷入专家负载崩塌或梯度数值扰动。为此,预训练(PT)阶段引入无辅助损失负载均衡(Bias-based balancing,以 \(\lambda_{lb}=10^{-3}\) 动态更新专家偏置项避免对抗性正则破坏语言建模)与 Router z-loss(\(\lambda_{z}=10^{-4}\)),采用 Sigmoid 独立门控替代竞争性 Softmax,全流程维持 FP32 门控精度与分组矩阵乘法(GMM)。预训练阶段在 6T 开源高质量代币上完成,以显著低于密集模型的 token 开销达到领先表现。第二阶段中期训练(MT)利用 500B 高质量精选知识、代码与数学数据退火,同时将上下文从 2,048 扩展至 8,192,强化专家专业化分工。第三阶段指令微调(SFT)覆盖 80M 指令样本,并切换为无丢弃(dropless)分发调度,防止结构化指令回答因容量截断而受损。第四阶段为确保模型适配手机端 DRAM,执行对称组级 INT4 量化感知训练(QAT,group size=32,激活动态量化至 INT8),在量化反向传播时将路由器权重冻结或维持在 FP32 计算,仅以 0.5% 的额外内存开销彻底稳固低比特下的路由决策准确率。

3. 面向移动 CPU/GPU 的端侧紧凑融合算子:转化理论稀疏为物理加速 标准端侧运行时(如 XNNPACK)仅对连续密集矩阵乘法进行了极致汇编优化,原生执行细粒度 MoE 将产生数十个小维度矩阵计算,导致严峻的内核发射开销与访存重排损耗。针对该瓶颈,论文在 ExecuTorch 框架内定制开发了一体化融合 MoE 算子:在分发阶段利用计数排序(Counting Sort)对当前 Batch 内各 Token 依据其指派的专家 ID 进行物理内存就地重排,使所有路由至同一专家的 Token 连续对齐,进而调用 torchao 高性能内核直接执行单次大块密集分组 GEMM 计算;同时将 Top-\(k\) 选择、门控计算、Gate/Up 投影拼接、SwiGLU 非线性激活、Down 投影以及加权发散写回(Weighted-scatter unpermute)完全融合进单一算子调用,最大程度摊薄中间张量搬运与动态量化代价。在真实端侧推理时,结合按需内存映射(mmap)机制,仅将触发激活的专家参数调入活跃缓存,使得 MobileMoE-S 的端侧实测峰值物理内存(Peak RSS)在长上下文场景下反而低于密集基线。

实验关键数据

主实验

在涵盖常识、知识、科学、阅读理解和推理的 14 个通用基准,以及由数学、代码、指令遵循和前沿推理组成的 8 个进阶基准上,MobileMoE 系列均展现出超越传统密集模型与前序 MoE 模型的性能边界。

模型 类型 激活/总参数 MMLU (5-shot) GSM8K (8-shot) HumanEval (0-shot) 14项基准总均分
Gemma 3 270M Dense 270M / 270M 26.5 5.8 12.8 32.5
SmolLM2 360M Dense 362M / 362M 25.9 10.0 0.0 35.8
Qwen3.5 0.8B Dense 749M / 749M 50.2 45.7 31.1 42.7
MobileLLM-Pro Dense 1.1B / 1.1B 38.7 31.8 59.8 40.0
Llama 3.2 1B Dense 1.2B / 1.2B 46.1 46.0 36.6 46.9
OLMoE-1B-7B MoE 1.3B / 6.9B 52.7 49.1 36.0 55.6
MobileMoE-S (SFT) MoE 272M / 1.3B 42.6 52.2 44.5 46.7
MobileMoE-M (SFT) MoE 528M / 2.8B 53.9 67.5 60.4 55.3
MobileMoE-L (SFT) MoE 922M / 5.3B 60.1 77.6 65.2 60.1

消融实验与端侧延迟分析

消融研究分别在缩放定律结构消融(500B Token 固定算力对比验证)以及两款真实旗舰智能手机(Samsung Galaxy S25,骁龙 8 Elite;iPhone 16 Pro,Apple A18 Pro)上展开端侧实测。

测试硬件与配置 对应模型 INT4 权重内存 首字时延 TTFT (1k 输入) 解码速率 (1k 输入) 8k 峰值 RSS (真实输入) 14项总评均分
Samsung S25 (4 CPU) MobileLLM-Pro (Dense) 0.55 GB 4.26 s 45.8 tok/s 1.91 GB 45.5%
Samsung S25 (4 CPU) MobileMoE-S (Ours) 0.68 GB 2.01 s (2.1×) 112.0 tok/s (2.4×) 1.49 GB (-22%) 44.0%
Samsung S25 (4 CPU) MobileMoE-M (Ours) 1.48 GB 4.32 s (1.0×) 64.6 tok/s (1.4×) 2.77 GB 52.5%
Samsung S25 (4 CPU) MobileMoE-L (Ours) 2.75 GB 8.99 s (0.5×) 36.6 tok/s (0.8×) 4.71 GB 57.8%
iPhone 16 Pro (Metal GPU) MobileLLM-Pro (Dense) 0.55 GB 1.20 s 59.2 tok/s - 45.5%
iPhone 16 Pro (Metal GPU) MobileMoE-S (Ours) 0.68 GB 0.32 s (3.8×) 151.3 tok/s (2.6×) - 44.0%

结构设计消融同样验证了各要素的必要性:在固定算力预算下,相比粗粒度基线(\(g=1\)),细粒度配置(\(g=8\))在保持参数量与计算量不变的前提下带来显著的验证损失下降;引入常驻共享专家(\(s=\checkmark\))不仅进一步降低了收敛损失,而且在相同硬件节点上的每步训练耗时缩短了约 12%,显著提升了并行通信效率。

关键发现

  • 稀疏化在端侧兼得计算与显存优势:MobileMoE-S 在匹配的 INT4 静态模型体积下,解码生成速率提升 2.2-3.4 倍,首字预填加速 1.8-3.8 倍;而在 8k 上下文下,由于其采用较小的主干隐藏层维度与较少层数,使得 KV Cache 与动态张量占用锐减,实测真实端侧峰值 RSS 反而比密集基线降低 22%(1.49GB vs 1.91GB)。
  • 知识与理解类任务收益最大:MoE 扩展总参数的红利在知识密集型(MMLU、TriviaQA)和复杂长文本阅读(DROP)上表现最为突出,如 MobileMoE-L 在 922M 激活参数下取得 60.1% 的 MMLU 成绩,大幅超越 1.3B 激活 / 6.9B 总量的 OLMoE(52.7%)。
  • 训练阶段各司其职:常识推理能力在 6T 预训练阶段即完全饱和;中期训练(MT)是知识强化与长上下文扩展的核心推手;而多步推理与数学代码求解能力(如 GSM8K 从 55.7% 跃升至 77.6%)绝大部分依赖无丢弃分发的指令微调(SFT)充分激发。

亮点与洞察

  • 端侧软硬件联合受限缩放法则:论文跳出了单纯以云端参数和算力为驱动的传统 Scaling Law 窠臼,首次将由 DRAM 限制的静态权重和可变 KV Cache 内存预算作为硬约束显式引入数学规划,给出了 sub-billion 级别 MoE 架构设计的明确理论抓手。
  • 端到端商用端侧落地闭环:不仅止步于仿真或 GPU 模拟,作者团队开发了基于 ExecuTorch 的重排与分组 GEMM 融合算子,在真实高通骁龙 8 Elite 与苹果 A18 Pro 手机上实机运行,彻底解决了 MoE 在端侧由于访存碎片化导致“算力低却跑得慢”的工业落地痛点。
  • 可复用的 QAT 路由稳定范式:在 4-bit 权重量化训练中,创新性地保持门控网络为 FP32 精度,仅引入约 0.5% 的可忽略显存开销,却有效杜绝了因梯度量化噪声造成的路由分配震荡,实现了模型在 4 倍压缩下保持 BF16 精度 95% 以上能力的极高保真度。

局限与展望

  • 作者承认的局限:端侧推理基准评测主要聚焦在移动 CPU 及 GPU 框架(XNNPACK/MLX),尚未全面适配各家芯片专有的移动 NPU 硬件加速器;此外后训练配方目前采用标准 SFT,尚未结合强化学习测试时计算(TTC)或思维链自我反思机制。
  • 自主审视的局限:在极端严苛的多轮复杂工具调用与纯多步逻辑推理上,超小尺寸的 MobileMoE 与两倍以上激活参数的紧凑型前沿密集模型(如 Qwen3.5 2B)仍存在一定表现差距,且多专家并发加载仍对移动存储持续读取带来隐性功耗开销。
  • 未来改进方向:探索针对移动端 NPU 的推测解码(Speculative Decoding)与 MoE 动态专家剪枝技术;将知识蒸馏与思维链强化学习融合至端侧 MoE 训练体系中。

相关工作与启发

  • vs MobileLLM / MobileLLM-Pro: 后者是深而窄(deep-and-thin)端侧密集小模型的代表,通过权重共享和极端架构比例榨取参数效率,但要提升能力必须线性增加推理 FLOPs;本文继承了其深而窄的主干长宽比先验,但通过引入细粒度稀疏专家将容量与计算解耦,在相同端侧内存占用下斩获数倍推理加速。
  • vs OLMoE-1B-7B: OLMoE 探索了百亿以下 MoE 的构建,但采用粗粒度专家且无共享专家,总参数量偏大(6.9B)难以在移动端运行,且缺乏端侧量化与专用推理内核;MobileMoE-L 凭借“细粒度 + 共享专家 + 缩放最优拓扑”,以减少 30% 激活参数和 23% 总参数的代价,在下游多任务上实现了全面超越。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次建立面向端侧 DRAM 与 FLOPs 双重约束的 MoE 缩放定律并打通真实手机端高效融合推理。
  • 实验充分度: ⭐⭐⭐⭐⭐ 包含 14 项通用与 8 项进阶评测、从 500B 到 6T 的完备消融轨迹,以及两款旗舰手机的全套实机延迟与物理内存剖析。
  • 写作质量: ⭐⭐⭐⭐⭐ 论证结构极其严谨,公式、理论推导、训练阶段分解与硬件部署细节自洽连贯。
  • 价值: ⭐⭐⭐⭐⭐ 为轻量化端侧大模型从密集结构迈向稀疏混合专家时代奠定了里程碑式的工程与理论基石。