BlockFFN: Towards End-Side Acceleration-Friendly Mixture-of-Experts with Chunk-Level Activation Sparsity

📄 arXiv: 2507.08771v2 📥 PDF

作者: Chenyang Song, Weilin Zhao, Xu Han, Chaojun Xiao, Yingfa Chen, Yuxuan Li, Zhiyuan Liu, Maosong Sun

分类: cs.LG, cs.CL

发布日期: 2025-07-11 (更新: 2025-07-30)

备注: 21 pages, 7 figures, 15 tables

🔗 代码/项目: GITHUB

💡 一句话要点

提出BlockFFN，通过块级激活稀疏性实现端侧加速友好的混合专家模型。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 激活稀疏性 端侧加速 块级稀疏性 可微路由

📋 核心要点

  1. 传统MoE模型路由方式不灵活，且块级稀疏性低，导致在端侧设备上难以高效加速。
  2. BlockFFN通过可微路由和块级稀疏性感知训练，提升模型性能和加速友好性。
  3. 实验表明，BlockFFN在token级和块级稀疏性上均有显著提升，并在端侧设备上实现了显著加速。

📝 摘要（中文）

为了减轻大型语言模型（LLM）的计算负担，具有激活稀疏性的架构，以混合专家模型（MoE）为代表，受到了越来越多的关注。然而，传统MoE的不可微和不灵活的路由会损害模型性能。此外，虽然每个token只激活少量参数，但这些稀疏激活的架构表现出较低的块级稀疏性，表明多个连续token的联合激活会激活很大比例的参数。这种稀疏模式不利于低资源条件（例如，端侧设备）下的加速，并且与主流加速技术（例如，推测解码）不兼容。为了应对这些挑战，我们引入了一种新的MoE架构BlockFFN，以及其高效的训练和部署技术。具体来说，我们使用集成了ReLU激活和RMSNorm的路由器进行可微和灵活的路由。接下来，为了提高token级稀疏性（TLS）和块级稀疏性（CLS），设计了CLS感知的训练目标，使BlockFFN更易于加速。最后，我们首次实现了结合激活稀疏性和推测解码的高效加速内核。实验结果表明，BlockFFN优于其他MoE基线，实现了超过80%的TLS和70%的8-token CLS。我们的内核在实际端侧设备上实现了高达3.67倍于密集模型的加速。所有代码和检查点均已公开发布。

🔬 方法详解

问题定义：现有混合专家模型（MoE）虽然具有激活稀疏性，但其路由机制通常是不可微的，限制了模型的性能。更重要的是，这些模型在块级别上的稀疏性较低，即多个连续的token会激活大量的参数，这使得它们难以在资源受限的端侧设备上进行高效加速，也与推测解码等主流加速技术不兼容。

核心思路：BlockFFN的核心思路是通过引入可微路由机制和块级稀疏性感知的训练目标，来提高模型的性能和加速友好性。可微路由允许模型在训练过程中学习更有效的路由策略，而块级稀疏性感知的训练则鼓励模型在处理连续的token时，尽可能地减少激活的参数数量。

技术框架：BlockFFN的整体架构基于MoE，主要包含以下几个模块：1) 输入层：接收输入token序列。2) 路由层：使用集成了ReLU激活和RMSNorm的路由器，根据输入token的特征，决定将token路由到哪些专家。3) 专家层：包含多个前馈神经网络（FFN），每个专家处理一部分token。4) 输出层：将各个专家的输出进行聚合，得到最终的输出结果。

关键创新：BlockFFN最重要的技术创新点在于其块级稀疏性感知的训练目标。传统的MoE模型只关注token级别的稀疏性，而BlockFFN则通过引入额外的损失函数，鼓励模型在处理连续的token时，尽可能地减少激活的参数数量。这使得BlockFFN在端侧设备上更容易进行加速，并且与推测解码等技术更加兼容。

关键设计：BlockFFN的关键设计包括：1) 使用ReLU激活和RMSNorm的路由器，实现可微且灵活的路由。2) 设计CLS感知的训练目标，包括稀疏性损失和负载均衡损失，以提高块级稀疏性。3) 实现高效的加速内核，结合激活稀疏性和推测解码，以在端侧设备上实现更高的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

BlockFFN在实验中表现出优异的性能，实现了超过80%的token级稀疏性（TLS）和70%的8-token块级稀疏性（CLS）。与传统的MoE模型相比，BlockFFN在端侧设备上实现了高达3.67倍的加速。这些结果表明，BlockFFN在提高模型性能和加速友好性方面具有显著优势。

🎯 应用场景

BlockFFN适用于资源受限的端侧设备上的大型语言模型部署，例如智能手机、嵌入式系统等。通过提高模型的加速友好性，BlockFFN可以降低计算成本和功耗，使得在这些设备上运行复杂的AI模型成为可能。此外，BlockFFN还可以应用于需要高吞吐量的场景，例如在线翻译、语音识别等。

📄 摘要（原文）

To alleviate the computational burden of large language models (LLMs), architectures with activation sparsity, represented by mixture-of-experts (MoE), have attracted increasing attention. However, the non-differentiable and inflexible routing of vanilla MoE hurts model performance. Moreover, while each token activates only a few parameters, these sparsely-activated architectures exhibit low chunk-level sparsity, indicating that the union of multiple consecutive tokens activates a large ratio of parameters. Such a sparsity pattern is unfriendly for acceleration under low-resource conditions (e.g., end-side devices) and incompatible with mainstream acceleration techniques (e.g., speculative decoding). To address these challenges, we introduce a novel MoE architecture, BlockFFN, as well as its efficient training and deployment techniques. Specifically, we use a router integrating ReLU activation and RMSNorm for differentiable and flexible routing. Next, to promote both token-level sparsity (TLS) and chunk-level sparsity (CLS), CLS-aware training objectives are designed, making BlockFFN more acceleration-friendly. Finally, we implement efficient acceleration kernels, combining activation sparsity and speculative decoding for the first time. The experimental results demonstrate the superior performance of BlockFFN over other MoE baselines, achieving over 80% TLS and 70% 8-token CLS. Our kernels achieve up to 3.67$\times$ speedup on real end-side devices than dense models. All codes and checkpoints are available publicly (https://github.com/thunlp/BlockFFN).