TEAM 完全拆解:时空一致性引导专家激活,MoE 扩散语言模型推理加速 2.2 倍(ICML 2026)

9103 字
46 分钟
TEAM 完全拆解:时空一致性引导专家激活,MoE 扩散语言模型推理加速 2.2 倍(ICML 2026)

AI 生成内容声明

背景:从自回归到扩散语言模型#

自回归生成的效率天花板#

过去几年大语言模型(Large Language Model, LLM)的生成几乎全部建立在自回归(autoregressive, AR)解码上:模型一次只能预测下一个 token,前一步的输出是后一步的输入,串行推进。这种串行依赖在硬件上很不友好——每一步只用到一个 token 的上下文,计算量小、显存访问量大,GPU 的算力大部分时间处于闲置,解码阶段(decoding phase)普遍是访存受限(memory-bound)的。投机解码(speculative decoding)、批量并行等都是围绕”如何让每一步多出几个 token”展开的修补。

扩散语言模型(diffusion language model, dLLM)提供了一条不同的路线:把生成过程建模为去噪(denoising)。LLaDA 系列(Large Language Diffusion Models,2025 年)是代表性工作:生成时先把整个回答初始化为一串 [MASK] 占位符,模型对所有位置同时做预测,给出每个位置是某个词的概率;置信度(confidence)超过阈值的 token 被”揭开”(unmask),其余位置重新盖上 [MASK] 进入下一轮去噪,直到所有位置都确定。由于一次前向传播(forward pass)可以同时处理整段序列,dLLM 天然支持并行解码,理论上每一步的产出是多个 token 而不是一个。

但第一代 dLLM 有一个致命伤:去噪过程使用的是全局双向注意力(bidirectional attention),每个 token 都要看到块内所有其他 token,导致解码过程中无法复用键值缓存(KV cache)——之前算过的 key/value 矩阵在下一轮因为注意力范围变化而失效,只能全部重算。这让 dLLM 的推理效率优势大打折扣。

块扩散:把自回归先验请回来#

块扩散(block diffusion)论文发表于 ICLR 2025(Oral) 用”分块”解决了这个问题:把待生成的序列切成若干块(block),块与块之间用因果注意力(causal attention),块内部才用双向注意力做并行去噪。这样一来,块与块之间的计算是自回归式的、可以复用 KV cache;块内部则是扩散式的并行解码。块大小(block size)从 1 到整段序列之间调节,模型就在”纯自回归”与”纯扩散”之间连续插值。

在此基础上,2025 年底出现了一批”自回归初始化 + 块扩散”的模型,代表是 SDAR(A Synergistic Diffusion-Autoregression Paradigm for Scalable Sequence Generation)与 LLaDA 2.0(LLaDA2.0: Scaling up Diffusion Language Models to 100B)。SDAR 的思路很巧妙:不去端到端训练一个扩散模型,而是把已经训练好的自回归模型做一次轻量”范式转换”(paradigm conversion)——通过少量数据适应训练,把 AR 模型的每个块内注意力改造成双向,从而既保留 AR 训练积累的语言能力,又获得块内并行解码的能力。论文的实验显示,这类模型在同等规模下准确率可以与最新的 AR 模型掰手腕,同时解码更快。

MoE 扩散模型:本文的主角#

混合专家(Mixture-of-Experts, MoE)是当下大模型扩展参数量的标准做法:把前馈网络(feed-forward network, FFN)替换为一组并行的专家(expert)网络,由一个路由器(router)为每个 token 挑选 top-k 个专家参与计算,其余专家不激活。总参数量大、每 token 激活的参数少,训练和推理都能省。DeepSeek-V3(Technical Report)、Mixtral(Mixtral of Experts)都是这套路线。MoE 扩散模型(MoE dLLM)自然也应运而生——SDAR 就有 30B-A3B 的 MoE 版本(30B 总参数、每 token 激活 3B),LLaDA 2.0 也扩到了 100B 规模。

到这里,两条技术线——“并行解码的扩散范式”和”稀疏激活的 MoE 架构”——各自都很高效,合在一起却出现了问题。这正是本文要讲的技术要解决的。

本文技术出处#

TEAM(Temporal-Spatial Consistency Guided Expert Activation for MoE Diffusion Language Model Acceleration,时空一致性引导的专家激活)由北京大学李萌(Meng Li)团队(北大 SEC Lab)的 Linye Wei、Zixiang Luo、Pingzhi Tang 等人完成,发表于 ICML 2026,论文在 arXiv:2602.08404,代码开源在 PKU-SEC-Lab/TEAM-MoE-dLLM。论文的定位是一个无需改权重、无需重新训练的即插即用(plug-and-play)推理加速框架:平均 1.94 倍、最高 2.2 倍的端到端加速,精度几乎无损。下面从问题出发,一步步拆解它。

问题:MoE 与扩散解码的组合错配#

解码过程的形式化#

先给块扩散解码一个精确的数学描述。设回答长度为 N=B×LN = B \times L,被切成 BB 个块,每块 LL 个 token,第 ii 个块记为 Yi=[yi0,yi1,,yiL1]Y_i = [y_i^0, y_i^1, \cdots, y_i^{L-1}]。给定提示词(prompt)PP,模型对整段回答做块级因式分解:

pϑ(Y^P)=i=1Bpϑ(Y^iP,Yi)p_{\vartheta}(\widehat{Y} \mid P) = \prod_{i=1}^{B} p_{\vartheta}\left(\widehat{Y}_i \mid P, Y_{\leq i}\right)

其中 YiY_{\leq i} 表示前 ii 个块的已知内容,ϑ\vartheta 是模型参数。这个式子的含义是:每个块 YiY_i 的生成条件于”提示词 + 之前所有块的最终结果”,块间严格因果,所以前面的块一旦确定就可以缓存——这是 KV cache 能用的数学基础。

在一个块内部,每一轮去噪迭代(decoding iteration)里,模型对块内每个位置 kk 做一次预测,得到 argmax 词与置信度:

yik^=argmaxvV  pϑ(yik=vP,Yi),ck=pϑ(yik=yik^P,Yi)\widehat{y_i^k} = \operatorname{argmax}_{v \in V} \; p_{\vartheta}\left(y_i^k = v \mid P, Y_{\leq i}\right), \qquad c_k = p_{\vartheta}\left(y_i^k = \widehat{y_i^k} \mid P, Y_{\leq i}\right)

其中 VV 是词表(vocabulary),ckc_k 是位置 kk 上预测词的 softmax 概率,即模型对该预测有多”自信”。随后按固定阈值 τ\tau 决定哪些位置被揭开:

yik={yik^,if ck>τ[MASK],otherwise,k[0,1,,L1]y_i^k = \begin{cases} \widehat{y_i^k}, & \text{if } c_k > \tau \\ [\mathrm{MASK}], & \text{otherwise} \end{cases}, \quad k \in [0, 1, \cdots, L-1]

置信度超过 τ\tau 的位置本轮接受(accept),其余位置盖上 [MASK] 进入下一轮。直到整块全部揭开,缓存该块,再处理下一块,如此往复直到生成 [EOS]。TEAM 实验里取 τ=0.95\tau = 0.95、块大小 L=32L = 32

错配在哪里#

扩散解码和 MoE 架构各自单独看都高效,但组合起来出现一个根本性错配(fundamental mismatch):每一轮去噪都要对块内所有 token 做一次完整的前向计算,每个 token 各自独立做专家路由;但每一轮真正被接受的 token 只有一小部分。 换句话说,每轮激活了一大堆专家,产出的却只有几个 token。

论文用 SDAR 30B-A3B 实测了这个错配的严重程度。下面这张图是全文的第一张图,横轴是解码迭代轮数,纵轴是每轮激活的不同专家数(distinct experts),蓝线是 vanilla(原始)MoE dLLM,橙线是加上 TEAM 之后:

图 1:每个前向传播中激活的专家数与接受的 token 数对比(SDAR 30B-A3B)。TEAM 用更少的专家激活解码出更多的 token
图 1:每个前向传播中激活的专家数与接受的 token 数对比(SDAR 30B-A3B)。TEAM 用更少的专家激活解码出更多的 token

图片来源:TEAM 论文 Figure 1(arXiv:2602.08404)。

SDAR 30B-A3B 每层有 128 个专家、每个 token 名义上只路由 top-8 个专家(8 experts per token),但论文在 Table 1 中实测:vanilla 模型每一轮前向平均激活超过 50 个专家,接近 128 个专家的一半;而每轮平均只接受约 3 个 token。 把激活摊到每个被接受的 token 上(论文定义的指标 APT,见后文),平均每个 token 要付出 17.66 次专家激活——是名义路由成本 8 的两倍还多,最差的 GSM8K 上高达 18.71。也就是说,MoE 的稀疏激活优势在扩散解码中被彻底稀释,几乎退化成稠密模型(dense model)的激活模式。

为什么会这样?两个原因叠加:

  1. 双向注意力迫使所有 token 每轮都参与计算。 块内是双向注意力,任何 mask token 的去噪都依赖块内所有其他 token(包括已经确定不再改变的 decoded token)的表示,因此已解码 token 必须每轮都出现在前向计算里。它们虽然不再变化,却依然各自独立触发专家激活。
  2. 每个 token 独立路由放大了激活集合。 块内 LL 个 token 的专家选择互不相同,一轮下来被触发的专家是这些选择的并集,这个并集随着 token 数增加而迅速膨胀。单请求(batch size = 1)场景下没有任何并发请求可以摊薄这份开销。

论文指出,这种开销在云上多请求并发时可以被摊薄,但对于解码速度和尾延迟(tail latency)高度敏感的场景——比如流式对话、边缘设备(edge platform)上的受限硬件——它就是实打实的瓶颈。

关键动机:块解码的三条一致性规律#

TEAM 的出发点不是直接去优化 kernel,而是先做了一次系统的行为分析:把一个块从初始化到全部揭开的完整解码轨迹(decoding trajectory)连同专家激活模式一起画出来,找到了三条规律。论文 Figure 2 是这三条规律的实证(测量自 SDAR 30B-A3B 在 GSM8K 上的一条回答,图中展示了 47 层中的第 0、24、47 层):

图 2:SDAR 30B-A3B 上专家激活与解码的时空特性。(a)各解码迭代激活的专家数;(b)第 6 步(共 11 步)由已解码 token 与 mask token 激活的专家分布;(c)每轮接受 token 的位置,以及相对下一轮迭代的隐藏状态相似度
图 2:SDAR 30B-A3B 上专家激活与解码的时空特性。(a)各解码迭代激活的专家数;(b)第 6 步(共 11 步)由已解码 token 与 mask token 激活的专家分布;(c)每轮接受 token 的位置,以及相对下一轮迭代的隐藏状态相似度

图片来源:TEAM 论文 Figure 2(arXiv:2602.08404)。

规律一:时间一致性(Temporal Consistency)#

块扩散要求同一个块在多轮去噪迭代中被反复处理,token 在这个过程中逐渐被接受并固定下来。已接受的 token 内容不再改变,它们的作用只是给剩余 mask token 提供上下文,但每一轮它们依然走完整的前向计算、独立触发专家激活。 图 2(a) 展示了三层上每轮激活的专家数:随着解码推进、固定 token 越来越多,已解码 token 贡献的专家激活占比不断上升,到后期成为主导。这些激活纯粹是冗余计算。

一个更微妙的观察:某个 token 的隐藏状态(hidden state)变化最大的时刻,恰好是”它被接受的那一轮”与”下一轮”之间;一旦它被接受并再经过一次前向传播,它的表示就基本稳定了。这个现象与先前工作 dKV-Cache(dKV-Cache: The Cache for Diffusion Language Models)等对扩散模型 KV 漂移的观察一致。

规律二:空间一致性(Spatial Consistency)#

再看未被接受的 mask token。它们之间几乎没有输入差异:所有 mask token 共享同一个 [MASK] 词嵌入(token embedding),唯一的区别来自位置编码(positional encoding),而同一块内位置编码的差异很小。因此空间上相邻的 mask token 在不同层上的专家路由模式高度一致——它们大概率路由到同一个专家子集。

图 2(b) 给出了第 6 步(共 11 步)的专家激活分布:已解码 token 激活的专家集合非常多样,接近均匀分布;而 mask token 的专家选择高度集中在少数几个专家上。论文 Table 4 用数字量化了这一点:块内所有 mask token 隐藏状态的两两余弦相似度(pairwise cosine similarity),在所有层里最高达 0.98–0.99,平均也有 0.84–0.86。也就是说,“少数专家服务了几乎所有 mask token 的解码”,集合外的专家要么几乎不被用到、要么只被个别 token 调用。

这条规律还有一个推论,解释了图 2(a) 的一个细节:既然 mask token 的专家选择高度集中,为什么每轮激活的专家总数还是那么多?因为已解码 token 激活的专家集合与 mask token 激活的集合重叠很小——已解码 token 的上下文各不相同,路由分散,贡献了大部分专家激活。这提示了优化方向:砍已解码 token 的计算,收益最大。

规律三:时空局部性(Temporal-Spatial Locality)#

图 2(c) 同时画出了每一轮被揭开的位置(红色框标注)和相邻两轮隐藏状态的相似度。观察结果:

  1. 接受顺序近似自回归。 一个块内的 token 大体上是从前往后被揭开的——毕竟模型由自回归模型初始化,自然语言的生成本身有因果结构,模型”心里”依然倾向于先确定靠前的 token。
  2. 每轮新接受的 token 在位置上聚集。 新揭开的位置往往紧挨着已解码的 token,而不是随机散布在整个块里。
  3. 部分 token 早期绝不会被接受。 既然接受顺序近似自回归,那么距离已解码位置很远、置信度又低的 token,在早期迭代里几乎不可能被揭开。

第三条直接导出”热/冷”分类的可行性:离已解码位置近、置信度高的 mask token(热 token,hot tokens)更可能马上被接受;离得远、置信度低的(冷 token,cold tokens)短期内不会被接受,为它们投入与热 token 同等的计算是浪费。

TEAM 总览:三类 token,三种策略#

基于这三条规律,TEAM 在每一轮前向传播之前,把块内 token 分成三类,分别采用差异化策略,最后合并成一次前向计算:

token 类别定义策略
已解码 token(decoded tokens)已被接受且已确定DCD:延迟缓存(Delayed Caching),接受后再算一轮,此后彻底退出计算、KV 与专家激活全部走缓存
热 mask token(hot tokens)置信度高于 τh\tau_h,或距已解码位置小于 LhL_hSEH:投机探索(Speculative Exploration),额外接受 top 置信度候选构造多个分支并行验证,提高每轮接受率
冷 mask token(cold tokens)其余 mask tokenLAC:受限激活(Limited Activation),第二轮路由时把专家选择限制在”必要专家集合”内

图 3 是 TEAM 的整体架构图,一图看全三路策略如何作用在一个解码块上:

图 3:TEAM 总览。对块内 token 采用差异化的专家激活与解码策略:已解码 token 通过一步延迟缓存削减冗余计算;热 mask token 采用激进的多分支投机探索利用闲置算力提高接受率;冷 mask token 通过两轮路由机制约束不必要的专家激活
图 3:TEAM 总览。对块内 token 采用差异化的专家激活与解码策略:已解码 token 通过一步延迟缓存削减冗余计算;热 mask token 采用激进的多分支投机探索利用闲置算力提高接受率;冷 mask token 通过两轮路由机制约束不必要的专家激活

图片来源:TEAM 论文 Figure 3(arXiv:2602.08404)。

注意三个策略的定位差异:DCD 是”省”——把确定不会再变的东西从计算里拿掉;SEH 是”赚”——利用省出来的算力并行验证更多候选,多收 token;LAC 是”堵”——把注定近期用不上的 token 的专家激活限制在最小必要集合。三者没有重叠:DCD 管已解码 token,SEH 和 LAC 分别管热/冷两类 mask token。

下面逐策略拆解。

策略一:DCD——已解码 token 的延迟缓存#

机制#

既然已解码 token 的隐藏状态在”接受 + 再经过一轮前向”之后近似稳定(规律一),自然想到把它们缓存起来,避免每轮重复计算。DCD 的具体流程是:

  1. 每一轮只对两类 token 做完整计算:当前仍是 mask 的 token,以及上一轮刚被接受的 token;
  2. 更早被接受的 token 完全退出计算,它们的键值对(key/value pairs,KV)和对应的专家激活直接从缓存(cache)中取出复用;
  3. 每轮前向结束后,把本轮新接受的 token 的 KV 插入缓存,供后续迭代使用。

“延迟一步”(delayed)是关键设计:不是接受当下就缓存,而是让新接受的 token 再多参与一轮完整计算,从第二轮开始才走缓存。 原因是隐藏状态在接受后的第一轮前向里才会”吸收”新的上下文变得稳定(图 2(c) 的观察),提前缓存会引入表示漂移。这一轮额外计算换来的是后续所有轮次的完全免除,稳赚不赔。

为什么”不刷新”也安全#

dKV-Cache 等先前工作为了对抗双向注意力下的 KV 漂移(KV drift),需要周期性做全局缓存刷新(global cache refresh)——每隔 NN 轮把所有 token 重算一遍。TEAM 则完全不做刷新(refresh-free)。论文给了两个理由:

  1. 块扩散下刷新本身收益有限。 块扩散的并行计算范围被限制在单个块内(L=32L=32),这个并行规模远达不到现代 GPU 的计算-带宽平衡点(compute-bandwidth balance point),解码阶段整体是访存受限的。在这种 regime 下,块内细粒度的 KV 缓存优化空间很小,刷新操作省下的计算微乎其微,反而要付出重算成本。

  2. MoE 下的收益被放大了。 与稠密模型不同,MoE 的已解码 token 每轮激活的是一大堆”与 mask token 几乎不重叠”的专家(图 2(b)),缓存它们意味着直接砍掉很大比例的专家激活与相应的参数访存。DCD 的收益主体是专家激活,不是 KV 计算。

  3. 自回归先验保证稳定。 由 AR 模型初始化 + 近似自回归的接受顺序,已解码 token 的表示在扩散迭代中保持高度稳定,因此缓存不需要刷新也不会漂移。

论文 Table 3 专门对比了刷新频率的影响(在加上 DCD 与 LAC、不加 SEH 的设置下):每 4 轮刷新(Refresh-4)、每 8 轮刷新(Refresh-8)与完全不刷新(Refresh-free,即 TEAM 的做法)。四个基准上,Refresh-free 的准确率与 Refresh-4/8 基本持平甚至更高(平均 77.57 分对 77.84 / 77.10 分),而激活专家数(APF)和加速比明显更好(平均 26.48 对 33.32 / 29.73,1.47× 对 1.29× / 1.38×)。也就是说,在块扩散 + AR 初始化的模型上,周期性刷新是纯粹的浪费

效果#

在 Table 3 的配置下(DCD 与 LAC 叠加、不含 SEH),APF 从 vanilla 的约 55 压到 26–33,带来 1.3–1.6 倍加速。它是三个策略中削减专家激活最猛的一刀,因为已解码 token 在后期迭代中贡献的专家激活占比最大(规律一)。

策略二:SEH——热 token 的投机探索#

动机:专家利用不足与热 token 识别#

DCD 解决的是已解码 token 的浪费,但 mask token 侧也有两个效率问题:

  1. 算术强度过低。 每个激活的专家平均只服务块内少数几个 token(块才 32 个 token,128 个专家被激活一半以上),专家计算量小、参数访存量大,GPU 的计算资源大量闲置。
  2. 存在注定”白算”的 token。 接近自回归的接受顺序意味着部分 token 在早期迭代绝不会被接受(规律三),为它们激活专家、跑完前向,下一轮又被重新 mask,纯属浪费。

于是 mask token 被进一步分成热、冷两类。论文给出热 token 的形式化定义:

yik-hot={yik(ck>τh)  or  (j,kj<Lh)}y_i^{k\text{-hot}} = \left\{ y_i^k \mid (c_k > \tau_h) \;\text{or}\; (\exists j, |k - j| < L_h) \right\}

其中 τh\tau_h 是判断”热”的置信度阈值,jj 遍历当前已解码 token 的位置,LhL_h 是”距已解码位置的最大距离”。含义:一个 mask token 是热的,当且仅当它本轮的置信度 ckc_k 超过 τh\tau_h,或者它落在某个已解码 token 的 LhL_h 邻域内。 前者对应”已经快够着阈值、下一轮很可能接受”的 token,后者对应”位置靠前、按自回归顺序马上要轮到”的 token。两条理由都指向同一个结论:它们比一般 mask token 更可能在接下来的迭代中被接受,值得投入计算。

机制:多分支投机探索#

对热 token,TEAM 做的是”投机”:不等置信度真的超过接受阈值 τ\tau,先额外接受置信度最高的候选 token,构造出多个分支(branch),在同一个前向传播里并行解码、并行验证。具体地(论文 §5.4 给出细节),在原始分支之外额外构造三个分支:

  1. 额外接受置信度第 1 高的候选 token(即使 ck<τc_k < \tau);
  2. 额外接受置信度第 2 高的候选 token;
  3. 同时额外接受这两个候选。

加上原始分支一共 4 条分支并行验证(论文默认 branch = 4)。下一轮看哪条分支的候选真的被接受,接受最多的分支胜出,其他分支作废。这样一轮前向可能比 vanilla 多接受好几个 token,总解码步数(迭代轮数)随之下降。

为什么分支在 MoE 上”近乎免费”#

这是 SEH 最核心的设计洞察。在双向注意力下,修改块内任何一个 token 都会通过注意力影响整块所有 token 的表示,因此每多一个候选分支,在稠密模型里就相当于多一整块的前向计算。先前的扩散模型加速工作(如 LoPA,Lookahead Parallel Decoding)之所以要上多卡,就是因为分支把计算强度(arithmetic intensity)推到稠密模型无法承受的高度——解码瓶颈从访存受限翻转成计算受限(compute-bound),单卡算不过来。

MoE 模型的结构正好让这件事变便宜:

  1. 计算天然分散在专家上。 推理延迟主要由 FFN 层主导,而 FFN 的计算被拆到各专家上。分支并行意味着同一批专家要服务更多 token,专家内的算术强度被填满——闲置算力被利用起来,而不是被要求”多激活新专家”。
  2. 分支之间高度相似,几乎不激活新专家。 额外接受的候选 token 与原始预测在表示上高度接近(它们本来就是同一个位置的 top 候选),路由结果也接近,所以新分支引入的专家激活增量很小。

图 4 直观展示了这一点:在 GSM8K 的一条回答上,第 24 层(共 47 层)的专家激活对比——加上探索分支后,每个专家服务的 token 数明显增加,而新激活的专家很少:

图 4:SDAR 在 GSM8K 一条回答上的投机探索专家激活情况(第 24 层,共 47 层)。加入探索分支后,原有专家服务的 token 变多、新激活的专家很少
图 4:SDAR 在 GSM8K 一条回答上的投机探索专家激活情况(第 24 层,共 47 层)。加入探索分支后,原有专家服务的 token 变多、新激活的专家很少

图片来源:TEAM 论文 Figure 4(arXiv:2602.08404)。

为什么是”对齐候选”而不是”更多样候选”#

一个自然的疑问是:既然要多验证候选,为什么只加 top-1、top-2 两个候选的三种组合,而不是继续加 top-3、top-4 让分支更”多样”?论文用图 6 的数据回答了这个问题——它统计了当前轮 top-3 置信度候选在下一轮被真正接受的概率:

图 5:SEH 中 top-3 置信度候选 token 在下一轮迭代中的接受概率
图 5:SEH 中 top-3 置信度候选 token 在下一轮迭代中的接受概率

图片来源:TEAM 论文 Figure 6(arXiv:2602.08404)。

第 1 候选的接受概率最高,第 2 候选已经明显偏低,第 3 候选几乎不会被直接接受。换句话说,构造”第 3 候选分支”这种追求多样性的分支,绝大多数时候是白跑一趟——既多占了计算,又几乎没有新增收益。相比之下,top-1、top-2 及其组合支持链式验证(chained verification):本轮先接受 top-1,下一轮 top-2 在 top-1 已被接受的上下文里变成”更可能被接受”的候选,两个分支可以接力推进。这就是”对齐的候选组合”(aligned token combinations)优于”多样的候选”的原因。

效果#

SEH 是提升”每轮接受 token 数”(TPF)的主力:Table 1 中 TPF 提升 1.49×–1.74×(GSM8K 3.16→4.79,HumanEval 2.91→5.07)。代价是分支多出来的一小部分专家激活——但正如上面分析的,这个增量很小,完全被接受率提升盖过。

策略三:LAC——冷 token 的受限激活#

机制:两轮路由#

冷 token 短期内不会被接受(规律三),为它们独占激活专家是浪费;但 mask token 的路由又高度集中(规律二),所以存在一个”能服务几乎所有 mask token”的共享专家子集。LAC 的策略是:把冷 token 的专家选择限制在”必要专家集合”内,而不是取消它们的计算。 论文用 Algorithm 1 描述(这里用伪代码复述):

输入:已解码 token 集合 D,mask token 集合 M,候选专家池 E0
输出:激活专家集合 E_A,路由权重 W
// 第 1 步:token 分类
找出新接受的 token D_a ⊆ D
按公式 (4) 找出热 token H ⊆ M
冷 token C ← M \ H
// 第 2 步:第一轮路由(必要激活)
W1 ← Router(D_a, H, E0) // 新接受 token + 热 token 正常 top-k 路由
E_A ← top-k(W1) // 取这些路由的专家并集 = 必要专家集合
// 第 3 步:第二轮路由(受限激活)
W2 ← Router(C, E_A) // 冷 token 的路由被限制在 E_A 内
// 第 4 步:合并
W ← Concat(W1, W2)
返回 E_A, W

拆开来看:

  • 第一轮路由只跑”必须准确”的 token:新接受的 token 和热 token。它们要么已经确定、要么马上要被揭开,路由必须精确,否则会污染下一轮的解码质量。它们 top-k 路由结果(按 token 并集)构成必要专家集合 EAE_A
  • 第二轮路由跑冷 token,但候选池被钳制在 EAE_A:冷 token 的 gate 分数只在 EAE_A 上计算并重新归一化(相当于把非必要专家的分数屏蔽掉),再选 top-k。这样冷 token 的专家激活被严格限制在必要集合内,不会引入任何”只为个别冷 token 服务”的独占专家。

为什么是”limited”而不是”cancelled”#

注意 LAC 的全称是 Limited Activation(受限激活),不是 Cancelled Activation(取消激活)。为什么不能直接把冷 token 的 FFN 计算整个跳过?两个原因:

  1. 保留未来被接受的可能。 冷 token 只是”近期”不会被接受,不是永远。如果完全取消它们的计算,等到它们终于轮到时,表示缺乏足够的去噪历史,可能造成质量损失。限制在 EAE_A 内则保留了完整的计算路径——只是专家选择被收窄,而收窄是有依据的:空间一致性表明 mask token 本来就会路由到共享子集,收窄几乎不改变它们的表示。
  2. 保持一次前向的完整性。 块内双向注意力要求所有 token 都参与,冷 token 的注意力部分照常计算,只有 FFN 的专家选择被限制。这样块内表示的相互依赖不被打断,避免引入系统性偏差。

论文指出,限制到 EAE_A 后冷 token 的表示质量几乎不受影响,同时消除了”token 专属专家激活”——这是 MoE dLLM 独有的浪费来源。

效果#

LAC 在 DCD 和 SEH 的基础上进一步压低 APT(每个解码 token 的等效专家激活),是三步中最精细的一刀。它对加速比的贡献在消融实验里可见(下文)。

实验:2.2 倍加速从哪里来#

实验设置#

  • 模型:SDAR 30B-A3B(MoE 块扩散模型,128 专家/层、每 token top-8 路由、47 层)。选用它是因为它提供官方开源评测协议;LLaDA 2.0 虽然也是 MoE dLLM,但官方没有 Hugging Face 格式的评测管线,因此不作为主实验平台。
  • 硬件:单张 NVIDIA A100 80GB。
  • 基准:代码生成 HumanEval、MBPP;数学推理 GSM8K、Math-500(全部 0-shot)。
  • 超参:接受阈值 τ=0.95\tau = 0.95,块大小 L=32L = 32;热 token 判定 τh=0.7\tau_h = 0.7Lh=3L_h = 3;投机探索 4 条分支。

核心指标#

论文用三个指标度量效率:

  • APF(Activated experts Per Forward pass):每轮前向实际激活的专家总数;
  • TPF(accepted Tokens Per Forward pass):每轮前向被接受的 token 数,即解码并行度;
  • APT(Activated experts Per decoded Token):APT=APF/TPFAPT = APF / TPF,把激活摊到每个产出的 token 上,是”解码一个 token 的实际成本”。

APF 衡量稀疏性,TPF 衡量并行度,APT 是两者综合后的真实单位成本——vanilla 的 APT 高达 17.66,意味着名义上”8 专家/token”的模型实际每解码一个 token 要激活 17 个专家。

主结果#

Table 1 是全文的核心数据表:

基准方法得分APF ↓TPF ↑APT ↓加速比
HumanEvalVanilla79.2753.342.9118.331.00×
HumanEvalTEAM79.88 (+0.61)34.48 (−35%)5.07 (1.74×)6.80 (−63%)2.20×
MBPPVanilla65.7649.592.7418.101.00×
MBPPTEAM65.76 (+0.00)30.92 (−38%)4.56 (1.66×)6.78 (−63%)2.08×
GSM8KVanilla90.6059.113.1618.711.00×
GSM8KTEAM90.30 (−0.30)36.20 (−39%)4.79 (1.52×)7.56 (−60%)1.83×
Math-500Vanilla76.0057.903.7415.481.00×
Math-500TEAM75.40 (−0.60)36.31 (−37%)5.57 (1.49×)6.52 (−58%)1.64×
平均Vanilla77.9154.993.1417.661.00×
平均TEAM77.84 (−0.07)34.48 (−37%)5.00 (1.59×)6.92 (−61%)1.94×

逐项读这张表:

  • 准确率几乎无损:平均 −0.07 分,HumanEval 甚至 +0.61。四个基准上最差也只掉 0.6 分(Math-500)。这是”插拔式”框架的底气——它不动模型权重,只改运行时的专家选择与缓存策略。
  • 专家激活砍掉 35–39%:APF 从平均 54.99 降到 34.48。注意这个数字已经包含 SEH 分支新增的专家激活——在三路策略叠加之后仍净减 37%。
  • 每轮接受 token 提升 49–74%:TPF 从 3.14 升到 5.00,SEH 的多分支验证是主要来源。
  • 单位成本跌破名义值:APT 从 17.66 降到 6.92——低于模型名义的”每 token 8 专家”路由成本。也就是说 TEAM 让 MoE 扩散模型实现了比它的”设计成本”更省的解码,这是消融里三路策略叠加的合力。
  • 端到端加速 1.64×–2.20×:平均 1.94×,峰值在 HumanEval 达到 2.20×。加速比与 APT 的降幅高度相关:APT 降得越多、加速越大,说明瓶颈确实在专家激活的访存与计算上。

消融:三路策略各自的贡献#

图 5 展示了逐步叠加三路策略的消融结果(纵轴:APT 与相对 vanilla 的加速比):

图 6:消融研究——相对 vanilla 模型,逐步叠加三路策略后的 APT(每解码 token 激活专家数)与加速比
图 6:消融研究——相对 vanilla 模型,逐步叠加三路策略后的 APT(每解码 token 激活专家数)与加速比

图片来源:TEAM 论文 Figure 5(arXiv:2602.08404)。

论文对消融结果的解读:

  • SEH 贡献最大的一步:它通过减少解码步数(每轮多收 token)大幅压低 APT——每轮接受率翻倍,激活摊薄。同时分支引入的专家增量很小(分支相似性),净收益显著。
  • DCD 其次:消除已解码 token 每轮重复触发的一大块专家激活(后期迭代中占比最大的部分)。
  • LAC 收尾:把冷 token 的专家激活严格限制在必要集合内,APT 进一步下降、加速比再上一个台阶。

超参敏感性:热 token 怎么判#

热/冷分类的超参是 (τh,Lh)(\tau_h, L_h) 的组合。论文在代码生成基准上扫描了五组(Table 2,DCD + LAC 配置下):

(τh,Lh)(\tau_h, L_h)HumanEvalMBPP平均分平均 APF
(0.4, 6)78.0566.9372.4923.35
(0.5, 5)81.0965.7673.4323.27
(0.6, 4)79.2768.0973.6823.07
(0.7, 3)79.2766.9373.1022.37
(0.8, 2)77.4462.6570.0521.33

趋势:τh\tau_h 从 0.4 升到 0.7、LhL_h 从 6 收到 3 的过程中,APF 单调下降(23.35 → 22.37),准确率基本保持(最差组合的 72.49 与最佳组合的 73.68 差距在 1.2 分内)。但继续收紧到 (0.8,2)(0.8, 2) 时,虽然 APF 进一步降到 21.33,准确率却明显跳水(平均 70.05,MBPP 从 66.93 掉到 62.65)。原因很直观:热 token 识别得太保守,就把大量本可投机探索的 token 划成了冷 token,SEH 的并行收益被削弱,同时冷 token 的表示质量开始受损。 (0.7,3)(0.7, 3) 是准确率与效率的最佳平衡点,论文定为默认配置。

两个佐证实验#

  • mask token 的相似性(Table 4):四个基准上,块内 mask token 隐藏状态两两余弦相似度最高 0.98–0.99、平均 0.84–0.86。这直接支撑”空间一致性”假设——mask token 输入同质化(同一个 [MASK] 嵌入 + 微小的位置编码差异),路由自然集中,LAC 的”共享专家子集”因此成立。
  • 刷新频率(Table 3):Refresh-4 / Refresh-8 / Refresh-free 三档对比,Refresh-free 准确率不降、APF 最低(26.48 vs 33.32 / 29.73)、加速比最高(1.47× vs 1.29× / 1.38×),验证了 DCD 的免刷新设计。

工程实现:插拔式加速#

TEAM 的代码库 PKU-SEC-Lab/TEAM-MoE-dLLM 展示了一个刻意设计的”侵入面最小”工程形态:

  • 改一个文件:仓库基于 SDAR 的官方评测环境,使用方式是把作者提供的 modeling_sdar_moe.py(或带观测输出插桩的 modeling_sdar_moe_mark.py)替换到从 Hugging Face 下载的 SDAR-30B-A3B-Chat-b32 模型 目录里,模型权重一行不改。
  • 推理即评测:TEAM 直接嵌在 OpenCompass 评测管线里,四个基准(GSM8K、Math、HumanEval、MBPP)各有现成配置,单卡即可运行:
Terminal window
CUDA_VISIBLE_DEVICES=0 python run.py configs/eval_sdar_hf_gsm8k.py
  • 三类逻辑的落点:DCD 的 KV 缓存与”新接受 token 多算一轮”逻辑、SEH 的分支构造与验证、LAC 的两轮路由,都发生在模型前向的运行时层(routing 与 cache 管理),与注意力/专家 kernel 解耦,这也是它能不改权重、不重训就插到任意 MoE 块扩散模型上的原因。

这种形态的启示:MoE 扩散模型的开销大头不在 kernel 细节,而在”决策层”——每轮让谁算、算几次、路由到哪。 把决策逻辑做成运行时可插拔的组件,一套策略就能跨模型复用。

局限与讨论#

  1. 验证面窄:主实验只有 SDAR 30B-A3B 一个模型(LLaDA 2.0 因缺官方 HF 评测管线未被纳入)。三路策略依赖的”AR 初始化 + 近似 AR 接受顺序”在这类模型上成立,但 SDAR 之外的 MoE 块扩散模型(如 LLaDA 2.0)是否同样成立,还有待验证。
  2. 单卡场景:TEAM 的目标是单卡/延迟敏感场景。云端大规模部署下,与 dInfer(dInfer: An Efficient Inference Framework for Diffusion Language Models)这类做专家并行(expert-parallel)的框架是互补关系——TEAM 管”每卡内激活哪些专家”,dInfer 管”专家分布到哪些卡上”,两者可以叠加。
  3. 超参需要微调τh\tau_hLhL_h、分支数都是需要按模型/任务选择的超参,表 2 显示它们对效率-质量平衡有明显影响。作为训练无关方法,它没有自适应地调节这些阈值。
  4. 度量口径:加速比来自端到端实验(A100 上),不是 kernel 级理论收益;对 FP8/量化等不同精度设置下的行为,论文未展开。

小结#

TEAM 的价值在于把”MoE 与扩散解码的组合为什么慢”这件事讲清楚了:双向注意力迫使每轮全块计算,而每轮只接受少量 token;每个 token 独立路由让激活专家集合爆炸。 论文通过三条实测规律(时间一致性、空间一致性、时空局部性)定位了浪费的具体位置,再用三个互补策略各管一段:DCD 缓存已解码 token、SEH 让热 token 多分支投机、LAC 限制冷 token 的专家集合。合起来的效果是 APT 从 17.66 压到 6.92(低于名义路由成本 8),平均 1.94×、最高 2.2× 的端到端加速,准确率几乎不变。

对理解大模型推理系统而言,这篇工作的方法论比数字更有价值:当两种高效机制组合出人意料地低效时,先做行为测量定位浪费,再按”数据分布的特征”设计差异化策略——这与 PagedAttention 从内存碎片测量出发、FlashAttention 从访存占比测量出发是同一套思路。MoE 扩散模型正在快速成熟(LLaDA 2.0 已到 100B 规模),专家激活层面的优化空间,值得继续关注。

参考资料#

  1. TEAM: Temporal-Spatial Consistency Guided Expert Activation for MoE Diffusion Language Model Acceleration(arXiv:2602.08404)
  2. TEAM 代码仓库(PKU-SEC-Lab/TEAM-MoE-dLLM,GitHub)
  3. TEAM ICML 2026 会议页面
  4. 李萌老师主页上的 TEAM 论文页面(北京大学)
  5. SDAR: A Synergistic Diffusion-Autoregression Paradigm for Scalable Sequence Generation(arXiv:2510.06303)
  6. SDAR 代码仓库(JetAstra/SDAR,GitHub)
  7. SDAR-30B-A3B-Chat-b32 模型权重(Hugging Face)
  8. LLaDA2.0: Scaling up Diffusion Language Models to 100B(arXiv:2512.15745)
  9. dKV-Cache: The Cache for Diffusion Language Models(arXiv:2505.15781)
  10. dInfer: An Efficient Inference Framework for Diffusion Language Models(arXiv:2510.08666)
  11. Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models(arXiv:2503.09573,ICLR 2025 Oral)
  12. Large Language Diffusion Models(LLaDA,arXiv:2502.09992)
  13. LoPA: Scaling dLLM Inference via Lookahead Parallel Decoding(arXiv:2512.16229)

文章分享

如果这篇文章对你有帮助,欢迎分享给更多人!

TEAM 完全拆解:时空一致性引导专家激活,MoE 扩散语言模型推理加速 2.2 倍(ICML 2026)
https://pinghaoyang.com.cn/aigc/posts/team-moe-dllm/
作者
平昊阳
发布于
2026-08-24
许可协议
CC BY-NC-SA 4.0

评论区

Profile Image of the Author
平昊阳
乘长风,破巨浪, 展鸿图于未央!
--
总访问量
--
访客数
公告
欢迎来到我的个人博客!欢迎关注交流吖!
更多相关公告,见
社交-留言」。
音乐
封面

音乐

暂未播放

0:000:00
暂无歌词
站点统计
文章
66
分类
16
标签
93
总字数
477,284
运行时长
0
最后活动
0 天前

文章目录