$\textit{Read-ME}$: Refactorizing LLMs as Router-Decoupled Mixture of Experts with System Co-Design
Authors: Ruisi Cai
Yeonju Ro
Geon-Woo Kim
Peihao Wang
Babak Ehteshami Bejnordi
Aditya Akella
Zhangyang Wang
Venue: NeurIPS 2024 poster (2024)
Keywords: MoE
Machine Learning
Efficient Inference
Mixture-of-experts
LLM
MoE Inference
Efficient serving
LLM serving
Primary area: deep_learning_architectures

Abstract: The proliferation of large language models (LLMs) has led to the adoption of Mixture-of-Experts (MoE) architectures that dynamically leverage specialized subnetworks for improved efficiency and performance. Despite their benefits, MoE models face significant challenges during inference, including inefficient memory management and suboptimal batching, due to misaligned design choices between the model architecture and the system policies. Furthermore, the conventional approach of training MoEs from scratch is increasingly prohibitive in terms of cost. 
In this paper, we propose a novel framework $\textit{Read-ME}$ that transforms pre-trained dense LLMs into smaller MoE models (in contrast to ``upcycling" generalist MoEs), avoiding the high costs of ground-up training. Our approach employs activation sparsity to extract experts. 
To compose experts, we examine the widely-adopted layer-wise router design and show its redundancy, and thus we introduce the pre-gating router decoupled from the MoE backbone that facilitates system-friendly pre-computing and lookahead scheduling, enhancing expert-aware batching and caching.
Our codesign therefore addresses critical gaps on both the algorithmic and system fronts, establishing a scalable and efficient alternative for LLM inference in resource-constrained settings.
$\textit{Read-ME}$  outperforms other popular open-source dense models of similar scales, achieving improvements of up to 10.1\% on MMLU, and improving mean end-to-end latency up to 6.1\%. 
Codes are available at: \url{https://github.com/VITA-Group/READ-ME}.

Read-ME: Refactorizing LLMs as Router-Decoupled Mixture of Experts with System Co-Design
Ruisi Cai*!, Yeonju Ro Geon-Woo Kim Peihao Wang', Babak Ehteshami Bejnordi', Aditya Akella', Zhangyang Wangb
'The University of Texas at Austin, "Qualcomm Al Research fruisi.cal) gukim, peihaowang, atlaswang/Hutexxs.. edu, fyro, akella) utexas. edu, behteshadqti quw..omm..cm.mmmmmmm com
Abstract
The proliferation of large language models (LLMs) has led to the adoption Mixture-of-Experts (MoE) architectures that dynamically leverage specialized subnetworks for improved efficiency and performance. Despite their benefits, MoE models face significant challenges during inference, including inefficient memory management and suboptimal batching, due to misaligned design choices between the model architecture and the system policies. Furthermore, the conventional approach training MoEs from scratch is increasingly prohibitive in terms of cost. In this paper, we propose novel framework Read-ME that transforms pre-trained dense LLMs into smaller MoE models (in contrast to "upcycling" generalist MoEs), avoiding the high costs of ground-up training. Our approach employs activation sparsity to extract experts. To compose experts, we examine the widely-adopted layer-wise router design and show its redundancy, and thus we introduce the pre-gating router decoupled from the MoE backbone that facilitates wystem-friendly pre-computing and lookahead scheduling. enhancing expert-aware aware batching and caching. Our codesign therefore addresses critical gaps on both the algorithmic and system fronts, establishing scalable and efficient alternative for settings. Read-ME outperforms other
LLM inference in essorcce-onnstrinee popular open-source dense models of similar scales, achieving improvements of up to 10.1% on MMLU, and improving mean end-to-end latency to 6.1 6.1%. Codes are available at: https://githut hwt.ggttth.ooo..o.eeeeoooooootee cww//om//////// VVTA -rr--Grroouure READ-ME
Introduction
The success of Mixture-of-EEpprts (MoE) 2 such as recently exemplified by the Mixtral model 3 in the era of large language models (LLMs)- lies in its remarkable ability to leverage the collective expertise specialized sub-networks, or "experts," " each proficient in handling specific subsets or aspects of the data. By dynamically routing data through these experts, MoE models effectively capture complex patterns, adapt to diverse data distributions, and offer superior predictive accuracy compared to traditional mingle-model approaches. In addition to performance promise, MoEs also have a natural appeal for resource- resource-limited devices due to their high sparsity, and therefore reduced
activated parameters per token, which can potentially save inference costs [4] 6 71.
However, MoE inference presents significant challenges for key system- level objectives:
Memory Management: Although Mořs activate only a subnetwork during inference, expert
selection determined on the fly by layerwise router, complicating efficient prefetching. This often forces reliance on naive prefetching algorithms. For example, prior work has
'Equal contribution: authors are listed alphabetically. Akella and Z. Wang also advised this work equally.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Pregattd spot
kpuda
Rouser
roulered Taken Excert
Expert 1
Reyuwa Gudve
Rkyww Qakut
Expert 2 Full
FiFFFIFFÍ
Takket rastened to Expoet
Pre-trained Dense Model
Refactoring
Depleyment
Inference Serving
Inference
Serving
time=0
at time11
Figure I: Overview of Read-ME. This figure shows the refactoring of pre-trained dense model (in yellow) into two experts (in red and green). After refactoring, the model deployed, and the serving timeline depicted. At time = 0, multiple inference requests (each a sequence of tokens) are queued, with expert assignment for each token undecided ("?") until processed by the router. Our router pre-gates tokens before inference, enabling expert-aware batching. Tokens are routed to their respective experts and batched accordingly: at = 0 for Expert I (red) and at t = for Expert 2 (green). New tokens enter the queue each time step, with routing computed only for incoming tokens marked "?".
predicted the next expert using hidden states from the previous layer and applied LRU cache replacement for recently used experts (8). While effective under certain conditions, such strategies depend on assumptions about expert locality and token predictability, which can become sub-optimal if those assumptions are violated (as shown in Table 4
Token Batching: Token batching techniques critical for efficient inference (e.g.. 9 are poorly suited to MoE architectures, where each batch contains tokens invoking different experts across layers, rendering batching strategies ineffective ($ 4.2).
Morcover, traditional MoEs are typically trained from scratch, which becomes prohibitively expensive as model scales increase. To mitigate this, some approaches, such as "upcycling" [10], reuse pre- trained dense LLMs to initialize experts in an MoE. While that efficiently scales MoEs by leveraging smaller, pre-trained models, it does not address the ffferencerelattee challenges mentioned earlier.
In this work, we tackle the opposite challenge: how to create a smaller MoE model from larger pre-trained models that enables resource-efficiens inference while minimizing training costs? Despite existing efforts 12 13 14], this problem remains underexplored Approaches like [1] 12 13
attempt MoE refactorization but still adopt unnfinl- layer-wise structures for inference. Similarly, [14] focuses on dynamically selecting "important" neurons during pre-filling and pruning others during generation, but this limited to long-content generation and requires neuron importance identification for each sequence.
To address both training and inference challenges, we introduce holistic MoE framework dubbed Read-ME. To minimize training costs, we "refactorize" a ppr--triied dense LLM into specialized experts through activation sparsity and optimize the routing policy ($ 3. For efficient inference. we examine the redundancy of layer-wise routers ($ 2.1 $ 2.2) and propose decoupling the router from the MoE backbone ($ 2.3). This allows us pre-gate all requests (token sequences) before inference and apply lookahead scheduling based on the experts to which tokens will be dispatched. Consequently, we propose expert-aware batching algorithm ($ 4.2) and an optimal expert caching
strategy inspired by Belady's 's offline caching algorithm [1
4.1
Figure illustrates framework. Our key contributions are:
We transform large pre-trained LLMs into Mixture-of-Experte (MoE) models with fewer activated parameters and small additional training cost (1 billion tokens). Our approach outperforms popular open-source models and compression techniques of similar scale on downstream tasks like MMLU [16].
We analyze the widely adopted layer-wise routers in existing MoEs and reveal design redundancies. Current caching policies and batching algorithms poorly suited to layer- wise MoEs. We propose a novel pre-gating router, decoupled from the MoE backbone, enabling better system-level optimization.

---

Our system achieves 6.1% reduction in mean latency and 10% improvement in tail latency compared to state-of-the-art systems. Our caching algorithm is both provably and empirically optimal, thanks to our algorithn-system co-design.
2
Pre-gating Sparse Mixture of Experts
In this section, we introduce our motivation and design of pre-gating MoB which enables system-level acceleration by sharing and precomputing expert selection for each layer.
2.1 System Drawbacks of Conventional Sparse MoE Design
An Mfxture-of-Expert (MoE) 17 layer consists of a routing network G and set of N expert networks { [Fi.... Fy J. In the forward pass, the routing network will first process input sequences and generate the gating weights. Then size-K K subset of experts will be dynamically activated and their outputs will be combined as final outputs according to the gating weights. In LLMs, MoE typically adopted the Forw-Forward Networks (FFN) within each transformer block 2 Suppose an LLM has L layers, the output of the 1-th layer can be formulated as:
N
(x), (z)ii) S (x),
(1)
where the superscripts indicate the layer indices, GO), F0) are point-wise functions operating tokens individually, and T(.) is the indicator function which filters experts with top- gating weights.
For shorthand, we denote 15") I(H E [N] G (x); 2 (z)))) S K).
As shown in Eq. D conventional MoEs assign a separate router to each layer. While this is commonly used by open-source MoEs like Mixtral K and OpenMoF [18), we highlight its system inefficiency. Layer-wise gating makes difficult to predict which expert to load until runtime ($ 4.1). and complicating request batching ($ 4.2). Specifically, layer-wise routers select the (-th layer expert i n based on the 1)-th layer outputs, which prevents pre-scheduling and pre-loading of data or model weights. This issue is especially problematic for billion-level parameter MoEs, where experts are usually distributed across devices (GPUs and CPUs in a machine) or even machines; in such situations, layer-wise selection accentuates high overheads of data I/O and communication among servers in the critical path inference.
2.2 Redundancy of Layer-wise Router
In this section, we demonstrate that layer-wise gating patterns are redundant in an MoE. In particular, we empirically find that expert selections between two adjacent layers are highly correlated.
Token
Roder Stat Laree
MISO
Expert
Expert
(b)
(c)
Figure (a) Visualization transition matrix between (l-1)-th layer and the 1-th layer, where each coordinate [ss, t), (i, represents P(S()) [i,j) = (s, f)). The row-wise sparse pattern suggests
that the router decision becomes almost deterministic given the previous layer's decision. (b) Mutual information ((8(0)00-)) SU-D), which indicates the learned knowledge shared by two neighboring layers is high. (c) Overview figure of router tuning and router distillation loss.
We use Mixtral-8x 7B (N 8, K = 2) B as a study case and analyze router decisions among its layers. Define the random variable (())
{ie [N] 1 1} the pair of experts selected
3

---

for each layer (1S4) 2). We are interested in the conditional probability of S()) between two consecutive layers: P(S = (i.j)]stt = ).). The transition matrix of the last two layers from Mixtral-8xx x7B is depicted in Figure 2 (a). The row-wise sparse pattern implies that the expert selection is almost deterministic given the previous layer' choices. For example, for tokens choosing expert-3 and expert-5 in the 30th layer, over 70% will select expert-1 and expert-5 at the 31st layer.
To further validate this observation, we plot the mutual information between expert choices of every two neighboring layers: S())). As reflected in the right sub-figure of Figure 2 knowing expert pairs used in the last layer significantly reduces the uncertainty of the next layer. Thus, the implicit knowledge learned by each router is extensively shared across layers.
2.3 Pre-Computed Routing Policy
The above observations suggest that among the many (x) routing paths, only a few are used during the inference. Therefore, layer-wise routing decisions are unnecessary for MoEs. Instead, we can separate the routerfrom the MoE backbone and pre-compute the routing path all at once.
First of all, we assume the indices of experts handling one domain of tokens are aligned, i.e. (E .... r!!")) } always forms routing path. We defer our approach to the construction of aligned experts to $ Next, we let a singleton network G generate gating weights for all layers. In particular, adopt one transformer block with causal attention as the model architecture of G. Gating weights computed in this way not only leverage the states of the current token but also take the information from the past tokens into consideration. Thus, tokens will have expert selections similar to the recent tokens, which ensures cache- friendly inference (see more details in $ $15.)).
Suppose the input sequence is (I1)t=1,-... ,T, the output for the t-th token the 1-th layer is:
N
y. = Σ e ( 2 G(as))d) 1=1
()
where Zgz = (. xx)) represents all the tokens before the t-th token. We note that G is independent of layer index l. Despite subtle change, it brings profound benefits enable system- level optimization. In brief, by separating the gating network from the transformer layers, expert selection can be determined at the outset and used to schedule the data-loading procedure for each layer. We defer more details on system co-design to $
3 Re-factoring Language Model with Pre-Gating MoE
In this section, we introduce the main technique to re-use a dense pre-trained model to construct our pre-gating MoE proposed in 2 In short, our approach first initializes each expert by structured pruning of a dense model on the corresponding data domains. Afterward, we instantiate gating network shared across layers and continue joint training of the router and experts.
Domain- Aware Expert Construction. We construct set of small experts by pruning the dense model with different data domains. To begin with, we point out that public language corpora often contain metadata indicating the domain of each subset. For example, the training dataset of LLaMA family [19 can be split into scientific articles [20), novels (21), and QAs (22), etc. We utilize this metadata to group data entries in the training corpus into sub-domains (DI.... Dx J. Observing that feature channels of a channel on each subset and keep top activated neurons to form the domain expert. Formally, let the number of experts equal to the number of sub-domains, and assume the dense model is two-layer FFN with hidden size D: (x) = W W2o(W1a)) then the i-th experts with hidden size d are initialized as: F((x) = W₂ M, a M W æ), Vi E [W]. in which M, is obtained by:
each subset sparsely activated [23, we compute the average magnitude
arg max EN-D, MW æ=|| M0
s.t. M1, 1, M 1a M
where M is constrained to be selection matrix without replacement. The mask for each layer jointly optimized so that the resultant experts are aligned layerwise and dedicated to the same data distribution. In our experiments, we set d se D/2. In addition, observe that certain

---

subset of channels is essential for all data, potentially due to the system prompt and the presence of commonsense knowledge. Therefore, we isolate the corresponding neurons as the permanent expert. which will be activated for all tokens, similar to previous designs [8, 24
Continual Training Objective. After initializing experts via structured pruning, perform joint training of randomly initialized gating networks and expert subnetworks via causal language modeling. In addition, we propose routing distillation loss to enhance the alignment between expert choice in pre-gating MoE and the activation sparsity in the original dense model.
We illustrate the training of our router in Fig. (c). Suppose the predicted token has embedding #++1- We feed æ++1 into the original dense model Fo and get a sparse selection matrix M, that indicates neurons with top 50% magnitude similar to Eq. Then we penalize this loss function:
CRD DKL (sottmax((((34))) | (o|t))/((mmma((MmMaaaaa M, 11:: ... M.MMMN)
(4)
1. M,M, 1a computes the Hamming distance between two masks induced by Mo M, We apply softmax to normalize these scores as the estimated selection probability of each expert for predicted token t++1-
Here, --1.(-1l-) represents Kullback-Leibler divergence. || [M
4 Expert-aware Inference System
We demonstrate how our refactoring and pregating concepts enable a novel, high-performance, and efficient MoE inference method. We address two key challenges in existing MoE models' inference: inadequate memory management and limited support for batched inference. Our problem setting is broad, aiming to serve multiple requests using an MoE model, each comprising sequence of tokens. This differs from previous systems, which focused on optimizing performance for individual requests.
4.1 Pre-gating Optimized Expert Prefetching and Caching
MoE models promise reduced memory usage during inference by loading only the parameters of required experts, skipping the rest. However, traditional layer-wise gating imposes significant loading costs. Previous approaches, such as on-demand loading 225 prefetching (26). and expert caching R, 27). attempt to address this. However, on-demand loading adds overhead to the critical inference path, and prefetching often loads unnecessary experts due to incomplete routing information, leading to suboptimal memory usage and performance (28). Additionally, caching strategies, based on request characteristics like temporal locality or activation sparsity, have mostly been evaluated in isolated single-request scenarios. In practice, expert caches are shared across multiple requests, making cache policies relying on per-request traits suboptimal. A global view across all requests is necessary for effective caching (see Table 4. Our work leverages pre-gating to develop more informed prefetching and caching strategies, resulting in significant system-level improvements.
Fine-grained Prefetching, By design, our pre-gating MoE architecture enables us to prefetch the exact expert layers needed for token or a request, avoiding guesswork. To further hide the latency in prefetching, we pipeline and thus overlap loading of experts and experts' computation at layer-wise wise granularity: specifically, while computing the ith layer's forward path in the compute stream, we load the i + 1st layer's is experts in separate loading stream.
Belady-inspine Caching. Prefetching can hide the loading latency of all but the first layer, which incurs significant cost. To mitigate this, we need a cache that stores relevant initial layers, and we argue that pre-gating enables an optimal caching strategy.
The classical Belady algorithm is known to be the optimal offline cache replacement algorithm, replacing the object that will be accessed farthest in the future. While impractical in real-world systems (due to unknown future accesses), our pre-gating architecture allows us to approximate By decoupling the router from the backbone MoE, requests in advance, enabling near optimml cache replacement.
compute future expert references across
- 1) = (ei, e2, es }. where the cache of size k and is filled with & experts e1...- F(e, () represents the next time after when expert will be requested. Then, our policy chooses the expert Cevict argmax.ecct-n
Suppose that the cache at time step / - 1 is follows: C(l-
for eviction.
5

---

Batchhig Happprs
Request
Queve
Roli
Expert
Roda
Expar
Roule
Expert
Batched!
Mol
UhigyAclvvsed Exports)
Batching Happern
Pre-garing
Lest
Expett
Gesrage maaners Exeat Bu butch
Figure 3: Challenges of MoE serving in current serving systems and Read-ME's hatching pipeline.
4.2 Expert-aware Batching
Current serving systems heavily rely batching improve inference efficiency, but effective batching for MoE models remains challenging. As shown in Figure (a), each token in MoE models of experts per layer, leading
multiple expert activations for batch of requests. For example, in a toy model with 4 experts per layer and batch of 3 tokens (one per request), 2/3/3 experts would be activated across the layers. In the Mixtral8x7B model [3] applied the chatbot arena dataset [29], we observed an average activation of 7.63 out of 8 experts, even with modest batch size of 56.8.
may invoke different
The core challenge is that while each token requires computation from only one expert per layer, must wait for all other tokens in the batch to complete their expert computations in the same layer (30) This bottleneck repeats at each layer, reducing the efficiency of batching. Ideally, a single loaded expert would serve multiple tokens in batch, but this is rarely achieved, affecting both performance and efficiency. For example, we observe linear increase in average per-token processing latency as the number of unique experts per batch grows (see Figure 3 (b)).
In contrast, pre-gating enhances inference performance by enabling the delayed creation of an optimal batch based required experts. For given set of tokens, pre-gate each and select subset for batching, depending on their identified expert requirements. The goal is to minimize the number of unique experts across all layers while maximizing the number of tokens in the batch. Morcover, as discussed in $ $2.3. our expert selection remains consistent across layers-iff a token is assigned to Expert 1, it will be routed Expert in every layer. This approach, combined with our batching strategy, ensures optimal efficiency. Algorithm I provides our batching pseudocode.
We note that in other Moïs, such batching isn't feasible because, shown in Figure 3 their expert selection each layer remains unknown until the request reaches the router. In Read-ME, experts are determined first, which allows batches to be created and submitted to MoE layers efficiently.
5 Evaluation
In this section, we start by describing the experimental details $5..T] Then we validate the refactorization effectiveness on downstream tasks in $ 5.2] In $ [.3] we evaluate the effec- tiveness of pre-gating and batching. 5.4 analyzes memory optimization techniques. In addition, provide experimental details in $ 5.1 and more experimental results in $. A
Table I: Details of router design. Fol-
lowing the standard Transformer archi- tecture [I), inserted router adds only 18 million additional parameters.
# Layers Heads
Vocab size
32000
5.1
Experimental Details
Embedding Dim.
512
Feature Dim.
MLP Intermediate Dim.
512
Model and Dataset We perform the MoE refactorization
Activation Function
SwiGi.U
32
based on Llama2-7B-chal [19] model, a popular open-source model pre- pre-trained on 2 trillion tokens. The training cor-
Positional Embedding Normalization
35 RoPE
RMSNorm
34
pus 35 involves the data collected from different resources: Arxiv 20), Books [21]. Common Crawl, C4 (36], Github,
Parums
18.0 M
Wikipedia [37] and StackExchange [22]. To generate experts, we collect 16 samples from each
data domain, with each sample consisting of 4096 consecutive tokens. During router tuning, we
use the subset of RedPajama dataet (35), with the same curation strategy. We present our detailed
router design in Table
We the standard Transformer [31] architecture with I-layer, 4-head

---

Algorithm I Read-ME Expert-aware Batching Algorithm (pseudocode)
Input NumExperts, RexQueuueyExpert, Max Tex Output ScheduledReq
1: for + to N NumErperts do
2:
len_regs_per_eeeeee cacss_per_epppeee _et__peer_erpps_lllpee_ (enn(eeeuucBEEEperrii
3: end for
4: while true do
5: E 4 orgmaa((en_ress__er_ceaa
6:
if faan_reg__ae let_rreggg_________ per__eperts_[]] <Marl Len - (en(SchedulledRe)) then
7:
Scheduled/Re + Scheduled/teq U BeqQueeeeByBxxerrf.
8:
ReqQueucByEExperle +-]
9:
else if pea__r___pe___peeeee +
2 then
n_available MazT'okenLen (en(ScheduleddRe))
12:
Seh + ScheduledReq U Byerpeerr n available)
13:
ReqQueeu t- BeeQQueeeeBBBEaperrrrr | :] EE] +- Ien Bew.RRQuueeBBBppeerrre
14:
len_regn pepperrerrsssseeeeeeeeepee___sse
15:
break
16:
else break
17:
end if
end while
design. The router is lightweight, consisting of 18 million additional parameters, and incurs negligible computational overhead. We use 8 A100 GPUs with 80GB of memory for all tuning experiments.
Continual- Tuning Details To co-optimize the router and expert networks, we iteratively tune each model component. Specifically, we first optimized the router by LRD. as detailed in $3. for 100 steps. We use the batch size of 64 in this router tuning stage. During this router tuning stage, we freeze the expert weights and solely tune the router weights. Then, during the ex- pert tuning stage, fix the router weights and modify the expert weights via language model- ing loss, for 200 steps, with batch size of 1 28. We set sequence length to 4096 for all stages, following the choice in the pre-training stage of Llama2 model [19]. This iterative training schedule conducted times. Detailed visualizations of the training dynamics are provided Section A.I For each round, the router tuning and expert tuning stages will cost 26 million and 105 million tokens, respectively. The whole continual-tuning process merely uses 1.04 billion tokens, negligible compared to the pre-training cost (2 trillion tokens). During each round of tuning, we use the cosine learning rate decay. At round 0, the initial learning rates are 5e 4 for router tuning and 5e 5 for expert tuning. The initial learning rate decays exponentially with a decay rate of 0.8 as the number of rounds increases.
Table 2: Hyper- Jyper-parameter choice during training.
Stage
Router Tuning
Expert Tuning
# Iteration per Round # Rounds
100
200
8
Initial LR at Round IR Decay within Round
5e: Cosine
Cosine
LR Decay type across Rounds
Exponential
Exponential
LR Decay rate
Rounds
0.8
0.8
Weight Decay
0.01
0.01
Batch Size
128
Sequence Length Tokens per Round
4096
4096
26 M
105 M
Tokens Total
1.04
Inference System Evaluation For our workload, we utilize the Chatbot Arena Conversation Dataset (29] to generate inference requests and replay conversation traces. Our setup employs single A100 GPU with 80GB of memory. The implementation is built on top of DeepSpeed inference primary metric, defined as the end-to-end latency
engine [38]. We use normalized lateney
divided by the generated token length, in line with previous works 99 39 38
5.2 Downstream Task Evaluations
We first validate the refactorization effectiveness on downstream tasks, as shown in Table comparing it to other models of similar scales, including the open-source models that trained from scratch, and
7

---

Thster EocostMUP Others
2.1
Read - ME
Mixtral-8x76
LatencyEmn
27
22.
10
20
30
40
50
60
Lateno 1msl
Temporal Distance
Figure 5: Latency evaluation and Temporal locality analysis. (Left)) Single inference latency measured on a 124 token generation task. (Center) Latency distribution measured on synthetic workload replaying Chatbot Arena
Dataset 29] ($ 5.]. (Right) Temporal distance measured on Arxiv dataset (20). and a subset of Redpajama [35]
the dense models pruned from larger pre-trained LLMs. We achieve the best average performance, outperforming all model variants from the Pythia 40 and Open-Llama-v2 4 families, as well as Sheared-Llama 42 We use just billion training tokens, considerably less than other models.
Table 3: Downstream task evaluation of our proposed method (Read-ME) compared open-source models, including dense models Pythia and Open-Llama-v2. the MoE model OpenMoE, and the compression method Sheared-Llama, The evaluation includes zero-shot performance on WinoGrande, ARC-Easy, LogiQA, CoQA; S-shot performance on MMLU; I 10-shot on Helllawag;; and 25-shot on presents in the form of (# Activated-Parameete Parameters # of tokens used. For compression methods like ours and Sheared- Llama, only tokens used for conversion are counted, excluding Llama-2 pre- pre-training costs.
The "#Param" column Training cost measured by the number
Method
#Param
Cost
MMLU
Hell.
Wino.
ARC-E
ARC-C
LogiQA
CoQA
avg.
Sheared- Llama
-lama
2.7B
50B
26.4%
70.8%
67.0%
67.0%
41.2%
28.3%
71.7%
53.2%
Pythia
2.8B
300B
26.9%
60.8%
59.7%
64.4%
36.4%
27.7%
61.9%
48.3%
Lpen-LLlama-22
3.4B
IT
25.7%
67.6%
63.5%
66.5%
39.0%
28.1%
54.4%
49.3%
OpenMoE
2.1B-8B
1.1T
26.2%
45.5%
60.3%
64. 14
30.3%
Read-ME
4.7B 17B
IB
38.9%
68.5%
67.7%
66.6%
42.3%
29.7%
74.8%
55.5%
Pythia
6.9B
300B
25.5%
67.1%
64.1%
67.3%
31.3%
25.3%
63.6%
49.2%
Open-Llama-v?
6.9B
IT
40.2%
66.7%
66.0%
63.0%
36.0%
27.6%
64.5%
52.0%
Llama-2 6.9B
2T
45.3%
78.6%
69.3%
76.4%
53.0%
31.0%
75.9%
61.4%
In Fig. we further provide a direct compari- son with other compression methods, which con- verts large LLM to small dense variant, on MMLU [16] benchmarks. Besides open-source models and Sheared-Llama 42) which are men-
45
MMLU Perfemmnnes
40-
Open-hrn
35-
tioned the previous table, we additionally in-
30-
clude recent compression techniques, including LLM-Pruner [43], SliceGPT (44). LaCo [45],
SisaPPT
Com Audoo
Qpareirr
ALLM
and Compresso [4), as our baselines. Read- MEachieves the best performance among the models with the number of activation parameters less than billion, and shows comparable per- formance with Dpen-Llama-v2-75 [41) More analysis is included in $ A.2
Activated Parameters (B)
Figure 4: Evaluation of Read-MEon MMLU [6 benchmark, compared to other open-source models and compression techniques performance numbers are col- lected from their respective papers.
5.3 Pre-gating and Expert-aware Batching
Inference Latency Breakdown. We evaluate the impact of the auto-regressive router introduced by our refactoring of the dense MoE on per-request inference latency. Unlike conventional layer-wise routers, usually linear layers, our ruto-regressive router comprises multi-head attention layer and an MLP layer (see $ 22.), potentially raising its computational cost.
Fig. 5 (left) illustrates the average per-token latency breakdown of a single isolated inference request measured in OpenMoE [18] with conventional layerwise routers, our refactored model with pregating router, and the original dense Llama2-7b model [19] we refactored. We find that the computational

---

overhead of our auto- regrrssseeeev router is minimal its contribution of 0,4% is much less compared to the router's net contribution in other MoE models (3.95%). This is because we use single router unlike other models with gating for each MoF layer; also, our router design is compact with only 18M parameters (Table I. Compared the dense model, we achieve net 19% reduction in latency via refactoring the MLP to MoE.
Batched Inference. We now evaluate the efficacy of our expert-aware batching. Fig (center) shows the latency distribution and the 95-th percentile latency (p95) during batched inference. We compare with two widely used techniques Decoding Decodingg-riorrtizzed batching [38. and Prefill- Prefill-prioritized prioritized batching [39 47 These methods utilize distinct queues for decoding requests and prefill requests, prioritizing batching of tokens from decoding and prefill requests, respectively.
Prioritizing either decoding or prefill requests yields comparable performance. In contrast, our method of constructing batches based on activated experts enhances the mean latency by 5.0-6. 1%% and reduces the p95 latency by 9.5-10.0% compared to these approaches.
The primary reason for this improvement is that our batching approach directly reduces the average number of unique experts invoked per batch by leveraging pre-gated information. Specifically, for decoding- decoding-prriorttzee prioritized and prefill-prioritized batching, the average number of unique experts per batch was 5.08 and 5.21, respectively, whereas our method reduces this to 3.51.
We observed a significant performance impact as prefill requests invoke more experts per batch compared to decoding requests. Prefill requests require tokens to dispatched to different experts, making it impractical to batch tokens by shared experts due to attention operations. As a result, a sub- stantial number of experts are invoked for each batch, negatively affecting performance. Fortunately, our auto-regressive router design improves temporal locality in prefill requests, often allowing tokens within the same request to select the same or a small number of experts. We explore this locality in greater detail in the following section.
High Temporal Locality. To analyze the locality, we measure the temporal distance of the tokens in sequence (Fig. (c)) We define temporal distance as the distance between two tokens selecting the same expert within a sequence 48).. Our result shows that our router leads to a smaller distance, indicating a high degree of temporal locality. Specifically, out of 4096 tokens, 2921 tokens follow the choice of the last token, compared to 850 tokens in Mixtral-8x71h The locality attributed to the regr-segeee design of our router, where the router's decision is based on the current and all previous tokens. As a result, a given token is likely to have similar expert selections with its recent predecessor tokens. However, note that this temporal locality appears only within the token sequence of a single request and does not appear across different requests.
5.4 Memory-Efficient Inference
We evaluate how well our approach ensure good perfor- mance while improving memory efficiency. In particular, we constrain the expert cache capacity to & (that is, up to i experts can reside in accelerator memory). In this setup, if a requested expert is not in memory, it must be loaded from host mem- ory, potentially increasing loading latency. As explained in 4.1 this loading overhead can be mitigated with prefetching, provided that we know which expert will be needed in Read- ME. We compare the end-to-end latency of requests from the versus not
Ordiard
Caste
Figure 6: Latency impact of prefetch- ing: We measured end-to-end latency synthetic workload generated by replaying Chatbot Arena Dataset (2 (Appendix [5.]
prefetching our approach enables
leveraging prefetching (On (On-demand Loading) (25). Figure 6 shows that for varying cache capacities, we consistently out- perform On On-demand Loading, with up 30% better latency.
In addition to proactively loading experts into memory, our approach also retains experts in cache further use memory optimally. Table [oole|[compaee three representative caching policies' hit ratios across varying cache capacities, including the Belady- inspired approach that our architecture enables. As noted earlier, our approach accommodates multiple requests where each request has a token sequence, in contrast with prior works focusing on a single eeques//token-sqqeeeee g 27

---

When multiple requests share the expert cache, temporal local- ity within a single request cannot leveraged across requests, limiting its effectiveness. This explains why LRU, which works well in single-request scenarios, underperforms in our setup. In contrast, our Belady-based algorithm excels at all cache capacities by utilizing future expert information across requests, thanks to the pre-gating router. When cache capacity constrained by system memory, latency can be significantly reduced with an optimized cache policy. Our Belady approach notably improves latency, particularly under limited cache sizes, though we omit detailed results for brevity.
Cache Policy
Cache Capacity
Random
Belady
2
34.19%
33.90%
44.16%
50.14%
52.42%
61.82%
67.52%
66.95%
77.21 %
5
82.91%
83.48%
88.03 %
Table 4: Cache hit ratio measured in batched inference setup.
6 Related Work
MoE Refactorization. Recent "MoE-fication" methods
33 49 optimize or group channels
using graph-based techniques but still rely on system-inefficient layer-wise routers. contrast, we are the first to identify the redundancy in layer-wiss routers and propose pre-gating router that enables expert pre-fetching. Similar [50] 14 ST we leverage activation sparsity [23] construct
experts, adaptively identifying important neurons and evicting less-important ones during inference.
Efficient Inference Serving. To deal with the limited memory in coource-constrainel settings, prior LLM inference works focused on optimizations such as offloading parameters to host memory [22 quantization (54 55 56), sparsity (57 (58) and MoE architectures [4 59
53
25
26 However, while
token batching 9 has garnered significant attention for dense models 13 problematic and underexplored in the context of MoE models.
7
38
60 it remains
Pre-gated MoE 28 is related to Read-MEas they too fine-tune a router to pre-gate using the ith layer's hidden states to compute the + 1th layer's routing: but they still maintain a layer-wise architecture which constrains batching. SiDA-Man MME [6]] separates the router from the inference path. However, tokens cannot be batched together because they do not share routing decisions across all layers. In addition, the offline routing function of SiDA is an approximation that may incorrectly guess expert selection, especially when the model scales. In contrast, Read-MEht exact routing, ensuring no performance drop during inference.
Mixtral-offloading 8 introduces speculation to "guess" routing decisions, resorting costly on- demand loading speculation fails. Caching is commonly used [62 [2 63 53 64]. including in MoE systems 88 27). which typically focus single requests. Prior caching methods are limited by layer-wise routing and lack of foresight into future requests.
Conclusions and Limitations
We address the under-explored challenge of reusing pre-trained L.M create smaller MoE model that enables efficient inference with minimal training cost. By leveraging activation sparsity, we construct specialized experts and integrate them via a router. Upon analyzing the layer-wise router design used in all open-source MoEs, we identify its inefficiency and redundancy. To overcome this, we propose pre-gating router, decoupled from the MoE backbone, enabling system-level optimizations that were previously unattainable.
Limitations. Our serving system is designed for single accelerator, and extending it to distributed serving remains a non-trivial task for future work. Our method has no negative socictal impact, as uses publicly released data and model checkpoints. This work is foundational research and is not tied to specific applications.
Acknowledgements
The work of Z. Wang is in part supported by the US Army Research Office Young Investigator Award (W91INF2010240 and a Research Gift from Qualcomm. Ro and Akella are supported by NSF grants CNS-21 2105-90 and CNS-2232135 and by Cisco Research and Meta.
10

---

References
[1] Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto. Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems,
2021.
[2] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022.
[3] Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford. Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401 04008, 2024.
[4] Zhiwen Fan, Rishov Sarkar, Ziyu Jiang, Tianlong Chen, Kai Zou, Yu Cheng. Cong Hao, Zhangyang Wang. et MDvit: Mixture-of-expert vision transformer for efficient multi-task learning with model-accelerator co-design. Advances in Neural Information Processing Systems, 35:28441-284577 2022.
[5] Zitian Chen, Yikang Shen, Mingyu Ding, Zhenfang Chen, Hengshuang Zhao, Erik G Learned- Miller, and Chuang Gan. Mod-squad: Designing mixtures of experts as modular multi-task learners. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recog- nition, pages 11828-1 11837, 2023.
[6] Rishov Sarkar, Hanxue Liang, Zhiwen Fan, Zhangyang Wang, and Cong Hao. Edge-moe: Memory-efficient multi-task vision transformer architecture with task-level sparsity via mixture- of-experts. 2023 IEEE/ACM International Conference on Computer Aided Design (ICCAD). pages 01-09. IEEE, 2023.
[7] Rui Kong, Yuanchun Li, Qingtian Feng, Weijun Wang, Linghe Kong, and Yunxin Liu. Serving moe models on resource-constrained edge devices via dynamic expert swapping. arXiv preprint ar0iv:2308. 15030, 2023.
[8] Artyom Eliseev and Denis Mazur. Fast inference of mixture-of-experts language models with offloading. arXiv preprint arXiv:2312.17238, 2023.
[9] Gyeong-In Yu, Joo Seong Jeong, Geon-Woo Kim, Soojeong Kim, and Byung-Gon Chun. Orca: A distributed serving system for (Trannformer--aas Based) generative models. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 521-538, 2022.
[10] Aran Komatsuzaki, Joan Puigcerver, James Lee- Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. Sparse upcycling: Training mixture-of-experts from dense checkpoints. arXiv preprint arXiv:2212.05055 2022.
[11] Zhengyan Zhang. Yankai Lin, Zhiyuan Liu, Peng Li, Maosong Sun, and Jie Zhou Moefication: Transformer feed-forward layers are mixtures of experts. arXiv preprint arXiv:21 10.01786, 2021.
[12] Zhengyan Zhang. Zhiyuan Zeng, Yankai Lin, Chaojun Xiao, Xiaozhi Wang, Xu Han, Zhiyuan Liu, Ruobing Xie, Maosong Sun, and Jie Zhou. Emergent modularity in pre- trained transformers. arXiv preprint 2tt.::22955.8989000 2023.
[13] Mikolaj Piórezyński, Filip Szatkowski, Klaudia Bałazy, and Bartosz Wojcik. Exploiting transformer activation sparsity with dynamic inference. arXiv preprint arX1v:2310.04361, 2023.
[14] Harry Dong, Beidi Chen, and Yuejie Chi. Prompt-prompted mixture of experts for efficient I1m generation. arXiv preprint arXiv:2404.01366 2024.
[15] Laszlo A. Belady. A study of replacement algorithms for a virtual-storage computer. IBM Systems journal, 101, 1966.
[16] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.0 030000099330....0 2020.
11

---

[17] Yihua Zhang, Ruisi Cai, Tianlong Chen, Guanhua Zhang. Huan Zhang. Pin-Y Yu Chen, Shiyu Chang, Zhangyang Wang, and Sijia Liu. Robust mixture-of-expee training for convolutional neural networks. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 90-101, 2023.
[18) Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou, and Yang You. Openmoe: An early effort on open mixture-of-experts language models. arXiv preprint arXiv:2402.01739, 2024.
[19] 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 arXiv:2307.0922 009888 2023.
[20] Colin B. Clement, Matthew Bierbaum, Kevin P. O' Keeffe, and Alexander A. Alemi. On the use of arxiv as dataset, 2019.
[21] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020.
[22] Stack Excahnge. Stack exchange data dump, 2024.
[23] Zonglin Li, Chong You, Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank J Reddi, Ke Ye, Felix Chern, Felix Yu, Ruiqi Guo, et al. The lazy neuron phenomenon: On emergence of activation sparsity in transformers. arXiv preprint arXiv:2210.06:: 2022.
[24] Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. Deepseekmoe: Towards ultimate expert specialization in mfiture-off-epeere language models. arXiv preprint arXiv:2401. 2024.
[25] Accelerate huggingface.co. https: http:p///wgggggggccoco/o///////s cessed 22-05-2024).
[Ac-
[26] Liang Shen, Zhihua Wu, WeiBao Gong, Hongxiang Hao, Yangfan Bai, HuaChao Wu, Xinxuan Wu, Jiang Bian, Haoyi Xiong, Dianhai Yu, and Yanjun Ma. Se-moe: A scalable and efficient mixture-of-experts distributed training and inference system, 2023.
[27] Leyang Xue, Yao Fu, Zhan Lu, Luo Mai, and Mahesh Marina. Moe-infinity: Activation-aware expert offloading for efficient moe serving, 2024.
[28] Ranggi Hwang, Jianyu Wei, Shijie Cao, Changho Hwang, Xiaohu Tang, Ting Cao, Mao Yang, and Minsoo Rhu. Pre-gated moe: An algorithm- system co-design for fast and scalable mixture-of-expeer inference. arXiv preprint arXiv:2308. 12066, 2023.
[29] Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric. P Xing, Hao Zhang, Joseph E. Gonzalez, and lon Stoica. Judging Ilm- as-adg with mt-bench and chatbot arena, 2023.
[30] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma. Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the- art natural language processing. In Proceedings of the 2020 Conference Empirical Methods in Natural Language Processing: System Demonstrations, pages 38-45 45, Online, October 2020.
Association for Computational Linguistics.
[31] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[32] Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05222.200
12

---

[33] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, Neurocomputiing. 568:127063, 2024.
[34] Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[35] Together Computer. Redpajama: An open source recipe reproduce llama training dataset, 2023.
[36] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified lext-to-text transformer. arXiv e-prints, 2019.
[37] Wikimedia Foundation. Wikimedia downloads.
[38] Reza Yazdani Aminabadi, Samyam Rajbhandari, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng. Olatunji Ruwase, Shaden Smith, Minjia Zhang. Jeff Rasley, and Yuxiong He. Deepspeed- inference: Enabling efficient inference of transformer models at unprecedented scale. In SC22: International Conference for High Performance Computing, Networking. Storage and Analysis, pages 1-15, 2022.
[39] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pages 61 1-166,6,222 2023.
[40] Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pages 2397-2430. PMLR, 2023.
[41] Xinyang Geng and Hao Liu. Openllama: An open reproduction of Ilama, May 2023.
[42] Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared Ilama: Accelerating language model pre-training via structured pruning. arXiv preprint arXiv:2310.06694 2023.
[43] Xinyin Ma, Gongfan Fang, and Xinchao Wang. LIm-pruner: On the structural pruning of large language models. Advances in neural information processing systems, 16:21702-21720, 2023.
[44] Saleh Ashkboos, Maximilian L Croci, Marcelo Gennari do Nascimento, Torsten Hoefler, and James Hensman. Slicegpt: Compress large language models by deleting rows and columns. arXiv preprint 40X5v:2401.15024, 2024.
[45] Yifei Yang, Zouying Cao, and Hai Zhao. Laco: Large language model pruning via layer collapse. arXiv preprini 2 1187, 2024.
[46] Song Guo, Jiahang Xu, Li Lyna Zhang, and Mao Yang. Compresso: Structured pruning with col- laborative prompting learns compact large language models. arXiv preprint arX1v:2310.05015, 2023.
[47] Huggingface tgi inference engine. text- generation inference [Accessed 20-05-2024).
http:://ggituu..comm/uggggggg
[48] Vidyadhar Phalke and Bhaskarpillai Gopinath. An inter-reference gap model for temporal locality in program behavior. ACM SIGMETRICS Performance Evaluation Review, 23(1):291- 300, 1995.
[49] Haizhong Zheng, Xiaoyan Bai, Beidi Chen, Fan Lai, and Atul Prakash. Learn to be efficient: Build structured sparsity in large language models. arXiv preprint arXiv:2402.06116. 2024.
[50] Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivas- tava, Ce Zhang, Yuandong Tian, Christopher Re, et al. Deja vu: Contextual sparsity for efficient Ilms at inference time. International Conference on Machine Learning, pages 22137-22176. PMLR, 2023,
13

---

[51] Varun Yerram, Chong You, Srinadh Bhojanapalli. Sanjiv Kumar, Prateek Jain, Praneeth Netra- palli, et al. Hire: High recall approximate top-k estimation for efficient Ilm inference. arXiv preprint arXiv:2402.09360, 2024.
[52] Weihao Cui, Zhenhua Han, Lingji Ouyang, Yichuan Wang, Ningxin Zheng, Lingxiao Ma, Yuqing Yang, Fan Yang. Jilong Xue, Lili Qiu, et al. Optimizing dynamic neural networks with brainstorm. In 17th USENIX Sympostum on Operating Systems Design and Implementation (OSDI 23), pages 797-815, 2023.
[53] Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. Zero-offload: Democratizing billion-scale model training. In USENIX Annual Technical Conference (USENIX ATC 21), pages 551-564, 2021.
[54] Yilong Zhao, Chien-Yu Lin, Kan Zhu, Zihao Ye, Lequn Chen, Size Zheng, Luis Ceze, Arvind Krishnamurthy, Tianqi Chen, and Baris Kasikci. Atom: Low-bit quantization for efficient and accurate Ilm serving. arXiv preprint arXiv:2310.19102 2023.
[55] Elias Frantar, Salch Ashkboos, Torsten Hoefler, and Dan Alistarh. Gplq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210. 17323, 2022.
[56] Lin Zhao. Awrq: Activation-awarc weight reformulation quantizer for large language models.
[57] Ajay Jaiswal, Shiwei Liu, Tianlong Chen, Zhangyang Wang, et al. The emergence of essential sparsity in large pre-trained models: The weights that matter. Advances in Neural Information Processing Systems, 36, 2024.
[58] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024.
[59] Zhixu Du, Shiyu Li, Yuhao Wu, Xiangyu Jiang. Jingwei Sun, Qilin Zheng, Yongkai Wu, Ang Li, Hai "Helen" Li, and Yiran Chen. Sida-moe: Sparsity-inspires data data-aware serving for efficient and scalable large -ffeerr-of-expeets models, 2024.
[60] GitHub - -NVIIIAATTTTTTTORLLLLLLL TensorRT-LLM provides users with an easy-to-use Python API define Large Language Models (LLMs) and build TensorRT engines that contain state-of- the-art optimizations to perform inference efficiently on NVIDIA GPUs. TensorRT-LLM also contains components create Python and C++ nuntimes that execute those TensorRT engines. github.com. htttp:://githue com/
-LM [Accessed 20-05-2024).
[61] Zhixu Du, Shiyu Li, Yuhao Wu, Xiangyu Jiang. Jingwei Sun, Qilin Zheng, Yongkai Wu, Ang Li, Hai "Helen" Li, and Yiran Chen. Sid: pparsity-inspiee data-awarg serving for efficient and scalable large mixture-of-experts models, 2024.
[62] Chijun Sima, Yao Fu, Man-Kit Sit, Liyi Guo. Xuri Gong, Feng Lin, Junyu Wu, Yongsheng Li, Haidong Rong, Pierre-Louis Aublin, et al. Ekko: A (Large-Scale) deep learning recommender system with [Low-Latency] } model update. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 821 839, 2022.
[63] Jaehoon Jung, Jinpyo Kim, and Jaejin Lee. Deepum: Tensor migration and prefetching in unified memory. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, pages 207-221, 2023.
[64] Jie Ren, Jiaolin Luo, Kai Wu, Minjia Zhang, Hyeran Jeon, and Dong Li. Sentinel: Efficient tensor migration and allocation on heterogeneous memory systems for deep learning. In 202/ IEEE International Symposium on fiigh-Performacce Computer Architecture (HPCA), pages 598 611. IEEE, 2021.
[65] 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:2310.06825, 2023.
14

---

A More Experimental Results
A.1 Training Dynamics
3.5
As detailed in $5.T] we iteratively tune the router and experts for rounds. We visualize the validation loss during the first rounds out of the total 8 rounds of training. In Fig. 7 the router tuning stages are marked in gray, while the expert tuning stages are marked in orange. Two observations can be drawn from Fig- ure ? (1) The validation loss decreases during both router tuning and expert tuning stages. (2) The val- idation loss reduction from router tuning saturates after two rounds, while the validation loss continues to decrease during expert tuning.
LM Loss
2.5
200
400
600
800
1000
1200
Iteaaiinn
Figure 7: Visualization on training dynamics.
A.2 MoE Achieves Better Efficiency-Accuracyy Trade-off than Dense Models.
focus converting large dense pre-trained model into smaller dense model. However, we argue that a smaller MoE model (i.e. the MoE model with the smaller number of activation parameters)
Prior compression-basee works 42
44 45 46
better target architecture. To ensure fair comparison, we (1) derive small dense model with 4.7B parameters, matching the size of a single expert network, using the same amount of data, and (2) fine-tune the obtained dense model for an equivalent number of steps. As shown in Table 5 refactorizing the pre-trained model into an MoE structure, rather than smaller dense variant, leads significant performance improvement. The models are evaluated based on performance on the MMLU [16], and perplexity across seven data domains included in RedPajama [35]
Table 5: We compare the Read-M SouudMMEperformaccc with dense model, and report the MMLU performance and perplexity on 7 data domains. By adopting an MoE as the target structure instead of dense model, our model achieve significantly better overall performance.
Evaluation
Arxiv
Books
C4
Common Crawl
Github
StackExchange
Wikipedia
MMLU
Dense
5.63
1.94
11.78
9.68
3.75
13.42
6.24
27.1%
Read-ME
4.18
1.31
10.57
7.72
2.39
12.52
3.94
38.9%
A.3 Read-ME Remains Effective without Prior Knowledge of the Training Domain
We additionally use the Mistral (65) model as the pre-trained dense model, and convert it to the MoE structure, with the proposed method. The task is challenging because we do have prior knowledge on the Mistral original training data, and our experiment in Table 6 shows that our method remains effective without the prior knowledge of the original training domain.
Table 6: Ablation study Mistral 65 pre-trained model.
Method
Pre-trained
Domain
Fine-tune
Domain
#Param
MMLU
Hell.
Wino.
ARC-E
ARC-C
LogiQA
CoQA
avg.
Read-ME. Llam-22
Red pajama
Red pajama
4.78-178
38.9%.
68.5%
67.7%
66.6%.
42.3%
29.7%
74.8%:
55.5%
Llama-2
Red -ajaapajma
6.9B
45.3%
78.6%
69.3%
76.4%
53.0%
31.0%
75.9%
61.4%
Misstal
N/A
Red-pajama
4.78-178
39.2%
19.1%
68.2%
77.1%
49.3%
30.9%
76.2%
60.0%
Misaral
NIA
6.9B
62.1%
84.5%
79.3%
82.7%
63.7%
33.5%
80.3%
69.4%
A.4
Computational Cost of Auto-regressive Router
For a detailed cost analysis auto-regressive router that we introduced, we added: (1) FLOPs com- parison, (2) latency, and (3) latency breakdown with larger batch size (high-throughput scenarios) of a Traditional Router (TR) and an Autoregressive Router (AR). To focus solely on the router's 's impact on latency, we controlled other variables (e.g., the number of activated parameters) to be the same.
Note that the computational cost of both the traditional router and the autoregressive router is theoretically linear batch size. Therefore, when the batch size is high (in tigh-throughhut
15

---

Table 7: Flops comparison between Traditional rouler and Auto-regressive router
Traditional Router
Auto-regressive Router
Flops/sample
4.7 KFLOPs
3 KFLOPs
Table 8: Latency [ms] comparison between Traditional router and Auto-rogressive router
bsz=5
bsz=5
bsz=10
bsz=10
10
bsz=20
bsz=20
bsz=30
bsz=30
TR
AR
TR
AR
TR
AR
TR
AR
Router
1.76
0.61
1.80
0.61
1.78
0.61
1.93
0.61
Attention
18.13
18.18
18.28
18.13
18.49
18.36
19.59
19.66
Expert/MLP
22.43
21.75
24.59
22.53
24.97
22.99
30.17
28.31
Sum
42.31
40.55
44.67
41.27
45.23
41.96
51.69
48.59
Table 9: Latency breakdown comparison between Traditional router and Auto-regressive
bsz=5
bsz=5
bsz=10 TR
bsz=10
bsz=20 TR
bsz=20
bsz=30
bsz=30
TR
AR
AR
AR
TR
AR
Router
4.15%
1.50%
4.02%
1.48%
3.93%
1.46%
3.74%
1.26%
Atlention
42.85%
44.85%
40.92%
43,92%
40.87%
43.75%
37.90%
40.47%
Exper/MLLP
53.01%
53.65%
55.06%
54.59%
55.20%
54.80%
58.36%
58.27%
Sum
100.00%
100.00%
100.00%
100.00%
100.00%
100.00%
100.00%
100.00%
scenarios), the cost increases linearly. In both cases, the computation can be parallelized, so this remains negligible in end-to-end latency even in high-throughput scenarios. In fact, we would like to clarify that the bottleneck in high-throughput scenarios is actually the expert layers, seen in Table Expert/MLP row. This issue can be addressed by the methods discussed Section Traditional layerwise routers do not allow for efficient system design, which underscores the need for carefull co-design of routers.
16

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justifiiaton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]", ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] " or "[NA]" is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: We explained the detailed claims in section 2-4 with proper observation figures (Fig 2) and added supporting evaluations Sec 5.
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 the paper and important assumptions and limitations. A No or
NA answer to this question will not be 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 to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
17

---

Justification:
We discussed the limitation of our method in Section 7.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model sell-specifiitiinn 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 be 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-texi system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't "te'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 set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification:
The paper does not include theoretical results.
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 in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch 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
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide the detailed experimental information including model configura- tions, syper--arameters and inference system setup choices Appendix Section B. Codes will be publicly released.
Guidelines:
18

---

The answer NA means that the paper does not include experiments.
If the paper includes experiments, a 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 a 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 it 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 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, other means that appropriate to the research performed.
also be provided via detailed
While NeurlPS 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) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily new 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., to 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 the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We use public datasets for both training and evaluation. Codes will be publicly released.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "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
/ //niss.ccpuulic//guus http://////iiiee///deuuubiiiiiii
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.
19

---

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 is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide experimental settings and details from both training and evaluation in Appendix Section B.
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 a 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, or 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: [Yes]
Justification:
We report error bars in Figure 6.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, at 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 run with given experimental conditions).
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 the standard error of the mean.
It is OK to report 1-sigma error bars, but 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 to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in 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 include the type of compute resources we used for experiments in Appendix Section B.
Guidelines:
20

---

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, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well 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 make 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: Our method does not have negative societal impacts, as we use publicly released data and model checkpoints. Our work is foundational research and is not tied to specific applications.
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 circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration 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 discuss broader impacts in Section 7.
Guidelines:
The answer NA means that there no socictal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious 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.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML.)
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
21

---

Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no 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 require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.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 cite the original paper//webxitellicensse of existing assets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.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 datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
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 the dataset/code/moodll as part 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 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
22

---

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 as details about compensation (if any)?
Answer: [NA]
Justification: Not Applicable
Guidelines:
The answer NA means that the paper does not involve crowdsourcing 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 to 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 equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Not Applicable.
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 NeurIPS 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": "0WGwUDMZjM", "reviewer": "Reviewer_wHVe", "summary": "The paper proposes an inference-aware method to convert a pretrained dense model into a Mixture-of-Experts architecture, where each expert is smaller than the original dense model. To extract the different experts, they use a dataset from given subdomain to identify the top activated channels. To route among the experts, the authors first break down the limitations of current per-layer routing schemes (waiting for all tokens at the layer to finish, redundancy across layers). They propose to decouple the router from the base model, and instead train a 1-layer transformer block as the router, which predicts token expert assignment autoregressively. \n\nOverall, these changes allow for inference friendly deployment, by enabling careful batching of examples that share similar experts, and with better caching of experts in resource constrained settings. The authors show that their approach strikes a good trade-off between efficiency and performance.", "questions": "1. If you keep the top activated neurons, should equation 3 be an argmax ?\n2. What is the impact of the routing distillation loss ? Can you provide an ablation experiment where it is not used to access its impact ?\n3. How would figure 3 look with READ-ME instead of standard MoEs ? This would be a good visualization to compare the methods. \n4. How were the baseline numbers in Figure 4 obtained ? Did the authors rerun the baselines ? Can you confirm that all these baselines start from the same base model ?\n5. The \"cost\" column of table 1 is misleading; given that your model starts from llama-2, you must either include the compute to create llama-2 in your analysis, or only monitor the additional compute starting from llama 2 (in which case the cost for llama 2 would be 0).", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors do a good job at breaking down the current bottlenecks in deploying MoE models \n2.  The proposed external routing approach is simple and effective\n3.  The latency evaluation and experimentation is well designed and clear", "weaknesses": "1.  I am a bit surprised that the gains in latency compared to the seed llama models are somewhat small (19%) given that the number of active parameters is reduced by 30%, and that there is also a decent performance drop from the seed llama model. For resource constrained settings, how would the approach of loading llama layer-by-layer fair against Read-ME ? Like Read-ME, this approach can know in advance what layers to load, and could hopefully retain the performance of the full lama model. \n2. Some parts of the analysis require further clarification (please see questions)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**[W1 - Comparison with layer-by-layer loaded Llama ]**\nThanks for the question. We compared the latency of the layer-by-layer loaded llama and Read-ME model in the following table. \n\nTable A. Latency comparison\n\n|Method| Latency [ms] |\n|---------|--------------|\n| Layer-by-layer Llama-7b |     111.909 |\n| Read-ME |      91.531 |\n\nThe difference arises from the size of the parameters to be loaded. The size of Read-ME's expert layer (including both the top-1 expert and residual expert) is approximately 25.0% that of the MLP counterpart in the LLaMA baseline. Regarding peak memory usage, Read-ME exhibits a 10% lower profile, though both models consume an insignificant amount of memory overall.\n\nOur method converts a dense pre-trained model into a MoE for efficient inference. Although this compression is generally not lossless, please note that our approach offers a better accuracy-cost trade-off than all baseline compression methods.\n\n**[Q1 - Understanding of Equation 3]** Sorry for the confusion. We found this is a typo in our original submission., As we are selecting the mask $\\boldsymbol{M}$ to maximize the magnitude of activated channels, the operator should indeed be $\\arg\\max$ in Equation 3. \n \n**[Q2 - Ablation on Routing distillation loss]** Thank you for the suggestion. We have ablated the routing distillation loss by removing the router training step and using a random routing mechanism instead, tuning only the expert weights using the language modeling loss. To ensure a fair comparison, we maintained the same number of training tokens and the same training schedule. The resultant performance is provided in Table B. The results show that with the routing distillation loss, the average downstream task accuracy increased from 51.5% to 55.5%, validating the necessity of the routing distillation loss. \n\nTable B: Ablation study on routing distillation loss. We compare the performance with and without the routing distillation loss, while keeping the number of training tokens, to validate the necessity of routing distillation loss.\n| Method | MMLU | Hell. | Wino.  | ARC-E | ARC-C | LogiQA | CoQA | avg.|\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| Read-ME (w/ routing distillation loss)| 38.9% | 68.5%  | 67.7%  | 66.6% | 42.3% | 29.7% | 74.8% |  55.5% |\n| Read-ME (wo routing distillation loss)| 30.6% | 65.1%  | 65.8%  | 64.2% | 39.9% | 24.9% | 69.7% |  51.5% |\n\n**[Q3 - Read-ME compared to Figure 3]** \nThis is a great suggestion! Please check Fig 3 of the PDF uploaded.  \nIn Fig 3-left, Read-ME batches tokens directed to the same expert, so all arrows point to a single expert at each layer. Also, there is only one router instead of three, resulting in all arrows following the same path from layer 1 to layer 3, leading to a unique activated expert count of 1.\nIn Fig 3-middle, Read-ME shifts the distribution to the lower left, where most points stay within the range of x < 4.5 and y < 55.\nFig 3-right already compares the traditional approach with Read-ME. \n\n**[Q4 - Baseline numbers in Figure 4]** Thanks for asking.\n- MMLU is the common evaluation benchmark for LLMs, so most of the numbers in Figure 4 are collected from their original paper, including OpenMoE, Llama-2, LaCo, and Compresso; For Sheared-Llama, Open-Llama and Pythia, as they do not report MMLU in their original paper, we test their publicly released checkpoints with the lm-eval-harness library[1]; LLM-Pruner and SliceGPT didn’t report MMLU nor release checkpoints, so we use the numbers reported in LaCo [3].\n- Yes, all of the methods reported in Figure 4 use the same base model - Llama-2.\n\n**[Q5 - Training Cost Calculation]** Thanks for the suggestion. First, our training cost calculation follows Sheared-Llama[2] (see Table 1 of the paper), a representative post-training method to generate a small model out of a large one by pruning. Second, the training cost here measures the computational resources needed by the LLM deployer to obtain a new model of an auxiliary architecture (a smaller model in the Sheared-Llama case, and a MoE in our case), given all the publicly available resources. Using the publicly released checkpoints (e.g. Llama-2) will not incur additional training costs. We appreciate the suggestion and will mention this in the caption. Please see the PDF for our tentative Table 1.\n\n[1] https://github.com/EleutherAI/lm-evaluation-harness\n\n[2] Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning\n\n[3] LaCo: Large Language Model Pruning via Layer Collapse", "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": "caR9HzdcEB", "reviewer": "Reviewer_ZdPD", "summary": "This paper proposes Read-ME, a novel framework for pruning large LLMs into smaller MoEs with minimal training cost. Read-ME separates the gating routers from the critical paths of the inference process and trains an individual expert subnetwork to perform offline pre-gating. With this refactorization of MoEs, the paper further demonstrates the effectiveness of Read-ME by designing optimal expert pre-fetching and caching and expert-aware batching for low-latency and high-throughput serving. Extensive experiments show that Read-ME produces small MoEs that outperform existing small dense models with significantly higher inference performance, lower latency, and higher memory efficiency.", "questions": "Pleease refer to the Weaknesses.", "limitations": "Pleease refer to the Weaknesses.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well-written, well-organized, and easy-to-follow.\n- The idea of pruning LLMs into smaller MoEs is quite interesting.\n- Experiments are thorough and extensive.", "weaknesses": "- Overall, this is a very interesting work. My biggest question is: what's the motivation behind pruning LLMs into smaller MoEs? If I need a smaller model, why don't I just prune an LLM and have a small but dense model, if the small dense model has approximately the same (or even less) number of activated parameters?\n\n- The observation of Figure 2 is interesting. However, what's the dataset (or input tokens) you use for plotting this figure? Does this observation still hold if the dataset (or input tokens) changes significantly?\n\n- Section 2.3, \"The above observations suggest that among many routing paths, only a few of them are in use during the inference...\" with only mutual information between adjacent layers may not justify this assumption. It would be better to visualize the end-to-end routing paths (from the first to the last layer) and explicitly show that only a few paths are used frequently.\n\n- The idea of separating routing logic from the inference process is not new. For example, [1] tries to distill the knowledge of gating networks and performs offline routing. How do you compare your expert subnetwork with [1]? Does training the expert subnetwork cost more than KD-based methods?\n\n- To follow up on the previous comment, the evaluation does not compare Read-ME with any existing offline routing approaches like [1]. If the major benefits of Read-ME come from separating routing logics, then the paper's contributions would decrease.\n\n- In evaluation, comparing Read-ME 4.7B with those baselines with fewer parameters may be unfair. Since Read-ME is pruned from much larger LLMs and still holds more activated parameters than baselines, it wouldn't be surprising to see that Read-ME has the best performance. This performance increase comes at the price of requiring more GPU memory. Once again, this question goes back to the motivation: why would someone need to prune LLMs into smaller MoEs instead of smaller dense models? Comparing Read-ME MoEs with smaller dense models pruned from the same LLM may help address this question.\n\n- This pruning process may need a theoretical justification on the performance guarantee, i.e., smaller MoEs are guaranteed to not suffer large performance degradations.\n\n\n[1] SiDA: Sparsity-Inspired Data-Aware Serving for Efficient and Scalable Large Mixture-of-Experts Model. MLSys'24\n\nMinor issues:\n\n- Figure 2(c) is never mentioned.\n\n- Figure 3 does not have indices for sub-figures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for all the interesting questions. Please see below. \n\n**[W1 - MoEs Motivation]** We validated that MoE achieves better cost-accuracy trade-off than small dense models acquired by pruning, and provided the results in Appendix C.1.\n\nWe mentioned that prior compression efforts focus on converting large dense pre-trained models into smaller dense models. However, we argue that a smaller MoE model, with fewer activation parameters, is a better target architecture. To ensure a fair comparison, we (1) create a 4.7B parameter-dense model matching a single expert network size, and (2) fine-tune it for the same number of steps. \nTable 5 shows refactorizing the pre-trained model into an MoE structure, rather than a smaller dense variant, leads to significant performance improvement. \n\nTable 5: By adopting an MoE as the target structure instead of a dense model, our model achieves better overall performance.\n|Eval|Arxiv|Books|C4|CC|Github|Stack.|Wiki.|MMLU|\n|-|-|-|-|-|-|-|-|-|\n|Dense|5.63|1.94|11.78|9.68|3.75|13.42|6.24|27.1%|\n|Read-ME|4.18|1.31|10.57|7.72|2.39|12.52|3.94|38.9%|\n\n**[W2 - Figure 2 Details]** For Figure 2, we used the Arxiv dataset, a subset of Red-pajama. The observation holds with Wikipedia and Github subsets as well. Please see the uploaded PDF for visualization results; we will add these in a future version.\n\nThe visualized high mutual information between two adjacent layers’ expert selection is sufficient to support our claim. Since $H(S_1, …, S_L) \\le \\sum_l H(S_l) - \\sum_l I(S_{l+1}; S_l)$, high layer-wise mutual information implies low-entropy (deterministic) path selection.\n\nRegarding visualizing end-to-end routing paths, firstly we would like to mention that since the model has 32 layers and each layer performs Top-2 selection out of 8 experts, there will be ${8 \\choose 2}^{32}$ (approximately $2 \\times 10^{46}$) possible paths. \n\nInstead, we calculated routing statistics for 180k tokens, finding only 8.4k paths out of $2\\times 10^{46}$ possible paths activated at least once. The top 20 paths are selected by 2805 tokens, and the top 40 paths by 4947 tokens. The observation validates that only a small fraction of routing paths are frequently selected. Please see the uploaded PDF for more visualization.\n\n**[W3 - Comparison with SiDA]** SiDA is an alternative approach that separates the router from the inference path. However, Read-ME offers better opportunities to improve inference efficiency because it is designed to enable ***expert-aware batching.*** Additionally, we would like to emphasize that, by design, the Read-ME router provides ***exact expert selection,*** while SiDA and prior works offer approximate selection [1,2]. In detail,  \n\n(1) Batching and latency:\nWith SiDA, *tokens cannot be batched together* because they do not share routing decisions across all layers, unlike Read-ME (Since SiDA is distilled from the Switch Transformer, it is not possible to change decisions of routers). This vastly magnifies the expert space when  batching and makes it difficult to compose an efficient batch that is aware of expert selection.\nIn effect, SiDA activates more experts for each batch at each layer, leading to increased latency. Table A compares the inference latency between SiDA and Read-ME. (SiDA did not release a checkpoint, so we used SwitchTransformer-8, the teacher model from which SiDA was distilled, as a proxy)\nTable A. Inference latency of SiDA and Read-ME.\n| |SiDA|Read-ME|\n|-|-|-|\n|Latency[ms]|62.34|48.59|\n|Avg # of Activated Experts|5.51|3.51|\n\n(2) Router accuracy:\nSiDA’s offline routing function is an approximation method that is distilled from original layerwise routers. Thus, the accuracy of routing is not 100%. Table B shows the “failure rate” of the SiDA router (defined as prediction miss rate on the expert activation of the trained router), and Table C shows the resulting degradation in task performance. With an incorrect guess, the entire expert selection can fail, leading to a collapse in inference, especially as the model scales. In contrast, *our method is exact, ensuring no performance drop during inference*, regardless of model size.   \n\nTable B. SiDA Router failure rate\n|#experts|SST2|MRPC|MultiRC|\n|-|-|-|-|\n|8|-1.00%|-2.59%|-8.26%|\n|128|-1.22%|-1.35%|-9.51%|\n\nTable C. Accuracy drop due to SiDA’s Router Failure\n|#experts|SST2|MRPC|MultiRC|\n|-|-|-|-|\n|8|-1.75%|-2.51%|-1.05%|\n|128|-6.98%|-7.41%|-7.44%|\n\n(3) Training cost:\nNote that SiDA is based on an MoE model and only distills the router, whereas our model is based on a dense model and builds both the expert network and the router, training them together. This means that SiDA's training time only accounts for the router training cost, while our training time includes both the router training cost and the expert specialization cost. Additionally, SiDA did not report the training cost in terms of tokens, making a meaningful comparison of training costs impossible. \n\nIn addition, our work introduces novel contributions to caching and prefetching, which SiDA lacks. SiDA relies on on-demand expert loading and FIFO-based expert eviction, which can negatively impact performance. Please refer to Section 4 for a detailed discussion of our contributions in these aspects. We will add discussion of SiDA in the revised version.\n\n[1] Fast inference of mixture-of-experts language models with offloading\n\n[2] SiDA: Sparsity-Inspired Data-Aware Serving for Efficient and Scalable Large Mixture-of-Experts Model\n\n**[W4 - Theoretical Analysis]** In this work, we empirically examine the redundancy of layer-wise routers and propose a system-oriented method to convert a pre-trained dense model to an MoE with minimal additional training costs. While our focus is on empirical validation, providing a theoretical justification is beyond the scope of our paper and something we hope to explore in future study.\n\nThanks for catching the missing references and indices. We will definitely correct it in our revised version.", "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": "KwhTSL6pLR", "reviewer": "Reviewer_EGvf", "summary": "The paper proposes a novel framework to enhance the efficiency of pre-trained LLMs by transforming them into post-hoc Mixture-of-Experts (MoE) models. The key innovation lies in decoupling the router from the MoE backbone, which facilitates pre-gating and lookahead scheduling, thereby improving memory management and batching during inference. The proposed method, Read-ME, demonstrates significant improvements in both latency and task performance compared to existing models.", "questions": "See weakness. Overall, this seems to be reasonably solid LLM co-design work, but my main concern is its limited evaluation to Llama2 only. Reporting addition experiments on Gemma or Mistral would be strongly encouraged.", "limitations": "Limitations are adequately addressed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper nicely bridges the gap between algorithmic advancements and system-level optimizations. By addressing both fronts, the proposed framework ensures that the improvements in model architecture translate into real-world performance gains.\n2. To the best of my knowledge, the introduction of a pre-gating, shared router decoupled from the MoE backbone is a significant innovation. Although breaking down dense LLMs into MoEs alone has been done before, this new gating approach allows for pre-computation and lookahead scheduling, addressing inefficiencies in traditional layer-wise routing systems.\n3. The expert-aware batching and optimal expert caching algorithms are well-designed to leverage the pre-gating architecture, showing clear improvements in mean and tail latency.\n4. The paper provides comprehensive experimental results that validate the effectiveness of the proposed approach. The improvements in MMLU performance and latency reductions are significant and clearly presented. The comparison with various baselines, including both dense and MoE models, is thorough and demonstrates the superiority of Read-ME across different metrics.", "weaknesses": "1. The experimental subject is limited. While the paper demonstrates impressive results on the LLaMA-2 model, the scalability and generalization of the proposed method to other model types and larger scales are not extensively explored. It would be essential to include experiments on more diverse datasets and larger models.\n2. Although the paper claims minimal overhead for the auto-regressive router, the detailed analysis of its computational costs relative to traditional routers could be expanded. Can you provide a more detailed breakdown of the computational overhead introduced by the auto-regressive router compared to traditional routers? How does this impact overall inference latency, especially in high-throughput scenarios?\n3. The paper could benefit from a more detailed explanation of the pre-gating and batching algorithms, potentially with pseudocode or flow diagrams to aid reproducibility. Sometimes I cannot fully follow the authors’ texts.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**[W1 - Additional experimental results]** Thanks for the suggestion! To validate that our method remains effective in other scenarios, we use the Mistral model as the pre-trained dense model, and convert it to the MoE structure, with the proposed method. The task is challenging because we do not have prior knowledge on the Mistral original training data, and our experiment in Table A shows that our method remains effective without the prior knowledge of the original training domain. \n\nTable A: Ablation study on Mistral[1] pre-trained model.\n|Method | Pre. domain | FT. Domain | # Param | MMLU | Hell. | Wino. |ARC-E | ARC-C | LogiQA | CoQA | avg. |\n|-|-|-|-|-|-|-|-|-|-|-|-|\n| ReadME-mistral | N/A | Red-Pajama | 4.7B - (17B) | 39.2% | 79.1% | 68.2% | 77.1% | 49.3% | 30.9% | 76.2% | 60.0% |\n| mistral | N/A | - | 6.9B | 62.1% | 84.5% | 79.3% |82.7% | 63.7% |33.5% | 80.3%| 69.4%|\n| ReadME-Llama-2| Red-Pajama | Red-Pajama | 4.7B - (17B) | 38.9% | 68.5% | 67.7% |66.6% |42.3% | 29.7%| 74.8% | 55.5%|\n| Llama-2 | Red-Pajama | - | 6.9B | 45.3% | 78.6% | 69.3% | 76.4% | 53.0% | 31.0% | 75.9% | 61.4%|\n\n\n[1] Mistral 7b. \n\n\n**[W2 - Computational cost of Auto-regressive Routers]**\n\nThanks for the suggestion! \n\nFor a detailed analysis, we added: (1) FLOPs comparison, (2) latency, and (3) latency breakdown with a larger batch size (high-throughput scenarios) of a Traditional Router (TR) and an Autoregressive Router (AR). To focus solely on the router’s impact on latency, we controlled other variables (e.g., the number of activated parameters) to be the same.\n\n\n(1) flops comparison\n\nTable B. Flops comparison between Traditional router and Auto-regressive router\n\n|              | Traditional Router | Auto-regressive Router |\n|-|-|-|\n| flops/sample | 4.7 KFLOPs         | 3 KFLOPs               |\n\n\n(2) latency [ms]\n\nTable C. Latency comparison between Traditional router and Auto-regressive router\n\n|            | bsz=5 | bsz=5 | bsz=10 | bsz=10 | bsz=20 | bsz=20 | bsz=30 | bsz=30 |\n|-|-|-|-|-|-|-|-|-|\n|            | TR    | AR    | TR     | AR     | TR     | AR     | TR     | AR     |\n| Router     |  1.76 |  0.61 |   1.80 |   0.61 |   1.78 |   0.61 |   1.93 |   0.61 |\n| Attention  | 18.13 | 18.18 |  18.28 |  18.13 |  18.49 |  18.36 |  19.59 |  19.66 |\n| Expert/MLP | 22.43 | 21.75 |  24.59 |  22.53 |  24.97 |  22.99 |  30.17 |  28.31 |\n| Sum        | 42.31 | 40.55 |  44.67 |  41.27 |  45.23 |  41.96 |  51.69 |  48.59 |\n\n\n(3) latency breakdown [%]\n\nTable D. Latency breakdown comparison between Traditional router and Auto-regressive router\n\n|            | bsz=5   | bsz=5   | bsz=10  | bsz=10  | bsz=20  | bsz=20  | bsz=30  | bsz=30  |\n|-|-|-|-|-|-|-|-|-|\n|            | TR      | AR      | TR      | AR      | TR      | AR      | TR      | AR      |\n| Router     |   4.15% |   1.50% |   4.02% |   1.48% |   3.93% |   1.46% |   3.74% |   1.26% |\n| Attention  |  42.85% |  44.85% |  40.92% |  43.92% |  40.87% |  43.75% |  37.90% |  40.47% |\n| Expert/MLP |  53.01% |  53.65% |  55.06% |  54.59% |  55.20% |  54.80% |  58.36% |  58.27% |\n| Sum        | 100.00% | 100.00% | 100.00% | 100.00% | 100.00% | 100.00% | 100.00% | 100.00% |\n\n\n\nNote that the computational cost of both the traditional router and the autoregressive router is *theoretically linear to batch size*. Therefore, when the batch size is high (in high-throughput scenarios), the cost increases linearly. In both cases, the computation can be parallelized, so this remains negligible in end-to-end latency even in high-throughput scenarios.\nIn fact, we would like to clarify that *the bottleneck in high-throughput scenarios is actually the expert layers*, as seen in table D – Expert/MLP row. This issue can be addressed by the methods discussed in Section 4. *Traditional layerwise routers do not allow for efficient system design*, which underscores the need for a careful co-design of routers.\n\n**[W3 - Reproducibility and Pseudocode]**\n\nIn Appendix A, we provide the pseudocode for the batching algorithm and pre-gating. In summary, at each scheduling step, we find the expert with the most requests and select that expert for the current step. We then check whether the scheduled tokens exceed the maximum token length or the maximum number of requests that can fit. This process is repeated until no more requests can be scheduled. We will release the code publicly and refine the text to improve understanding as well.", "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": "ZMZt1A4c3T", "reviewer": "Reviewer_AX77", "summary": "This paper proposes to perform data-specific pruning from a regular LLM and, with a small amount of continual training, builds a set of smaller LLMs. Then they are used as experts, and a router network is trained to route requests to them. Note that the routing decision is one per query and not per token or per layer, and hence the proposal is very different from regular MoE. Empirical results are presented with Llama-7B as the starting point and show that the resulting ensemble has good accuracy-latency tradeoff.", "questions": "Please see above.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Please see the summary. The results seem sound.", "weaknesses": "The main results of Table 1 does not support that the proposed method is at a better accuracy vs inference-cost trade-off. The average accuracy of the proposed method is only marginally better than Sheared-Llama yet its effective inference size of 4.7B is much larger than 2.7B. In the other direction and comparing against the Llama-7B starting point, the overall accuracy cost of pruning and ensemble is quite substantial from 61.4% to 55.5%. \n\nIn additional to the tasks in Table 1, it would help to add perplexity comparison.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your time and effort for the review.  We answered your questions as follows. \n\n**[W1 - Effectiveness of Read-ME]** Thank you for your question. \n- We would like to first clarify that our method significantly outperforms Sheared-Llama. For example, Sheared-Llama nearly performs at random guess levels on the MMLU task (achieving only 26.4% accuracy on a 4-class classification), whereas Read-ME achieves a much higher accuracy of 38.9%. \n- Second, our method is also more training-efficient, using only 1 billion tokens compared to Sheared-Llama's 50 billion tokens (*x50* times higher than ours). To further back up our claim, we perform the Sheared-llama continued tuning using the same amount of tokens as ours. We show the downstream task performance in Table A. By using 1B tokens for tuning, the performance of Sheared-llama reduces to 48.2%, which is significantly lower than our method (55.5%). Note that we use the publicly released Sheared-llama pruning checkpoint [1] as the initialization.\n- Additionally, our method aims to convert a dense pre-trained model into a MoE for efficient inference, and such a compression method is typically not lossless. However, our method achieves a superior accuracy-cost trade-off compared to all baseline compression methods. To validate our approach, we plot the MMLU performance of our method alongside a large number of baseline methods in Figure 4. This data is also presented in tabular form in Table B. Our method outperforms all baseline models/methods of similar sizes.\n\nTable A: Our method is more training-efficient, using only 1 billion tokens compared to Sheared-Llama's 50 billion tokens (x50 times higher than ours). If running sheared-llama with only 1 billion tokens, the performance drops significantly.\n| Method | Cost | MMLU   | Hell. | Wino. | ARC-E | ARC-C | LogiQA | CoQA  | avg.  |\n|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| Sheared-Llama  | 50B  | 26.4%  | 70.8%  | 67.0% | 67.0% | 41.2% | 28.3%  | 71.7% | 53.2% |\n| Sheared-Llama-efficient  | 1B  | 25.4%  | 59.4%  | 61.9% | 65.8% | 35.8% | 25.1%  | 64.2% | 48.2% |\n| Read-ME  | 1B | 38.9% | 68.5% | 67.7% | 66.6% | 42.3% | 29.7% | 74.8% | 55.5% | \n\nTable B: Evaluation of Read-ME on MMLU benchmark, compared with other representative open-source models and compression techniques.\n| Method | OpenMoE | Sheared-Llama | Pythia | Open-Llama-v2 | LLM-Pruner | Compresso | **Read-ME** | LaCo | SliceGPT | Pythia | Open-Llama-v2 | Llama-2 |\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| #Param | 2.1B | 2.7B | 2.8B | 3.4B | 4.5B | 4.5B | **4.7B** | 4.7B | 4.8B | 6.9B | 6.9B | 6.9B |\n| MMLU | 26.2% | 26.4% | 26.9% | 25.7% | 23.9% | 25.9% | **38.9%** | 26.5% | 28.9% | 25.5% | 40.2% | 45.3%|\n\n\n[1] https://huggingface.co/princeton-nlp/Sheared-LLaMA-2.7B-Pruned\n\n\n\n**[W2 - Perplexity Comparison]**\nThank you for your suggestion. We tested the perplexity on Wikipedia, with the results presented in Table C. We set the sequence length to 1024, which is within the sequence limit of all models. \nTable C demonstrates that our method shows negligible performance degradation (from 3.91 to 3.94) compared to Llama-2, the pre-trained dense model.\nFurthermore, our method achieves better performance compared to most of the baseline methods of similar sizes, including Sheared-Llama, and Pythia. We will add the column to Table 1 in the future version, and we have also provided a tentative Table 1 in the uploaded PDF.\n\nTable C: Comparison on Wikipedia perplexity.\n| Model | # Params | Wikipedia PPL |\n|:---:|:---:|:---:|\n| Sheared-Llama | 2.7B | 6.77 |\n| Pythia | 2.8B | 6.02 | \n| Open-Llama-v2  | 3.4B | 3.69 |\n| Read-ME  | 4.7B  | 3.94 |\n| Pythia | 6.9B | 5.49 |\n| Open-Llama-v2 | 6.9B | 2.85 | \n| Llama-2 | 6.9B | 3.91 |", "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:
We thank all reviewers [R1(wHVe), R2(ZdPD), R3(EGvf), R4(AX77)] for their thoughtful and constructive feedback. We are grateful that the reviewers found our approach interesting and effective [R1,R2,R3], the paper well-written and well-organized [R2], and the experimental results thorough and extensive [R2,R3].

We have thoroughly addressed all of the concerns raised by the reviewers. As a model compression effort, we emphasize our superior accuracy-cost tradeoff and training efficiency compared to all baseline methods. We provided additional experiments on routing distillation loss (R1), dense counterpart comparisons (R1,R2, R4), other evaluation metrics (R4), cost analysis of the auto-regressive router (R3), and generalizability to other model families (R3). 

For other questions, we added detailed discussions for each reviewer. Please check our PDF if it is mentioned in the answer.

Meta-review: The manuscript has been reviewed by four reviewers. After the rebuttal, three reviewers gave positive ratings, while one gave borderline rejection.

Overall, the reviewers praised the submission for the motivation, the presentation, and experimental validation. 

The AC agrees with consensus from the reviewers, and believes that the manuscript, despite not flawless, should be presented to a large audience.

Please, however, do account for the comments in the final version.

Congrats!
Decision: Accept (poster)