BlockFFN: Rumo a uma Mistura de Especialistas Amigável à Aceleração no Lado do Terminal com Esparsidade de Ativação em Nível de Blocos

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

July 11, 2025

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

cs.AI

Resumo

Para aliviar o fardo computacional dos grandes modelos de linguagem (LLMs), arquiteturas com esparsidade de ativação, representadas pela mistura de especialistas (MoE), têm atraído crescente atenção. No entanto, o roteamento não diferenciável e inflexível do MoE tradicional prejudica o desempenho do modelo. Além disso, embora cada token ative apenas alguns parâmetros, essas arquiteturas esparsamente ativadas exibem baixa esparsidade em nível de bloco, indicando que a união de múltiplos tokens consecutivos ativa uma grande proporção de parâmetros. Esse padrão de esparsidade é desfavorável para aceleração em condições de baixos recursos (por exemplo, dispositivos de borda) e incompatível com técnicas de aceleração predominantes (por exemplo, decodificação especulativa). Para enfrentar esses desafios, introduzimos uma nova arquitetura MoE, o BlockFFN, bem como suas técnicas eficientes de treinamento e implantação. Especificamente, usamos um roteador que integra ativação ReLU e RMSNorm para um roteamento diferenciável e flexível. Em seguida, para promover tanto a esparsidade em nível de token (TLS) quanto a esparsidade em nível de bloco (CLS), objetivos de treinamento conscientes de CLS são projetados, tornando o BlockFFN mais amigável à aceleração. Por fim, implementamos núcleos de aceleração eficientes, combinando esparsidade de ativação e decodificação especulativa pela primeira vez. Os resultados experimentais demonstram o desempenho superior do BlockFFN em relação a outras baselines MoE, alcançando mais de 80% de TLS e 70% de CLS para 8 tokens. Nossos núcleos atingem até 3,67 vezes de aceleração em dispositivos reais de borda em comparação com modelos densos. Todos os códigos e checkpoints estão disponíveis publicamente (https://github.com/thunlp/BlockFFN).

English

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.67times speedup on real end-side devices than dense models. All codes and checkpoints are available publicly (https://github.com/thunlp/BlockFFN).