Learn To be Efficient: Build Structured Sparsity in Large Language Models
Authors: Haizhong Zheng
Xiaoyan Bai
Xueshen Liu
Zhuoqing Mao
Beidi Chen
Fan Lai
Atul Prakash
Venue: NeurIPS 2024 spotlight (2024)
Keywords: LLM inference efficiency
Moefication
Contexual Sparsity
Primary area: natural_language_processing
TL;DR: We propose a novel training algorithm to train efficiency-aware LLMs that have more structured contextual sparsity for fast inference.

Abstract: Large Language Models (LLMs) have achieved remarkable success with their billion-level parameters, yet they incur high inference overheads. The emergence of activation sparsity in LLMs provides a natural approach to reduce this cost by involving only parts of the parameters for inference. However, existing methods only focus on utilizing this naturally formed activation sparsity in a post-training setting, overlooking the potential for further amplifying this inherent sparsity. In this paper, we hypothesize that LLMs can learn to be efficient by achieving more structured activation sparsity. To achieve this, we introduce a novel training algorithm, Learn-To-be-Efficient (LTE), designed to train efficiency-aware LLMs to learn to activate fewer neurons and achieve a better trade-off between sparsity and performance. Furthermore, unlike SOTA MoEfication methods, which mainly focus on ReLU-based models, LTE can also be applied to LLMs like LLaMA using non-ReLU activations. Extensive evaluation on language understanding, language generation, and instruction tuning tasks show that LTE consistently outperforms SOTA baselines. Along with our hardware-aware custom kernel implementation, LTE reduces LLaMA2-7B inference latency by 25% at 50% sparsity.

Learn To be Efficient: Build Structured Sparsity in Large Language Models
Haizhong Zheng
Xiaoyan Bait
Xueshen Liu"
Z. Morley Mao
Beidi Chen
Fan Lai
Atul Prakash
University Michigan *Cannegie Mellon University
University of Illinois Urbana- Champaign
(hzzheng, , smallyan, liuxs, zmao, aprakash) Sumi ch. edu beidi cmu. edu, fanlai@illimois.
Abstract
Large Language Models (LLMss) have achieved remarkable success with their billion-level parameters, yet they incur high inference overheads. The emergence activation sparsity in LLMs provides a natural approach reduce this cost by involving only parts of the parameters for inference. However, existing methods only focus on utilizing this naturally formed activation sparsity in post-training setting, overlooking the potential for further amplifying this inherent sparsity. In this paper, we hypothesize that LLMs can learn to be efficient by achieving more structured activation sparsity. To achieve this, we introduce novel training algorithm, Learn-To-be Co-bbe-Effiiiett (LTE), designed to train efficiency- to learn to activate fewer neurons and achieve a better trade-off between sparsity and performance. Furthermore, unlike SOTA MoEfication methods, which mainly focus on Rel.U-based models, LTE can also be applied to LLMs like LLaMA using non-ReLU activations. Extensive evaluation on language understanding, language generation, and instruction tuning tasks show that LTE consistently outperforms SOTA baselines. Along with our hardware-aware custom kernel implementation. LTE reduces LLaMA2-78 AB inference latency by 25% 50% sparsity. We make our code publicly available at GitHut
LLMs
Introduction
The exponential growth data volumes and model sizes has catalyzed significant breakthroughs in large- scale models, enabling a wide range of applica- tions 22] 47 53
activations
Traisins
30]. Among them, large lan- guage models (LLMs), like GPT-3 (2). OPT [47] and LLaMA 38 39), have demonstrated impressive natural language ability. However, the skyrocketing number of model parameters [2]] [36] and dataset has introduced challenges fur- ther scaling those models. The exponential growth in model size has not only inflated the deployment costs of LLMs, due their significant computational and
48
FFN
FFN with structiond sparsity
size [5T 45 52 44]
Figure 1: Learn-To-be- Effic-ffient (LTE) per- forms efficiency-aware training to construct structured model contextual sparsity for fast inference. Activated neurons are in orange.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

memory demands during inference, but also affected the user experience in many latency- satency-sensitive
applications, such as chatbots [T] and autonomous driving [7 35]. Recent advances have been
leveraging sparsity to improve L.M inference efficiency, including model weight quantization [
and pruning [22] token sparsity (00 46 and activation sparsity 21 49).
During LLM inference, Feed Forward Network (FFN) layers are the primary efficiency bottleneck, accounting for over 60% of the FLOPs and I/O operations [2I]. but they exhibit high activation For example, in a 175B OPT model, more than Recent advances leverage this activation sparsity via MoEfication [49 25]: they convert FFN layers to MoE layers by grouping neurons in the FFN intermediate layer into n experts and then applying a router select k most important experts for each input token. Unlike training the MoE model from scratch, MoEfication converts FFN layers in pretrained L.MS into MoF layers with relatively small training resources.
sparsity, especially in ReLU-based LLMS 18
95% of activations in FFN layers are zeros 21
However, existing advances only focus on manipulating the pre-trained activation sparsity of neurons, overlooking the potential for further increasing this inherent sparsity. Moreover, they are limited to the historically ReLU-based LLMs, whereas emerging advanced models use soft activation functions for better model quality (e.g.. SwiGLU in LLaMA BT and GeGLU in Gemma (77,, exhibiting much lower natural activation sparsity. While existing works suggest replacing the model activation with ReLU to enable sparsity 49 25,,
we show that this can hurt model performance (Section (5.))
In this paper, we hypothesize that LLMs can learn to be efficient and achieve more structured activation sparsity. As shown in Figure our key insight is that, without compromising model quality, we can also train I.Ms to be efficincyy-wwae by developing more structured sparsity, which is more friendly for achieving hardware speedup. However, creating structured sparsity in pretrained LLMs is non-trivial due to two practical challenges:
1. How to train routers more stably? accuracy drop (Sectior How to jointly train the model and routers in a MoEfication setting is still open-ended.
? The widely-used Top-k Softmax routing can lead to a severe
2. How select the right experts in serving? The number of experts needed in MoE layers depends on specific inputs and layers, implying a trade-off between model efficiency and quality.
To address these challenges, we introduce Learn-Tow
(LTE), novel training algorithm to train efficiency-awand models. LTE integrates an efficiency loss penalty, encouraging models to activate fewer neurons in their FFN layers while keeping good task performance. Additionally, LTE adopts a threshold-based Sigmoid routing strategy to select experts and employs a two-stage training mechanism to improve training stability. LTE achieves a more flexible selection of experts instead of selecting fixed number of experts for all layers and inputs. Same as Deja Vu (2 and moefication [49. LTE provides very structured sparsity. We further develop a custom CUDA kernel with Triton, a Python-like CUDA programming language, to enable wall-clock time speedup from such structured sparsity (Section 4.3
We evaluate LTE on both encoder-based models (RoBERTa RoBERTa (0 and decoder-based models (GPT2-Medium (88), LLaMA2-7B (38) with the HuggingFace's 's transformers. Our extensive experiments on Natural Language Understanding (NLU) tasks, downstream Natural Language Generation (NLG) tasks, and instruction tuning tasks, show that LTE consistently outperforms state- of-the-art designs. For instance, LLaMA with LTE provides a 1.83x- 2.59x FL.PPs speed-up on NI.G tasks without compromising model quality. After integrating our hardware-efficien implementation of sparse matrix multiplication kernel, LTE reduces LLaMA2-7B 25% wall-clock time latency at around 50% sparsity. Our evaluation results demonstrate that LTE effectively builds more structured sparsity even on LLMs with soft activation functions.
2 Related Work
Mixture of Experts (MoE). MoE was proposed by [4] few decades ago to build an integral system with subset networks. Recently. MoE layers have been used in the Transformer architecture as a substitute for the MLP block 32 with the recent Mistral-7B model [55 as a prominent example. In MIP blocks of MoE transformers, instead of using an FFN layer for calculating output, MoE layers construct multiple smaller FFN layers and employ a router to choose a subset of experts to do conditional computation. Even though transformers can have billions of parameters, only subset of

---

experts are activated for each token, thus effectively increasing model activation sparsity (i.e., the number of skipped neurons in execution).
LLM Contextual Sparsity. Recent studies
18 21 show that trained ReI.U-hased transformers
naturally show a great sparsity in the activation of FFN layers. For example, in 175B OPT model, more than 95% activations in FFN layers are zeros [21]. While SOTA models are increasingly using diverse soft activations, like GeLU and SwiGLU, existing work shows that replacing soft activations with ReLU and fine-tuning the model can increase the activation sparsity without greatly hurting model quality 49 25 18 Yet, we find that this benefit not consistent for all datasets (Section (.2) This emerging sparsity indicates that a large portion of neurons are unnecessary in LLM inference (i.c.. sparsity), which we can leverage to improve LLM inference efficiency like using MoEfication [49].
Recent papers [21 49 show that single FFN layer in pretrained LLMs can be converted to an MoE
layer, by splitting the matrices of FFN layers into exclusive sub-matrices and employing a router to activate only neurons in subset of experts. Similar to MoE. models, MoFfication can effectively
accelerate LLM inference by only using part of the parameters 49 21
Weight Sparsity. Another orthogonal direction regarding model sparsity is static model weight sparsity 8 (3), which often prunes the model weights and keeps them static during inference. For instance, Wanda [36] estimates model weight importance and prunes unimportant weights, so all inputs will use the same subset of weights. In contrast, contextual sparse models select a different subset of weights for different inputs. Nonetheless, model pruning can be applied to contextual sparse models too, meaning a complementary optimization contextual sparsity.
Model
Background and Motivation
In this section, we first briefly introduce the MoEfication of FFN layers in transformers. Then, we present a study on the limitation of applying noisy top-K Softmax routing the Moefication setting.
3.1 Background: MoEfication
MoEfication 49 is a way to group neurons in FFN layers, thereby converting FFN layers to MoE layers for better execution speedup. For transformer whose hidden states and FFN intermediate dimension are dmodel and dRFN, respectively, the FFN layers process the input as follows:
h = xW1 + b2
FFN(x) = (1))
where E Rdmmodei W1, W2 are weight metrics, b2, by are biased term and o is activation function. MoEfication aims to group deF N neurons in the FFN intermediate layer into n experts. Then, a router is trained for each FFN layer to select & experts (k < <),, thus reducing activation load, to speed up inference. A recent work, Deja Vu [1], uses a similar setting but treats each neuron as an expert.
However, SOTA MoEfication methods overlook the potential for further optimizing the activation sparsity of L.S Besides, for LLMS employing soft activations such as GeLUU [22 and SwiGLU (I), MoEfication [49] proposes to replace soft activations with ReLU and fine-tune the model to improve activation sparsity. Although recent works [49 25) show that this replacement has a marginal impact on model performance, we notice that this hypothesis is not widely applicable (Section 5.2).
3.2 Limitations of Noisy Top-K Softmax Routing
One potential solution to train and convert the model into MoF. layers is noisy top-K Softmax routing 33). It selects k experts based on the
MRPC (RoBERTa base)
QNLI (RoBERTa base)
Accuracy.
highest router outputs and then calculates Soft-
max values for these selected outputs. The out- put of the MoE layer is determined by summing the products of these softmax values with their corresponding experts' outputs. Additionally, a small Gaussian noise added to the router's outputs during training to encourage exploration across different experts, which addresses the is- sue of unselected experts being non-differentiable.
Ful Model
Put Model
Softmax with neiry tos
Sortmax with noiy top-K
de
Sparsity
Sparsity
Figure 2: Models trained with noisy top-K Soft- max routing experience severe accuracy drops. even at very low levels of sparsity.
3

---

To study the effectiveness of noisy top-K Softmax routing, we follow the standard MoEfication setting [49 to form experts in FFN layers: we fine-tune a model on a downstream dataset and group neurons into experts via parameter clustering. Subsequently, we integrate Softmax routers into models and jointly train both the model and routers. As shown in Figure 2 we observe noticeable accuracy drops even at very low sparsity levels. When diving into the expert scores produced by routers, we observe a very biased expert score allocation: only one expert in an MoE layer has an expert score close to I, while other experts receive nearly 0 scores. This results in only one expert contributing to the inference in each MoE layer, leading to performance drops.
A potential reason behind this biased allocation is that the sum of expert scores in Softmax routing is constrained to Unlike traditional MoE models, which typically select 1-2 experts [32 7. MoEfied models need to select a much larger number of experts 49), but the sum of experts score is still 1; this can complicate the allocation of the Softmax budget and cause a biased allocation.
4 Methodology: Learn To be Efficient
In this section, we present LTE (Learn To be Efficient) to learn more structured MoE in FFN layers for fast inference. Specifically, we tackle two key challenges toward practical MoEfication for diverse LLMs: (1) How to train routers more stably? (Section 4.2); and (2) How to select the right experts in serving? (Section 0.cction4.2.22 4.2.2
4.1 Experts Grouping
The first step of LTE is to group neurons in FFN layers to form experts. Expert grouping aims to group neurons that are often activated together, thereby improving efficiency by reducing the total number of activated neurons. Here, constructing too many experts (e.g., each neuron as an expert) can significantly increase the cost, while too few can lead to poor model quality. Moreover, we need to group neurons without significant performance drop the pre-trained model.
Following the previous work 49 clustering to group MLP neurons into different experts. Fundamentally, in FFN layers, cach neuron associated with a column of W1 (which is a dmodel vector). Parameter clustering first treats the corresponding vector in W1 as the feature for neuron, then applies the balanced K-Means algorithm (33 to cluster neurons into n clusters. Each cluster has 32 neurons by default and will be treated as an individual MoE expert We also conduct an ablation study on other alternative grouping
we group 32 neurons an expert in FFN layers, and use parameter
strategies Section 5.5
4.2
Adaptive Expert Routing
4.2.1 Router Design
After constructing the experts, we need to decide the right routing strategy. Similar to existing MoE designs, we use a fully-connected layer as the router network. The input to the router is the output of the preceding self-attention block, and the router output is the expert score for each expert.
Routing funccion. To address the biased expert score issue (Section 3.2), we employ the Sigmoid function router outputs get expert scores. Unlike the Softmax function, the Sigmoid function computes each expert score independently, which circumvents the biased expert score due to the constraint the sum of outputs in the Softmax function:
G(x); Sigmoid(x: Wg.s) 1 + W
where r is the input to the corresponding FFN layer, is the index for the expert, and Wg. is the router network weight for the i-th expert.
Threshold- based experts selection. SOTA methods 49
select fixed number of experts for all layers and inputs. However, the necessary number of experts may differ depending on inputs and layers. Here, propose select experts more adaptively for different inputs and layers for
Forr FFN layers using SwiGLU (e.g.. LLaMA), there are two linear networks to calculate neuron values: SwiGLU (x) (Swish(æW) @ 2V 22) We use the weight of the gate network, W. to cluster neurons.

---

better trade-off between model sparsity and quality, by leveraging threshold- based method to enable adaptive selection. The MoEfied FFN layers become:
(FN ( = Σ (())))))
where E(x); is the output for the i-th expert, 1()) is the indicator function, and 7 is predefined threshold. Only experts with score larger than 7 will be activated. Consequently, within the same MoE layer, as the router generates larger expert scores, more experts are selected.
4.2.2 Two-stage LTE Training
Training the router poses three practical challenges. Firstly, as we consider training experts indepen- dently to mitigate bias (e.g., using Sigmoid routing function), it is difficult to tease apart important experts from the rest. Secondly, the con-differentibiiiiyy of threshold based expert selection further complicates router training. Lastly, setting thresholds for each MoE layer is non-trivial.
To tackle these challenges, we next propose novel two-stage training algorithm with efficiency loss penalty to MoEfy LLMs. Our training algorithm consists of two training stages: a) model-router training stage and b) model adaption stage.
Stage 1: Model-router training. In this stage, we jointly train routers and model parameters to capture the importance of experts for given inputs in the Sigmoid routers. To address the non- differentiability issue, we switch the expert selection to a "soft" mode:
FoN = Σ G(r).((x).
1=1
Instead of selecting a subset of experts, we always select all experts and multiply expert outputs E(x); with the corresponding expert score G(z), make both router and model differentiable. Since there is no discrete selection in MoE layers, all parameters are trained for each iteration.
Compared to Softmax routers, Sigmoid routers score each expert independently and do not introduce any competition on expert scores among experts. This makes Sigmoid routers alone cannot identify more important experts in MoE layers. As such, we design an efficiency loss penalty, Lef ficiency, to introduce competition among experts for Sigmoid routers:
Expert Score Distribution (Wlo Seperability Loss)
Derett
Threshod
Expert Score
Expert Score Distribution (wl Separability Loss)
Denait
Threshoid
feef ficieney
XDX Σ G 11
(4)
LN
Expert Score
where L is the number of layers, and N is the number of experts in each layer. The efficiency loss calculates the mean of expert scores across all layers. This efficiency loss penalizes the output magnitudes of routers, driving routers selectively allocate lower expert scores to less important experts, which helps distinguish experts with different importance. Moreover, instead of assigning the same sparsity for all layers, the efficiency loss also induces inter-layer orchestration, allowing LTEE-trined models to allocate adaptive sparsity for different layers. We show that LTE models are capable of adaptively allocating sparsity across different layers (Figure 16).
Figure 3: Expert distribution w/ and w/o sep- arability loss. The expert scores the model trained with the separability loss are much more separable than the model trained with- out the separability loss.
For the threshold in Equation 2 instead of choosing a threshold for each router, we propose to select predetermined fixed threshold and then train models to fit this threshold, i.e., train models to be threshold aware. Specifically, we use a threshold separability regularizer to drive models to make expert scores more separable for this given threshold:
L
N
L Cseparabiiity
1
Σ
(5)
LN
(Gi(z); 7)
1
5

---

where т is predefined threshold for all routers (we set T 0.5 for all experiments). The separability loss encourages router outputs to diverge from the threshold т. which makes the router outputs more separable. An example is illustrated in Figure 3 The only difference between the models in the two figures if the model is trained with a separability loss. We find that the expert scores are not separable when training without the separability loss, making it hard to decide the threshold. However, training with the separability loss markedly improves the separability of expert scores, allowing us to choose experts with the predefined threshold.
Combined with the task performance loss Leask, our training loss in model-router training stage is:
L.1 Lrack + n XL separrabbis
(6)
The efficiency coefficient, 7), is used control the trade-off between inference efficiency and task performance, and 17 can be used to control the sparsity of trained models. A higher indicates a larger inference cost penalty, leading
more sparse model. À the separability loss, we set 0.5 for our evaluation. We conduct an ablation study on both hyperparameters in Section 5.5
Stage 2: Model Adaptation. To accelerate model inference, need switch routers discrete selection mode (Equation 2), and the model trained in Stage needs further training to adapt to the changes in router outputs. In the model adaptation stage, freeze the router parameters, switch routers to discrete selection mode (Equation 2), and fine-tune the model to adapt to the discrete routing with the task performance loss: L.2 Ltask-
4.3 Hardware-efffcient Implementation
illustrated in Figure same as Deja Vu [2I] and moefication 44). LTE pro-
T
vides very structured sparsity. After se- lecting neurons in intermediate layers, CUDA kernel needs only to load the rele- vant columns and rows of W Wup and Wdown from GPU global memory to SRAM compute dot product, thus reducing memory-I/O as well as computational over- heads. Thanks to the structured spar- sity provided by LTT,, the kernel avoids non-coalesced memory access by storing column-mapi major Wep and row-major W Wdown. tomized MLP layer translate the sparsity wall-clock time latency reduction. The evaluation of our custom kernel is presented in Section 5.3
Xout
Ww
Wdown
GeLU(xX
Figure Structural sparsity provided by LTE in FFN layer. After selecting neurons, unselected rows and columns won't loaded and used.
In this paper, we use Triton 2.3.0 to implement a cus-
Experiment
In this section, we conduct extensive evaluationS verify the effectiveness of LTE. For instance, in Section 5.2 LLaMA with LTE provides a 1.83x - 2.59x FLOPs speed-up on NLG tasks. Along with our hard ware custom kernel, LTE reduces 25% LLaMA2-7B inference latency at 50% sparsity.
5.1 Experimental Setting
Models. We implement LTE on both encoder-based models (RoBERTatae- RoBERTa aajgg)) and decoder-based models (GPT2-Medium, and LLaMA-7B) with the HuggingFace's transformers. Following previous work 49) we set each expert to contain 32 neurons FFN layers for all models.
Datasets: 1) Natural Language Understanding (NLU): we evaluate on eight tasks from the GLUE
dataset 40 SST-2 (3,, RTE [5), CoLA (2), MNLI (43], QNLI [29]. QQP [3] STS-B 33, and
MPRC 6 We report the Pearson correlation coefficient for STS-B, BB, the Matthews correlation coefficient for CoLA, and the accuracy for the other six datasets. 2) Natural Language Generation
(NLG)) we evaluate LTE on E2E 227, XSum 26). and 10k 24 For both E2E and XSum, we report the ROUGE-L [19 to measure the generation text quality and report the perplexity performance for the Wikitext dataset. 3) Instruction Tuning: besides downstream tasks, also

---

evaluate LTE on instruction tuning tasks to evaluate LTE's generalization capabilities. We use Tulu dataset [1 to perform instruction tunning with LTE, and we evaluate models with MMLU benchmark [T]. MMLU benchmark is a comprehensive evaluation dataset designed to evaluate the generalization capabilities of L.MSs and consists of 15908 questions from 57 distinct tasks.
Baselines. We compare two baselines with LTE in our paper: (1) MoEfication [9]: MoEfication proposes to replace soft activations with ReLU and then fine-tune the model on downstream datasets. In our evaluation, we report the number of MoEfication with parameter clustering and ground-truth expert selection. (2) Deja Vu 21:: Deja Vu employs predictors to estimate the values of neurons in intermediate layers, pruning neurons with absolute magnitudes (absolute values for non-ReLU activations). The models used in Deja Vu evaluation are models fine-tuned on specific tasks as well.
5.2 End-to-end Performance Comparison
LTE, MoEfication, and Deja Vu provide the
QNLI (RoBERIa-larga)
MNU (RoBERIa large)
same type of structured sparsity in FFN lay-
ers. For the convenience of comparison, we use
Full Mode
FFN neuron sparsity as a measure of efficiency in this section. We will further discuss how our hardware-efficient implementation trans- lates this sparsity into wall-clock time latency reduction in Section [.3] For our method (LTE), we get models with different sparsity by tuning efficiency loss 17 in Equation 6
Digavu
PuBl Mode
Dejovu
Rail Notiicatoo
RoLL & Mutfication
Sparsily
Sparsily
MRPC (RoBERTe-large)
SST-2 (RoBERTo-large)
Drjavo
Dajavu
No performance drop with 80-95% sparsity NLU tasks. We first evaluate our methods on eight NL. Tasks in GLUE datasets (40) with two encoder-based models (RoBERTabee and RoBERTatagee Due to the space limit, we report four datasets (MRPC, SST2, QNLI, and MNLI) on RoBERTaagge in Figure 5 and other results in Figure 14 in Appendix C
RellI + Mocfcation
Rell MnEfication
LTE (Ours)
UTE (Ours)
Sparsity
Sparsity
Figure 5: LTE consistently outperforms other baselines on 4 NLU datasets.
As shown in Figure and Figure [4] our method (LTE) consistently outperforms the base- lines on all datasets. MoEficatiin exhibits per- formance drop when sparsity reaches 70%-80% and Deja Vu experiences performance drop at relatively low sparsity level (e.g., 30%), while LTE maintains good performance even at higher sparsity levels (>90%)). Moreover, we notice that while ReLU-based models achieve high sparsity with a slight decrease in performance, replacing GeLU with ReLU in RoBERTa can af- fect performance on some datasets. For instance, replacing GeLUU with ReLU in a RoBERTa model results in around 10% accuracy drop on the MRPC dataset.
Table 1: GFLOPs per token for LLaMA-7B on different datasets with permitting quality drops. FLOPs of routers are also included in our method. N/A stands for the method failing to achieve the expected performance.
XSum
E2E
Wiki
Full Deja Vu MoEfication R-LLaMA+MoE
12.06
12.06
12.06
7.92
6.42
7.92
10.45
N/A
N/A
8.27
7.39
11.1
LTE (Ours)
5.38
4.65
6.59
50% FLOPs saving on NLG tasks. We next evaluate LTE with decoder-based models on three types of generation tasks. Similar to NLU tasks, 6 that LTE outperforms all other baselines. NLG tasks are more challenging: all methods start to have a performance drop at lower sparsity level compared to NLU tasks. Yet, LTE- -rr-nne models have better resilience to the sparsity increase. Especially at a high sparsity level, while other baselines fail to generate meaningful content, LTE still generates content with a marginal quality drop. Similarly, we also observe the performance drop caused by replacing soft activation with ReLU.
Additionally, to better understand the computation saving of LTE, we compare the FLOPSs per token of different methods on ILLMMA-BB, allowing an up to 0.05 ROUGE-L decrease for XSum and E2E, and an up to 0.5 perplexity (PPL) increase for WikiText- 103. It is noteworthy that the FLOPs for routers are included in our method, constituting approximately 1% of the total FLOPs of FFN layers. The results are presented in Table I The evaluation results show that LTE provides a 1.83x 2.59x FLOPs speed-up up and gives the most FLOPs savings among all methods.
7

---

XSum (GPT2-M)
EZE (GPT2-M)
WikIText-103 (GPT2-M)
Model
638
Deja
MoEtication
ROUGE4
ROUGE:
(Ours)
0-
Full Model
Full Model
Deja Vi
Daja
MoERication
LTF (Ouns)
ITF (Ours)
Sparsity
Sparsity
Sparsity
XSum (LIaMA.78) 78)
E2E (LLaMA.78)
WikIText- 103 (LLaMA 78)
Full Model
Daja
6.3
15.0
Molfication
ROUGE-
6-3-
13.8
LLLMM RLLMMA MoErication
Full
8
(Ours)
Deja Vl
Deja Wu
MoEfication
MoEfication
MLNs-MMEticttoo LTE (Ours)
MLLEca-MEEncaton LTE (Ours)
1
Sparsity
Sparsity
Sparsity
Figure Performance comparison across three NLG datasets (each column). We compare LTE (Ours)
with other baselines on two decoder-based models: GPT2-M and LLaMA-7B (each row).
LTE Outperforms SOTA on instruction tuning. To
0.6
further evaluate the effectiveness of LTE, we evaluate LTE as chatbot and verify model performance in a few-shot learning scenario. We fine-tune LTT on the Tulu instruction tunning dataset [41]. and then evaluate the few-shot performance on MMLU bench- mark [1I] We use LLaMA2-7B
MMLU Accuracy (5-shot)
0.4
LLaMA2-Tulu
LLaMA2
the base model
Deja Vu
for LTE training. For a fair comparison, we also fine-tune base models with the Tulu dataset for other baselines. Specifically, for Deja Vu, we use Tulu-fine- tuned LLaMA2-7B the base model. For Moefi- cation, we use mulu-fine-tumed ReLULLaMA-7B as the base model ensure the best possible baseline performance. We present the 5-shot MMLU accuracy comparison in Figure 7 The evaluation results show that LTE outperforms other baselines across all sparsity levels. Yet, LTE shows a weaker sparsity- performance trade-off compared to downstream tasks. Our hypothesis is that LTE routers change the structure of the original LLaMA model, which influences general language patterns learned by models pretraining. Instruction tuning alone is not enough fully recover this ability (considering instruction tuning has much smaller datasets). We believe that training with a wider variety of data (like pretraining data) will further improve LTE performance.
R-LLaMA Moefication
(Ours)
oh
Sparsity
Figure 7: LTT outperforms other baselines on 5-shot MMLU benchmark.
5.3 Translate Sparsity to Wall-clock Time Speedup
25% wall-clock time latency re- duction. As discussed in Sec-
FFN Inference Latency
End-to-end Generation Latency
ITE
tion 4.3 the sparsity introduced by
Latency (Us)
Latency (ms)
LTE is highly structured and can be casily translated latency re-
duction. This section demonstrates that our custom Triton kernel ef- fectively translates the LTE spar- sity wall-clock time speed up. We evaluate wall-clock time speed up using LLaMA2-7B on single 3090Ti. We also reported vanilla Pytorch indexed matrix multiplica- tion. Due to the requirement for
100-
Our Kemei
Pyterch Incex
326
siz
3004
Sparsity
Sequence Lengthh
Figure 8: Wall-clock time latency comparison dense FFN blocks and our custom Triton kernel (Left). End-to-end gener- ation latency comparison between dense model and LTE with 50% sparsity (Right).

---

Table 2: Performance comparison with model pruning. We apply Wanda [36] on WikiText103 fine-tuned LLaMA-7B with the author-ima inthhorrmpleeeetedd code. Besides using C4 the calibration data suggested by the Wanda paper, we also try to use WikiText for calibration to improve this baseline. The sparsity reported here the overall sparsity the entire model, which different from the FFN sparsity. The evaluation results show that LTE achieves a lower perplexity than Wanda.
Wanda (2:4)
Wanda
Wanda (Unstructured)
Wanda (2:4)
Wanda (4:8)
Wanda (Unstructured)
LTE
Method
(4:8)
(1) 2)
Overall Sparsity
50%
50%
50%
50%
50%
50%
52%
Calibration Data
C4
C4
C4
Wiki 10.82
Wiki
Wiki
PPL
11.45
8.39
7.04
8.17
6.87
5.95
new memory allocations, slicing in Pytorch is very time-consuming. This makes indexed matrix multiplication even slower than dense matrix multiplication at most sparsity levels. We report the wall-clock time latency comparison in Figure (noting that the router inference latency is already included in LTE latency). Compared to a dense FFN block, our custom Triton kernel achieves nearly linear speed-up with respect to sparsity (Left). At around 50% sparsity, LTE reduces end-to-end generation latency by approximately 25%.
5.4 Addtional Comparison
Model pruning. As discussed in Section besides dynamic contextual sparsity provided by LTE, static sparsity provided by model pruning also common method to accelerate model inference. In this section, we present performance comparison hetween Wanda [36] and LTE in Table 2 The sparsity reported here is the overall sparsity of the entire model, which is different from the FFN sparsity reported in Seciton 5.2 For LTE, we recalculate the overall sparsity based on the FFN sparsity. The evaluation results show that LTE achieves lower perplexity than Wanda. The evaluation results show that, given a similar sparsity, LTE achieves better perplexity.
5.5
Ablation Study
Full Model
Full Model
Random
Co-sctivation
Kmmans
Wic WicThet Liabt
11
122 LLaM5
Abumn LL.MG
Sparsity
Sparsity
Figure 9: Performance compar- Figure 10: The relation between ison between different experts and sparsity on different mod- grouping algorithms.
Figure 11: Performance com- parison between different À on WikiText- 103 (GPT2-M).
els and tasks.
Experts grouping algorithms. Three different expert grouping algorithms are explored in (49 random grouping, parameter K-means clustering, and co-activation graph split. The co-activation graph split algorithm constructs graph where each neuron represents node and the edges denote the frequency of co-activation between neuron pairs. This graph is subsequently divided into subgraphs using the graph split algorithm [16), with each subgraph representing group of neurons. To study the effectiveness of these grouping algorithms, we compare their LTE performance by fine-tuning GPT2 on the WikiText- 103 dataset. As shown in Figure 9 both co-activation graph split and parameter K-means grouping have very similar performance (similar to finding in 49 and outperform random grouping. Given that parameter K-means grouping is simpler to implement, as co-activation graph split requires collecting activation data, we use K-means grouping our paper.
The relation between efficiency loss hyperparameter and sparsity. As discussed in Section efficiency loss hyperparameter 7) is used balance the trade-off between task performance and sparsity (efficiency). We illustrate the relation between r) and sparsity across different models and tasks in Figure 10 For a given task and model, increasing results in greater sparsity.

---

Ablation study on separability loss. Another hyperparameter LTE train- ing, X, is used to encourage the sepa- rability of router outputs, which facili- tates us to choose the threshold to pick neurons (as shown in Figure 3). We compare models trained with different lambda Figure IT When increasing À from 0 (i.e., separability loss) to 0.5, we observe a constant perplexity drop for all sparsity. However, if we keep increasing >, perplexity does not further decrease.
MRPC (RoBERTa-basee
WikiText- 103 (GPT2-M)
Full Model
Accuracy
(Ours)
Random Router
Full Modal
Random Router
Sparsity
Sparsity
Figure 12: Performance comparison between random routers and routers trained with LTE. Models with LTE-trained routers consistently outperform models with random routers.
Effectiveness of LTE-trained routers. We next demon- strate the effectiveness of LTE's first training stage, the model-router training stage, in helping LLMs gain bet- ter sparsity. To build baseline, instead of using routers trained in the model-router training stage, we use randomly initialized networks as routers and train LLMs adapt to random routers. Since we cannot control the sparsity of LLMs with threshold in a random router, decide the sparsity for MoE layers by picking K top experts: we first fix K decide the sparsity and then start train the model. Figure 12 compares the performance of models with different routers. We observe that routers trained with LTE significantly contribute to achieving better trade-off between sparsity and model performance, which proves the effectiveness of the model-router training stage.
WikIText-103 (GPT2-M)
Model
MoE
Deja Mu+Fine-tuning
05
02
04
5.6
1.
FFN Sparsity
Figure 13: Performance comparison
with new baselines: 1) a-MoE 4 and Deja VU (21 with further fine-tuning. We evaluate two baselines by fine-tuning GPT2 medium with WikiText103. LTE outperforms the other two baselines across different sparsity.
Other baselines with further fine-tuning, A key distinc- tion between LTE and other baselines is that LTE requires additional fine-tuning to make the model efficiency-aware. In this section, we assess whether further fine-tuning can enhance the performance of other baselines, aiming to understand the contribution of additional training to LTE's effectiveness. Specifically, we fine-tune two baselines on GPT2-Medium using the WikiText103 dataset: we first fine-tune GPT2-M with Wikitext and apply Deja Vu the fine-tuned models. Then, we further fine-tune MoEfied models with WikiText103. 2) We implement o-ooE A on GPT2-Medium model and fine-tune models with the WikiText103 dataset. For both baselines, we train models for the same training time as we train LTE models. The evaluation results shown in Figure show that LTE outperforms the other two baselines across different sparsity levels.
Conclusion
In this paper. we explore how to enable better activation sparsity for fast LLM inference. We introduce novel algorithm. LTE, for the model automatically learn to activate fewer neurons, by grouping neurons into different experts and then adaptively selecting important ones for specific inputs and layers. Our extensive evaluations, spanning four LLMs and eleven datasets, show that LTE can achieve 1.83x 2.59x FLOPs speed-up up on LLaMA, thus execution efficiency, outperforming ffate-of-the-as designs. We believe that our work can inspire more research on designing more efficiency-awant training methods, making LLMs more accessible to the broad community.
Acknowledgement
My work was partially supported by DARPA Grant HR001 12020008, National Science Foundation Grant No. 2039445, and Cisco Research Grant. Any opinions, findings, and conclusions recom- mendations expressed this material are those of the author(s) and do not necessarily reflect the views of our research sponsors.
10

---

References
[1] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
[2] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1900 1901, 2020.
[3] Daniel Cer, Mona Diab, Eneko Agirre, Inigo Oopez-Gazpio, and Lucia Specia. Semeval- 2017 task 1: Semantic textual similarity multilingual and crosslingual focused evaluation. In Proceedings of the 1 11th International Workshop on Semantic Evaluation (SemEval-2017). Association for Computational Linguistics, 2017.
[4] Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. Approximating two-layer feedforward networks for efficient transformers. In Findings of the Association for Computational Linguistics: EMNLP 2023, pages 674 692, 2023.
[5] Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entail- ment challenge. In Proceedings of the First International Conference Machine Learning Challenges: Evaluating Predictive Uncertainty Visual Object Classification, and Recognizing Textual Entailment, ML.CW05, page 177-190, Berlin, Heidelberg, 2005. Springer- ppringer-Verlag..
[6] William B. Dolan and Chris Brockett. Automatically constructing a corpus of sentential paraphrases. In Proceedings of the Third International Workshop on Paraphrasing (IWP2005), 2005.
[7] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1):5232 222-22200,0 2022.
[8] Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot. In International Conference on Machine Learning, pages 00323-10337. PMLR, 2023.
[9] Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17323: 2022.
[10] Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages
2021.
[11] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint urXiv:2009.03300 2020.
[12] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint 1rXiv:1606.08415, 2016.
[13] Shankar Iyer, Nikhil Dandekar, Kornél Csemnai, et al. First quora dataset release: Question pairs. data. quora. com, 2017.
[14] Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adaptive mixtures of local experts. Neural Computation, 3(1):79-87) 87, 1991.
[15] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:23 2023.
[16] George Karypis and Vipin Kumar. fast and high quality multilevel scheme for partitioning irregular graphs. SIAM Journal scientific Computing, 20(1):359-392, 1998.

---

[17] Xin Li, Yeqi Bai, Pinlong Cai, Licheng Wen, Daocheng Fu, Bo Zhang, Xuemeng Yang, Xinyu Cai, Tao Ma. Jianfei Guo, et al. Towards knowledge-driven autonomous driving. arXiv preprint arXiv:2312.043106 2023.
[18] Zonglin Li, Chong You, Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank Reddi, Ke Ye, Felix Chern, Felix Yu, Ruiqi Guo, et al. The lazy neuron phenomenon: On emergence of activation sparsity in transformers. In Conference on Parsimony and Learning (Recent Spotlight Track), 2023.
[19] Chin-Yew Lin. Rouge: A package for automatic evaluation of summaries. In Text summarization branches out, pages 74-81, 2004.
[20] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach, 2019.
[21] Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivas- lava, Ce Zhang, Yuandong Tian, Christopher Re, et al. Deja vu: Contextual sparsity for efficient Ilms at inference time. In International Conference on Machine Learning, pages 22137-22176. PMLR, 2023.
[22] Xinyin Ma, Gongfan Fang, and Xinchao Wang. language models. arXiv preprint arXiv:2300.. /627,, 2023.
On the structural pruning of large
[23] Mikko Malinen and Pasi Fränti. Balanced k-means for clustering. In Structural, Syntactic, and Statistical Puttern Recognition: Joint IAPR International Workshop, S+ SSPR 2014, Joensuu, Finland, August 20-22, 2014. Proceedings, pages 32-41. Springer, 2014.
[24] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint 169.07460 2016.
[25] Iman Mirzadeh, Keivan Alizadeh, Sachin Mehta, Carlo C Del Mundo, Oncel Tuzel, Golnoosh Samei, Mohammad Rastegari, and Mehrdad Farajtabar. Relu strikes back: Exploiting activation sparsity in large language models. arXiv preprini arXiv:2310 10.05533 2023.
[26] Shashi Narayan, Shay Cohen, and Mirella Lapata. Don't give me the details, just the sum- mary! Topic-aware convolutional neural networks for extreme summarization. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium, 2018.
[27] Jekaterina Novikova, Ondrej Dušek, and Verena Rieser. The E2E dataset: New challenges for end-to-end generation. In Proceedings of the 18th Annual Meeting of the Special Interest Group on Discourse and Dialogue, Saarbrücken, Germany, 2017. arXiv:1706.09254
(28) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 188):9 2019.
[29] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text, 2016.
[30] Joseph Redmon, Santosh Divvala, Ross Girshick, and Ali Farhadi. You only look once: Unified, real-time object detection. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 779-788, 2016
[31] Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.052000 2020.
[32] Noam Shazeer, Youlong Cheng. Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanan- takool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, et al. Mesh-tensorflow: Deep learning for supercomputers. Advances neural information processing systems, 31, 2018.
[33] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations, 2016.
12

---

[34] Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over sentiment treebank. In David Yarowsky, Timothy Baldwin, Anna Korhonen, Karen Livescu, and Steven Bethard, editors, Proceedings of the 2013 Conference on Empirical Methods Natural Language Processing pages 1631-1642, Seattle, Washington, USA, October 2013. Association for Computational Linguistics.
[35] Jiachen Sun, Haizhong Zheng, Qingzhao Zhang, Atul Prakash, Zhuoqing Mao, and Chaowei Xiao. Calico: Self-supervised camera-lidar contrastive pre-training for bev perception. In ICLR, 2024.
[36] Mingjie Sun, Zhuang Liu, Anna Bair, and Zico Kolter. A simple and effective pruning approach for large language models. arXiv preprint WIXi::2306.1165, 2023.
[37] Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rívière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based on gemini research and technology. arXiv preprint arXiv:2403.08295, 2024.
[38] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.139771 2023.
[39] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint urXiv:2307.0922 2023.
[40] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: 1804.v1804.074000 2018
[41] Yizhong Wang, Hamish Ivison, Pradeep Dasigi, Jack Hessel, Tushar Khot, Khyathi Chandu, David Wadden, Kelsey MacMillan, Noah A Smith, Iz Beltagy, et al. How far can camels go? exploring the state of instruction tuning on open resources. Advances in Neural Information Processing Systems, 36, 2023.
[42] Alex Warstadt, Amanpreet Singh, and Samuel Bowman. Neural network acceptability judgments. arXiv preprint arXiv: 1805. 180..1241, 2018
[43] Adina Williams, Nikita Nangia, and Samuel R. Bowman. A broad-coverage challenge corpus for sentence understanding through inference, 2018.
[44] Yue Wu, Zhiqing Sun, Huizhuo Yuan, Kaixuan Ji, Yiming Yang, and Quanquan Gu. Self-play preference optimization for language model alignment. arXiv preprint arXiv:2405.00675, 2024.
[45] Mengzhou Xia, Sadhika Malladi, Suchin Gururangan, Sanjeev Arora, and Danqi Chen. Less: Selecting influential data for targeted instruction tuning. arXiv preprint arXiv:2402.044 2024.
[46] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309. /7453, 2023.
[47] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint 2rXiv:2205.00000 2022.
[48] Xiaokuan Zhang, Haizhong Zheng, Xiaolong Li, Suguo Du, and Haojin Zhu. You are where you have been: Sybil detection via geo-location analysis in osns. In 2014 IEEE Global Communications Conference, pages 698-703. IEEE, 2014.
[49] Zhengyan Zhang, Yankai Lin, Zhiyuan Liu, Peng Li, Maosong Sun, and Jie Zhou. Moefication: Transformer feed-forward layers are mixtures of experts. In Findings of the Association for Computational Linguistics: ACL 2022, pages 877-890, 2022.
13

---

[50] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H _ Heavy-hitter oracle for efficient generative inference of large language models. arXiv preprint arXiv:2306. 14048, 2023.
[51] Haizhong Zheng, Rui Liu, Fan Lai, and Atul Prakash. Coverage-centric coreset selection for high pruning rates. In ICLR, 2023.
[52] Haizhong Zheng, Jiachen Sun, Shutong Wu, Bhavya Kailkhura, Zhuoqing Mao, Chaowei Xiao, and Atul Prakash. Leveraging hierarchical feature sharing for efficient dataset condensation. In ECCV, 2024.
[53] Haizhong Zheng, Minhui Xue, Hao Lu, Shuang Hao, Haojin Zhu, Xiaohui Liang, and Keith Ross. Smoke screener or straight shooter: Detecting elite sybil attacks in user-review social networks. In NDSS, 2018.
14

---

Broader Impact
This paper presents the Sein-TTTbe--fiicee (LTE) algorithm to accelerate the LLM inference. A more efficient LLM inference can significantly enhance the accessibility and sustainability of LLMs. By reducing computational demands, it addresses environmental concerns through lower energy consumption and helps democratize access advanced Al technologies. We believe that work can better help smaller organizations, educational institutions, and researchers, who previously faced barriers due resouree limitations, access LLMs more easily, To sum up, the LTE algorithm not only advances the field technically but also broadens the scope of AI's benefits to wider, more inclusive community.
Limitations
While LTE outperforms other baselines, LTE needs further fine-tuning to get contextual sparse models. (Yet, training one-time investment that yields long-term benefits for LLM serving.) future study is to combine LTE with other parameter-efficient fine-tuning techniques to reduce I.TE training overhead. Another limitation our computational resource, which prevents us from training LTTE on pretraining or RLHF. As discussed in Section 5.2 we believe that training with a wider variety of data will further improve LTE performance. We leave this for future exploration.
A Code and Datasets License
Codebase. Our model implementation is based on Huggingface transformer repo: Apache- 2.0 license Datasets. We list the license of used datasets follow:
GLUE dataset 40
Custom License
WikiText103
24
CC BY-SA 3.0
XSum
26
MIT License
E2E
CC4.0-BY-SA
Tulu
ODC-BY
MMLU
Custom License
B
Experiment Settings
We presented all training hyperparameters in Table 3 5 and For hyperparameters presented in (). select the best hyperparameter for each task. We follow the settings in RoBERTa paper (20) to fine-tune RoBERTa on GLUE datasets. We set the coefficient for the separability loss (A in Equation 6)) to be 0.5 for all stage training. We use different 17 control the sparsity of trained models (Figure 1O). Hardware: All models are trained and evaluated on A A100, A40, and 3090Ti, depending memory usage and availability.
C Additional Evaluation Results
C.1 NLU Performance Comparison
In Figure 14] we present the evaluation results of the rest four NLU datasets. We have similar findings as we discussed in Section [.2] LTE achieves a better trade-off between sparsity and task performance. However, we notice that, LTE outperforms MoEfication on the CoLA dataset, but has worse performance than KLA at low sparsity level. The size of CoLA is relatively small and may not be efficient for models learn good routers.
C.2 Further Analysis on Sparsity Provided by LTE
15

---

RTE (RoBERI2-Iorge)
CoLA (RoBERTD-large
STSB (RoBERTD-large
QOP (RoBERta-large)
Accurazy
Accuracy
Accurecy
Model
C4
Mode
Nosel
Putl Made
Dujavu
(Ours)
(Ours)
(Ours!
(Ours)
Sparsity
Sparsity
Sparsity
Sparsity
MRPC (RoBERTB-OaS) -baEE)
S5T-2 (RoBERTB-base)
QNLI (RoBERTa-base)
MNLI (RoSERIa-base)
Accuracy
AccurDyy
Accura 1
Full Mode
Fill Model
Ful Meddl
Full Mode
Desevu
Dejave
Dejevu
RelU Nocfication
Rewi + MoEiication
REW MoERcation
ReLU MoEniation
LIE (Durs)
UE rOurs!
LE (Ours)
LIE (Durs)
57
25
Sparsity
Sparsity
Sparsity
Sparsity
RTE (RoBERIa-base)
COLA (RoBERTa-base)
STSB (RoBERIa-base)
QQP (RoBERta-base)
Accuracy
Accuracy
Accuracy
Full Mosel
Ful Model
Ful Mode
Ful Model
Diparu
Digars
Dijeru
Dejavu
Relu Motfication
Relli Mottication
Rmill Motfication
Relli McEtication
tOurs!
LTE
Sparsty
Sparsity
Sparsity
Sparsity
Performance comparison across four NLU datasets from GLUE dataset. The evaluation results show that LTE consistently outperforms other baselines in each dataset for both models.
Figure 14:
Contextual Sparsity for Larger Batches. Similar to previous contextual sparsity work (21), we observe that as the batch size increases, more neurons are activated, leading to a reduction in union sparsity. These dynamics are illustrated in 15 Similar to Deja Vu [21]. the number of activated neurons does not grow linearly with the batch size, indicating a non-uniform distribution of parameter access from different inputs. This property can extend LTE for a nrightthroughuu batch setting by batching inputs activating similar neurons together to achieve a high union sparsity.
0-0.45
Union Sparsity
0-5.5
0
Batch Size
Figure 15: Union sparsity changes w.r.t. batch size.
Adaptive Sparsity across Layers. Instead of assigning a fixed sparsity each layer, LTE uses the efficiency loss to introduce competition across different layers to allocate the computation budget to more important layers. Here, we conduct study on the distribution of sparsity across layers in models trained with LTE. As shown in Figure 16 we observe that LTE achieves different sparsity levels across different layers. For example, the MRPC (RoBERTagge) ) model with 0.46 sparsity has 0.2 sparsity at layer 2, but has more than 0.8 sparsity at layers 10 and 11.
Another intriguing finding is that RoBERTa and GPT2-M MM have very different sparsity pat- terns. In RoBERTa, the sparsity of MoE layers significantly increases with depth, contrasting with GPT2-M, where MoE layer sparsity de- creases as layer depth increases. This difference can stem from differences in the design of these two transformers and tasks. Given that MRPC a sentence classification task, incoder-based transformers like RoBERTa primarily rely on the "[CLS]" token for classification. Consequently, other tokens become less important in deeper layers, allowing sparser FFN layers. In contrast, GPT2, a decoder-based transformer designed for language generation, requires all tokens to generate the next token. Furthermore, recent research [00] on interpretability suggests that deeper FFN layers store more complex semantic information, which may drive deep FFN layers to be less sparse.
MRPC (RoBERTa Base)
WikiText: 103 (GPT2.M)
Sparsity
Sparsity 0.46 Sparsty
Sparsity. 0.52 Spanityy
Layer index"
Layer Index
Figure 16: Sparsity across different layers. We present the sparsity of different layers in an encoder-based model (RoBERTanae Tanse)) for MRPC and a decoder- based model (GPT2-M) for Wiki- Text103. Different layers learn to have different sparsity, and RoBERTa and GPT2-M have dif- ferent sparsity patterns across layers.
16

---

Table 3: RoBERTa-base hyperparameter for LTE training on GLUE dataset.
Hyperparameter
LTE-Stage (2e-5, 5e-5)
LTE-Stage
Learning rate
(5e-5)
Training batch size
[16, 32)
(16, 32)
Training epochs
10) 10)
(1,3, 10]
Weight decay
0
Warm up ratio
0.06
0.06
Table 4: RoBERTa-large hyperparameter for LTE training on GLUE dataset.
Hyperparameter
LTE-Stage
LTE-Stage
Learning rate
(2e-6, 5e-6, 3e-5)
(2e-6, 5e-6, 3e-5]
Training batch size
[16. 32/
(16. 32)
Training epochs
[1, 10)
(3, 10)
Weight decay
0
Warm up ratio
0.06
0.06
Table 5: GPT2-Medium hyperparameter for fine-tuning and LTE training on XSum, E2E, and WikiText- 103 datasets. (Numbers in the bracket correspond to three datasets order.)
Hyperparameter Learning rate Training batch size
Fine-tuning
LTE-Stage
LTE-Stage
5e-5
5e-5
5e-5
8
Training epochs
3
(1, 3)
Weight decay
Warm up ratio
0.06
0.06
0.06
Table 6: LLaMA hyperparameter for fine-tuning and LTE training on XSum, E2E, and WikiText-103 datasets.
Hyperparameter
Fine-tuning
LTE-Stage
LTE-Stage
Learning rate Training batch size
1.5e-5
1.5e-5
1.5e-5
16
16
16
Training epochs
3
Weight decay
Warm up ratio
0.06
0.06
0.06
Table 7: Hyperparameters for instruction tunning with Tulu dataset.
Hyperparameter
Fine-tuning
LTE-Stage
LTE-Stage
Learning rate
2e-5
2e-5
2e-5
Training batch size
128
128
128
Training epochs
3
4
Weight decay
Warm up ratio
0.06
0.06
0.06
17

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
The main claims in the abstract and introduction accurately summarize the paper's contributions and scope.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discussed the limitations of our work in the Appendix.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [NA]
18

---

Justification:
No theory is included in the paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We introduced detailed implementation details and hyperparameters used in the Appendix.
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
19

---

Answer: [Yes]
Justification:
We upload the code in the Supplementary files.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (httpss:/niis.cc.
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We introduced detailed implementation details and hyperparameters used in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Training LLMs incurs significant costs. To ensure the reliability of our evalua- tion and avoid the potential bias of Furthermore, the consistency of our conclusions across multiple models and benchmarks can also be treated as form of repeated evaluation and support the robustness of our findings.
results, we employ entirely random seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
20

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out of range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We introduced compute resources used in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Yes, we check that our paper conform with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We included broader impacts in the Appendix.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
21

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
No model or dataset is released in this paper.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We include the license for the used code and datasets in the Appendix.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
22

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
23


Reviews:
[{"review_id": "DWRJQ8q6Dq", "reviewer": "Reviewer_zH4X", "summary": "This work aims to introduce structured sparsity to large language models (LLMs) to improve their execution efficiency. To achieve this, it enhances previous MoEfication methods by employing a sigmoid-based non-competitive routing function and a threshold-based expert selection, allowing for adaptive expert numbers. Experiments across various models and different language understanding and generation tasks validate the effectiveness of the proposed method.", "questions": "My questions have been included in the weakness section. I'm willing to adjust my scores if my concerns are properly addressed.", "limitations": "This work does not suffer from notable negative societal impacts as it aims to improve the efficiency of LLMs.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-motivated and easy to follow.\n\n2. The proposed method exhibits good soundness and can achieve real-device speed-up with the developed CUDA kernel.\n\n3. The proposed method has been evaluated across both encoder and decoder language models, achieving a consistently improved accuracy-efficiency trade-off.", "weaknesses": "1. The major concern is that the technical contribution and novelty of this work are somewhat limited. The use of MoE with sigmoid functions to avoid competition among experts has been adopted in previous works, such as [1][2], and an extension to MoEfication will intuitively work. The authors are expected to analyze the key differences that make the proposed method particularly suitable for MoEfication.\n\n[1] \"Approximating Two-Layer Feedforward Networks for Efficient Transformers,\" R. Csordás et al., EMNLP'23.\n\n[2] \"SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention,\" R. Csordás et al., arXiv'23.\n\n2. One missing baseline is structured weight pruning [3][4]. Considering this work mainly targets a per-task fine-tuning setting, structured weight pruning can already achieve decent sparsity with non-trivial generation speed-up without performance dropping [3][4]. This comparison can inform the community whether weight sparsity or context sparsity is more cost-effective or if they can be applied together.\n\n[3] \"A Simple and Effective Pruning Approach for Large Language Models,\" M. Sun et al., ICLR'24.\n\n[4] \"Fluctuation-based Adaptive Structured Pruning for Large Language Models,\" Y. An et al., AAAI'24.\n\n3. Task-specific fine-tuning for each task is too costly. The authors are highly encouraged to perform continued pretraining and evaluate across different tasks to validate the generalization capability of the proposed method.\n\n4. I wonder how the task-specific fine-tuning is performed for the baseline methods. Specifically, is Dejavu/MoEfication performed on top of the fine-tuned model, or is the model fine-tuned after applying these techniques, or are there any smarter strategies? This question arises because the proposed method simultaneously updates both model weights and expert selection strategies, which may be a key reason why it outperforms the baselines.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are glad that the reviewer found that our work is well-motivated and sound. We thank the reviewer for the constructive feedback and appreciate the opportunity to address the points you have raised.\n\n---\n\n> **Q1:** Novelty concerns: A sigmoid router was proposed in previous MoE works like [1][2]. The authors are expected to analyze the key differences that make the proposed method particularly suitable for MoEfication.\n\n**A1:** We thank the reviewer for pointing out the missing references, and we are happy to clarify the key differences between LTE and MoE with a sigmoid router, and the novelty of our paper.\n\nFirst, we would like to clarify that, compared to traditional MoE, the MoEfication problem has three unique challenges: 1) Router designing for a larger number of selected experts; 2) Router training on pretrained models;  3) Adaptive sparsity in different layers of the pretrained model. *(We kindly refer to the General Response Q1 for more details.)* \n\nEven though the Sigmoid-MoE[1] can handle the first challenge, it ignores the second and third challenges, which can cause an inferior performance. Instead, LTE also addresses two other challenges: For the second challenge, LTE adopts an indicator function (Eq. 2) to avoid the expert output scale and a two-stage training algorithm to solve the non-differentiability issue of the indicator function.  For the third challenge, LTE employs an efficiency loss (Eq. 4) to introduce competition across different layers, leading to more adaptive sparsity in different layers (Figure. 15 in the paper). Those novel designs make LTE better address the challenges in MoEfication.  *(We kindly refer to the General Response Q2 for more details.)* \n\n**Sigmoid-MoE Evaluation.** To better understand how Sigmoid-MoE works on MoEfication tasks, we implement the Sigmoid-MoE[1] on the GPT2-M and fine-tune it with the WikiText103. We train Sigmoid-MoE models with the same training time as we train the LTE models. The comparison results are in Figure 1 in the rebuttal PDF: Sigmoid-MoE overcomes the collapse issue of the softmax router, but still underperforms LTE. \n\n>**Q2:** One missing baseline is structured weight pruning [3][4]…\n\n**A2:** We thank the reviewer for bringing the model pruning for discussion. As discussed in lines 94-99 of our paper, while both LTE and structured weight pruning provide structured sparsity for inference acceleration, they provide two different types of sparsity. The contextual sparsity offered by LTE is more flexible and adaptive compared to the static sparsity offered by model pruning. \n\nTo provide a clearer comparison between these methods, we apply Wanda to a Wikitext-103 fine-tuned LLaMA-7B model and report the results in Table 1 in the rebuttal PDF. The evaluation results show that LTE achieves better performance than Wanda given the same level of sparsity.\n\n\n> **Q3:** Task-specific fine-tuning for each task is too costly. The authors are highly encouraged to perform continued pretraining and evaluate across different tasks to validate the generalization capability of the proposed method.\n\n\n**A3:** We agree with the reviewer that validating the generalization capability can better demonstrate the effectiveness of the proposed method. In our submission, we already included this evaluation in Figure 7 (Section 5.2). We use the Tulu dataset to supervise fine-tune LTE models and evaluate the few-shot performance on the MMLU dataset (which is a large comprehensive dataset consisting of 15,908 questions from 57 distinct tasks.) The evaluation results show that LTE still outperforms other baseline methods in this supervised fine-tuning setting, which demonstrates the generalization capability of LTE. \n\n> **Q4:** I wonder how the task-specific fine-tuning is performed for the baseline methods. Specifically, is Dejavu/MoEfication performed on top of the fine-tuned model, or is the model fine-tuned after applying these techniques, or are there any smarter strategies? \n\n**A4:** Deja Vu and MoEfication are proposed as post-training methods in their papers. Following this, we implemented Deja Vu or MoEfication in a post-training manner in our paper: we first fine-tuned the model with the specific datasets and then applied Deja Vu or MoEfication to the fine-tuned model to MoEfy models. (This is also the implementation suggested in the MoEfication paper).\n\nHowever, we understand the reviewer's concern that further fine-tuning those MoEfied models (all parameters) can further improve performance. We conduct an evaluation on fine-tuning Deja Vu models. We first fine-tuned the model with Wikitext-103 and applied Deja Vu. Subsequently, we further fine-tuned the MoEfied Deja Vu model with Wikitext-103 (for the same training time as LTE training). Note that the first fine-tuning is necessary for Deja Vu to collect data to train predictors.\n\nThe evaluation results are reported in Figure 1 of the rebuttal PDF. We find that while the additional fine-tuning does improve the performance of Deja Vu, LTE still outperforms Deja Vu with additional fine-tuning.\n\n---\n\nThanks for the attentive reading of the manuscript and constructive feedback. We will incorporate these changes into our final version.\n\nWe hope our response addresses all the concerns and that the reviewer will consider raising the rating accordingly. We are more than glad to answer any further questions.\n\n---\n\n[1] \"Approximating Two-Layer Feedforward Networks for Efficient Transformers,\" R. Csordás et al., EMNLP'23.\n\n\n[2] \"SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention,\" R. Csordás et al., arXiv'23.\n\n[3] \"A Simple and Effective Pruning Approach for Large Language Models,\" M. Sun et al., ICLR'24.\n\n\n[4] \"Fluctuation-based Adaptive Structured Pruning for Large Language Models,\" Y. An et al., AAAI'24.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "k69ZFVSmR0", "reviewer": "Reviewer_Jye3", "summary": "This article introduces a novel training algorithm, LTE, designed to train large language models (LLMs) to achieve more structured activation sparsity during inference. Thus, it enhances their efficiency without compromising performance until a very high sparsity.", "questions": "1. Have you tried LTE also on the attention layers?", "limitations": "The article includes a limitation section.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. LTE performs excellently across all datasets: In multiple natural language understanding (NLU) tasks, LTE shows no significant performance degradation at high sparsity levels (80-95%) and maintains good performance even at sparsity levels exceeding 90% of FFN.\n\n2. This article develops a CUDA kernel to speed up inference by reducing memory and computational overheads.", "weaknesses": "1. Stage 1 of LTE will train all the model's parameters. How do you implement Dejavu? As far as I know, Dejavu is a post-training method that freezes the model's parameters. I'm not sure whether this is a fair comparison.\n\n2. The LTE algorithm introduces a two-stage training process, which might be complex and computationally intensive to implement. This can be a barrier to practical adoption in resource-constrained environments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are glad that the reviewer found that our work has excellent performance. We thank the reviewer for the constructive feedback and appreciate the opportunity to address the points you have raised.\n\n---\n\n> **Q1:** Stage 1 of LTE will train all the model's parameters. How do you implement Dejavu? As far as I know, Dejavu is a post-training method that freezes the model's parameters. I'm not sure whether this is a fair comparison.\n\n**A1:** Deja Vu is proposed as a post-training method in their paper. Following this, we implemented Deja Vu in a post-training manner in our paper: we first fine-tuned the model with the specific datasets and then applied the Deja Vu to the fine-tuned model to MoEfy models.\n\nHowever, we understand the reviewer's concern that fine-tuning those MoEfied models (all parameters) can further improve performance. We conduct an evaluation on fine-tuning Deja Vu models. We first fine-tuned the model with Wikitext-103 and applied Deja Vu. Subsequently, we further fine-tuned the Deja Vu model with Wikitext-103 (for the same training time as LTE training). Note that the first fine-tuning is necessary for Deja Vu to collect data to train predictors.\n\nThe evaluation results are reported in Figure 1 of the rebuttal PDF. We find that while the additional fine-tuning does improve the performance of Deja Vu, LTE still outperforms Deja Vu with additional fine-tuning.\n\n> **Q2:** The LTE algorithm introduces a two-stage training process, which might be complex and computationally intensive to implement. This can be a barrier to practical adoption in resource-constrained environments.\n\n**A2:** Compared to the two post-training baselines, even though LTE introduces additional training overhead, we argue that LTE training is a one-time effort. Since LTE significantly saves the inference overhead, this approach is beneficial in the long run for serving LLM more efficiently.\n\nMoreover, as we discussed in Q1. Even if we further fine-tune the entire model MoEfied with Deja Vu for the same training time as LTE models, the Deja Vu models fail to achieve comparable performance to LTE models, which indicates that, besides training, other LTE components also contribute to the performance.\n\n>**Q3:** Have you tried LTE also on the attention layers?\n\n**A3:** We focus on the MoEfication for FFN layers in this work, but we do think applying MoEfication on attention layers can be an interesting and promising exploration direction. A potential solution is to set each attention head as an expert and set a sigmoid routing to decide if a head in the attention layer should be used. We believe MoEfication on attention layers can further increase the contextual sparsity of the paper, and we plan to leave this for our future study.\n\n---\n\nThanks for the attentive reading of the manuscript and constructive feedback. We will incorporate these changes into our final version.\n\nWe hope our response addresses all the concerns and that the reviewer will consider raising the rating accordingly. We are more than glad to answer any further questions.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0XDCaaN7hf", "reviewer": "Reviewer_aQHS", "summary": "The paper presents a new approach (LTE) aimed at improving the inference efficiency of large language models by developing structured activation sparsity. The method trains LLMs to activate fewer neurons in FFN layers while attempting to maintain task performance. The approach works by grouping neurons into experts and using a routing strategy (based on Sigmoid instead of Softmax) to select experts adaptively. The authors evaluate the method on RoBERTa, GPT2, and LLaMA models across various NLP tasks. They report that LTE outperforms existing baselines and provides FLOPs and a latency reduction by utilizing sparsity through a custom CUDA implementation.", "questions": "1. Since you are introducing a different training approach, how does the total training time compare to the baseline?\n2. Can you provide more insights into the limitations of Softmax routing?", "limitations": "1. Complexity overhead of the training.\n2. Many of the presented insights are already in the MoE literature, except for the Sigmoid routing.\n3. Limited number of baselines (using only two baselines, there is a lot of work about sparsity in LLMs). Imo, the paper would be much stronger if it compares against more (sparsity) baselines.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper presents a new method for inducing sparsity in LLMs. It is based on several MoE concepts but uses Sigmoid-based routing instead of the traditional Softmax.\n2. The authors tested their method on multiple models, datasets, and task types.\n3. The paper includes a custom CUDA kernel implementation to strengthen the applicability of the method.", "weaknesses": "1. The two-step training increases the complexity of applying the approach\n2. Dependency on multiple hyperparameters/thresholds.\n3. The paper keeps claiming that existing methods focus on existing sparsity in pre-trained models, but inducing sparsity in training (even in LLMs) is not new and has been around for a while.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are glad that the reviewer found that our work has extensive evaluation and our custom kernel increases the applicability. We thank the reviewer for the constructive feedback and appreciate the opportunity to address the points you have raised.\n\n---\n\n> **Q1:** The two-step training increases the complexity of applying the approach.\n\n> **Q4:** Complexity overhead of the training. Since you are introducing a different training approach, how does the total training time compare to the baseline?\n\n**A1&A4:** Compared to the two post-training baselines, even though LTE introduces additional training overhead, we argue that LTE training is a one-time effort. Since LTE significantly saves the inference overhead, This approach is beneficial in the long run for serving LLM more efficiently.\n\nMoreover, as we discussed in Q4 of the General Response. Even if we further fine-tune the entire model MoEfied with Deja Vu for the same training time as LTE models, the Deja Vu models fail to achieve comparable performance to LTE models\n\n>**Q2:** Dependency on multiple hyperparameters/thresholds.\n\n**A2:** We agree with the reviewer that LTE training introduces additional thresholds and hyperparameters. However, we argue that those additional thresholds and hyperparameters do not increase the training complexity of LTE. \n\n**1)** To avoid manually selecting thresholds, we design the separability loss (Eq. 5 in the paper) to make models threshold-aware. This separability loss encourages router outputs to diverge from a predefined threshold (Figure 3 in the paper), which enables us to use the same threshold for all evaluations. In our paper, we set the threshold to 0.5 for all models and tasks, as 0.5 represents the midpoint for the sigmoid output. \n\n**2)** For the hyperparameter ($\\lambda$ in Eq. 6) for the separability loss, our ablation study (Figure 11 in the paper) shows that this parameter is not sensitive to performance. Once it is increased to a certain value. It won’t introduce performance gains. We set this hyperparameter to 0.5 for all models and tasks. \n\n**3)** The only sensitive hyperparameter is $\\eta$ in Eq.6, which is necessary to control the sparsity level of the model (like the number of selected expert $k$ in traditional MoE models).\n\n> **Q3:** The paper keeps claiming that existing methods focus on existing sparsity in pre-trained models, but inducing sparsity in training (even in LLMs) is not new and has been around for a while.\n\n**A3:** We thank the reviewer for pointing out the confusion in our claim. The existing methods that we discussed in the paper stand for *MoEfication methods transforming a pretrained dense model into MoE models* rather than traditional MoE training.\n\nAs we discussed in the General Response Q1, MoEfication presents unique challenges, and our evaluations show that the MoE training baselines (such as noisy-top K softmax, sigmoid-MoE, and fine-tuning Deja Vu models) all underperform LTE. As far as we know, our paper is the first work to explore the sparse training upon a pretrained model. \n\n> **Q5:** Can you provide more insights into the limitations of Softmax routing?\n\n**A5:**  The limitations of Softmax routing in MoEfication come from the fact that the sum of the softmax output is one. In traditional MoE models, only a few experts are selected (typically fewer than four), resulting in expert scores that are not too small. However, in MoEfication, a much larger number of experts can be chosen. When the sum of the expert scores is distributed among all experts, each expert may receive a very small score (almost zero in many cases). Since the expert outputs are multiplied by these scores, extremely small scores will scale the expert outputs to very small values, affecting inference performance. \n\n> **Q6:** Many of the presented insights are already in the MoE literature, except for the Sigmoid routing.\n\n**A6:**  Although some aspects of LTE design have been discussed in previous work, MoEfication introduces unique challenges (as detailed in General Response Q1). Directly applying existing MoE techniques to the MoEfication scenario is non-trivial. \n\nOur evaluation shows that directly applying MoE training methods, such as noisy-top K softmax, sigmoid-MoE, and fine-tuning Deja Vu models, does not outperform the proposed LTE methods. This result indicates that LTE is more effective for the MoEfication scenario and provides a strong baseline for this area.\n\n> **Q7:** Limited number of baselines (using only two baselines, there is a lot of work about sparsity in LLMs). Imo, the paper would be much stronger if it compares against more (sparsity) baselines.\n\n**A7:**  In addition to the Deja Vu and MoEfication baselines reported in Section 5, we also compare the performance of the noisy top-k softmax routing, which is the most common MoE training strategy, in Section 3.2. Due to the training collapse issue of the softmax router, we did not conduct large-scale evaluations on other tasks.\n\nMoreover, as suggested by Reviewer #1 (zH4X), we evaluated and compared **three additional baselines**: Sigmoid-MoE and Deja Vu with fine-tuning (Figure 1 in the rebuttal PDF); Model Pruning Wanda (Table 1 in the rebuttal PDF). The evaluation results show that LTE still outperforms these baselines. A potential reason for LTE's better performance is that MoEfication presents unique challenges compared to traditional MoE training, which are better addressed by LTE. *(we kindly refer to the General Response Q1 for more details).*\n\nWe hope those new baselines can address your concerns.\n\n---\n\nThanks for the attentive reading of the manuscript and constructive feedback. We will incorporate these changes into our final version.\n\nWe hope our response addresses all the concerns and that the reviewer will consider raising the rating accordingly. We are more than glad to answer any further questions.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8wsZ3DKRjB", "reviewer": "Reviewer_BZ5G", "summary": "This paper proposes a method to increase the structured sparsity of models through training, called Learning-to-be-efficient (LTE). It introduces a new training loss that guides the model to activate fewer neurons while maintaining original performance. Simultaneously, LTE employs a threshold-based sigmoid routing strategy, allowing flexible expert selection instead of a predefined fixed number. To achieve acceleration, the authors further implement an efficient version based on Triton. Compared to previous sparse acceleration methods, LTE can be applied to activation functions other than ReLU, offering better generality.", "questions": "- Regarding the experimental analysis in Line 135: Since MoEfication approximates dense model computation, expert scores are only used for expert selection and do not scale the expert outputs. This ensures that when all experts are selected, the computation is entirely equivalent to the original dense model. Therefore, it should not be possible to have a situation where two experts are selected but only one contributes. Could the authors further elaborate on this phenomenon?\n- Concerning the expert grouping method, is there room for further improvement in clustering W_1? LLaMA uses GLU, where the intermediate representation is obtained by element-wise multiplication of two vectors, which is different from the vanilla FFN studied in the original MoEfication.\n- What principles guide the selection of thresholds, and is there any transferability to other LLMs?", "limitations": "Yes, the authors have adequately addressed the limitations.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The method demonstrates good generality and can be applied to various activation functions.\n- The designed separability loss is intuitive and effective.\n- Experiments on diverse types of tasks prove the effectiveness of the method.\n- Implementation of a Triton kernel achieves computational time acceleration.", "weaknesses": "- Due to the mainstream dense models having an FFN to attention ratio of approximately 2:1, although good sparse acceleration is achieved in FFN, the overall speedup is not high. If experiments were conducted on models with a higher proportion of FFN, the acceleration effect would likely be better, potentially enhancing the impact of the work.\n- Some experimental and methodological details lack sufficient discussion. See the questions section for more details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are glad that the reviewer found that our work has extensive evaluation and shows strong empirical performance. We thank the reviewer for the constructive feedback and appreciate the opportunity to address the points you have raised.\n\n---\n\n\n> **Q1:** Due to the mainstream dense models having an FFN to attention ratio of approximately 2:1, although good sparse acceleration is achieved in FFN, the overall speedup is not high. If experiments were conducted on models with a higher proportion of FFN, the acceleration effect would likely be better, potentially enhancing the impact of the work.\n\n**A1:**  We agree with the reviewer that only improving sparsity for FFN layers does have a speed-up limit. Applying LTE in a more FFN-intense model can gain more speed-up. Another potential solution is to apply MoEfication in attention layers. We will leave this for future exploration.\n\n>**Q2:** Regarding the experimental analysis in Line 135: Since MoEfication approximates dense model computation, expert scores are only used for expert selection and do not scale the expert outputs. This ensures that when all experts are selected, the computation is entirely equivalent to the original dense model. Therefore, it should not be possible to have a situation where two experts are selected but only one contributes. Could the authors further elaborate on this phenomenon?\n\n**A2:** The reason multiple experts are selected but only one contributes (Line 135) is due to the adoption of the noisy top-K softmax routing strategy in Section 3.2. With this routing strategy, the expert outputs are multiplied by their corresponding expert scores. Given the large number of experts in MoEfication and the fact that the sum of the softmax expert scores is 1, many of the selected expert scores are nearly zero. When those expert outputs are multiplied by the nearly zero scores, the outputs are scaled to zero, resulting in no contribution to the inference.\nWe hope that this explanation can address your concerns, we will polish the writing of Section 3.2 to make this point more clear.\n\n>**Q3:** Concerning the expert grouping method, is there room for further improvement in clustering W_1? LLaMA uses GLU, where the intermediate representation is obtained by element-wise multiplication of two vectors, which is different from the vanilla FFN studied in the original MoEfication.\n\n**A3:** For LLaMA, our current design uses the parameter clustering on the gate matrix (whose output feeds into the activation function) to group the neurons. In our early evaluation, we also tried to use another matrix (up-matrix) to group the neurons. However, we observed no significant difference in performance between the two approaches. As we discussed in Figure 9 in the paper, another strategy (co-activation) also shows similar performance. Our hypothesis is that, since LTE updates the model weights and routers, the model can adjust the grouping to some extent, even though the initial grouping is not optimal. However, if the initial grouping is too bad (like random), the model may not be able to make such a large adjustment.\n\n>**Q4:** What principles guide the selection of thresholds, and is there any transferability to other LLMs?\n\n**A4:**  In our evaluation, we set the threshold to 0.5 for all models and tasks, as 0.5 represents the midpoint for the sigmoid output.  As shown in Figure 3 in the paper, we observed that the separability loss creates a significant gap between the pre-set thresholds. we believe that a threshold around 0.5 can also work well.\n\n---\n\n\nThanks for the attentive reading of the manuscript and constructive feedback. We will incorporate these changes into our final version.\nWe hope our response addresses all the concerns and that the reviewer will consider raising the rating accordingly. We are more than glad to answer any further questions.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
## General Response

Dear reviewers,

We thank all the reviewers for their constructive reviews towards improving our work. 

We are pleased that reviewers found our paper’s advantages: “LTE constantly achieves better performance-sparsity trade-off across multiple models, datasets, and task types.” (All reviewers); “The customized kernel improves the soundness of the work.” (All reviewers).

For the comments and concerns discussed in the reviews, we write this general response to address some common concerns and separate responses to each individual review. 

---

> **Q1:** The challenges of MoEfication compared to traditional MoE training.

**A1:** MoEfication presents three unique challenges:

**1) Router designing for a larger number of selected experts:** Unlike traditional MoE designs, which typically select a small number of experts, MoEfication can choose a much larger number of experts, which leads to softmax routing collapse (as discussed in Section 3.2), which lead to commonly used Softmax routing collapse in our empirical evaluation.

**2) Router training on pretrained models:** MoEfication is based on a pretrained model, unlike traditional MoE models that are trained from scratch. Before MoEfication, the outputs of the grouped experts are not scaled in the pretrained model, but typical MoE designs need to scale outputs with expert scores to make the router differentiable. The scaling of expert outputs in a pretrained model can hurt performance, even with fine-tuning.

**3) Adaptive sparsity in different layers of the pretrained model.** Recent work [2] shows that different layers in a pretrained model have different levels of sparsity. Traditional MoE designs use a predefined k to set the sparsity for each layer, which prevents adaptive sparsity across different layers.

Those challenges make it non-trivial to directly apply MoE training for the MoEfication problem.

> **Q2:** The novelty and contribution of our paper.

**A2:** While some MoE-related concepts have been discussed in the MoE literature, as far as we know, LTE is the first method to address all three aforementioned challenges in MoEfication. 

**1) Efficiency-aware Sigmoid router:** We adopt a sigmoid router and use an efficiency-aware loss to introduce competition among experts, thereby avoiding the collapse issue associated with softmax routing, which address the first challenge.

**2) Indicator function with two-stage training:** We use an indicator function to select experts without scaling expert outputs. To address the non-differentiability issue of the indicator function, we proposed a two-stage training algorithm to jointly train the model and routers. This address the second challenge.

**3) Threshold-based router with efficiency loss:** We utilize a threshold-based router, which allows for a more adaptive selection of experts in each layer. The efficiency loss introduces competition across different layers, leading to adaptive sparsity in different layers (see the ablation study in Figure 15 in the Appendix), which addresses the third challenge.

Besides handling those three challenges, another novelty of LTE is the introduction of inference efficiency as an optimization goal. This approach encourages models to only use necessary parameters for inference by training, which is different from the common practice of predefining a sparsity level in typical MoE.

> **Q3:** Additional baseline comparison: Sigmoid-MoE[1]

**A3:** As suggested by Reviewer #zH4X, we implemented the Sigmoid-MoE to test its performance on MoEfication tasks. We applied Sigmoid-MoE to the GPT-2 Medium model and fine-tuned it with the WikiText-103 dataset. The Sigmoid-MoE model was trained for the same training time as the LTE models. The comparison results are presented in Figure 1 in the rebuttal PDF. While Sigmoid-MoE overcomes the collapse issue of the softmax router, it still underperforms LTE. 

> **Q4:** Additional baseline comparison: Deja Vu fine-tuning

**A4:**  Deja Vu and MoEfication are proposed as post-training methods in their papers. Following this, we implemented Deja Vu or MoEfication in a post-training manner in our paper: we first fine-tuned models with the specific datasets and then applied Deja Vu or MoEfication to the fine-tuned modesl to evaluate performance.

However, we understand the reviewers’ concern that fine-tuning those MoEfied models (all parameters) could further improve the performance of the baselines. To better study this problem, we conduct an evaluation on fine-tuning the Deja Vu model (all parameters) to compare the performance. We first fine-tuned the model with Wikitext-103 and applied Deja Vu. Subsequently, we further fine-tuned the Deja Vu model with Wikitext-103 for the same training time as LTE training. Note that the first fine-tuning is necessary for Deja Vu to collect data to train predictors.

The evaluation results are reported in Figure 1 of the rebuttal PDF. We find that while the additional fine-tuning does improve the performance of Deja Vu, LTE still outperforms Deja Vu with fine-tuning.

---

To sum up, MoEfication presents unique challenges, and it is non-trivial to directly apply MoE techniques to MoEfication. Considering the novel design and strong empirical performance of LTE, we believe that LTE presents a strong baseline for future study. 

We agree with the reviewers that the paper writing can be further improved, and we believe that those issues can be well addressed in our final version.

---

[1] "Approximating Two-Layer Feedforward Networks for Efficient Transformers," R. Csordás et al., EMNLP'23.

[2] "The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers." Li, Zonglin, et al., ICLR’23.

Meta-review: Reviewers agree that the paper delivers a significant real-world improvement in LLM inference efficiency by training a structured sparsity which means that even relatively low sparsity (say 50%) is beneficial and leads to wall-clock improvements in latency and runtime.

The approach combines existing ideas (sigmoid MoE, MoEfication, sparsity-encouraging losses) but in a novel way, and has the potential to open new areas of related work.

Reviewers are concerned with complexity of the training process, but the rebuttal argues effectively that this complexity is justified by the improvement in inference efficiency.

Reviewers also note that applying to the FFN layers only does limit the opportunity for end-to-end improvement, but as the rebuttal notes, this is potential future work.   Publishing the paper now will allow other researchers to benefit from the ideas.
Decision: Accept (spotlight)