BlockFFN: Towards End-Side Acceleration-Friendly Mixture-of-Experts with Chunk-Level Activation Sparsity
Authors: Chenyang Song
Weilin Zhao
Xu Han
Chaojun Xiao
Yingfa Chen
Yuxuan Li
Zhiyuan Liu
Maosong Sun
Venue: COLM 2025 (2025)
Keywords: mixture-of-experts
activation sparsity
inference acceleration
TL;DR: We propose BlockFFN, an effective MoE architecture more friendly for end-side acceleration, as well as its sparsity-aware training objectives and efficient acceleration kernels.

Abstract: 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 at https://github.com/thunlp/BlockFFN.

Published as a conference paper at COLM 2025
BlockFFN: Towards End-Side Acceleration-Friendly Mixture- of-Experts with Chunk-Level Activation Sparsity
Chenyang Song, Weilin Zhao, Xu Han, Chaojun Xiao, Yingfa Chen, Yuxuan Li, Zhiyuan Liu, Maosong Sun
Dept. of Comp. Sci. & Tech., Institute for AT, Tsinghua University, Beijing, China (scy22, manl23)@mails..... @mails. tsinghua.edde edu.cn,
edu.
Abstract
To alleviate the computational burden of large language models (LLMs), architectures with activation sparsity, represented by mixture-of-experts (MoB), have attracted increasing attention. However, the non-differentiable and inflexible routing of vanilla MoE hurts model performance. Moreover, while each token activates only few parameters, these sparsely-active architectures exhibit low chunk-level sparsity, indicating that the union of multiple consecutive tokens activates large ratio of parameters. Such sparsity pattern is unfriendly for acceleration under low-resource condi- tions (e.g., end-side devices) and incompatible with mainstream accelera- tion techniques (e.g,, speculative decoding). To address these challenges, we introduce novel MoE architecture, BlockFFN, as well its efficient training and deployment techniques. Specifically, we use a router integrat- ing ReLU activation and RMSNorm for differentiable and flexible routing. Next, to promote both token-level sparsity (TLS) and chunk-level spar- sity (CLS), Cl-r-aware training objectives are designed, making BlockFFN more acceleration-friendd kernels, combining activ ation 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 3.67 speedup real end-side de- vices than dense models. All codes and checkpoints are available publicly!.
Finally, we implement efficient acceleration
1 Introduction
To reduce the high costs of training and deploying large language models (LLMs), various efficient LLM architectures are proposed (Wan al., 2023). A popular paradigm is designing architectures with activation sparsity, indicating that considerable part of LLM parameters contribute weakly LLM outputs given specific inputs, and thus can be skipped (i.e., not activated) in the forward and backward computation. Mixture-of-expent outstanding representative and has been adopted by many recent models such as Mixtral- 8x22B (Jiang et al., 2024) and DeepSeek-Ve V3 (Liu et al., 2024b). Based on MoF, techniques including load balancing (Wang et al., 2024a) and expert parallelism (He et al., 2021) are adopted to achieve remarkable efficiency on cloud-side servers.
(MoE) is an
However, few efforts explore sparsely-activater architectures under low- rwworeource conditions (e.g., end-side devices), where it is challenging to deploy huge MoE models and highly- distributed frameworks with expert parallelism. For end-side MoE models, which generally serve only few users, some typical issues not required to be considered (e.g., load balancing, see Appendix A), while raising the following two challenges.
Performance compromise caused by imperfect routing. Existing MoE models generally compromise performance due to two significant routing drawbacks: non-differentiability
"Equal Contributions.
*Corresponding Authors.
'https://gi thwbb http:///tthuu..cmmm/u//////laaaa

---

Published as a conference paper at COLM 2025
31
Vanilla Topk DeepSeekMot GRIN
10
Decoding Speed (talentord)
Chunkc oven Sparsty (55
29
RoMoE
BlockFFN
50
50
Train Loss
2.
50
27
20
43
30
20
BlockFFN-11
DeepSenkMMoF166
10
ProSparse LLaM/2-70
Massingaes
our (LBA) (32
120
Length
Decoding Method
(a)
(b)
(c)
Figure 1: (a) For models with high TLS (except BlockFFN-1.2BB) CLS quickly collapses to
lower level as a chunk contains more consecutive tokens. (b) Smoothed training curves of the BlockFFN-1 1.2 and other MoE baselines with the same total and activated parameters. (c) The speed of different decoding methods, where "Ours (1-Tok)" and "Ours (32-Tok)" are our leven and chunk-level sparsity- acceleration kernels, respectively.
and inflexibility. Specifically, most mainstream MoE models adopt a TopK router (Fedus et al., 2022; Jiang et al., 2024) with discrete and non-differentiable computation. Conse- quently, only activated parameters have complete gradients and are well updated at each step, which harms the convergence efficiency of MoE models (Liu et al., 2024c). Moreover, Topk makes each token activate the same number of experts, enforcing an inflexible activa- tion pattern, which may weaken model performance (Huang et al., 2024). Few works can well alleviate both drawbacks, see Section 2.1.
Acceleration unfriendliness caused by low chunk-level sparsity (CLS). To make sparsely-activated architecture more friendly for acceleration, just increasing the ratio of weakly-contributed experts for each token, namely the token-level sparsity (TLS), is not enough. Instead, the ratio of weakly-contributte experts for multiple consecutive tokens, namely the chunk-level sparsity (CLS), critical for practical acceleration.
Specifically, low CLS can eliminate the value of activation sparsity when combined with speculative decoding, a mainstream acceleration method that requires LLMs to process multiple consecutive tokens at the same time (Leviathan et al., 2023). Besides, other im- portant resource-saving techniques, such as offloading, also become more challenging to implement due to the large differences in activation patterns within specific chunk, leading to frequent GPU-CPU communication overheads. Unfortunately, existing works mainly focus improving TLS (Mirzadeh et al., 2023; Song et al., 2024; 2025), but low CLS values still exist in most sparse architectures (see Figure 1a).
To address the above challenges, introduce the following novel MoE architecture, well as its training techniques and efficient end-side deployment.
For model architectures, we propose BlockFFN, novel MoE paradigm that minimizes performance compromise. Specifically, its router module integrates ReLU and RMSNorm. ReLU computes differentiable and flexible activation patterns, enabling each token deter- mine the number of activated experts adaptively. RMSNorm generates learnable magnitudes of activation values. This separation of activation patterns and magnitudes alleviates the disturbance on activation magnitudes induced by regularization (e.g., the shrinkage of activation magnitudes caused by L1 (Rajamanoharan et al., 2024)), as regularization applies solely to the ReLU activation pattern. Through experiments, we demonstrate the better performance of BlockFFN compared other MoE variants (Figure 1b, Table 2 and 3).
For training techniques, we introduce CLS-aware training objectives to improve the CLS of BlockFFN, which can make BlockFFN more friendly for acceleration, including the activation locality loss and the chunk sparsification loss. The former is to increase the similarity of activation patterns between neighbor tokens, helping reduce the gap between TLS and CLS. The latter is to increase the overall sparsity level. While existing sparsification objectives such as L1 (Song et al., 2025) are applied to each token independently, our chunk sparsification loss directly minimizes the probability that specific expert is activated by
2

---

Published as a conference paper at COLM 2025
Model
Activation values A(x)
Single expert outputs E,(x)
Neuron-level activation sparsity
o(W
Non-gated FFN
w
Wap RAXX,,
Wdown R
Gated FFN
<WW) (
(wijx) W down
(Shazeer, 2020)
quie RAXN,
Wape Waoon RR
dRcLU
X) ( gu
ReLU(wap Weono
(Song et al., 2024)
W quie RAXN,
Wap. W down E IR
Block-level activation sparsity
Vanilla TopK MoE (Jiang et al., 2024)
TSpf(Suft WT W router Rd.M
xo)ter
DeepSeek-V1/91 /V2 (Dai et al., 2024)
routtr x)]] Wrntter (N, Nate))
[Tstarer (Sok(Softmax(MM (WT
DeepSeck-V3
1 Norm (
(W x))]
(Liu al., 2024b)
W Winter [R,x(N. NAmr))
WT [cr( (enc)) Wilxo]
down
DynamicMoF:
(Sof(Sof WT
nowter x))
Wgater Wap E Rexod
(Huang et 2024)
W Wotttrr E R4.x
W down E Rd.xdy
GRIN
OparseMixer(Softmmme (WT
router)))
(Liu et al., 2024c)
router RdLNNG
ReMoB
ReLU rudg
(Wang et al,, 2024b)
W Wrutter E R
BlockFFN (Ours)
RMSNorm(ReLUUM ())
WT doomn" o7 (W.y)))
Wrnater R4
Wap RRdxxd, W Wown Rdxxd
Table 1: Comparison between different architectures with activation sparsity. o denotes an activation function (e-g, ReLU, Swish, GELU). de is the intermediate dimension each block-level expert. For simplicity, the expert index i is omitted in notations of (x
at least one token within the chunk. In experiments, we obtain average TLS values higher than 80% and 8-token CLS values higher than 70% (Table 2).
For end-side deployment, we implement efficient acceleration kernels for BlockFFN, com- bining activation sparsity and speculative decoding for the first time, and demonstrate its practical effectiveness real end-side devices such as NVIDIA Jetson Orin NX. To enhance the efficiency of verifying multiple tokens in speculative sampling, we leverage the high activation similarity across multiple tokens induced by the high CLS level. This enables merging memory accesses to the same expert across different tokens, reducing the memory access volume to the union of experts activated by these tokens (Figure 2). Additionally, we implement the kernels based on CUTLASS (Thakkar et al., 2023) and utilize tensor cores to boost computational efficiency. Overall, the kernel achieves an acceleration ratio of 3.67 x, compared to the baseline auto-regressive (AR) decoding (Figure 1c and Table 6).
2 Preliminaries and Related Works
2.1 Architectures with Activation Sparsity
To reduce the computation expenses of LLMs, various inference acceleration methods are proposed. Quantization (Xiao et al., 2023; Yao et al., 2023; Shao et al., 2023) and distil- lation (Gu et al., 2023; Hsieh et al., 2023) compress LLMs by using low bit-widths and transferring knowledge into smaller models, respectively. Weight pruning (Ma et al., 2023; Sun et al., 2023; Frantar & Alistarh, 2023; Xia et al., 2023) reduces FLOPs by removing weakky-contributed parameters (regardless of inputs). Speculative decoding (Li et al., 2024a; Cai et al., 2024; Zhao et al., 2024) uses a smaller model to generate multiple candidate tokens and lets the L.M itself verify these tokens in parallel.
3

---

Published as a conference paper at COLM 2025
Besides the above post-training methods, efficient architectures with activation sparsity can also effectively reduce the computation overhead of LLMs (Xue et al., 2024; Zhang et al., 2024b; Liu et al., 2024b). Specifically, in aparsely-activatee architectures, considerable part of the parameters contribute weakly to the model outputs given specific inputs.
In this work, we mainly focus on the activation sparsity within FFN layers. Typically, sparsely-activated FFN with hidden dimension dj, can be written in a unified MoE format:
FxN Σ (,(x) Ex(),, A [A)(x), (x(x AN, (x)).
(1)
where x E Rdh is the input hidden state, and N₂ is the number of experts. (x(x) E R and (x) E R denote the i-th activation value and expert outputs, respectively. If some (x) is zero
low value, the corresponding expert E, is weakly-contributed. Token- level sparsity TLS denotes the average ratio of weakly-contributed experts for a single token, while chunk-level sparsity CLS² is the ratio of experts contributing weakly to all tokens within a consecutive chunk of length L. Based on the granularity of experts, LLM architectures with activation sparsity can be divided into the following two categories.
Neuron-level activation sparsity commonly exists in mainstream LLMs (Li et al., 2022), where each expert is composed of single neuron, i.e., single columns or rows within the FFN parameter matrices. For example, LLaMA2-7B is estimated to have about 70% TLS (Zhang et al., 2024a). However, as every single neuron is generally too small to be memory access unit (i.e., < 1MB in BF16), most neuron-level sparse LLMs have bad memory locality. This makes it difficult realize practical acceleration due to large TO overheads.
Block-level activation sparsity indicates that each expert is composed of multiple neurons or MLP modules. Represented by MoE (Fedus et al., 2022), such architectures are currently the mainstream solution for sparsity-based acceleration due to their better memory locality.
Nevertheless, the routing strategies of many MoE models, especially TopK routers, have non-differentiable and inflexible activation patterns, which limit model performance (Luo et al., 2024). Works such as GRIN (Liu et al., 2024c), ReMoE (Wang et al., 2024b), and DynamicMoE (Huang et al., 2024) try to address these issues. We list and compare several representative architectures of activation sparsity in Table 1. There are also special methods based on direct expert merging (e.g., SMEAR (Muqeeth et al., 2023) and Lory (Zhong et al., 2024)) that cannot be naively expressed as Equation 1.
In this work, BlockFFN absorbs the merits of both categories. With good memory locality of block-level experts, it adopts a Reltivacte router (common in neuron-level settings), with activation values scaled by RMSNorm (i.e., the architectural difference from ReMoE).
2.2 Acceleration with Activation Sparsity
For neuron-level architectures, due to the relatively bad memory locality, designing tailored acceleration frameworks is complicated. Deja Vu (Liu et al., 2023) and Powerinfer (Song et al., 2023) utilize actívation predictors to forecast the actívation values, thus reducing 10 overheads. P'owerlnfer22 (Xue et al., 2024) introduces complex 1O pipelines and neu- ron caches to promote higher speedup on specific smartphones. However, these all risk potentially inaccurate inference due to the imperfect performance of activation predictors.
Block-level architectures have relatively more available frameworks. FastMoE (He al., 2021) and Tutel (Hwang et al., 2023) mainly focus on distributed training or inference with multiple GPUs working concurrently, while MegaBlocks (Gale et al., 2023) emphasizes the large-batch training of MoE. However, few of them are tailored for deploying MoE on end-side devices, where it is generally impractical to adopt a distributed implementation, and the service requirements shrink to mmall-batch inference for individual users. Under end-side conditions, sparsity-based acceleration will face different challenges.
As far as know, we present the first work to address the acceleration combining activation sparsity and speculative decoding. Specifically, we improve the chunk-level sparsity of models through Cl-S-aware training objectives, making BlockFFN more friendly
4

---

Published as a conference paper at COLM 2025
for sparsity-based acceleration and speculative decoding. Moreover, our acceleration kernels are well applicable to end devices and have remarkable effectiveness.
3
Methodology
In this section, we first introduce the overall architecture of BlockFFN (Section 3.1) and CLS-aware training objectives (Section 3.2). Then, the acceleration kernels are introduced in Section 3.3, combining activation sparsity and speculative decoding for the first time.
3.1 BlockFFN Architecture
Expert modules Considering the better memory locality of block-level activation sparsity, we make each BlockFFN expert an MLP with an activation function:
(x)
W down
Swish (i) ( x),
(2)
where i is the expert index, and w(f up E Rdaxd, wii
down E R& (d, are learnable weights.
Following DeepSeekMoE (Liu et al., 2024b), we use fine-grained expert segmentation to increase flexibility, namely de << dn. We specifically add a Swish activation (Ramachandran et al., 2017) to increase the honlinearity. Notably, we choose a vanilla non-gated MLP for experts instead of the more popular gated variant (Dauphin et al., 2017; Shazeer, 2020), as we find that a gated MLP can destroy the router sparsity (See Appendix I).
Router module BlockFFN adopts a linear router with ReLU activation instead of TopK. As common activation function in neuron-level sparse LLMs, ReLU is fully differentiable and can generate sparser activation patterns than other common activations (e.g,, Swish) (Luo al., 2024). Moreover, ReLU allows each token adaptively activate different numbers of experts. This alleviates the inflexibility issue of conventional TopK routing.
On the other hand, as one major difference from ReMoE (Wang et al., 2024b), we add an RMSNorm layer (Zhang & Sennrich, 2019) after ReLU:
A (x) W
A (x) Re)U(A((x)), A(x) RMSNomM( (x)),
(3)
where Wrouter is learnable parameters. Such design makes the magnitude of activation values adaptively learned through RMSNorm, indicating better flexibility than vanilla softmax. Besides, RMSNorm separates the ReLU activation pattern A (x) from the final activation value A (x). This alleviates the disturbance activation magnitudes by direct regularization, which may hurt performance (Rajamanoharan et al., 2024) (Section 4.1.4).
3.2 CLS-Aware Training Objectives
The low chunk-level sparsity (CLS) is one important obstacle to fully leveraging activation sparsity in practical acceleration, especially under conditions where multiple consecutive tokens are processed in parallel (e.g., speculative decoding). The improvement of CLS involves two imporlant aspects: (1) how to promote activation locality; (2) how to promote higher overall sparsity. We propose two respective training objectives.
Activation locality loss Activation locality refers to the similarity of activation patterns between neighbor tokens, which also indicates the gap between TLS and CLS. To promote this property, we introduce the activation locality loss as an additional training objective:
A LeftShift(A (x)), Lat BCE|a(a (x)) (((aaa(()))))
(4)
where o and denote the sigmoid function and the sharpness hyper-parameter, respectively. We approximate the activation pattern through sharp sigmoid function applied on A (x). LeftShift operator left-shifts a tensor in the sequence dimension, and finally, the binary cross entropy BCE minimizes the gap between the soft activation patterns of neighbor tokens.

---

Published as a conference paper at COLM 2025
Real Activation
Union Activation
Irrelemant Activation
TLS Mean 13233
7
6'6'6'6)
12
[nnion]
3
No
6
A(X) E R
Inion (2 3
Sparse (E Uniom)
R,4.
Mask
Immlesant Activation
4
Nr expert number
decexpert interm ediate size
x. R
Wap E Rdx0DG,d))
dy hidden size
1. number of tokens in chunk
Figure 2: The overall framework of our acceleration kernels (the up projection part), which combines speculative decoding and chunk-level sparsity for higher efficiency. The down projection part has similar implementation.
Chunk sparsification loss cannot be achieved without considerable reduction in computation, which relies on high sparsity level. Conventionally, L1 (Song et al., 2025) and router entropy (Huang et al., 2024) are both effective methods to improve sparsity, but they are applied independently to each token and cannot directly optimize the chunk-level sparsity.
Despite the increase in activation locality, practical acceleration
Therefore, we design the chunk sparsification loss, which directly minimizes the chunk-level sparsity of chunk with L consecutive tokens. Suppose Pik is the probability of the i-th expert activated by the k-th token, while E Pix = The loss is the average probability that the i-th expert is activated by at least one token within this chunk (i.e., Paat):
N
N [Pul 11 = Norm (Al (x)),) Pir =1 exp Σ In(1
PR)),,
Los
p
(5)
Ne
act
where A (x) E N specifically denotes the ReLU activation pattern of the k-th token, and Norm operator normailizes it in the expert dimension.
The overall training objectives are computed by: Liotal = Lim + Aat + Acshes where À and Acs are corresponding factors. We introduce an adaptive factor scheduler adaptively determine Ass according to the dynamics of Les, see Appendix B.
3.3 Acceleration Kernels
We implement acceleration kernels for BlockFFN, which are applicable to end-side devices and effectively combine chunk-level sparsitt and speculative decoding.
Specifically, during the speculative sampling process, the draft model proposes draft tokens. When BlockFFN verifies these tokens, the router activation values are denoted as A(x) R, while the index union of activated experts for these tokens is Union(x)) xx). Due to BlockFFN's high CLS level, the size of Union (x) only accounts for small ratio to the total expert number. Therefore, by only involving the experts in Union (x) for computation, memory access is reduced and sparsity- acceleration can be achieved in verification.
However, different experts may be activated by different subsets of tokens, which not friendly for hardware parallelization. To address this issue, we leverage the characteristic that CLS and TLS values of BlockFFN are relatively close, indicating that each expert in Union( (x) is activated by the vast majority of tokens. Therefore, we only need to precompute all tokens for every activated expert for better GPU utilization, and subsequently discard computations induced by irrelevant activations. Specifically, for up projection, the hidden states of all и tokens participate in the matrix multiplication with the experts in Union (x), yielding intermediate result mid, illustrated in Figure 2. Finally, apply mask based on the sparse pattern of A(x) to remove irrelevant activations. Similar sparse computation is also conducted for down projection.

---

Published as a conference paper at COLM 2025
Small
Medium
Large
XLagge
Setting
TLS
CLSy 1 PPL.
TLS
CLSs
PPL1
TLS
CLSy
PPL
TLS CLSg t PPLJ
Dense
14.90
10.03
9.29
8.49
TopK
79.17
49.18
16.22
85.00
62.25
10.58
83.33
59.40
9.96
82.14
61.05
8.87
DSMoE
79.17
49.27
15.53
85.00
66.22
10.69
83.33
62.06
9.89
82.14
60.28
8.86
GRIN
79.17
50.45
15.50
85.00
61.48
10.40
83.33
59.08
9.72
82.14
60.89
9.03
ReMoE
78.33
42.44
14.60
84.43
52.00
10.42
82.80
50.79
9.60
81.93
51.01
8.78
BlockFFN
80.54
71.38
14.88
84.25
75.87
10.23
84.05
73.79
9.52
81.81
72.78
8.69
Table 2: The average perplexities (PPL) and chunk-level sparsity for consecutive tokens (CLSg) the validation data under close TLS. "Dense" is the upper bound setting, which involves vanilla Transformers with the same parameter numbers as MoE settings.
Small
Medium
Large
XLarge
Setting
C.R.t
R.C.F
C.R.1
R.C.
C.R.t
R.C.1
C.R.t
R.C.t
Dense
45.15
32.72
52.68
44.85
54.89
48.87
57.30
49.53
TopK
43.73
28.99
50.90
42,47
52.84
43.05
54.98
49.29
DSMoE
44.35
31.90
50.46
44.89
52.54
46.45
55.28
50.57
GRIN
43.69
31.61
50.96
42.56
52.93
45.04
54.65
49.68
ReMoE
45.22
32.78
51.42
43.05
53.77
47.33
55.56
47.12
BlockFFN
44.80
33.93
51.75
43.74
54.44
51.60
56.42
50.73
Table 3: The average evaluation scores on two groups of benchmarks: commonsense reasoning (C.R.) and reading comprehension (R.C.). "Dense" is the upper bound setting.
The matrix multiplication kernel is modified based on CUTLASS GEMM (Thakkar et al., 2023), where we modify the outer loop of the up projection and the inner loop of the down projection to only scan through those activated experts in Union (x), see Appendix J. To match the requirements of CUDA Tensor Core, we set the number draft tokens to 32.
4 Experiments
4.1
Architecture Rationality
4.1.1
Overall Results
To demonstrate the rationality of our architecture, we conduct experiments by comparing BlockFFN with multiple sparsely-activated architectures: Vanilla TopK MoE, DeepSeekMo (DSMoE) (Dai et al., 2024), GRİN (Liu et al., 2024c), and ReMoE (Wang et al., 2024b) (see Appendix C). To ensure fairness, we keep consistent settings for attention layers and MoE experts (i.e., the number and intermediate dimension of experts) throughout baselines and BlockFFN. Besides, all settings (within each scale) have close parameter numbers, training token numbers, and token-level sparsity. We involve four parameter scales: Small (0.1B), Medium (0.5B), Large (0.8B), and XLarge (1.2B). See Appendix D for model settings.
We adopt two comparison metrics: perplexity (PPL) on validation datasets and evaluation scores on benchmarks. Benchmarks include two groups: commonsense reasoning (C.R.) and reading comprehension (R.C.). See Appendix E for details about data and benchmarks.
The PPL and evaluation scores are shown in Table 2 and 3, respectively. The training curves of the "XLarge" settings are drawn in Figure 1b. We can draw the following observations:
(1) Performance: Under close parameter numbers, all the settings (except for Small ReMoB and BlockFFNN cannot match the "Dense" setting, due to the performance compromise of sparsification. However, under close TLS values (i.e., identical average FLOPs for each token), BlockFFN outperforms other MoE baselines in terms of validation PPL, train loss, and scores downstream tasks, showing less performance compromise.
7

---

Published as a conference paper at COLM 2025
35
BoccFFFN
Activation Magnttdd
Wloce N wib RMSMorrm
1.0
RoMoE (L3 objectives)
Tokee Frequuency
33
25
Experts (84
Layer
Figure 3: For each token in vocabulary, we calculate its frequencies and average ratios of activated experts, which show bimodal distribution of expert allocation.
Figure 4: The layer-wise distributions av- erage activation magnitudes on the "Small" settings. While BlockFFN uses CLS-aware objectives, ReMoE adopts L1 regularization.
(2) Sparsity: Under close TLS values, BlockFFN always has considerably higher CLS values than other baselines. Attributed to CLS-oriented training objectives, this property makes BlockFFN more friendly for acceleration.
4.1.2 Expert Selection Stability
Low-resource conditions often require the implementation of memory techniques, such as offloading, where the weights of experts are loaded into memory only when they are activated. Such technique calls for higher expert selection stability. Specifically, the distribution of selected experts should be as similar as possible across consecutive tokens, so that the costs of expert IO can be saved.
In this section, we demonstrate that BlockFFN has significant expert selection stability, which measured by the reuse ratio, namely, within the activated experts of one token, the average ratio of experts that are also activated by its next token. Within sequence with L I tokens, the set of experts activated by the i-th token is denoted by S, The reuse
ratio of this sequence is calculated by + -/- [S,S Sj
As shown in Table 4, the high reuse ratios of BlockFFN models over 85% ensure satisfactory memory efficiency and good adaptability to offloading.
Scale
Small
Medium
Large
XLarge
Reuse Ratio (%)
90.28
89.69
87.13
89.57
Table The average reuse ratios of BlockFFN models on the validation data.
4.1.3 Analysis of Expert Allocation
While ReLU-based routing is intrinsically differentiable, in this section, we examine how the router allocates experts and whether the inflexibility of activation patterns truly addressed. Based on the validation data, we calculated the frequencies and average ratios of activated experts for each token, which are shown in Figure 3. These results demonstrate bimodal distribution of expert allocation.
Specifically, the smaller peak lies between the activation ratio interval between 10% and 15,, which involves tokens such as numbers (e.g., "ø", "1"), single characters (e.g. "a", "b"), and reserved words programs (e.g., "import", "return"). These tokens have more deterministic meanings and thus require fewer experts for processing. By contrast, the larger peak between 20% and 25% mainly involves those tokens with more complex or diverse meanings, such as English pronouns and Chinese characters. Therefore, they need more experts to understand. Such a bimodal allocation of experts demonstrates that ReLU activation can truly address the routing inflexibility and allocate resources more wisely.
8

---

Published as a conference paper at COLM 2025
Direct Ablation
Substitute Objectives
Setting
TLS
CLSg
PPL1
Setting
TLS
CLSg t
PPL1
AL+CS
80.54
71.38
14.88
AL+L1
79.86
65.16
15.50
CS
81.67
67.56
15.66
AL+Ent
81.05
69.13
15.97
AL 63.55
52.59
14.89
79.35
45.13
15.01
Null
48.56
14.89
14.85
Ent
79.04
45.78
15.12
Table 5: Ablation studies the training objectives. "AL+CS" is our standard setting. "AL", "CS", "L1" and "Ent" indicates activation locality, chunk sparsification, L1 norm, and router entropy, respectively. "Null" is the setting without any additional training objectives.
4.1.4 RMSNorm and Activation Magnitude Disturbance
In this section, we examine the effectiveness of the RMSNorm in our router module. First, conduct an ablation study the "Small" setting. After removing the RMSNNrm layer, the validation PPL rises from 14.88 to 15.04, indicating the effectiveness of RMSNorm.
Moreover, to inspect the effects of RMSNorm, we calculate the average activation mag- nitudes (computed by L2 Norm) on the validation data. As shown in Figure 4, under all settings, higher layers (closer to output) generally have larger activation magnitudes. Without RMSNorm, the magnitudes of activation values in BlockFFN considerably rise with worse performance. By contrast, ReMoE, which has similar architecture BlockFFN without RMSNorm, suffers from significantly smaller activation magnitudes. This is at- tributed the activation shrinkage issue induced by L1 regularization directly imposed on activation values (Rajamanoharan et al., 2024). We assume that RMSNorm, preventing activation values from direct regularization (Section 3.1), potentially alleviates activation magnitude disturbance and maintains more stable and appropriate magnitude level.
Besides the above issues, expert granularity also has an important influence. Through experiments, we find that the validation loss generally decreases with finer experts, but the marginal benefits quickly diminish with > 40 experts for BlockFFN Medium. However, the relationship between sparsity and expert granularity is nonmonotonic, with the best setting of 40 experts achieving the highest sparsity. See Appendix F for more details.
4.2 Training Objective Rationality
In Section 3.2, we introduce activation locality (AL) loss and chunk sparsification (CS) loss training objectives. In this part, mainly discuss whether such a practice is reasonable and better than other potential substitutes.
First, we conduct direct ablation studies by removing either loss CS loss. As shown in the left part of Table 5, without AL (Setting "CS"), ), the model suffers from lower CLS and considerably higher PPL. On the other hand, without "CS" (Setting " "AL" and "Null"), the sparsity (both TLS and CLS) can be extremely low. These demonstrate the division of labor: CS is mainly responsible for global sparsification, while AL is to promote the CLS with less performance compromise (compared with the direct application of large CS loss). A possible explanation for why " "AL+CS" performs better than pure "CS" is the competing relationship between "AL." and "CS". Specifically, the introduction of AAL" weakens the sparsification effect of "CS", producing lower TLS but higher CLS, and the better performance is attributed to the lower TLS level.
Next, we explore other potential substitute training objectives for sparsification. This includes the L1 norm (L1) (Song et al., 2025) and router entropy loss (Ent) (Huang et al., 2024). As shown the right part of Table 5, replacing CS with L1/Ent (Setting "AL+L1" and "AL+Ent") cause considerable drop in performance. Besides, due to the absence of AL, "L1" and "Ent" cannot reach satisfactory CLS, either. Therefore, CS is a more competitive sparsification partner of AL with less performance compromise.

---

Published as a conference paper at COLM 2025
Setting
MT.
Trans.
Summ.
QA
Math
RAG
Average
Huggingface
7.61
7.32
7.54
7.81
7.76
5.96
7:33 (0.57
Baseline AR EAGLE-2
12.61
13.32
12.40
13.04
12.80
12.86
12.84 (1.00x)
23.70
20.26
20.74
21.99
25.01
22.47
22.36 (1.74x)
Ours (1-Tok)
38.41
40.39
37.37
42.05
44.54
39.38
40.36 (3.14x)
Ours (32-Tok)
49.43
39.18
42.70
46.09
59.85
45.76
47.17 (3.67x)
Table 6: Decoding speeds (token/sec) and average speedup ratios on NVIDIA Jetson Orin NX. "Ours (1-Tok)" is our token-level acceleration kernel purely dependent on sparsity, while "Ours (32-Tok)" is our efficient chunk-level acceleration kernels that combine EAGLE- 2 and chunk-level sparsity. The speedup ratios are relative "Baseline AR".
4.3 Practical Inference Acceleration
Speedup experiment To demonstrate the efficacy of our acceleration kernels, we conduct speedup experiments on Spec-Bench (Xia et al., 2024), a comprehensive benchmark for speculative decoding, with NVIDIA Jetson Orin NX 16GB. To ensure comparison fairness, except for the vanilla Huggingface auto-regressive decoding, all baseline methods are implemented within the framework of FR-Spec (Zhao et al., 2025), which applies CUDA kernels to reduce IO overheads and is much faster than Huggingface. These baselines include Baseline AR (i.e., faster FR-S TR-Spec rego-regressvve implementation), and EAGLE- 2 (Li et al., 2024b). Moreover, since acceleration effects are more significant on larger models, specifically train 2.8B BlockFFN model as the base of our efficiency experiment.
From Table 6, we have the following observations: (1) "Baseline AR" is considerably faster than "Huggingface" indicating that our experimental framework is efficient enough and can alleviate the influence of potential IO overheads. (2) "Ours (32-Tok)), our acceleration kernel combining speculative decoding and activation sparsity, achieves the highest decoding speed, with 3.67 x speedup. Meanwhile, it is faster than the pure sparsity setting "Ours (1-Tok)" and the pure speculative decoding "EAGLE-2". This demonstrates the value of such combination and reveals the significant value of utilizing activation sparsity in end-side device inference acceleration. See Appendix H for speedup data on independent datasets.
Upper bound analysis Moreover, conduct further experiments and find that both "Ours (1-Tok)" and "Ours (32-Tok)" can reach the theoretical upper bound of FFN speedup ratios induced by the corresponding token-level sparsity and the average union sparsity of tokens contained in an EAGLE-2 draft tree, see Appendix G.
5 Conclusion
In this work, we propose BlockFFN, novel MoE architecture equipped with ReLU-based differentiable and flexible routing strategy, which enables BlockFFN to outperform existing MoE counterparts, Next, we advocate more attention chunk-level sparsity (CLS), and introduce the CLS-aware training objectives to promote the 8-token CLS to over 70%, offering BlockFFN better activation locality and more friendliness for end-side device acceleration. Finally, our efficient acceleration kernels achieve up to 3.67x speedup on NVIDIA Jetson Orin NX than the baseline auto-regressive decoding, reaching the sparsity-induced upper bound of FFN acceleration.
Acknowledgments
This work is supported by the National Key R&D Program of China (No.2022Z01163312) Beijing Municipal Science and Technology Plan Project (224110001124025) and a grant from the Guoqiang Institute, Tsinghua University. Our research is also supported by Huawei and can be carried out using the Huawei Ascend AT technology stack.
10

---

Published as a conference paper at COLM 2025
References
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. PIQA: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020. URL https: //jss./j.../////.//.... aaai wwwoorrinnnooeo g/index.
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D Lee, Deming Chen, and Tri Dao. Medusa: Simple LLM inference acceleration framework with multiple decoding heads. In Forty-first International Conference Machine Learning, 2024. URL https://arxiv org/pdf/ 20/.//401.100 10774.
Christopher Clark, Kenton Lee, Ming-We Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 1924-29366 2019. URL https: aatpss///lclantoooggyyy or 0tg/N19-110 pdf.
Jonathan Clark, Eunsol Choi, Michael Collins, Dan Garrette, Tom Kwiatkowski, Vitaly Nikolaev, and Jennimaria Pallomaii. TyDi QA: A benchmark for information-seeking question answering in typologically diverse languages. Transactions of the Association for Computational Linguistics, 8:454-470, 2020. URL https: http:/:////aaaaooooooooooo org/ 2020. tacl-1. 30. .d..
Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. DeepSeekMoE: Towards ultimate expert specialization in mixture-of-experts language models. CoRR, 2024. URL http://arxiv. org//46f/2401.00666 06066.
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. International Conference Machine Learning, pp. 133-941. PMLR, 2017. URL https: httpeeiiiig........rrs//////aaaa
mwr.
Ning Ding, Yulin Chen, Bokai Xu, Yujia Qin, Zhi Zheng, Shengding Hu, Zhiyuan Liu, Maosong Sun, and Bowen Zhou. Enhancing chat language models by scaling high- quality instructional conversations. ar preprint orXie:2305.1223, 2023. URL https:
//arxiv.org//df//33333 /2385. 14233. pdf.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Llama herd of models. arXiv preprint arXin.2407.2178, 2024. URL http:tps://axxxi.org//f//////// 21783.
William Fedus, Barret Zoph, and Noam Shazeer. Switch Transformers: Scaling to trillion
parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022. URL https: //ww://wwwwwwww wwwpm.orr/////o/oommmmooo/oooerl
pdf.
Elias Frantar and Dan Alistarh. SparseGPT: Massive language models can be accurately pruned in one-shot. In International Conference on Machine Learning, pp. 10323-10337. PMLR, 2023. URL https::/prooceedingg ml mllr.press/v222/frrmms arrreesvveee/rraaaa22aaaaaaaaaaa pdf.
Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. MegaBlocks: Efficient sparse training with mixture-of-experts. Proceedings of Machine Learning and Systems, 5:288-304, 2023. URL https::/procceedingg mlsys org/paper www.se/apppeeeea////// pdd.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The Pile: An 800GB dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020. URL https://art / wwwppdfg///ooopppoopopopppp...pp / 2001. 00027. .df.
Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. Knowledge distillation of large language models. Xiv preprint arXio:2306.08543 2023. URL https://aaxii: wwwpaxiwi...ooo/oooooooo org/pdff// /2066 2306.08543. pdf.
11

---

Published as a conference paper at COLM 2025
Jiaao He, Jiezhong Qiu, Aohan Zeng, Zhilin Yang, Jidong Zhai, and Jie Tang. FastMoE: A fast mixture-of-expert training system. arXiv preprint ar arXi::2103.13262, 2021. URL http://axxvvvooopddf/////// /21./ 13262.
Cheng-Yu Hsieh, Chun-Liang Li, Chin-Kuan Yeh, Hootan Nakhost, Yasuhisa Fujii, Alexan- der Ratner, Ranjay Krishna, Chen-Yu Lee, and Tomas Pfister. Distilling step-by-stepe outperforming larger language models with less training data and smaller model sizes. arXiv preprint arXiv:2305.02301 2023. URL https://arxi:: otg/ /2005 02301 pdff.
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, et al. MiniCPM: Unveiling the potential of small language models with scalable training strategies. arXiv preprint ar Xiv:2400 203.:2204.003995, 06395, 2024.
Quzhe Huang, Zhenwei An, Nan Zhuang, Mingxu Tao, Chen Zhang, Yang Jin, Kun Xu, Liwei Chen, Songfang Huang, and Yansong Feng. Harder tasks need more experts: Dynamic routing in MoE models. arXiv preprint ar 2rXit:2403.07652, 2024. URL https: //arxiv.org// /2403. 24662.7652..
Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang, Ze Liu, Han Hu, Zi- long Wang Rafael Salas, Jithin Jose, Prabhat Ram, et al. Tutel: Adap- tive mixture-of-experts at scale. Proceedings of Machine Learning and Systems, 269-287, 2023. URL https://proceeet 561 555644488f7334225555222222222823 http:://reeccaaa222222---eeanaaa
www.mslog/ysp....oo..ooo.. g/paper...... Tiles//appr//222/fille
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint ar arXio:2401.04088 2024. URL https:
//arxiv.org/ /401 04088.
Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michal Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygózdz, Piotr Sankowski, et al. Scaling laws for fine-grained mixture of experts. arXiv preprint arXio:2402.07871, 2024. URL https: //txxv/aavvv///////////////aaavl wwwprwwwwwww...oooo........ooerr 2200224022220 07871.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from Transformers via speculative decoding. In International Conference on Machine Learning, pp. 19274-19286. PMLR, 2023. URL https://proceedings mlr, aeesss//2///eeiiaaaaa2222a/eeeee pdf.
Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Marc Marone, Christopher Akiki, Jia Li, Jenny Chim, et al. Star- Coder: may the source be with you! arXiv preprint arXiv:2305.06161, 2023. URL
http://aaxxi..org/df/////
06161.
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang, EAGLE: Speculative sampling requires rethinking feature uncertainty. In Forty-first International Conference on Machine Learning, 2024a. URL https://arxiv.. http:///axxxvoooppdd///// /2tp//ff//1111111 2201/ 15077.
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. EAGLE-2: Faster inference of language models with dynamic draft trees. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, Pp. 7421-7432, 2024b. URL https:
//aclanthology, www.aaattoooogogg/go/eee
mmg.:l/....rr/22222emmmammmmaaa ennlp-main. 422.pdf..
Zonglin Li, Chong You, Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank J Reddi, Ke Ye, Felix Chern, Felix Yu, Ruiqi Guo, et al. The lazy neuron phenomenon: On emergence of activation sparsity in Transformers. In The Eleventh International Conference on Learning Representations, 2022. URL htttps://operreviees www.nnt/tddffidddJnmmmmm YCk-.
Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, et al. DeepSeek-V2: A strong, economical, and efficient eixpurs language model. ar preprint ar Xivv:2405.04434 2024a. URL https://arxit www.:oogg//////ee/m /240. 04434.
12

---

Published as a conference paper at COLM 2025
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. DeepSeek- V3 technical report. arxio preprint ar Xio:2412.19437, 2024b. URL
http://aaaxiv..o/pfff.... /2412. 19437.
Liyuan Liu, Young Jin Kim, Shuohang Wang, Chen Liang, Yelong Shen, Hao Cheng, Xi- aodong Liu, Masahiro Tanaka, Xiaoxia Wu, Wenxiang Hu, et GRIN: Gradient-informeed MoE. Xre preprint ar Xr1::2209.12136 2024c. URL https: / wwt.:ooogg///////////// 12136.
Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shri- vastava, Ce Zhang, Yuandong Tian, Christopher Re, et al. Deja Vu: Contextual sparsity for efficient LLMs at inference time. In International Conference on Machine Learning, pp. 22137- 22176. PMLR, 2023. URL https: https://proceedings sma.eeddinnnmm........... press//202//iuueee
liu23am. pdf.
Yuqi Luo, Chenyang Song, Xu Han, Yingfa Chen, Chaojun Xiao, Zhiyuan Liu, and Maosong Sun. Sparsing Law: Towards large language models with greater activation sparsity. arXiv preprint ar Xir:2411.00233, 2024. URL http::////aaxxv..r..as/ee
Xinyin Ma, Gongfan Fang, and Xinchao Wang. LLM--runee: On the structural pruning of large language models. Xrv preprint arXir:2305.11127, 2023. URL https://arxiv.orr pdf/2305. 10117.ppp..
Iman Mirzadeh, Keivan Alizadeh, Sachin Mehta, Carlo Del Mundo, Oncel Tuzel, Golnoosh Samei, Mohammad Rastegari, and Mehrdad Farajtabar. ReLU strikes back: Exploiting activation sparsity in large language models. arviv preprint ar arXiv:2310.04564 2023. URL https://arxiv. http://awwx..ooo//dp.//// org/pdf /2310. 04564. pdf.
Mohammed Muqeeth, Haokun Liu, and Colin Raffel. Soft merging of experts with adaptive routing. arxiv preprint ar arXio:2306.03745, 2023. URL //atps://arxiv.... wtt.w//og//////////////// 22222336677445
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc-Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1525-1534, 2016. URL https: /ttp//aaaaaaana...a.aocoo wrg/p16 11411611114111 1.4....f..f
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter] J Liu. Exploring the limits of transfer learning with a unified text-to-text Transformer. Journal of machiie learning research, 21(140):1-67, 2020. URL https: //wp::jwww........../......r www. httg:/apprss/vlllloo
200 074/20-072220 pff.dd.ff.
Senthooran
Rajamanoharan, Arthur Conmy, Lewis Smith, Tom Lieberum, Vikrant Varma, János Kramár, Rohin Shah, and Neel Nanda. Improving dictionary learning with gated sparse autoencoders. ar preprint Xiw:2404.16014 2024. URL pttpp://arrxe org/ / 2404.16014.
Prajit Ramachandran, Barret Zoph, and Quoc V Le. Searching for activation functions. arviv preprint ar arXic:1710.05941 2017. URL https://arxivv org. 01945591.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan Le Bras, and Yejin Choi. SocialIQA: Commonsense reasoning about social interactions. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 4463-4473, 2019. URL https:// aclanthology. org/D19- 1454. .ff.
Wenqi Shao, Mengzhao Chen, Zhaoyang Zhang, Peng Xu, Lirui Zhao, Zhiqian Li, Kaipeng Zhang, Peng Gao, Yu Qiao, and Ping Luo. Omniquant: Omnidirectionally calibrated quantization for large language models. In The Twelfth International Conference on Learning Representations, 2023. URL /tttp:///arivv wrg/pdf
Noam Shazeer. GLU variants improve Transformer. Xrv preprint xxXiv:2002.05202,200 2020. URL https://arxiv.orr wwtp:/ivwwrggdd/////// /2002222.//000/.e//// 22wpp///.222222020... pdf.
13

---

Published as a conference paper at COLM 2025
Luca Soldaini, Rodney Kinney, Akshita Bhagia, Dustin Schwenk, David Atkinson, Russell Authur, Ben Bogin, Khyathi Chandu, Jennifer Dumas, Yanai Elazar, et al. Dolma: An open corpus of three trillion tokens for language model pretraining research. arviv preprint arXiv:2402.00159 2024. URL https: http//axivvvooro/ee org. httgg/df//2a /dd0/2202.00 .0109..0009
Chenyang Song, Xu Han, Zhengyan Zhang, Shengding Hu, Xiyu Shi, Kuai Li, Chen Chen, Zhiyuan Liu, Guangli Li, Tao Yang, and Maosong Sun. ProSparse: Introducing and enhancing intrinsic activation sparsity within large language models. In Proceedings of the 31st International Conference Computational Linguistics, pp. 2626-2644, January 2025. org hwwg/2225.cco coling-mmmiiiii 180. pdf.
URL
Yixin Song, Zeyu Mi, Haotong Xie, and Haibo Chen. PowerInfer: Fast large language model serving with consumer-grade GPU. arXiv preprint ar xrxo:23322.12556 2023. URL
https://arxiv wttgpdf///// 2312. 12456. 2455..ddf.
Yixin Song, Haotong Xie, Zhengyan Zhang, Bo Wen, Li Ma, Zeyu Mi, and Haibo Chen. Turbo Sparse: Achieving LLM SOTA performance with minimal activated parameters. arxiv preprint ar arXix:246..05955 2024. URL httpp:///arxis org/ /406 05955.
Mingjee Sun, Zhuang Liu, Anna Bair, and Zico Kolter. A simple and effective pruning approach for large language models. Xiv preprint ar Xit::2306.1166 2023. URL https: //arxiv.org www.voorg/pppoo pdf 2206. 1695...... pdf.
Vijay Thakkar, Pradeep Ramani, Cris Cecka, Aniket Shivam, Honghao Lu, Ethan Yan, Jack Kosaian, Mark Hoemmen, Haicheng Wu, Andrew Kerr, Matt Nicely, Duane Merrill, Dustyn Blasig, Fengqi Qiao, Piotr Majcher, Paul Springer, Markus Hohnerbach, Jin Wang, and Manish Gupta. CUTLASS, Jan 2023. URL https:///itthei wtt.tub.ccm//DDDIA/tla..aaaaaaa
Zhongwei Wan, Xin Wang, Che Liu, Samiul Alam, Yu Zheng, Jiachen Liu, Zhongnan Qu, Shen Yan, Yi Zhu, Quanlu Zhang, et al. Efficient large language models: A survey. Transactions on Machine Learning Research, 2023. URL https:///peerreiiiw.eeee bsCCJHb08A.
Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, and Damai Dai. Auxiliaryloss--ree load balancing strategy for mixture-of-experts. arXiv preprint arXio:2408.15664, 2024a. URL
15664.
Ziteng Wang, Jianfei Chen, and Jun Zhu. ReMoE: Fully differentiable mixtrrro-ptttttuuuueeeeextt with ReLU routing. arXiv preprint arXit:2412.147111, 2024b. URL htt.httss///rrww.og/////a 2412.14711. 14711.
Yuxiang Wei, Zhe Wang, Jiawei Liu, Yifeng Ding, and Lingming Zhang. Magicoder: Em- powering code generation with OSS-Instruct. In Forty-first International Conference on Machine Learning, 2024. URL https: /ttp::////xxvv.or/ ottpd////2///////////////// /2312 02120.
Heming Xia, Zhe Yang, Qingxiu Dong, Peiyi Wang, Yongqi Li, Tao Ge, Tianyu Liu, Wen- jie Li, and Zhifang Sui. Unlocking efficiency in large language model inference: A comprehensive survey of speculative decoding. In Findings of the Association for Computa- tional Linguistics ACL 2024, pp. 7655-7671, 2024. URL https: htttps:/accaanthol..g/ 022/2024. findings- acl, 456. pdf.
Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared LLaMA: Accelerating language model pre-training via structured pruning. Xrv preprint arXio:2310.06694 2023, URL /tttps://arxiv /tggpdd/23111 0w6.0.6669000 pdf.
Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao WU, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language mod- els. In International Conference Machine Learning, pp- 38087-380000 38099. PMLR, 2023. URL https://prooceedingg mlr.
Can Xu, Qingfeng Sun, Kai Zheng, Xiubo Geng, Pu Zhao, Jiazhan Feng, Chongyang Tao, and Daxin Jiang. WizardLM: Empowering large language models to follow complex instructions. arxie preprint arii:22304.122 2023. URL nttpps://arxer /ttp://rxxxvvooppp////////////// 1 1224.
14

---

Published as a conference paper at COLM 2025
Zhenliang Xue, Yixin Song, Zeyu Mi, Le Chen, Yubin Xia, and Haibo Chen. Powerlnfer-2: Fast large language model inference on a smartphone. ar preprint arXiv:2406.06285 2024. URL https: /arxiiv org/pdf htggpdff/4405 /206 06282.
Greg Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs V: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXio:2203.03466, httpp//ff//.....2.2 03466.
2022. URL https: www.aaiwww....ggrooooooeo 2203.
Zhewei Yao, Cheng Li, Xiaoxia Wu, Stephen Youn, and Yuxiong He. A comprehensive study on post-training quantization for large language models. ar preprint ar erXtv:2303.08322 2023. URL https://arxixs wwtp://oro///////// 2303. 08302. pdf.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can machine really finish your sentence? In Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics, pp- 4791-4800, 2019. URL https://aclanthoology
1rg/P19-1472.ppt
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019. URL https://pproeeedingg hetprdiinnnnnnniniissseeee newrips.cc// cc/pperr/ 772
pdf.
Zhengyan Zhang. Yixin Song, Guanghui Yu, Xu Han, Yankai Lin, Chaojun Xiao, Chenyang Song, Zhiyuan Liu, Zeyu Mi, and Maosong Sun. ReLU2 wins: Discovering efficient activation functions for sparse LLMs. ar ar preprint ar xrXi::2402.03804, 2024a. URL https: //arxiv.org/pddf 2402. 03804. pdf.
Zhengyan Zhang, Chaojun Xiao, Qiujieli Qin, Yankai Lin, Zhiyuan Zeng, Xu Han, Zhiyuan Liu, Ruobing Xie, Maosong Sun, and Jie Zhou. Exploring the benefit of activation sparsity in pre-training. In Forty-first International Conference on Machine Learning, 2024b. URL nttps://openrewee/ wwtpee//pf/dddddff//eefeeee id=KFXXPCcobh.
Weilin Zhao, Yuxiang Huang, Xu Han, Wang Xu, Chaojun Xiao, Xinrong Zhang, Yewei Fang, Kaihuo Zhang, Zhiyuan Liu, and Maosong Sun. Ouroboros: Generating longer drafts phrase by phrase for faster speculative decoding. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 13378 13393, 2024. URL https: //aclanthology . org/2022/2 emnlp-maii calp-mmain... pdf.
Weilin Zhao, Tengyu Pan, Xu Han, Yudi Zhang, Ao Sun, Yuxiang Huang, Kaihuo Zhang, Weilun Zhao, Yuxuan Li, Jianyong Wang, et al. FR-Spec: Accelerating large-vocabulary lan- guage models via frequency-ranked speculative sampling. arXiv preprint arXiv:2502. 14856, 2025. URL https: wwtpxiiivoog//piieee wtt/ /2002.202222 14856.
Zexuan Zhong, Mengzhou Xia, Danqi Chen, and Mike Lewis. Lory: Fully differen- tiable mnixture-of-expert for autoregressive language model pre-training. arXiv preprint 03133.
XxXin:2405.03133, 2024. URL https:
15

---

Published as a conference paper at COLM 2025
A Influence of Load Balancing
As a common practice of training MoE models, load balancing is to make the activation frequency of each expert balanced possible so that the model can be more friendly for the distributed deployment with expert parallelism, where experts are separately deployed on different devices and work concurrently. However, most end-side devices (which this work mainly focuses on) do not contain so many computation devices or cores, and thus cannot well support expert parallelism. Instead, is more important to reduce global computation costs and promote activation locality, which is critical for end-side deployment techniques such as offloading and speculative decoding.
Setting
TLS
CLSs
PPL.
AL+CS
80.54
71.38
14.88
AL+CS+LB
78.06
69.68
15.26
Table 7: Load balancing is less important for end-side deployment and can cause potential performance degradation. "LB" indicates the load balancing with auxiliary loss.
Therefore, in this work, we do not consider load balancing and only focus on sparsification and the activation locality issue. Besides, load balancing can potentially cause performance degradation (Wang et al., 2024a). Specifically, as shown in Table 7, under similar TLS, the PPL. suffers from a considerable increase after adding the load-balancing auxiliary loss.
B Adaptive Factor Scheduler for Chunk Sparsification Loss
With the language modeling loss Lim, the training objective is:
Lioaa Lim + Aat Lat + Acshess
(6)
where a and Ass are corresponding factors.
Considering the difficulty of tuning scheduler for Aes, which controls the overall sparsity level. Concretely, this scheduler keeps Acs constant the initial value A2g for the first Na: steps. Next, for every Nudj steps, the scheduler adjusts Acs according to the change of Les, increasing the factor when Les increases and vice versa. Formally, the behavior of this scheduler at step m = (i+ 1) is:
we introduce an adaptive factor
if Not else if Ycs $ 1 otherwise
AvglCts (1+1)Nan
Tes A/s
Yes Avgl ( (((i-)))
(NNa
max(7mmin Tc) A/2
where C/g denotes the loss value at step f, and Vmin is the minimum magnification ratio.
C Experimental Settings
First, we give detailed introduction to baseline MoE architectures used in our experiment:
(1) Vanilla TopK MoE is currently the most common MoE implementation, adopted by works such as Switch Transformer (Fedus et al., 2022) and Mixtral (Jiang et al., 2024). Their routers are composed of the softmax and TopK functions.
(2) Dee (DSMoE) (Dai et al., 2024) adopts similar TopK MoE architecture but introduces shared experts for improvement, which are consistently activated by each token.
(3) GRIN (Liu et al., 2024c), also using the TopK activation, adopts innovative routing strategy called SparseMixer-v2. This alleviates the non-differentiable issue through an approximation of the missing gradients.
(4) ReMoE (Wang et al., 2024b) makes the router differentiable by introducing a ReLU-based router module. Though similar to our design, ReMoE does not apply RMSNorm after ReLU,
16

---

Published as a conference paper at COLM 2025
Scale
dn
de
Ne
Nayyer
Nio
Aat
A
batch size
Mppe
0.1B (Small)
768
64
48
16
1.19
10a
2e
5e 2
1.57 x 106
10000
0.5B (Medium)
1280
128
40
27
4.95 x 10
2e
3
5e
2
3.15 x 10
20000
0.8B (Large)
1536
128
48
32
8.17
10a
le
3
5e
2
2.36 106
30000
1.2B (XLarge)
1792
128
56
35
1.19 x 109
2e -
3
5e
3e-2 2
3.15 x 106
40000
2.8B
2048
128
128
36
2.80
109
2e
1e
3.15 106
25000
Table 8: The major structural settings and hyper-parameters of our experimental models. Mayer Niot, and prr denote the number of layers, the number of non-embedding parameters, and the pre-training steps, respectively.
and more importantly, its L1 regularization directly imposed on activation values can cause activation magnitude disturbance (Section 4.1.4) and harm performance.
Next, we list the syper-parameters in Table 8. We adopt jul' parametrization (Yang et al., 2022) to promote training stability and reduce the influence of hyper-parameters. Therefore, we can adopt the same setting for the following parameters: peak learning rate Ir = 0.01, ß1 = 0.9, $2 0.95, weight decay 0.1. We use the WSD scheduler to adjust the learning rates in the training process (Hu et al., 2024; Dubey et al., 2024). As for the adaptive factor scheduler, under all BlockFFN settings, we adjust the factor every Nadj = 100 steps, with
Nst 1000 and 7min 1.025.
D Model Settings
For all settings, which include BlockFFN, the upper bound "Dense", and the other MoE baselines, we maintain the close number of total parameters, activated parameters (i.e., TLS), and training tokens. Moreover, the number and the intermediate dimension of experts are also exactly the same, following the fine-grained expert segmentation of DeepSeekMoE (Liu et al., 2024b). As for the attention layer, apply the multi-latent attention (MLA) (Liu et al., 2024a) for models from O.1B to 1.2B, while adopting group query attention (GQA) for BlockFFN-228BB to make the acceleration implementation easier. Therefore, we ensure that the differences between different settings only lie in the routing strategy and training objectives, which the key improved points of our work. The detailed structural settings of our models are listed in Table 8.
E Datasets and Benchmarks
Training data The pre-training data of BlockFFN is comprehensive mixture of multiple corpora across various categories. This includes C4 (Raffel et al., 2020), Pile (Gao et al., 2020), Dolma (Soldaini al., 2024), CommonCrawl, StarCoder (Li et al., 2023), and other collected raw corpus. Besides, to obtain reasonable evaluation results, we perform decay stage before evaluating models on benchmarks (Hu et al., 2024; Dubey et al., 2024). For this stage, instruction-tunni tunnin data are added, including Evollnstruct (Xu et al., 2023), UltraChat (Ding et al., 2023), OssInstruct (Wei al., 2024), and other collected SFT datasets.
Validation data The validation data has the same distribution the pre-training data. Deduplication is conducted to alleviate the intersections between pre-training and validation data, so that the validation data cannot be easily over-fitted.
Evaluation benchmarks The task-specific benchmarks used in our experiments can be divided into two groups: commonsense reasoning (C.R.) and reading comprehension (R.C.). The former group includes PIQA (Bisk et al., 2020), SIQA (Sap et al., 2019), and HellaSwag (Zellers et al., 2019). The latter group includes LAMBADA (Paperno et al., 2016), Dy QA (Clark et al., 2020), and BoolQ (Clark et al., 2019). For both groups, the evaluation metric is Shoho accuracy.
17

---

Published as a conference paper at COLM 2025
Small
Medium
Lange
XLarge
Setting
PIQA
SIQA
Hella.
PIQA
SIQA
Hella.
PIQA
SIQA
Hella.
PIQA
SIQA
Hella.
Dense
65.51
38.23
31.71
69.80
40.53
47.71
70.62
42.27
51.77
71.98
42.32
57.61
TopK
63.98
36.28
30.93
68.44
40.43
43.82
69.53
41.45
47.55
71.60
41.86
51.49
DSMoE
65.13
36.95
30.96
68.28
40.02
43.09
69.31
40.99
47.33
70.84
41.61
53.40
GRIN
64.20
36.08
30.78
68.28
40.02
44.59
69.31
41.35
48.13
69.86
41.76
52.34
ReMoE
64.91
38.18
32.56
69.31
40.38
44.57
71.38
40.38
49.55
71.06
42.22
53.40
BlockFFN
64.96 37.62
31.82
69.80
39.56
45.89
71.49
41.40
50.43
71.16
42.84
55.26
Table 9: The evaluation scores on the three benchmarks of commonsense reasoning (C.R.).
Small
Medium
Large
XLarge
Setting
LAM.
TyDi.
BoolQ
LAM.
TyDi.
BoolQ
LAM.
TyDi.
BoolQ
LAM.
TyDi.
BoolQ
Dense
30.41
13.41
54.34
43.80
36.59
54.16
49.04
39.77
57.80
52.22
40.23
56.15
TopK
29.03
5.91
52.02
42.40
34.55
50.46
44.85
41.59
42.72
49.16
40.23
58.47
DSMoE
28.37
17.05
50.28
43.74
33.41
57.52
48.52
44.77
46.06
50.30
47.95
53.46
GRIN
28.41
12.27
54.16
42.77
35.00
49.91
44.87
38.18
52.08
50.44
44.55
54.04
ReMoE
29.44
15.68
53.21
42.58
32.05
54.53
45.90
39.55
56.54
51.35
30.23
59.79
BlockFFN
30.53
21.14
50.12
42.75
37.50
50.98
45.97
46.59
62.23
50.30
43.41
58.47
Table 10: The evaluation scores on the three benchmarks of reading comprehension (R.C.).
Performance on independent benchmarks In Table 3, we only list the average evaluation scores of two benchmark groups. In this section, we provide the evaluation results on independent benchmarks, as shown in Table and 10.
F Effect of Expert Granularity
Sparsity Value (%6
2.35
Validation
2.34
2.38
TLS
2.32
CLSa
10
160
320
10
160
320
Number of Experts (Ne)
Number of Experts (Ne)
Figure 5: The validation loss of BlockFFN Medium with different expert granularities.
Figure 6: The TLS and CLS of BlockFFN Medium with different expert granularities.
Expert granularity has long been demonstrated to influence the performance of MoE mod- els (Krajewski et al., 2024). Specifically, given fixed computation budget (assumed proportional to the parameter scale), what the best trade-off between the expert number N₂ and expert dimension d,? To solve this problem, we conduct experiments on BlockFFN Medium with different expert granularities. These models are evaluated from four aspects: the validation loss, the token-level sparsity, the chunk-level sparsity, and memory locality,
First, as shown in Figure 5, while the loss drops considerably with coarse granularities (i.e., small Ne), the marginal benefits of granularity increase gradually diminishes with > 40 experts. Moreover, Figure further displays a nonmonotonic relationship between sparsity and expert granularities. The setting with 40 experts, which we adopt in our main experiments (Section 4.1.1), achieves both the highest TLS and CLSg. Finally, as larger memory access units generally have better memory locality and hardwane-friendllnnss
18

---

Published as a conference paper at COLM 2025
do not expect extremely fine granularity. To sum up, 40 experts is the best setting for BlockFFN Medium. We leave more quantitative analyses for future studies.
G Upper Bound Analysis of Acceleration Kernels
Setting
FFN Time (ms)
TLS
Setting
FEN Time (ms)
CLS
Baseline AR Ours (1-Tok)
61 7.8 (12.8%x)
EAGLE-2
88 27 (30.7%xx
1 12.45%
30.86%
Ours (32-Tok)
Table 11: The upper bound analysis of our kernels. TLS and CL.S values evaluated on Spec-Bench decoding tokens, which are close to the FFN time consumption ratios of "Ours (1-Tok) Baseline AR" and "Ours (32-Tok) EAGLE-2", respectively.
To delve deep into the ability of our acceleration kernels, we conduct an upper bound analysis by inspecting the time consumption of FFNs separately. As shown Table 11, "Baseline AR" and "EAGLE-2" can be viewed as the sparsity-ablation setting of "Ours (1-Tok)" and "Ours (32-Tok)", respectively, with only about 12.8% and 30.7% FFN time consumption. Surprisingly, we find that these two time consumption ratios are quite approximate the TLS and CLSspec, respectively. Note that "Ours (32-Tok)" adopts a draft tree size of 32, and CLS e is calculated by the average ratio of experts activated by the union of all the tokens contained the EAGLE-2 draft tree. This phenomenon indicates that both kernels can reach the theoretical speedup upper bound in FFN acceleration induced by the corresponding token-level sparsity and the union sparsity of tokens in a draft tree.
Notably, although our CLS-aware training objectives do not directly optimize the tree-level union spansity, these objectives tailored for consecutive chunks are effective for tree patterns, since each path from the root node to the leaf node is still composed of consecutive tokens.
H Inference Acceleration on Independent Datasets
MT.
Trans.
Summ. Tokens/s
QA
Setting
Tokens/s
Speedup Tokens/s Speedup
Speedup
Tokens/
Speedup
Huggingface
7.61
0.60
7.32
0.55
7.54
0.61
7.81
0.60
Baseline AR EAGLE-2
12.61
1.00
13.32
1.00
12.40
1.00
13.04
1.00
23.70
1.88
20.26
1.52
20.74
1.67
21.99
1.69
Ours (1-Tok)
38.41
3.05
40.39
3.03
37.37
3.01
42.05
3.22
Ours (32-Tok)
49.43
3.92
39.18
2.94
42.70
3.44
46.09
3.53
Table 12: Detailed speedup results on NVIDIA Jetson Orin NX (1st part).
Math
RAG
Average Speedup
Setting
Tokens/s
Speedup
Tokens/s
Speedup
Tokens/s
Huggingface
7.76
0.61
5.96
0.46
7.33
0.57
Baseline AR EAGLE-2
12.80
1.00
12.86
1.00
12.84
1.00
25.01
1.95
22.47
1.75
22.36
1.74
Ours (1-Tok)
44.54
3.48
39.38
3.06
40.36
3.14
Ours (32-Tok)
59.85
4.68
45.76
3.56
47.17
3.67
Table 13: Detailed speedup results on NVIDIA Jetson Orin NX (2nd part).
In Table 6, we provide the decoding speeds each dataset contained in Spec-Bench and the average speedup ratio. In this section, we list the speedup ratios on each independent dataset, shown in Table 12 and 13.
19

---

Published as a conference paper at COLM 2025
MT.
Trans.
Summ.
QA
Math
RAG
Average
2.73
2.17
2.38
2.67
2.83
2.57
2.66
Table 14: The acceptance lengths on each independent dataset of Spec-Bench.
On most datasets, "Ours (32-tok)" achieves the best inference efficiency. However, there exists an exception, "Translation" (Trans.), where "Ours (32-tok)" underperforms "Ours (1- tok)". This indicates that the combination of chunk-level activation sparsity and speculative decoding has worse performance than utilizing token-level activation sparsity alone. After careful examination, we find this is attributed the shortest EAGLE-2 acceptance length on this dataset (Table 14), which hurts the efficiency of speculative decoding. Therefore, the spnrsity speculative decoding is more reasonable when speculative decoding works efficiently, generally with longer acceptance lengths and larger models.
I Ablation Studies on the Gated Expert Variant
For the design of BlockFFN expert modules, we choose the non-gated MLP instead of the more widely adopted gated variant. To support this choice, we conduct an ablation study on BlockFFN (Small). As shown in Table 15, using gated MLP for expert modules can cause extremely low sparsity, which is quite surprising result worth further study. A possible explanation may lie in the "competition" of sparsity between the router module and the expert module, indicating that the router sparsity and the expert sparsity vary in the opposite direction. Therefore, the higher sparsity of gated MLPs can significantly weaken the sparsity of the router module.
Expert Design
TLS f
CLSg
PPL/
Non-gated Gated
80.54
71.38 25.44
14.88
28.39
15.17
Table 15: The ablation results of the gated variant for BlockFFN expert modules.
J More Details about Acceleration Kernels
Figure 7 shows more details about our efficient acceleration kernels. These include the execution details of the two-loop structure of the kernels, and how the weights of activated experts are transferred between different memory hardware (e.g,, SRAM and HBM).
20

---

Published as a conference paper at COLM 2025
Ne: expert number
de: expert intermediate size
d hidden size
n number of tokens in chunk
Real Activation
Input Hidden
A (x) E R
Ruxda
Inner Loop
N₂
Copy to SRAM M
Union = 236 6
Bx
W
Inner Loop
B
Outer Loop (E Union)
(2)
de
Compute on SRA M
W (3)
O utput to HBM
W
up
n
(6)
mid E (xx(Nd)e)
dn
W up IR xdd
Inner Loop (E Union)
Mask
Inrelevant A ctivation
Copy to SRAM M
W down E (Nede)
W W 2) W() W(4) (5)
(6)
de
down down down down down down
dg
Outer Liter-Loop Loop
Compute on SRA M
O utput to HBM
FFN(x) E RRdx
Inner Loop (E Union)
Figure 7: The detailed framework of our efficient acceleration kernels.
21


Reviews:
[{"review_id": "aJzqJhfmdb", "reviewer": "Reviewer_SU7d", "summary": "This paper proposes a chunk-sparse MoE design that is better compatible with acceleration techniques. The results show improvement in block sparsity in routing decisions with equal or better PPL than other MoE methods. With the accompanying CUDA kernel and speculative decoding, the decoding speed is significantly higher. The paper is well-written, providing motivations and empirical evidence to substantiate the design decisions.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. Realistic problem and solution: The paper tackles a previously overlooked realistic challenge, applying inference speed-up techniques for MoE models. Also, the authors didn't stop at \"our method has higher block sparsity, making it possible for decoding speedup\", but went all the way to develop the kernels to prove the speed gain.\n2. Well-rounded experiments: This paper's experiments are extensive, providing sufficient evidence on both the model performance and speed-up. In addition, the method is supported by plenty of ablation and distribution analyses to help readers understand the proposed method and the effect of various components.", "reasons_to_reject": "1. Remaining question on task performance: The performance evaluation focuses on language modeling, reading comprehension, and multiple-choice commonsense; it is still unclear whether BlockFFN can perform well on tasks that require long generation, such as the translation and summarization tasks from this paper's speed evaluation.\n2. Complex method: The proposed method includes architecture design, loss regularizations, and particular scheduling for the introduced loss term that needs to be adopted together, rather than a single pluggable element. This increases the barrier to adoption and future extensions. It also makes it hard to find ways to convert a conventional pretrained model for the proposed method.", "questions_to_authors": "1. I still find it hard to understand why CS loss hurts PPL while AL+CS helps. I think AL is mostly implied from CS. Improving the block sparsity naturally discourages abrupt changes in the routing distribution, except for the block boundaries. Could you provide more thoughts on why both are necessary? And why does CS alone cause so much performance degradation?", "extra_scores": {}}, {"review_id": "FCJLbaGQZZ", "reviewer": "Reviewer_G7oy", "summary": "This paper proposes BlockFFN, a novel Mixture-of-Experts (MoE) architecture designed for end-side device acceleration. The authors aim to improve both token-level sparsity (TLS) and chunk-level sparsity (CLS), and present an architecture that combines ReLU-based differentiable routing with RMSNorm to achieve this. The paper introduces CLS-aware training objectives and deploys a set of custom CUDA kernels to demonstrate the practical acceleration benefits on consumer-grade hardware (Jetson Orin NX).\n\nThe problem is timely and important: bringing MoEs to edge and desktop environments requires careful engineering and novel solutions. The authors show a comprehensive set of experiments that evaluate both hardware-level performance (decoding speeds) and language modeling metrics (perplexity, evaluation scores).", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- **Timely and relevant**: Efficient MoEs for edge hardware are an important direction, and this work directly targets that gap.\n\n- **Strong experimental results**: The model is thoroughly benchmarked both in terms of quality (perplexity, reasoning, QA tasks) and real-world decoding speed.\n\n- **Hardware applicability**: Unlike many MoE papers, the authors test their system on actual consumer-grade hardware (Jetson NX), which gives weight to their claims.\n\n- **Novel technical contributions**: The use of ReLU + RMSNorm for routing and the introduction of chunk-level sparsity training losses are fresh and well-motivated.\n\n- **Batch size = 1 performance**: Table 5 shows substantial speedups even with single-token speculative decoding or no batching—this is especially important for local deployment, where batching isn’t practical.\n\n- **Content-aware routing**: The observed bimodal expert activation distribution suggests that the routing policy is content-aware and semantically meaningful.", "reasons_to_reject": "- **Speculative decoding is context-limited**: The paper includes speculative decoding as part of the proposed acceleration strategy, but this technique is primarily effective in server-side deployments where batch sizes >1 and longer contexts are common. On edge devices like Jetson Orin NX, where inference typically runs with batch size = 1, the gains from speculative decoding are less relevant.\n\n- **Memory concerns with dynamic routing**: On small devices, loading different experts per token (especially when using ReLU) might force all experts to be kept in memory, which limits scalability. It’s unclear if expert selection is stable enough across tokens to allow chunk-wise caching or offloading.\n\n- **Missing sparse TopK MoE baseline**: Table 5 shows that BlockFFN achieves significant speedups even for single-token inputs, which strongly suggests that activation sparsity itself, not speculative decoding, is the main driver of performance on the tested hardware. However, there is no comparison with a sparse TopK-based MoE that activates fewer experts per token (e.g., via adaptive thresholds). Without this, it’s unclear whether BlockFFN’s dynamic ReLU-based routing is essential.\n\n- **Undiscussed training cost**: The paper doesn’t discuss the training cost of BlockFFN relative to baselines. The use of dynamic routing and auxiliary sparsity losses likely adds overhead that may not be trivial when scaling further.\n\n- **Pretraining dataset scale unclear**: If BlockFFN is trained on a significantly larger or higher-quality dataset compared to the baselines, this could confound comparisons with other MoE variants. This needs to be controlled for.\n\n- **ReMoE comparison could be better**: The authors state that RMSNorm is the main difference between BlockFFN and ReMoE, but no explicit ablation confirms this. A controlled test isolating the effect of RMSNorm would be helpful. This could be done by adding RMSNorm to ReMoE itself and doing the evaluation.", "questions_to_authors": "- **Sparse TopK comparison**: In Table 5, your model shows strong speedups with single-token input, suggesting that activation sparsity is the primary source of improvement. Could a TopK MoE model with fewer-than-K activated experts (e.g., using thresholding) achieve similar results?\n\n- **Memory implications of dynamic routing**: Have you analyzed whether expert selection is stable across token chunks? Could experts be preloaded per chunk and reused to reduce memory overhead?\n\n- **RMSNorm vs ReMoE**: Is RMSNorm the only architectural difference between BlockFFN and ReMoE? If so, could you provide an ablation to confirm that it’s the primary contributor to performance improvements?\n\n- **Server-side evaluation**: Since speculative decoding is more useful in server scenarios, have you evaluated BlockFFN on server-class GPUs (e.g., A100 or H100) with larger batch sizes? It would be helpful to understand how well your method scales in that context.", "extra_scores": {}}, {"review_id": "5Vs3Yn0lXr", "reviewer": "Reviewer_ohTG", "summary": "This work proposes a variant of mixture of experts for feed-forward layers with simpler sparsity in routing and chunk-aware loss. The router is inspired by ReMoE with additional root mean squared normalization to scale the activation. The activation locality loss is designed to differentiate with the activation pattern in neighboring position and chunk sparsification loss is designed to take into account the activations in prior history so that the expert routing is not skewed, but balanced for efficient computation. Experiments measured by perplexity and end-task performance, e.g., commonsense reasoning, shows clear gains when compared with other variants measured on models differentiating the number of parameters with faster inference on average.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The proposed mixture of experts, BlockFFN, is inspired by ReMoE in that the routing is normalized by root mean square to preserve invariants considering other activation norms. The feed-forward layer is further simplified by using the architecture from the original Transformer, not gated variants with scaling of up projection. \n- The two loss functions for routing mechanism is inspired to trade the locality in neighboring position so that the activation pattern is clearly differentiated, while the chunk sparsitication loss considers locality when considering a history of activations. These losses are well motivated and ablation studies show gains.\n- It is well engineered and descriptions are good enough to reproduce the results.\n- Experiments are well designed, demonstrating the effectiveness of the proposed method.", "reasons_to_reject": "- The motivation of BlockFFN is not clear and the ablation studies are not comprehensive. If my understanding is correct, each expert is basically identical to the feed-forward layer in the original Transformer, but others are using the gated variant. I'd like to see why the proposed expert is employing the simpler variant and what is the impact of using the simpler one for other routing methods, e.g., GRIN.\n- The proposed loss functions are the major contribution of this work and they sound orthogonal to the routing method. I'd like to understand what is the impact when applying those losses to other routing methods.\n- This work should run experiments to measure the impact of the token length in chunk sparsification loss (CLS), in order to quantify the gains of the approach. I suspect longer token length of $L$ will hurt the performance in, e.g., perplexity, but with the more efficiency in computation.", "questions_to_authors": "Please check the reason for rejection.", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a simple but effective approach for making MoEs more efficient by changing the activation function (ReLU followed by RMSNorm) as well as adding in additional objectives to encourage chunk-level sparsity (which is helpful for speculative decoding). The reviewers generally found the approach to be well-motivated, and the empirical results to be strong. Showing practical speedups on edge devices (section 4.3) was especially appreciated. On the negative side, there were some slight concerns with regard to baselines (e.g., comparison against a TopK baseline with a dynamic threshold), as well as whether this approach would work long-context tasks. However, most of these were addressed during the rebuttal, and I think this paper is a clear accept.
Decision: Accept