BlockFFN : Vers une architecture Mixture-of-Experts adaptée à l'accélération côté terminal avec une parcimonie d'activation au niveau des blocs

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

July 11, 2025

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

cs.AI

papers.abstract

Pour alléger la charge computationnelle des grands modèles de langage (LLM), les architectures à activation parcimonieuse, représentées par le mélange d'experts (MoE), ont suscité un intérêt croissant. Cependant, le routage non différentiable et rigide du MoE classique nuit aux performances du modèle. De plus, bien que chaque token n'active qu'un petit nombre de paramètres, ces architectures à activation parcimonieuse présentent une faible parcimonie au niveau des blocs, indiquant que l'union de plusieurs tokens consécutifs active une grande proportion de paramètres. Un tel schéma de parcimonie est peu adapté à l'accélération dans des conditions à ressources limitées (par exemple, les appareils en bout de chaîne) et incompatible avec les techniques d'accélération dominantes (par exemple, le décodage spéculatif). Pour relever ces défis, nous introduisons une nouvelle architecture MoE, BlockFFN, ainsi que ses techniques d'entraînement et de déploiement efficaces. Plus précisément, nous utilisons un routeur intégrant l'activation ReLU et RMSNorm pour un routage différentiable et flexible. Ensuite, pour promouvoir à la fois la parcimonie au niveau des tokens (TLS) et la parcimonie au niveau des blocs (CLS), des objectifs d'entraînement conscients de la CLS sont conçus, rendant BlockFFN plus adapté à l'accélération. Enfin, nous implémentons des noyaux d'accélération efficaces, combinant pour la première fois la parcimonie d'activation et le décodage spéculatif. Les résultats expérimentaux démontrent la performance supérieure de BlockFFN par rapport aux autres modèles de référence MoE, atteignant plus de 80 % de TLS et 70 % de CLS pour 8 tokens. Nos noyaux permettent une accélération allant jusqu'à 3,67 fois sur des appareils réels en bout de chaîne par rapport aux modèles denses. Tous les codes et points de contrôle sont disponibles publiquement (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).