UniMoMo:基于专家合并的大规模推荐模型MoE加速 UniMoMo: Expert Merging-Based MoE Acceleration for Large Recommendation Models
用校准流量驱动的专家合并把推荐MoE压缩为更小的标准MoE,近无损且提速至2.21倍
前置知识
稀疏混合专家(Sparse MoE)
一种条件计算架构:每层包含 $E_\ell$ 个专家网络(这里是 SwiGLU 前馈块)和一个路由器。路由器为每个输入 token 计算打分 $z_\ell(x)$,只激活 top-$k$ 个专家,按 softmax 归一化权重 $a_{\ell,e}(x)$ 加权输出 $y_\ell=\sum_{e\in A_\ell} a_{\ell,e} f_{\ell,e}(x)$。它把存储容量与激活计算分离,用有限计算量支撑更大参数规模。
本文的全部操作对象就是训练好的稀疏 MoE 层:压缩目标是把每层专家从 $E_\ell$ 减到 $M_\ell$ 并同步缩小路由器,同时保持标准 top-$k$ 推理形态,不理解这套机制就无法读懂问题设定。
SwiGLU 门控前馈块
专家内部常用的激活结构:中间激活为 $\mathrm{SiLU}(XW_G^\top)\odot XW_U^\top$,再乘下投影 $W_D$。每个专家约含 $3dh$ 个参数(上投影 U、门控 G、下投影 D 三部分)。由于 SiLU 非线性门控的存在,先平均参数再前向与先计算输出再平均不可交换。
论文明确指出简单参数平均对非线性 SwiGLU 无效,这正是引入最小二乘激活修正(Eq.9–10)的原因,理解这一结构才能明白重建步骤为何必不可少。
KL 散度与亲和度评分
KL 散度 $q(i\|j)$ 衡量两个分布的差异且不对称。本文把每个专家在共享校准 token 上的输出概括为各向同性高斯(均值向量 $\mu_{\ell,e}$、标量方差 $v_{\ell,e}$),用两个高斯摘要间的对称化距离 $D_{ij}=\frac{1}{2}(q(i\|j)+q(j\|i))$ 定义亲和度 $s_{ij}=1/(1+D_{ij})\in(0,1]$。
这是构建专家亲和图、决定谁和谁合并的核心度量,替代了传统基于参数距离的判据,是方法部分最先出现的数学工具。
图粗化与平均链接聚类
图粗化指反复把图上最相似的节点对合并成超节点,直到节点数达到目标。平均链接(average-link)用两组间所有点对相似度的均值 $\mathrm{sim}(A,B)=\frac{1}{|A||B|}\sum_{i\in A}\sum_{j\in B} s_{ij}$ 衡量簇间相似度,避免了单链接的链式效应和全链接对离群对的敏感。
整篇方法被形式化为'预算约束下的图粗化':$E_\ell$ 个专家节点被贪心合并到 $M_\ell$ 个簇,保护规则也嵌入在合并的合法性判定里,这是方法主线。
Moore–Penrose 伪逆与最小二乘
求 $\min_{T}\|PT-Q\|_2$ 的闭式解为 $T=P^\dagger Q$,其中 $P^\dagger$ 是 Moore–Penrose 伪逆,即使 $P$ 秩亏也给出最小范数解。本文用它求线性映射 $T_m$,把合并后专家的实际中间激活 $P$ 修正到流量加权的目标激活 $Q$,再把 $T_m$ 折入下投影权重 $W_D\leftarrow \bar W_D T_m^\top$。
这是修正 SwiGLU 非线性失配的手段,修正被折叠进存储参数,因此推理时零额外算子开销,这是'输出仍是标准 MoE'这一卖点的技术保障。
路由熵与暴露保护机制
用校准集上的路由质量分布 $p_{\ell,e}$ 计算归一化熵亏 $S_\ell=1-H/\log E_\ell$:均匀路由时 $S_\ell=0$,流量一家独大时趋近 1。层越集中,按 $\gamma_\ell=\gamma_{\min}+(\gamma_{\max}-\gamma_{\min})S_\ell^\beta$ 保护越多专家,取路由质量 $u_{\ell,e}$ 最大的 $K_\ell=\lceil\gamma_\ell E_\ell\rceil$ 个为保护集,禁止其被反复合并。
高流量专家的重建误差会随残差流放大到后续层,这套层自适应保护是论文防止精度下降的关键,也是与 HC-SMoE、MergeMoE 等既有专家合并工作的本质差异。
采样评估指标 HR@10 / NDCG@10 / AUC
推荐排序常用协议:每个用户的正样本对 99 个采样负样本排序。HR@10 检查正样本是否进入前 10;NDCG@10 按命中位置以 $1/\log_2(r+1)$ 折扣,奖励排在靠前的命中;AUC 统计正样本得分高于负样本的比例。所有方法共享同一份划分和候选列表,五次运行取均值。
论文以五次运行平均 NDCG@10 作早停与主指标,所有核心结论(如 99.92%–102.30% 的源相对比率)都建立在该协议之上,读懂指标才能正确解读实验。
研究动机
工业推荐系统需要大容量模型,但最终能部署哪种容量由服务成本决定。近年 Wukong、DHEN、RankMixer、TokenMixer-Large 等工作不断扩大推荐骨干的特征交互深度与稀疏参数容量;稀疏 MoE 进一步用条件计算为长行为序列、异构交互频率和多任务学习扩容(MMoE、MoS、频率感知专家、SMES 等)。问题在于:一套 MoE 推荐检查点训练完成后,仍然要存储完整专家库、推理时对全部 $E_\ell$ 个专家打分路由,而架构设计层面的工作并不回答'训练完成后到底需要多少专家库来服务'这个检查点层面的问题。现有加速路线要么改变路由方式或执行内核(均衡路由、expert-choice、MegaBlocks、ScatterMoE),要么设计并从头训练新的稀疏架构,都保留训练好的专家身份,缺乏一种检查点级机制能在明确的服务预算下移除冗余专家容量。
本文的目标是本文要解决的是一个部署问题:给定一个已训练好的推荐 MoE 检查点和一个明确的每层专家预算 $M_\ell<E_\ell$,把它转换成一个更小的标准稀疏 MoE——保留原有 top-$k$ 路由形态、不引入任何压缩专用的在线模块,从而能直接落进现有部署管线;同时在给定预算下最小化对被路由专家功能的扰动。转换只允许使用未标注的校准交互(测试交互全程排除),转换后进行与源模型协议一致的短暂监督微调。作者特别强调 $M_\ell$ 是服务预算而非为刷指标调出来的超参,排序质量在转换之后度量,而不是作为算法约束被强加,因此论文把 NDCG、HR、AUC、参数量与延迟分开报告。
与已有工作不同的是,既有专家缩减方法依赖参数距离或输出重建误差,但这两个判据在推荐场景都不完整:参数距离无法反映两个专家在推荐器产生的具体隐藏状态上是否给出相似响应——训练后两个专家可能参数相距很远却行为相近,也可能参数接近而非线性门控导致输出大幅偏移;输出相似性同样不够,因为同样大小的重建误差落在处理大量 token 的高流量专家上,随残差流向后续层的危害远大于落在低流量专家上。因此部署划分必须在同一专家数约束下联合考虑行为兼容性与路由暴露度。UniMoMo 的独特切入是把该转换表述为一个以校准数据为条件、由推荐流量约束的图粗化问题:亲和度、可合并规则、重建权重全部来自同一条未标注校准流,而不是把每个专家对等同对待。
核心方法
直觉上,两个专家如果对相同的推荐状态给出相似响应、且都不承载太多路由流量,就可以安全合并成一个;反之,处理大量 token 的高流量专家一旦被反复修改,重建误差会沿残差流放大到后续层。技术路线分四步:(1) 函数画像——旁路路由器,让每层全部 $E_\ell$ 个专家在最多 20 批共享校准 token 上各前向一遍,把每个专家的输出概括为均值 $\mu_{\ell,e}$ 和标量方差 $v_{\ell,e}$(Eq.4),同时累计原 top-$k$ 的概率加权路由质量 $u_{\ell,e}$;(2) 用高斯摘要间的对称化 KL 距离构建专家亲和图(Eq.5–6),并用层内路由熵决定保护集;(3) 在保护约束下贪心地平均链接粗化到 $M_\ell$ 个簇(Eq.7);(4) 按路由质量加权平均簇内 U/G/D 参数,用最小二乘修正 SwiGLU 中间激活失配(Eq.9–10),路由行取算术平均初始化,最后对压缩后的完整模型做短暂的监督微调。
核心创新是把专家缩减定义为'预算约束下的部署划分决策',并让行为相似性与流量暴露度在同一条校准流里联合发挥作用。与 HC-SMoE、MergeMoE 等输出重建方法的本质区别有三点:其一,亲和度既不是参数距离也不是无条件的输出误差,而是所有专家在同一批共享输入上的响应比较——共享输入保证了成对差异只归因于专家行为而非不同的路由分布;其二,层自适应保护机制(Eq.13–14)把路由熵转化为确定性的可合并规则:两个受保护簇不能合并,受保护单例最多吸收一个非保护簇后退出后续合并,均匀层少保护、集中层多保护;其三,重建针对 SwiGLU 的非线性失配做最小二乘激活修正并折入存储参数,推理时保持标准 top-$k$ 计算、无任何 UniMoMo 专用模块。三者共同构成一个'推荐原生'的转换目标,输出可以直接被现有专家内核执行。
方法步骤详情
对每层的完整流程:输入为冻结的 $E_\ell$ 个 SwiGLU 专家、路由器与最多 20 批校准输入。第一步画像:旁路路由器,每个专家在共享 token 上计算输出并逐批累计 $\mu_{\ell,e},v_{\ell,e}$(存储仅 $O(E_\ell d)$),同时记录原 top-$k$ 的概率加权路由质量 $u_{\ell,e}$ 与分布 $p_{\ell,e}$。第二步定保护集:计算归一化熵亏 $S_\ell$,得 $\gamma_\ell=0.1+0.3S_\ell$、$K_\ell=\lceil\gamma_\ell E_\ell\rceil$(8 专家时 $K_\ell\in[1,4]$),取 $u$ 最大的 $K_\ell$ 个专家为保护集 $\mathcal{P}_\ell$。第三步粗化:每个专家初始化为单节点,边权 $s_{ij}=1/(1+D_{ij})$,每轮合并平均链接相似度最高的合法对,直到剩 $M_\ell$ 个簇;若保护簇数超过预算,仅在该轮禁用保护选最高相似对,下一轮恢复约束。第四步重建:簇内按 $w_e\propto u_e$ 加权平均 U/G/D 三组参数;用校准状态计算目标激活 $Q=\sum_e w_e\,\mathrm{SiLU}(XW_G^\top)\odot XW_U^\top$,与合并后实际激活 $P$ 做最小二乘得 $T_m=P^\dagger Q$,折入 $W_D\leftarrow\bar W_D T_m^\top$;路由行取簇内算术平均。全部层替换后按源协议监督微调,按验证 NDCG@10 选检查点,测试交互全程隔离。
技术新颖性
技术新颖性体现在四点。第一,问题定义新:明确把'训练后专家库缩减'定义为检查点转换,输出必须是标准稀疏 MoE,从而与改路由/改内核的加速工作和从头训练的架构工作正交,也让同一检查点可以在多个服务预算下导出。第二,度量新:各向同性高斯摘要只需 $O(E_\ell d)$ 存储,即可在全协方差需要 $O(E_\ell d^2)$ 的宽推荐层场景下逐批刻画专家行为;均差惩罚持久的输出偏移、方差项区分响应尺度,$s_{ij}$ 只用于排序合并候选,不需要概率解释或全局阈值。第三,约束新:暴露保护被表述为确定性合并优先级而非硬约束——保护数超预算时单轮回退,保证任意 $M_\ell$ 可行。第四,理论交代完整:Eq.16 给出最终输出误差上界,拆成最小二乘残差项(恰为 Eq.10 最小化的量)乘以合并下投影范数,加上源下投影离散度项,说明初始化何时精确、误差何时出现。相比 MergeMoE 虽也用路由统计加最小二乘,但缺乏推荐流量条件下的暴露规则与共享输入协议。
实验结果
主实验在 Amazon Beauty(过滤前 2,023,070 次交互、52,374 用户、121,291 物品)、KuaiRec 大矩阵(7,176 用户、10,728 物品、12,530,806 次交互)和 TenRec(约 500 万用户、1.4 亿交互)上覆盖 2/4/6 个 MoE 块和两个部署预算,共九个配对设置,五次运行取均值。六块主表(Table 1)中,8→4 专家的 UniMoMo 在 Beauty 上 NDCG@10 达 0.2686(Origin MoE 0.2649、MergeMoE 0.2682),延迟 3.127ms 对源 4.833ms;KuaiRec 上 0.6963 对源 0.6893,延迟 4.210ms 对 6.860ms;TenRec 上 0.5987 对 0.5938,延迟 7.451ms 对 9.566ms。跨九个设置,4 专家检查点取得源相对五次均值 NDCG@10 比率 99.92%–102.30%、实测 A100 加速 1.28×–1.63×;2 专家 top-1 激进点为 98.36%–104.24%、加速 1.47×–2.21×。相对同预算同适配预算的 MergeMoE,NDCG 增益为 0.0004/0.0028/0.0051(Beauty/KuaiRec/TenRec)。消融(Table 2)显示:去掉双信号 NDCG 0.2662、HR 0.4401,完整方法 0.2686、0.4445,合计提升 0.0018/0.0034。专家数敏感性(Table 3):8→4 延迟降 40.36% 保留 99.37% NDCG(单位延迟 NDCG 提升 66.61%),8→2 再省 15.77% 延迟但 NDCG 掉 1.68%。深度敏感性(Table 4):2→6 块 NDCG 涨 6.04%、HR 涨 6.29%,6→8 仅涨 0.63% 而延迟涨 145.70%,支持 6 块为质量-效率工作点。补充实验(Table 5)在 Mamba4Rec/TransAct/TWIN/SDIM 四骨干 × MicroVideo/KuaiVideo/Ebnerd 三数据集上 24 项 AUC/GAUC 全部第一,领先最强替代 MoS 约 0.51–1.37 个百分点。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 排序推荐(Amazon Beauty,6 MoE 块,8→4 专家,top-2) | NDCG@10 | UniMoMo 0.2686(延迟 3.127ms,参数 20.25M) | Origin MoE 0.2649(4.833ms,24.97M);MergeMoE 0.2682(3.108ms) | 相对源检查点约 +1.4%,相对同预算 MergeMoE +0.0004,加速约 1.55× |
| 排序推荐(KuaiRec,6 MoE 块,8→4 专家,top-2) | NDCG@10 | UniMoMo 0.6963(延迟 4.210ms,参数 10.82M→6.10M) | Origin MoE 0.6893(6.860ms);MergeMoE 0.6935(4.250ms) | 相对源 +1.0%,相对 MergeMoE +0.0028,加速 1.63× |
| 排序推荐(TenRec,6 MoE 块,8→2 专家 top-1,UniMoMo-A) | NDCG@10 / HR@10 / AUC | 0.6054 / 0.7894 / 0.9042(延迟 6.515ms,参数 155.46M) | Origin MoE 0.5938 / 0.7692 / 0.8966(9.566ms,162.54M) | 相对源 +2.0% NDCG,加速 1.47×;是激进操作点下全面超源的案例 |
| 推理延迟(A100 80GB,BF16,batch 256,九个数据集-深度配对) | 实测加速比 | 4 专家 1.28×–1.63×;2 专家 top-1 1.47×–2.21× | Origin MoE(8 专家 top-2) | 全部配对设置为正收益,且排名指标保持源相对 98.36%–104.24% |
| 跨骨干可移植性(Mamba4Rec/TransAct/TWIN/SDIM × MicroVideo/KuaiVideo/Ebnerd) | AUC / GAUC(共 24 项) | 24 项全部第一,例如 Mamba4Rec+MicroVideo AUC 70.68 | MoS(最强替代,同例 69.46)及 Vanilla/DSelect-k/GShard/Expert Choice | 领先最强替代 0.51–1.37 个百分点 |
| 效率-质量权衡(Amazon Beauty,专家数扫描) | NDCG@10 保留率 / 延迟降幅 | 8→4:保留 99.37% NDCG,延迟 -40.36%(3.127ms vs 5.243ms) | 8 专家对照(NDCG 0.2703,5.243ms) | 单位延迟 NDCG 提升 66.61%;8→2 再省 15.77% 延迟但 NDCG -1.68% |
局限与改进
作者承认的局限写在第 5 节:实验证据限于 RankMixer 骨干、每层 8 个源专家、2–6 个 MoE 块、采样排序和单卡 A100 推理;划分完全依赖有代表性的校准流量,校准流中未出现的行为无法影响亲和度与暴露估计;路由质量本身识别不出'低流量但编码了稀有有用行为'的专家;虽然不采集新数据,校准日志仍受底层推荐系统的隐私与公平约束。我的补充观察:其一,与 MergeMoE 的差距(如 Beauty 上 0.0004)远小于五次运行的波动量级,作者也只称'匹配的五次运行增益'而非统计优越,工程上说服力主要来自延迟与参数的确定性下降;其二,评估采用 1 正 99 负的采样排序而非全库检索,与工业全量排序场景存在差距;其三,2 专家激进点在 $K_\ell>2$ 时必须禁用保护做回退,此时'保护高流量专家'的核心主张实际失效,作者也承认不把该点作为硬保护证据,而高于 100% 的比率部分来自微调而非压缩本身;其四,$E_\ell=8$ 很小,亲和图 $O(E_\ell^2 d)$ 构建与 $O(E_\ell^3)$ 合并成本在更大专家库(LLM 级 MoE)下的扩展性未验证;其五,$\gamma_{\min}=0.1,\gamma_{\max}=0.4,\beta=1$ 全程沿用未调参,跨域稳健性存疑。
独立分析的弱点
独立分析有四个弱点。第一,校准代表性是单点依赖:合并计划是一次性离线产物,若线上流量分布漂移(新类目、季节性、新用户群体),被合并的专家组合可能迅速过时,改进方向是把转换做成可周期重算的流水线,或引入流量漂移检测自动触发重合并。第二,逐层独立划分浪费跨层冗余:不同层的专家可能对相似行为模式重复建模,做跨层联合划分或按路由轨迹分组(如 MoE Pathfinder 的思路)有望进一步压缩预算。第三,最小二乘只修中间激活:Eq.16 的第二项(源下投影离散度 $\sum_e w_e\|H_e\|_F\|W_{D,e}-\bar W_{D,m}\|_2$)没有被直接最小化,且修正目标与最终排序损失脱节,改进方向是对层输出乃至端到端 NDCG 做加权蒸馏式重建。第四,预算分配是外生给定的:所有层共用同一 $M_\ell$,而不同层的熵亏 $S_\ell$ 不同、可压缩性天然不同,可以按 $S_\ell$ 自动分配各层预算,让集中层少压、均匀层多压,把'层自适应保护'进一步升级为'层自适应预算'。
未来方向
作者指出的延伸:把检查点转换确立为部署原语,让同一训练好的检查点在多个服务预算下导出(文中 4 专家/2 专家两端点即是示例),并在更多推荐骨干上验证可移植性(附录已在 Mamba4Rec、TransAct、TWIN、SDIM 四骨干上给出初步支持)。基于成果可以继续延伸的方向包括:与量化、剪枝等正交压缩手段组合,把专家参数量的下降折算成端到端收益;把熵驱动的暴露保护迁移到 LLM 级 MoE(Mixtral、DeepSeek-MoE)的专家合并,检验其在 $E_\ell$ 更大、路由更稀疏时是否依然有效;面向多任务推荐 MoE 的任务感知合并,避免合并掉只服务稀有任务的小流量专家;在分布式训练/推理中量化专家库缩减带来的通信与设备放置收益;以及为保护机制建立近似保证,把 Eq.16 的误差界与路由质量联系起来,形成可联合优化的目标函数。
复现评估
复现难度中等偏低。数据全部公开:Amazon Beauty、KuaiRec(big matrix)、TenRec,附录另用 MicroVideo/KuaiVideo/Ebnerd;论文给出完整协议(保留至少 5 次交互的用户、按时间戳把最后两次交互留作验证/测试、更早前缀做训练、1 正 99 负采样评估、所有方法共享同一划分与候选列表)。方法细节完备:Eq.1–17 覆盖全部公式,超参固定($\gamma_{\min}=0.1$、$\gamma_{\max}=0.4$、$\beta=1$、校准 20 批、AdamW 训练学习率 $10^{-3}$、适配阶段 $5\times10^{-4}$ 至多 5 epochs 且 patience 为 3)。算力需求很低:压缩只需前向校准(一次性成本 $O(NE_\ell dh)$ 画像、$O(E_\ell^2 d)$ 建图、$O(E_\ell^3)$ 合并、$O(Nh^2+h^3)$ 最小二乘),骨干是小型 RankMixer(hidden 128、序列长 16、每层 8 专家 top-2),延迟测量单卡 A100 80GB、BF16、batch 256、80 次同步前向。未提及开源代码或检查点,贪心粗化与最小二乘重建需自行实现,这是主要门槛;五次运行、固定种子偏移等评测协议均写明,指标对齐应无歧义。
论文图表
分两部分展示方法的输入输出:(a) 离线阶段,UniMoMo 以训练好的 MoE 推荐器、未标注校准交互和目标专家预算(专家计数)为输入,输出一个压缩后的 MoE 推荐器;(b) 在线阶段,压缩后的模型接收用户交互历史和候选物品,输出排序分数或 top-K 推荐列表,与普通推荐器使用方式完全一致。
一图界定问题:输入是检查点加预算,输出是标准 MoE,没有任何压缩专用在线模块。理解了这个输入输出约定,才能明白'检查点转换'这一部署问题与改架构、改路由方法边的区别。