你在小厨房里做一桌菜,常用调料放在手边,其他都收在柜子里。如果每做一步就换一套调料,时间多半耗在开柜门、拿取和归位上。混合专家模型也有类似问题:它明明每次只调用少数“专家”,却可能处理一个词就换一批专家,设备只好反复搬运参数。
StickyMoE想解决的正是这笔隐形开销。它不等模型部署后再优化缓存,而是在预训练时就教路由器养成一种习惯:语义连贯的一段文字,尽量连续调用相同的专家。论文在小型和中型模型上报告,专家切换率最多下降59%,中型模型的困惑度还同时改善。不过,目前证据全部来自作者Ali Kayyam的这篇论文,尚无独立复现;论文表格中的部分具体数值在供稿文本里缺失,因此下文只采用能够确认的结果。
MoE省了计算,为什么还会卡在搬运上?
混合专家模型(Mixture-of-Experts,MoE)可以理解为一个内部设有多个专科组的模型。每个文字单元,也就是token,只交给少数专家处理。论文介绍,常见设置会从8至128个专家中启用1至2个。这样可以扩大模型的总参数容量,又不用每次计算全部参数。
负责分配任务的是路由器(router)——它像分诊台,根据当前token的内部表示选择专家。问题在于,标准路由器会逐个token独立判断。相邻token即使属于同一句话,也可能被分给完全不同的专家。
在显存或设备内存有限时,所有专家权重无法同时待在高速内存里。模型需要把参数从容量较大但较慢的存储,搬到容量较小但较快的内存。刚装入专家A,下一个token却要专家B,就得换出再装入。这叫权重搬运。论文指出,在受带宽限制的设备上,搬一次专家权重甚至可能比这一轮计算本身更耗时。
已有办法通常在部署阶段救火,例如缓存近期使用的专家,或者提前猜测接下来会调用谁。另一类方法则在模型训练完后,只微调路由器。它们都没有改变预训练期间形成的专家分工与路由习惯。
把“少换人”直接写进训练目标
StickyMoE的关键很简单:给训练目标加一项“路由一致性损失”。它比较相邻token面对各个专家时的概率分布;两次分配差得越远,惩罚越大。于是,路由器会倾向于让语义连续的位置沿用相近的专家组合。
这不是部署后的缓存补丁,也不需要修改MoE架构。论文称,Soft版本只新增一个控制惩罚强度的超参数 。 较小,模型更看重原本的语言建模任务; 较大,模型更积极地减少专家切换。
真正重要的不是“多加了一项损失”,而是加入的时机。StickyMoE从训练第一步就施加约束。梯度不仅更新路由器,也会传回更早的注意力层和词嵌入层。换句话说,分诊台在学习稳定分诊,专家和上游表示也同时适应这种安排。论文认为,这正是它相对训练后微调的优势:后者只能重新安排已经定型的专家,前者则让专家分工与路由决策一起长成。
这里有一个明显风险:如果所有token都塞给同一位专家,切换率当然最低,但MoE也就失去了意义。论文继续保留标准的负载均衡损失,用来避免少数专家包办全部工作。作者报告,在所测设置中,专家利用率熵始终接近四专家配置下的最大值,未观察到路由坍缩。
相邻一致,还不够
只约束相邻token仍可能出现“缓慢漂移”。每一步变化都很小,但走过一长段后,路由已经换到了另一组专家。
为此,论文还提出Soft-Hard版本。它把序列切成固定窗口,以窗口第一个token的路由分布作为锚点;越靠近窗口末尾,偏离锚点受到的惩罚越强。Soft约束控制每一步不要跳得太远,锚点约束则控制整段不要逐渐漂走。窗口结束后锚点重置,路由器可以重新选择。
这个设计仍然发生在训练阶段。模型部署时不需要增加新的前向计算,也不用更换路由结构,普通的LRU等缓存策略可以直接利用更稳定的专家访问序列。
59%的下降说明了什么?
论文在WikiText-2上训练了两种GPT式MoE语言模型:小型模型约880万参数,中型模型约2200万参数;两者均有4个专家,每个token启用其中两个。所有实验训练1万步,并在验证集上比较标准MoE、StickyMoE、硬窗口约束、模拟的ReMoE式训练后路由微调,以及简化版Oracle-MoE。
作者报告,Soft版本在两个规模上都把专家切换率最多降低59%。在中型模型的适中设置下,困惑度——衡量模型预测下一个token有多不确定,越低通常越好——不但没有恶化,反而优于基线。论文据此认为,一定强度的一致性约束可能兼具正则化作用。不过,更激进的约束仍会带来质量与局部性的取舍。
模拟的训练后微调在两个规模上的切换率变化都不到1%。硬窗口方法虽然也能减少切换,却付出了更高的困惑度成本。Soft-Hard版本则在中型模型上取得论文所测方法中的最低切换率和最高缓存命中率,同时维持较温和的质量代价。基于这些受控实验,作者称StickyMoE在“质量—局部性”前沿上优于训练后微调。
还有一个有意思的层间差异。最底层L0最难变得稳定,因为它面对的是尚未经过注意力上下文混合的token表示;L1和L2的切换下降更明显。这个结果支持了论文的直觉:有了上下文后,路由更容易按连续语义形成稳定分配。
为什么值得关注?
这项工作的意义不只是一种MoE正则项。它把服务系统里的成本,提前变成模型训练时可以优化的目标。过去的思路多是“模型给出怎样的访问序列,系统就设法缓存”;StickyMoE反过来问,能不能让模型从一开始就生成更适合缓存的访问序列。
这种模型设计与服务约束的协同尤其适合端侧场景:快速内存装不下全部专家,而每次从慢速存储加载权重都很贵。若路由本身更连续,现有缓存机制就更容易命中,无须在推理路径上额外修补。
但59%下降仍是作者在小规模实验中得到的最好结果,不代表所有模型、数据集或设备都能复现。更准确地说,论文证明了“训练时塑造内存访问习惯”值得继续验证,还没有证明它已经解决了大规模MoE的端侧部署问题。
局限与未知
- 实验仅覆盖WikiText-2上的约880万和2200万参数模型,且采用固定随机种子。结果能否扩展到大型生产模型,材料没有给出证据。
- 论文主要用专家切换率和模拟LRU缓存命中率作为代理指标,没有报告真实端侧设备上的权重搬运量、内存占用、延迟或能耗;因此“更省内存的推理”目前仍是由代理指标推导。
- 一致性损失会同等惩罚所有切换,包括句号、主题变化或对话轮次等自然边界。作者提出未来可让约束识别这些边界,但尚未完成验证。