Mixture of In-Context Experts Enhance LLMs' Long Context Awareness
Authors: Hongzhan Lin
Ang Lv
Yuhan Chen
Chen Zhu
Yang Song
Hengshu Zhu
Rui Yan
Venue: NeurIPS 2024 poster (2024)
Keywords: long context awareness
large language model
attention mechanism
Primary area: natural_language_processing
TL;DR: We propose MoICE, which enhances LLMs' context awareness. MoICE introduces a router in each attention head within LLMs, which dynamically directs the head's attention to contextual positions crucial for completing the head's function well.

Abstract: Many studies have revealed that large language models (LLMs) exhibit uneven awareness of different contextual positions. Their limited context awareness can lead to overlooking critical information and subsequent task failures. While several approaches have been proposed to enhance LLMs' context awareness, achieving both effectiveness and efficiency remains challenging. In this paper, for LLMs utilizing RoPE as position embeddings, we introduce a novel method called "Mixture of In-Context Experts" (MoICE) to address this challenge. MoICE comprises two key components: a router integrated into each attention head within LLMs and a lightweight router-only training optimization strategy:(1) MoICE views each RoPE angle as an 'in-context' expert, demonstrated to be capable of directing the attention of a head to specific contextual positions. Consequently, each attention head flexibly processes tokens using multiple RoPE angles dynamically selected by the router to attend to the needed positions. This approach mitigates the risk of overlooking essential contextual information. (2) The router-only training strategy entails freezing LLM parameters and exclusively updating routers for only a few steps. When applied to open-source LLMs including Llama and Mistral, MoICE surpasses prior methods across multiple tasks on long context understanding and generation, all while maintaining commendable inference efficiency.

Mixture of In-Context Experts Enhance LLMs' Long Context Awareness
Hongzhan Lin Ang Lv Yuhan Chen
Chen Zhu
Yang Song Hengshu Zhu? Rui Yan
Gaoling School of Artificial Intelligence, Renmin University of China
XiaoMi AI Lab Career Science Lab, BOSS Zhipin 4 NLP Center, BOSS Zhipin {linhongzhan, anglv, ruiyan/Sruc. edu. (chenyuhan5) com
Abstract
Many studies have revealed that large language models (LLMs) exhibit uneven awareness of different contextual positions. Their limited context awareness lead to overlooking critical information and subsequent task failures. While several approaches have been proposed enhance LLMSs context awareness, achieving both effectiveness and efficiency remains challenging. In this paper, for LLMs utilizing RoPE as position embeddings, we introduce novel method called "Mix- ture of In-Context Experts" (MoICE) to address this challenge. MoICE comprises key components: router integrated into each attention head within LLMs and lightweight router-only training optimization strategy: (1) MoICE views each RoPE angle ing the attention of a head to specific contextual positions. Consequently. each attention head flexibly processes tokens using multiple RoPE angles dynamically selected by the router to attend to the needed positions. This approach mitigates the risk of overlooking essential contextual information. (2) The router-only training strategy entails freezing LLM parameters and exclusively updating routers for only few steps. When applied to open-source LLMMs including Llama, Mistral and Qwen, MoICE surpasses prior methods across multiple tasks on long context understanding and generation, while maintaining commendable inference effi- ciency. Moreover, we also demonstrate the effectiveness of MoICE in pre-training a language model from scratch.
'in-context' expert, demonstrated be capable of direct-
Introduction
Although large language models (LLMs) have demonstrated impressive capabilities across diverse NL.P tasks, several studies [2 g 31 have pointed out that the contextual awareness of L.LMs
not as powerful as widely believed, constraining their application in tasks demanding extensive
contextual awareness, such as in-context learning (28] 46), coherent long text generation [49 26
and Retrieval-Augmentte Generation (RAG, [19 6 100 tasks necessitating in-context retrieval {
Liu et al. [27] identified a common issue termed the "Jost-in-middle" phenomenon, indicating that LLMs often exhibit weaker awareness of information situated in the middle of the long context compared to the beginning or end. Chen et al. [8] highlighted challenges arising from a mathematical property of RoPE 38 wide-used positional embedding in LLMs, which impedes attention to
'Equal contribution. Hongzhan Lin and Ang Lv proposed the idea of MoICE. Hongzhan Lin and Yuhan Chen designed the MoICE router architecture and implemented efficient code. Experiments conducted by Hongzhan Lin, while Ang Ly led writing. Code available httttp::/ggthhiii wttthuub.com/pnnnnoo//oooooo "Coresponding authors: Rui Yan truiyan@ruc. edu. en) and Yang Song fsongyang@kanzhun.com com
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

FIN
PWN
PEN
Hend
Head
Head
1
0
@
Oy
G
Gx
G
Byy
Matoi Buader
Male
(u) Attention Buckets
Ms- PoE
e MoICE
indinates taking Q and K and using RoPE angle 0, compute the attention scores.
to more information RoPR backgrounds.
Some methods developed to enhance LLMs' context awareness. (a) Attention Buckets g selects N different RoPEs and conducts N parallel inferences for each input. The outputs are then aggregated in the final layer. (b) Ms-PoE [49] employs a unique RoPE angle for each attention head. However, it needs additional forward pass for RoPE angle assignment. (c) MoICE integrates router within each attention head. This novel plug-in selects several of the most suitable RoPE angles for each token. The selected RoPE angles collectively contribute computing the attention scores. MoICE demonstrates superior memory efficiency and performance.
Figure I:
specific positions within the long context. Consequently, if critical information coincides with such positions, task performance suffers.
Many works [23 49 48 have attempted enhance the long-context awareness of LLMs. Central to these efforts is the enhancement of attention heads which serve as the linchpin for contextual awareness, given that FFNs in language models do not introduce token interaction. Chen al. B proposed an inference algorithm named Attention Buckets (AB), which enhanced the context aware- ness of LLMs by executing N inference instances, each with a distinct RoPE angle, and aggregated the outputs at the final layer. Zhang et al. 49 observed the varying awareness of attention heads to contextual positions. They proposed an inference algorithm named Ms-PoE. Ms-PoE enhances the utility of position-awar aware heads by re-scaling the positional embedding indices, equivalent to assigning each head a unique RoPE angle. Figure iilustrates these approaches. However, these approaches each come with their own drawbacks: AB conducts excessive redundant FFNs calculations, leading to high memory consumption. In Ms-PoE, determining distinct re-scale factor for every attention head needs additional forward pass. Meanwhile, each attention head still depends re-sealed static RoPE As highlighted by AB [8), this leads to limited awareness of certain contextual positions, thereby constraining its potential. Moreover, significant drawback of both AB and Ms-PoE lies in their static assignment of the RoPE angle for each attention head throughout the generation. However, the generation progresses, the positions of crucial tokens shift, necessitating corresponding adjustments in the required RoPE angles for each head.
single
In this study, we present Mixture of In-Contex Experts (MoICE), novel plug-in in of LLMs for enhancing context awareness. Specifically, We conceptualize a unique RoPE angle as an "in-context expert," as it can allocate a head's more attention to certain contextual positions [8] We integrate a router within each attention head, which discerns the potentially important tokens for the head and dynamically selects K RoPE angles that provide comprehensive awareness of these tokens for attention computation. Through the re-computation of only few query-key dot products, attention patterns computed with selected RoPE angles are aggregated to produce the final attention pattern. This approach yields two primary advantages: (1) It eliminates unnecessary computational overhead in AB, enhancing efffciency. (2) The dynamic expert selection of each head for arbitrary tokens introduces flexibility not attained in previous studies. This minimizes the risk of the initial RoPE angle assigned to head failing to work due to crucial token positions shifting during generation.
Consequently, MoICE not only surpasses AB's effectiveness but also achieves commendable effi- ciency. We name our approach as "Mixture of In-Context Experts" (MoICE) due to the aggregation of attention patterns calculated with different RoPE angles resembling the concept of "Mixture of Experts" (MoE, (37). When applying MoICE to open-source LLMs, freeze LLMs' parameters

---

and conduct lightweight training only on the MoICE routers. With only a few quick updates, MoICE surpasses many competitive baselines in tasks involving long-context generation and understanding.
In summary, our main contribution is the introduction of MolCE, a novel plug-in for enhancing L.MMs context awareness. It achieves head-and token-specific dynamic multiple RoPE angles assignment, outperforms previous methods across various tasks, and maintains commendable inference efficiency,
2 Background
We introduce some background of Mixture of In-Context Experts, including (1) the rotary position embeddings commonly used by mainstream LLMs, (2) the primary problem addressed in this paper: the limited context awareness of LLMs, (3) an explanation of the underlying reasons for this limitation.
Position embedding Positional embedding is crucial for Transformer 43 to perceive sequence order and compensate for the pesition-agnosiee nature of the attention mechanism. In this paper, we mainly focus on LLMs using Rotary Position Embedding (RoPE, (38 which is the prevalent position embedding in current LLMs. We discuss other position embeddings in Appendix [
In a Transformer layer with H attention heads employing RoPE, where d represents the hidden state dimension of each attention head, let d and kh. denote the query vector position n and key vector at position m in the hith head. To encode position information, RoPE initially applies rotary matrix the query and key vectors:
chh = Roy.. In E ERd,
R m = Rey.. km E IRd
0 0
r
cos n0 sin n0,
sin n0,,. ndl
0
Ra
where Fe
=
10..022-1.1
(2)
Here, Oj. B3 21/d ; E [0. d/2 - 1). is termed as the rotary angle of RoPE, and By is typically a fixed base. The subscript j serves to differentiate various RoPE angles +, each associated with a distinct By. a distinction necessary for discussions in Section This approach effectively incorporates relative position information between m and n in the query-key key product during attention computation:
TT In
kh (Re
A
(Re
AT -
Rey.--. mn k
@gTT
AA -Re, m-n
k/
Attn A mm = Softmax
= Softmax
Vd
Vd
Here, Attn am
denotes the attention score assigned by the h-th head at position n to position m.
Context awareness of LLMs LLMs struggle with limited context awareness, significantly impacting their performance in tasks like long-text generation [49), Retrieval-Augmennte Generation (RAG, 19 interactions
Artention before Softmax
B 10000
50000 8 100000
6
10 40). and multi-turn human-agent
g involving complex contexts. Liu et al. (27] iden- tified problem known as "I -os--ie-the-Middlee"" where LLMs process the beginning and end of the context well but have reduced awareness of the middle. Chen et al. ( observed that LLMs using RoPE exhibit uneven context awareness, favoring certain positions. Peysakhovich et al. (34] further highlighted that LLMs exhibit variable attention to document-level token segments based on their contextual positions. Lv et al. [29 observed that language models develop context awareness, es- pecially in their ability copy, through "grokking [5)." "M They
200
400
600
100
Relative tokee positions
Figure 2: Different θ, alter the up- per bounds of attention scores be- tween token and its x-distance neighbors. Each angle is distin- guished by its own base value B₁.
suggest pre-training models with increased regularization to enhance this capability.
3

---

Attention waveforms According to Chen et al. 8 LLM's uneven awareness of different contextual positions is due to RoPE's mathematical characteristics. Within RoPE, the attention score exhibits "waveforms" when retrieving the same token from the context, based on their relative positions. The troughs in these waveforms can impair task performance, especially when critical tokens are situated at these positions during generation. Different RoPE angles produce waveforms with troughs occurring at different positions. These phenomena are depicted Figure A detailed derivation of the depicted curves in Figure is provided in Appendix
3 Mixture of In-Context Experts
In this section, we first introduce the core component of MolCE, the MolCE router, detailed in Sc.tion[.1]] Subsequently, we delve into the optimization of MoICE in Section 3.2 Figure providess an overview of MolCE. The discussion in this section focuses solely on a single layer transformer for clarity, with the same principles applying to any other layers.
3.1 Architecture
We aim to design an enhanced attention mechanism in LLMs that dynamically attends to crucial information across various contextual positions required for complet- ing the head's 's function. As result, we can mitigate the performance drop caused by inadequate context awareness, Motivated by insights of Chen et al. [8), who demonstrated that a distinct RoPE angle + could direct the attention heads more focus on specific contextual positions, we pro- pose the integration anism. This routing mechanism is designed select the appropriate RoPE angles for processing a token. We im- plement the router with the SiLU activation function:
FFN A
MoICE
Softmax
Router
contextual aware routing mech-
SILU
0g
Multi-Layer Perceptron (MLP)
W
Wa
K
Router (q) == W3 (SiLU (W1q) @ (Weq)). (5)
Here, q is the query vector that encapsulates the contextual information for the task. This router input indicates the specific information for which the current token is "query- ing." W, W2 Rxx and W; E N are weight matrices, where N denotes the number of the number of RoPE angle candidates. Considering each head's distinct function [32 45 30). we integrate router into every at- tention head in the LLM. Notably, router's decision is independent of other heads and dynamic the context.
Figure The structure of MoICE, Only the router's parameters when plugged into an LLM. For clarity, the figure illustrates single head, with N=3 and K=2 as toy demonstration ex- amples.
trainable
As defined in Eq. [5] the router outputs an N-tuple distribution, indicating the weight it allocates for each RoPE angle. In each step in generation, the router selects K angles from a set of N angles vv) } for attention computation. Specifically, we first identify the K RoPE angles with the highest routing weights and normalize their relative weights using the Sof tmax function, resulting P R representing the probability distribution over the selected RoPE angles within the h-th head:
(O,, O2,
TopK-Indices" argsort (qh) K),
(6)
p Sof tmax [Router (a")) [TopK-Indices":)
where q represents the query at position n within the k-th attention head in a Transformer layer. Subsequently, aggregate the attention scores computed with these chosen K in-context experts based on their routing weights to derive the final attention scores for head h from position n to position
q . mn-n kb
Attn mm
Σ
P Softmas
Vd
JETpK Indices

---

Considering RoPE angles impact how attention heads allocate attention and focus specific contex- tual positions, we view each distinct RoPE angle as an in-context expert, in contrast traditional
in-weight experts [6 21 18 where the experts are learnable parameter weights. Given that these in-context experts together augment L.LMs' context awareness, we term this method Mixture of In-Context Experts (MoICE, Section
3 Figure @egureBillae Billustratess the overview of MolCE. Our proposed
MoICE has three major advantages:
(I) We only add additional computational overhead to the query-key dot products, resulting in minimal increase memory usage and negligible impact on inference speed (Section 4.2)
(2) MoICE dynamically selects suitable RoPE angles token- wise and wesd-wise, offering unprece- dented flexibility and unlocking the full potential of each attention head.
(3) Concerning LLMMs' context awareness enhancement, MoICE addresses longstanding issue: the relative position of the relevant information will shift during generation, leading to previous static modification of the attention heads 47 will be sub-optimal during practical generation. The contextual-aware dynamic routing MoICE is not bothered by this issue.
3.2
Router-only training
To train the newly incorporated MoICE routers in LLMs, the most straightforward way is simulta- neously update the LLMs' parameters alongside the routers. However, updating the original LLMs' parameters can result in catastrophic forgetting. Therefore, we propose a more effective and efficient strategy, the router-only training strategy, which freezes the LLMs' parameters and solely optimizing the routers.
Given an input sequence, we calculate the negative log-likelihood loss (Lnu) for language modeling. During backward propagation, only the parameters of MoICE routers are updated. To mitigate the possibility of a router favoring specific experts disproportionately [16] 44], we incorporate an auxiliary loss Cave following [16]. An ablation study on this auxiliary loss is in Table II Given an input of T tokens and N experts, we calculate the Lu by the scaled dot- dot-product between frequency vector F and probability vector P:
N = a N. Σ F Pj-
Eq. 8 avoids the router falling into sub-optimal solution favoring specific experts overwhelmingly, as its minimal is achieved when the routing probability is uniform. Here, a is the weighting factor for load balancing loss. F, denotes the proportion of the j-th expert selected across all positions and attention heads, while P, denotes the proportion of router weight assigned to expert j:
T
H
F,
1ij E TopK TopK-Indices,).
T
H
(9)
T
H
Pj
E
1i TopK-Indices ) P.D.
TXH
1=1
Our overall training objective is to minimize the following loss:
C = Lull + Laaz-
(10)
Experiment
4.1 Setup
To evaluate the efficacy of MoICE, we implement with open-source LLMs, which we will introduce later, and conduct lightweight training of MoICE routers small and general dataset. Subsequently, we evaluate the enhanced LLM's capability to zero-shot undertake multiple tasks in long context
understanding and generation, detailed in Section 4.2 and Section
5

---

Training data We use training datase which extracts the one thousand longest entries from OpenHermes AЛ] OpenHermes multi-ourree integrated dataset containing high-quality syn- thetically generated instruction and chat samples. A detailed analysis of other training data is in Section
Hyperparameters for MolCE-router-only training We froze all the original parameters of the open-source LLMs we used and only trained the MoICE router. Following Attention Buckets (8), employed the RoPE angle set of N 7 items. each assigned with base values follows: {1.0 x 10a, 1.75 x 10 1.8 x 10 1.9 x 10 2.0 x 10 2.25 x 10a 2.5 x 10*3 for Llama2-7B and Mistral-7B, {1.0 X 10", 1.25 x 10 1.4 x 10 1.8 x 10, 1.9 10 2.25 10% 2.5 x 10%} for Qwen1.5-7B. By default, unless otherwise specified, the attention head selects K=7 bases to ensure a fair comparison with [8] Section 4.2 introduces our baselines in detail. Section 5 delves into the impact of set size and the number of selected items.
We implement& warm-up strategy comprising 20% of the total steps, with a maximum learning rate of 0.0001. The batch size is 128, a set as 0.3. We train the MoICE routers for epoch (about minutes) on four A800-80G GPUs.
4.2 Long context understanding and generation
Following the L-Eval benchmark D. we evaluated the LLM with tasks categorized into two main groups: closed-ended and open-ended tasks. Closed-ended tasks primarily focus on the capacity for understanding and reasoning within long contexts, including tasks like multiple-chor questions from QuALITY [5). Coursera, TOEFL [TI], and True/False question answering from SFiction. F On the other hand, open-ended tasks include summarization generation and pormaformat question-answeriin answering tasks, requiring extracting information from lengthy in-context documents. The open-ended tasks comprise a subset of 181 questions drawn from 29 diverse long documents.
Baselines and open-source LLMs
In evaluating the efficacy of our proposed MoICE, we compare methods known for enhancing the capacity of LLMs to understand
it against several and generate long contexts. These baselines include two context extrapolation techniques: Positional Interpolation (PI, (77) and Dynamic NTK [15). Additionally, we consider two inference algorithms for awreeessswareness enhancement: Ms- 49 and Attention Buckets (8].
We evaluate all these methods alongside our MolCE on three representative open-source L.Mss that
utilize RoPE for positional embeddings: Llama2-7B-CCI 22) Misttral77B-Instruuc-vvo.. 20 and Qwen 1.wenl.5-BB-- 1.5------hhhh and 32,768, respectively. Mistral-7B employs sliding window attention (SWA) mechanism with window size of 4,096 tokens, enabling it to accommodate longer contexts than the default. Therefore, conduct experiments with a context length of 192 on Mistral-7B. using SWA the exclusive baseline for comparison. For PI and Dynamic NTK, we apply a scaling ratio of 1.5, while for the remaining baselines, we adhere to the hyperparameters specified in their original papers. All methods are tested on a single A800-80G GPU, except for applying AB to Mistral-7B -B-kk, which needs 2 GPUs due substantial memory requirements.
(2). Llama2-7B and Qwen (0wenl.5-7B support a pre-trained context length of 4,096
Evaluation metrics We adopt the exact match for closed-ended tasks. For open- open-ended tasks, we employ GPT-4-Turbo 33 as the judge to evaluate the effectiveness of various enhancement methods on open-source I.MMs. This evaluation compares their performance against GPT3.5 1ew..---Troooo6kkeee 16k-0613 across 181 questions.
Results and analysis We report experimental results in Table MoICE significantly enhances the overall performance of Alamm-2-7BB-le (with p-value < 0.02 in the t-test) in both closed-ender and open-ended tasks. On Mistral, MoICE outperforms all baseline models significantly (p- value 0.02). We also report the mean and standard deviation of MoICE in Table I0] Standard fine-tuning degrades the performance of original L.Mss, demonstrating catastrophic forgetting and proving that the improvement of MolCE does not stem from more training. These results underscore MolCE's
2.5- 1k-1k- longest
httpa://coursera.rgg http:://gttabbbooo///cctttt///tt

---

Table I: Experimental results on the L-Eval Benchmark [ Applying to various models, MoICE demonstrate superior performance compared previous competitive approaches. We emphasize the highest score in bold.
Closed Ended Task
Open Ended Task
Method
Coursera
QuALITY
TOEFL.
SFiction
Average
wins
ties
win-rate
Llama2-7B Hamma2-7B-Chal 42
36.77
38.12
55.02
60.16
47.52
68
117
34.94
Fine-tuning
32.85
30.20
$1.30
59.38
43.43
65
91
30.52
7
38.23
38.61
56.51
61.72
48.77
76
112
36.46
Dynamic NTK
15
39.11
55.76
62.50
49.41
112
MssPoE 40
39.24
40.10
55.76
63.28
49.60
110
40.41
41.09
56.88
61.72
50.02
85
114
39.23
MoICE (Ours)
39.83
42.08
56.13
64.84
50.72
89
118
40,88
Miistral-7--nnstruct--a 1B-Instrucc--8 -8k
20
4,.20
44.06
62.08
61.72
53.27
71
105
34.11
Finc-tuning
25.29
26.73
25.65
50.00
31.92
53
85
26.38
SWA
44.77
42.57
62.08
60.94
52.59
73
89
32.45
+I [2]
44.19
44.06
64.68
62.50
53.86
73
96
33.43
Dynamic
62.08
63.28
53,20
MssPoR 49
46.37
45.05
61.34
57.03
52.45
106
37.84
AB 8
46.08
42.57
62.08
62.50
53.31
87
110
39.22
MoICE (Ours)
47.82
46.53
64.68
62,50
55.38
85
117
39.36
Qwenl,5-7B-Chal
78.44
61,88
61.19
69.53
67.76
83
119
40.83
+ PT [
76.58
61.88
60.32
70.3
67.27
83
107
39.11
Dynnamic NTK 15
78.07
62,38
60.32
70.31
67.77
84
111
40,20
Mss-oo 49
75.47
60.89
60.47
71.88
67.18
OOM
OOM
N/A
78.44
0OM
OOM
OOM
N/A
OOM
OOM
N/A
MoMCEE (Ours)
78.44
62.87
61.77
71.09
68.54
105
41.59
Following win-rate = (win counts + 0.5 counts)
Table 2: Practical inference time (in minutes) / GPU memory costs (GB) on single A800-800 80G GPU for each method applied Llama2-7B-CHB (top) and Distral-7B-Instruce (bottom), respectively, Due to out-of-memory issues, AB can not accomplish many tasks, denoted as OOM in the table.
Method
Coursera
QuALITY +
TOFFL
SFiction J
Open-Ended |
Average |
78.7
18.1 83.162.. /62.5
19.9/56.5 56.5
5.0/33.2
45.9
20.0/ 61.8
Ms-PoE
4.1/2
6.0/27.8
6.7/
6.0/27.8
20.2/289
8.6 28.1
MoICF (Ours)
5.0/ 19.6
11.0/19.7
10.2/ 19.5
1.6/
34.2/
12.4 19.4
AB g
OOM 14.1 / 50.3 13.4/25.7
0OM
37.2/714
OOM
OOM
N/A
Ms-PoE A9
11.2/4844 7..9
9.8/25.4
50,3
72.8 / 62.4
22.5 /47.4
MolCE (Ours)
11.3//2.4
2.38
79.
22.5 24.2
efficacy in enhancing LLMs' ability to understand and generate long contexts, both of which require high context awareness. Furthermore, these results underscore the broad applicability of MoICE across different LLMs.
Regarding efficiency, we provide practical inference time and memory costs associated with AB, Ms-PoE, and MoICE in Table For fair comparison, we utilize Flash Attention [3] across all approaches. While achieving superior overall performance, MolCE remains at an inference speed similar Ms-PoE and notably excels in memory efficiency compared these two baselines.
We also perform further experiments on one additional long context benchmark LongBench 4 which are detailed in Appendix
4.3 Retrieval-augmmented generation (RAG)
Retrieval-augmeened generation (RAG) tasks involve retrieving numerous documents related to the current generation. The retrieved documents arranged in the context. RAG necessitates that LLMs have robust context awareness to pinpoint crucial documents, process the retrieved information effectively, and integrate it generate responses.
Following g 49, we employ the MDQA task to evaluate the efficacy of MoICE in enhancing LLMMs performance in RAG tasks. Mcanwhile, MDQA offers the bonus of allowing flexible control over
7

---

Table 3: The experiment results on the MDQA task. MoICE achieve superior average performance compared to previous competitive approaches. We emphasize the highest score in bold.
Method
3
5
7
10
Gap
Avg.
Llama2-77 1aaaa-2-B-Chtt
64, 14
65.95
64.97
62.67
67.53
4.86
65.05
MMsPPP
66.06
64.29
63.99
62.22
64.75
3.84
64,34
AB 7
66.36
66. 14
65.25
63.20
64.93
3.16
65.18
MoICE (Ours)
65.50
66.33
65.61
64.11
65.84
2.22
65.48
Method
1
8
15
23
30
Gap
Avg.
Mistral- -B- 1B-Inntruct--8 8kk
58.38
47.42
46.97
49.68
50.81
11.41
50.65
Ms-PoE 49
52.76
41.24
42.80
42.90
43.58
11.52
44.66
AB Y
58.57
47.57
47.12
49.83
50.96
11.45
50.81
+ MolCE (Ours)
61.81
52.54
52.43
50.36
49.34
12:47
53.30
the location of documents, enabling more precise evaluation of LLMs' context awareness across various contextual positions.
Our MDQA experiments leverage a subset of NaturallQuestions-Oper [25 24
consisting of 2,655
queries, following 44 (27) Each query is paired with context consisting of 10 or 30 documents (with an average of 1,722 or 5,046 tokens), depending on the model 1aama-2-7B-chie or Mistral-7B- Instruct-8k), tasked with answering based on this contextual information. Only one document among these comprises useful information for the given query. We compare Ms-PoE, AB, and MolCE, testing each method through iterations. For Llama, the relevant document positioned 1st, 3rd, 5th, 7th, and 10th within the context, while for Mistral, it is positioned 1st, 8th, 15th, 23rd, and 30th, respectively.
In Table MoICE on Llama demonstrates the highest average performance across most positions, showeasing remarkable stability. Its accuracy scores show minimal variation, with only a marginal difference of 2.22 points between its highest and lowest values. On Mistral, MoICE exhibits significant average improvement (p-value 0.02). Notably, when the relevant document is positioned at the end of the context, all methods on Llama exhibit decrease compared to the original model, although MoICE shows a minimal decline. This phenomenon also happens in the Mistral model. We posit that this decline may stem from the original model predominantly directing attention towards nearest However, as approaches enhance awareness of various contextual positions, the model's attention to the nearest documents is diffused by other positions, as its overall capacity for context awareness constant and limited. Nevertheless, MoICE consistently emerges as the superior-performen method overall across language models.
documents 27
5 Method analysis
In this section, we delve into a comprehensive analysis of the properties of MoICE. We illustrate how N, the total number of in-context experts (Section[5.1 and K, the specific number of selected in-context experts (Section 3.2),, influence MoICE. We further demonstrate that MoICE is robust to training data (Section 5.3 Additionally,
present case study demonstrating the dynamic selection of in-context experts for tokens during generation (Section 5.4). Finally, we verify the effectiveness of language model with MoICE architecture in pretraining stage (Section 5.5
5.1 The effect of expert total numbers N
We investigate the impact of the total number of experts. Employing the search algorithm proposed by Chen et al. B8, we obtain various sets of different sizes, each comprising complementary base values. The searched expert sets are detailed in Appendix []] We apply MoICE -lahh-t-BB-chat and test the model on L-Eval tasks. The results are presented in Table The results of the original Llama are denoted as (N=1) the table. As the table illustrates. MoICE demonstrates improvement to LLMs' context awareness with increasing N, with noticeable improvements even when N is as small as 3. However, as N reaches 9, the average performance close to N=7, indicating a performance plateau. This suggests that having ==7 experts is sufficient for general usage.

---

Table 4: The performance of Llama-2-7B chat enhanced by MoICE with N in-context experts. We show results marked with color to emphasize the improvements over the original model.
Method
Coursera
QUALITY
TOEFL.
SFiction
Avg.
Original (N=1)
36.77
38.12
55.02
60.16
47.52
N=3
37.65
40. 10
55.76
62.50
49.00
N=5
38.23
39.60
56.13
63.28
49.32
N=7
39.83
64.84
50.72
N=9
40.26
41.58
56.13
64.84
50.70
Table 5: The improvement of context awareness of Llama-2-7B-chat by MoICE, wherein each head dynamically selects diverse K experts (N=7). We show results marked with color to emphasize the improvements over the original model.
Method
Coursera
QuALITY
TOEFL
SFiction
Avg.
Original (N=l)
36.77
38.12
55.02
60.16
47.52
K=1
35.03
35.64
56.51
61.72
47.22
K=3
39.83
41.58
56.
64.84
50.60
K=5
38.52
39.60
56.13
64.84
49.77
K=7
39.83
42.08
56.13
64.84
50.72
Equal Weights
36.48
38.12
53.90
61.72
47.56
Random Weights
15.55
28.71
21.75
8.59
18.65
5.2
The effect of selected experts number K
With a fixed number of experts (7) we examine the effect of different numbers of chosen experts K with values of 3. 5, and We consider two additional setups where experts are selected with equal weights ("Equal Weights") and with random weights ("Random Weights") ), using Llama- 2-7B-chat as the case study. The results are presented in Table Setting K doesn't 't enhance or significantly degrade the model's performance. aligning with our assertion in the Introduction (Section I) assigning a single and unique RoPE angle each head inadequately explores the head's functionality. For K greater than 3, performance improvements become evident. This shows that the MolCE router in our method can select the appropriate combination of experts better aware the context. Randomly selecting experts ruins the model's language modeling ability, leading to aberrant outputs.
5.3 MoICE is robust to training data
We further analyze the impact of the data for training routers. We udditionally use three instruction fine-tuning datasets from different sources: self-instruct dataset, Airoboros [22]; and two datasets for LLM alignment with long context, Long-Alpace (9. and LongAlign 3 The hyperparameters
remain consistent as mentioned in Section 4 As presented Table 6
MoICE exhibits almost
identical scores when trained on different data, showcasing the robustness of our method.
Table The improvement of context awareness of Llama-2-7B-chat by MoICE trained on various data.
Training Data
Coursera
QuALITY
TOEFL
Avg. SFiction
OpenHermes 41
39.83
42.08
$6.13
64.84
50.72
Airoboros 22
39.68
41.58
56.13
64.84
50.56
Long-Alpaca
39.68
42.08
56.13
64.84
50.68
LongAlign 3
39.68
41.58
56.13
64.84
50.56
5,4 The visualization of dynamic routing states
provide a case study exemplifying the dynamic routing mechanism within MoICE during text generation. Depicted in Figure 4 in the Appendix, the MolCE router of each head independently selects distinct experts. each step of the generation process, these heads dynamically choose experts for each new token. This dynamic utilization of diverse RoPE angles within each attention

---

head maximizes the potential of attention heads across various inputs, a capability not attained in prior research, including both Attention Buckets and Ms-PoE.
5.5 Applying MoICE to the pretraining stage
We further evaluate the performance of language model with MoICE architecture in pretraining stage. Specifically, we train language model with a Llama architecture of 49M parameters, with and without MoICE respectively. We pretrain a small model from scratch and observe the effectiveness of MoICE. More experimental details be found in Appendix C We measure the model's context awareness on the Key-Valun Retrieval [27] task, which uses multiple randomly generated key-value string pairs as prompts evaluate the model's ability to extract the correct value corresponding to given query from the context. One prompt example can be found in Figure 6
Table 7: Experimental results on the Key- Value Retrieval task [4, We evaluate the model's 's awareness of different positions by controlling the position index of the key-value pair corresponding to the query among 10 key-value pairs in the prompt.
Position
3
5
7
Baseline
0.476
0.324
0.328
0.344
0.502
+ MolCE
0.652
0.762
0.634
0.622
0.814
From the results in Table 7 we can see that our model can significantly increase the contextual capabilities of the pretrained language model, which indicates the potential of scaling up our method in pretraining stage.
6 Conclusion
In this paper, we introduce a novel approach enhancing the context awareness of LLMs termed Mixture of In-Context Experts (MoICE). Through lightweight training, open-source LLMs such as I. Ilama and Mistral, enhanced by MoICE, demonstrate improved context awareness. Across numerous tasks demanding substantial context awareness, MolCE-enhanced LLMs consistently outperform competitive baselines, all the while maintaining commendable efficiency. A distinctive feature of MoICE is that it first implements head- and token-specific RoPE angles assignment for attention heads, pivotal factor contributing to its success. This paper underscores the need to address the inherent limitations in current LLMs and advocates for thorough exploration of their existing capabilities.
Acknowledgments and Disclosure of Funding
This work was supported by the National Natural Science Foundation of China (NSFC Grant No. 62122089), Beijing Outstanding Young Scientist Program NO. 1 2019100020098, and Intelligent Social Governance Platform, Major Innovation & Planning Interdisciplinary Platform for the "Double-Fine Class" Initiative, Renmin University of China, the Fundamental Research Funds for the Central Universities, and the Research Funds of Renmin University of China. Ang Lv is supported by the Outstanding Innovative Talents Cultivation Funded Programs 2024 of Renmin University of China.
References
[1] Chenxin An, Shansan Gong, Ming Zhong, Mukai Li, Jun Zhang, Lingpeng Kong, and Xipeng Qiu. L-eval: Instituting standardized evaluation for long context language models. arXiv preprint arXiv:2307. 1088, 2023.
[2] Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, Binyuan Hui, Luo Ji, Mei Li, Junyang Lin, Runji Lin, Dayiheng Liu, Gao Liu, Chengqiang Lu, Keming Lu, Jianxin Ma, Rui Men, Xingzhang Ren, Xuancheng Ren, Chuanqi Tan, Sinan Tan, Jianhong Tu, Peng Wang, Shijie Wang, Wei Wang, Shengguang Wu, Benfeng Xu, Jin Xu, An Yang, Hao Yang, Jian Yang, Shusheng Yang, Yang Yao, Bowen Yu,
10

---

Hongyi Yuan, Zheng Yuan, Jianwei Zhang, Xingxuan Zhang, Yichang Zhang, Zhenru Zhang, Chang Zhou, Jingren Zhou, Xiaohuan Zhou, and Tianhang Zhu. Qwen technical report. arXiv preprini arXiv:2309. 16609, 2023.
[3] Yushi Bai, Xin Lv, Jiajie Zhang, Yuze He, Ji Qi, Lei Hou, Jie Tang. Yuxiao Dong, and Juanzi Li. Longalign: A recipe for long context alignment of large language models. arXiv preprint arXiv:2401. 18058, 2024.
[4] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longhench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308. 14508, 2023.
[5] Samuel R Bowman, Angelica Chen, He He, Nitish Joshi, Johnny Ma, Nikita Nangia, Vishakh Padmakumar, Richard Yuanzhe Pang, Alicia Parrish, Jason Phang, et al. Quality: Question answering with long input texts, yes! NAACI. 2022, 2022.
[6] Jiawei Chen, Hongyu Lin, Xianpei Han, and Le Sun. Benchmarking large language mod- els in retrieval-augmented generation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 17754-17762, 2024.
[7] Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306. 15595, 2023.
[8] Yuhan Chen, Ang Lv, Ting-En Lin. Changyu Chen, Yuchuan Wu, Fei Huang. Yongbin Li, and Rui Yan. Fortify the shortest stave in attention: Enhancing context awareness of large language models for effective tool use. In Lun- Wei Ku, Andre Martins, and Vivek Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 11160-11174, Bangkok, Thailand, August 2024. Association for Computational Linguistics.
[9] Yukang Chen, Shaozuo Yu, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Long alpaca: Long-context wowncctioo-oolowiing models. https: thub. com/ dvlab- aawaabrreeaarc/ooogleeetoee
2023.
[10] Xin Cheng, Di Luo, Xiuying Chen, Lemao Liu, Dongyan Zhao, and Rui Yan. Lift yourself up: Retrieval-augmenned text generation with self-memory, Advances in Neural Information Processing Systems, 36, 2024.
[11] Yu-An Chung, Hung- Yi Lee, and James Glass. Supervised and unsupervised transfer learning for question answering. In Marilyn Walker, Heng Ji, and Amanda Stent, editors, Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume 1 (Long Pupers), pages 1585-1594, New Orleans, Louisiana, June 2018. Association for Computational Linguistics.
[12] Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models. arXiv preprint 2rXiv:2401.066 0666666,6 2024.
[13] Tri Dao. Flashattention-22 22 Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, 2024.
[14] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chupter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pages 4171-4186, Minncapolis, Minnesota, June 2019. Association for Computational Linguistics.
[15] emozilla. Dynamically scaled rope further increases performance of long context llama with zero fine-tuning. htttps:/www.. http//twwww.ddd..o.//// http://////oaaalLaaa////aaaattts 14mrgpr/ dymamicallly_scaledd further int.eeaeeenereeseeeeeeaeea

---

[16] 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.
[17] Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus. http:://Skylion007, htt.-kyylliiooooiiitttbbbb. 2o//peenneebTextCrrpus 2019.
[18] Zhuocheng Gong, Ang Lv, Jian Guan, Junxi Yan, Wei Wu, Huishuai Zhang, Minlie Huang. Dongyan Zhao, and Rui Yan. Mixture-of -f-ouulless Reinventing transformers as dynamic assemblies of modules, 2024.
[19] Yizheng Huang and Jimmy Huang. A survey on retrieval-augmeeted text generation for large language models. arXiv preprint arXiv:2404. 10981, 2024.
[20] 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 2rXiv:2310006855 2023.
[21] Albert 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. 044088, 2024.
[22] jondurbin. airoboros: using large language models to fine-tune large language models. https 7/github. gegithubbcom/jodduuree..... http:/jdduubnaarroooooooo
[23] He Junqing, Pan Kunhao, Dong Xiaoqun, Song Zhuoyang, Liu Yibo, Liang Yuxin, Wang Hao, Sun Qianguo, Zhang Songxin, Xie Zejian, et al. Never lost in the middle: Improving large lan- guage models via attention strengthening question answering. arXiv preprint a.X.1:2311.09198 2023.
[24] Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming- Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoe Le, and Slav Petrov. Natural questions: A benchmark for question answering research. Transactions of the Association for Compulational Linguistics, 7:452-466, 2019.
[25] Kenton Lee, Ming- Wei Chang, and Kristina Toutanova. Latent retrieval for weakly supervised open domain question answering. In Anna Korhonen, David Traum, and Lluís Mårquez, editors, Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 6086-6096, Florence, Italy, July 2019. Association for Computational Linguistics.
[26] Jia-Nan Li, Quan Tu, Cunli Mao, Zhengtao Yu, Ji-Rong Wen, and Rui Yan. Streamingdialogue: Prolonged dialogue learning via long context compression with minimal losses. arXiv preprint arXiv:2240..88 12 2024.
[27] Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts, 2023.
[28] Yao Lu, Max Bartolo, Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio, editors, Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 8096, Dublin, Ireland, May 2022. Association for Computational Linguisties.
[29] Ang Lv, Ruobing Xie, Xingwu Sun, Zhanhui Kang, and Rui Yan. Language models "grok" to copy, 2024.
[30] Ang Lv, Kaiyi Zhang, Yuhan Chen, Yulong Wang, Lifeng Liu, Ji-Rong Wen, Jian Xie, and Rui Yan. Interpreting key mechanisms of factual recall in transformer- based language models, 2024.
[31] Ang Lv, Kaiyi Zhang, Shufang Xie, Quan Tu, Yuhan Chen, Ji-Rong Wen, and Rui Yan. Are we falling in middle-intelligenee trap? an analysis and mitigation of the reversal curse, 2023.
12

---

[32] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Harfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovit, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. http://wwsooomrrriuuuuuuuu/b//// -eaaaiinnnnn-nnnnnittott heads/index.htm
[33] OpenAl. Gpt-4 technical report, 2024.
[34] Alexander Peysakhovich and Adam Lerer. Attention sorting combats recency bias in long context language models, 2023.
[35] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets, 2022.
[36] Ofir Press, Noah A. Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation, 2022.
[37] Noam Shazeer, * *Azalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gate mixture-of- experts layer. In International Conference on Learning Representations, 2017.
[38] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputiing 568:127063 2024.
[39] Yutao Sun, Li Dong, Barun Patra, Shuming Ma. Shaohan Huang, Alon Benhaim, Vishrav Chaudhary, Xia Song, and Furu Wei. A length-extrappoattbbe transformer. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki, editors, Proceedings of the 61st Annual Meeting of for
the Association for Computational Linguistics (Volume I: Long Papers), pages 14590-14604, Toronto, Canada, July 2023. Association for Computational Linguistics.
[40] Tao Tan, Yining Qian, Ang Lv, Hongzhan Lin, Songhao Wu, Yongbo Wang, Feng Wang, Jingtong Wu, Xin Lu, and Rui Yan. Pear: Dosition-embeedding-ggne enhances retrieval-augmennte generation with zero inference overhead, 2024.
attention re-weighting
[41] Teknium. Openhermes 2.5: An open dataset of synthetic data for generalist Ilm assistants, 2023.
[42] 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 preprini arXiv:2307.09228 2023.
[43] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[44] Haotao Wang, Ziyu Jiang, Yuning You, Yan Han, Gaowen Liu, Jayanth Srinivasa, Ramana Kompella, Zhangyang Wang, al. Graph mixture of experts: Learning on large-scale graphs with explicit diversity modeling. Advances in Neural Information Processing Systems, 36, 2024.
[45] Kevin Ro Wang. Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: circuit for indirect object identification in GPT-2 small. In The Eleventh International Conference on Learning Representations, 2023.
[46] Kaiyi Zhang, Ang Lv, Yuhan Chen, Hansen Ha, Tao Xu, and Rui Yan. Batch-ICL: Effec- tive, efficient, and order-agnostic in-context learning. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar, editors, Findings of the Association for Computational Linguistics ACI. 2024, pages 10728-10733 Bangkok, Thailand and virtual meeting, August 2024. Association for Computational Linguistics.
[47] Liang Zhang, Katherine Jijo, Spurthi Setty, Eden Chung, Fatima Javid, Natan Vidra, and Tommy Clifford. Enhancing large language model performance to answer questions and extract information more accurately, arXiv preprint axXiv:2402.01722..20 2024.
13

---

[48] Qingru Zhang, Chandan Singh, Liyuan Liu, Xiaodong Liu, Bin Yu, Jianfeng Gao. and Tuo Zhao. Tell your model where attend: Post-hoc attention steering for 1lms. arXiv preprint arXiv2231.00226,, 2023.
[49] Zhenyu Zhang, Runjin Chen, Shiwei Liu, Zhewei Yao, Olatunji Ruwase, Beidi Chen, Xiaoxia Wu, and Zhangyang Wang. Found in the middle: How language models use long contexts better 2024.
via plug-and and-play positional encoding. arXiv preprint
14

---

A Results on LongBench
LongBench is benchmark for bilingual, multitask, and comprehensive assessment of long context understanding capabilittes of large language models. We choose 16 tasks from LongBench, spanning five long-text application scenarios: Single-Doc QA (NarrativeQA, Qasper, and MultiFieldQA enn), Multi-Doc QA (HotpotQA, 2 2WikiMQA and Musique), Summarization (GovReport, QMSum and MultiNews), Few-shot Learning (TREC, TriviaQA, SAMSum and LSHT) and Synthetic Tasks (Passage Count, PassageRetrieval--e and PassageRettreeal--20
For evaluation, follow the setup of the LongBench benchmark. The input length of LLama2-7B is set to 4k, Mistral-7B to 8k, and Qwenl.5-7B to 32k. We report the average task scores for each scenario for all methods in Table
Table 8: Experimental results on the LongBench Benchmark bold.
We emphasize the highest score in
Method
Single-Doc
Multi-Doe
Summarization
Few-shot Learuing
Synthetic Tasks
Average
Llama2-7B-Chat 42
25.54
18.47
23.37
51.78
3.94
24.62
23.42
23.73
25.34
7.63
.TK 15
24.73
23.67
25.41
51.97
8.33
26.82
Ms-PoE
24.59
25.33
AB 8
27,06
22.94
25.52
52.84
8.62
27.40
MMMCEE (Ours)
23.70
9.71
20
27.20
19.89
24.22
$2.41
5.06
25.76
30.94
24.94
9.35
NTK LS
30.46
21.21
23.89
52.41
8.44
27.28
27.90
AB [
29.81
21.95
25.58
$4.42
7.89
27.93
(Ours)
31.09
Owen1.5-7B-Chul
34.66
35.91
25.77
56.89
33.83
37.41
28.28
17.08
57.51
32.67
32.03
NTK [5
31.35
23.98
24.95
56.64
32.50
33.88
Ms-PoE
OOM
0OM
0OM
AB
0OM
0OM
0OM
0OM
0OM
N/A
MoICE (Ours)
39.37
37.35
34,83
MolCE consistently shows improved average performance across models with 4k, 8k, and 32k context lengths, surpassing previous competitive approaches.
B Attention waveforms
In this section, we will elaborate on attention waveforms and the concept of complementarity. Assuming d Rhh is the attention score (before softmax) of the n-th position to mn-th position on the 1ithh attention head. The attention score can be formulated follows:
(Re 1T (Ro.mkm)
d/2-1
Re
Σ
an
[2 : 2j + 11jk 2ÿ (a
1=0
d/2-1
Anzy
kmnz + ah2z41 khnzs41) cos ((n m)0,)
km2j+1 Thzs41 kmn2))) sin ((n-((n)))) - m)0,),
where 8, B- B-///d, B is the rotary base of RoPE. Considering context-awareneet task, basic context awareness relies on attending to the same token and then copying its next token as outputs [32] To simplify the calculation, positions on the attention when retrieving the same token from the context. This impact (or the intensity of attention) is dubbed attention waveform W by [8].
set both q. and kt as all-one vectors to observe the impact of relative
15

---

d/2-1 Σ
W
2 cos ((n m))3)
As illustrated in Figure 2 the waveform exhibits two notable mathematical properties concerning attention scores: demonstrates fluctuations and undergoes a gradual decay with the increasing relative position (i.e., long-term decay).
Chen al. [8] observed that crucial information falling within the troughs of a waveform might diminish the performance of models employing RoPE. Mcanwhile, they pointed out the waveform, characterized by peaks and troughs, vary across RoPE bases. When leveraging the peaks of one attention wave to compensate for the overlook of the troughs in another, the model's capability to perceive and process information from diverse contextual positions can be enhanced. When a set of bases possesses this waveform characteristic, they are termed "complementary."
C Experimental details pretraining
In this section, we provide detailed experimental setup in Section 5.3 The model we has layers, 6 heads per layer, hidden layer dimension of 512 and intermediate size of 1280. We train the model using the WeenWeeText pretraining [17) dataset. We use four GTX A800-80Gs for training for 600k steps, with context window of 512. During pretraining, we use learning rate of 0.004 with a cosine annealing schedule and 6,000 warm-up steps.
D Discussions on more position embeddings
In this section, we discuss other position embeddings and demonstrate why they are not studied, e.g.. discarded in LLMs, do not exhibit attention waveform pattern, or are in the same family of RoPE: Firstly, the waveform pattern only exists in position embeddings constructed by cosine functions. Regarding the cosine embedding used in the original Transformer, it does exhibit long-term decay and periodic waveforms. However, this embedding disregarded in modern LLMs. Moreover, these embeddings are incorporated before the imitial model layer rather than during the attention computation, making hard to assess their impact on attention patterns. Secondly, the learned positional embeddings utilized in BERT [14] lack mathematical constraints to display periodic patterns. They are similarly added before the first model layer. Thirdly, Alibi [36] introduces linear bias to attention scores. The linear bias is devoid of wave patterns. The remaining popular positional embeddings used in LLMs such as xPos [39 are RoPE-based variants. These variants predominantly modified for long-context extrapolation rather than better context awareness. Therefore, they share the same shorteoming: tokens in attention trough are less focused on, thereby limiting context awareness, which is the study focus in our paper.
Details on expert sets E
Utilizing the RoPF-hase searching algorithm as proposed by Chen et al. g Table 9 resulting sets for different values of N.
illustrates the
Table 9: Searched Sets for Different N
N
Searched Set
3
{ {10,000, 18,000, 19,000)
5
[10,000, 17,500, 18,000, 19,000, 20,000) { |10,000, 17,500, 18,000, 19,000, 20,000, 22,500, 25,000]
7
9
(10,000, 13,500, 17,500, 18,000, 19,000, 20,000, 22,500, 24,000, 25,000)
16

---

165
0.155
150
0.16
24 58
24 500
T4
Taken
Figure 4: The routing weights across two distinct attention heads at the 27th layer in Llama-2-7B chat. The input tokens are randomly sampled from the training data, and the attention heads under observation also randomly selected. The horizontal axis depicts the input tokens, while the vertical axis represents experts with varying RoPE angles. Due to their distinct functions, each head dynamically chooses different experts process individual tokens. Input text can be found in Figure
5
F
Limitations
In this paper, we introduce plug-in module called MoICE, which is integrated into the attention heads of open-source LLMs to enhance their context awareness. One limitation that, due limited computational resources, we did not investigate the effectiveness of pretraining a language model with more parameters using the MoICE architecture. Furthermore, our proposed method exploits the potential for context awareness within LLMs, but it does not imbue the models with additional inherent context awareness abilities. Achieving this may necessitate more extensive data train all model parameters.
G Broader impacts and safety issues
Our novel lightweight plug-in approach efficiently enhances the context awareness of open-source LLMs. This advancement holds great promise for enhancing the effectiveness of LLMs across diverse scenarios characterized by extensive and complex contexts, such as RAG, tool utilization, and role-playing. The safety issue of our method mainly comes from the large language models used, as they might output toxic and biased texts, which is a common safety issue regarding LLM research.
Table 10: The mean and standard deviation of MoICE. We repeat L-eval [ experiments times with different random seeds. The randomness of MolCE results from the initialization of MolCE router when training, which causes slight differences in performance.
Open
Coursera
QUALITY
TOREL.
Average
win-rate
Liuma2 TB- Chat 12
36.77 4 0.00 38.12 0.00
55.02 0.00
60.16
0.00
47.52 a: 0.00
68.00 4: 0.00
117.00 = 0.00
34.94 0.00
MoICE
39.65 à 0.32
41.88
0.27
56.28
0.21
64.84
a: 0.00
50.66 1: 0.05
$9.00 4: 1.00
117.20
1.48
40.77 0.20
Mistral-78- St
M0
0.00
44.06
0.00
0.00 62.08
61.72
0.00
0.00
0.00
00.00-11
0.00
34.11
MclCE
48.08 4 0.24
46.73 a 0.27
65.35 081
62.18 a 1.19
55.59 4: 0.16
85.00 4: 1.10
115:20
2.05
39.39
17

---

Here's how retrieve the top tracks of artist:
"'java import
org-apache. wttpppehee.....oc.. httphhhtiiicc..... client.methods. Hethoos.HHtpteett import
wwwpphe.http.iige....e... www...lpp.ceettttmmmmmmmm Cetppp..ceeet.HHtCeeeeeeeeete ButpllieenBuilleer
public class SpotifyAPI private static final String
ARTIST _____ "spotify_ad ___aristt_ddd lisssssdd;i public static void main(String[] args) throws Exception { String accessToken
access "oocees"tokke"; HttpGet
httpGet -
hiiiss///ootoo-aaccc/ooottt--ss ARTIST ID));
N, "Bearer JSONO
response new
r.create() (aantt((.build(.eee. bww.dddd...ccceeeetteceeeeeeeeee for (Object trackObj response...
Figure 5: The input text in Figure To clearly display, we only show part of the input text, where the text with yellow background corresponds the decoded tokens.
"eb098018-bdb5":
"970cbed8-3665",
"0a9d957f-2256";
"be09fd63-4dfa""
"e2b49af9-d0e3"
"c5ed6251-085d"
"Bcce145105ee1"
"2d5932f7-acd8".
"eb2E4a8d-e0b7":
"e0acbe2c-d478".
0c8c0695-dd33":
"086d71ab-35c0",
79a1c002-4ba6""
"e69f5f62-250e"
"boclc9dffc13f"
"3ce6b12e-6223",
ee17ccc7-6342"
"41c410e1-776c"
483f6a4d-9aa4":
"3711356c-6df1"
"ee17cc77-6342";
"41c
Figure 6: The input prompt example in Section 5.5 We use 10 key-value pairs as examples in prompt, which includes query key. We insert the query key-value pair different positions of examples (In the prompt example above, the query key is inserted in the 9th position). The model's 's task is to find the value corresponding to query key and output it, which evaluates its ability of context awareness.
Table 11: The ablation study the auxiliary loss of MoICE. To assess the impact of this loss term, we perform an ablation experiment on two LLMS by removing from Eq. [0] The results show a significant drop performance, highlighting the positive impact of the auxiliary loss
Method
Coursera
QuALITY
TOFEL
SFiction
Average
Llama2-7B-chan
[42
MolCE w/o aux loss
39.83 39.83
41.58 42.08
56.13 56.13
62.50
50.01
MoICE w/ aux loss
64.84
50.72
Mistral-7B-IIntrucce
20
MoICE w/o aux loss
47.67
46.04
64.68
58.59
54.25
MoICE w/ aux loss
47.82
46.53
64.68
62.50
55.38
18

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer:
[Yes]
Justification: We propose effective and efficient approach for enhancing the context awareness of LLMs.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not 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.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We create separate "Limitations" section in Appendix.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section 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 well-specification, asymptotic approximations only holding locally). The authors should reflect 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 is low images are taken in low lighting. Or speech-to-text 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 fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
19

---

Justification:
Our paper does not introduce 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 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 or 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've shared the link the code. We promise open source.
Guidelines:
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 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either 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 to describe the particular way they provide for reprodueibility. the case of closed-source models, 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 have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
20

---

Answer: [Yes]
Justification:
We promise to open code. We have posted an anonymous code link.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https guwdes/dde/eeee Cooeeeoodeesssseeeeee 7/nips.com http://bbb///iiiidedeeee/eaaaaaa cassiicyncccyy 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. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
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 understand the results?
Question:
Answer: [Yes]
Justification:
See Section 4.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report the p-value in the t-test in the "Results and Analysis" paragraph of Section 4.2.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall 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).
21

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out 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 to reproduce the experiments?
Question:
Answer: [Yes]
Justification: See Section 4.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http://aaaar...cc//ba/a//hccccee
Answer: [Yes]
Justification:
We have conformed the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurlPS 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., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: See Section "Broader Impacts and Safeguards."
"
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
22

---

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

---

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


Reviews:
[{"review_id": "nhlUf6liGL", "reviewer": "Reviewer_BQEC", "summary": "After rebuttal: raised score by 1 point after discussion.\n\n---\nThe paper proposes a new strategy Mixture of In-Context Experts (MoICE) to increase the input context length of LLMs while allowing the model to function effectively on longer context inputs. Their key idea is to introduce a routing mechanism at each attention head of the transformer that allows selection of multiple positions (RoPE angles) dynamically to effectively process tokens at different parts of the input context.\nThey implement the proposed MoICE strategy on Llama-2-7B-chat and Mistral-7B-instruct-8k, and evaluate it on tasks in the L-Eval benchmark which consists of 4 close-ended tasks (Multiple choice questions, classification etc.) and ~181 questions on open-ended generation tasks.", "questions": "Similar to weaknesses.\n\n1. Have you tried tuning and evaluating on context length greater than 8k?\n2. Have you considered evaluations on other benchmark tasks for long context?", "limitations": "No limitations have been listed.\n\nClear limitations in terms of any memory usage or implementation details and challenges in training and datasets used for training would be of value to the community.", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "* The paper proposes an interesting idea and explains it reasonably well.\n* The implementation on 2 open source LLMs Llama-2-7B-chat and Mistral-7B-instruct-8k and analysis are valid.", "weaknesses": "1. A major weakness is the training of the router for context lengths of just 8k. While increasing from 4k to 8k is valuable. Having an experiment or model with larger input context lengths (perhaps atleast 16k) will be of great value.\n\n2. Evaluations of long context abilities on other benchmarks. While evaluations on L-Eval are reasonable, it would have been valuable to report on atleast one other popular benchmark such as ZeroScrolls [1]\n\n[1] Shaham, Uri, Maor Ivgi, Avia Efrat, Jonathan Berant, and Omer Levy. \"Zeroscrolls: A zero-shot benchmark for long text understanding.\" arXiv preprint arXiv:2305.14196 (2023) -- EMNLP 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your valuable feedback and suggestions! We hope our response could address your concerns.\n\n---\n### 1. To test MoICE with LLMs whose context length is greater than 8k. (Weakness 1 & Question 1)\n\nThanks for your valuable suggestions. We have implemented MoICE on Qwen1.5-7B-Chat, whose pre-training context length is **32k** . The results on LongBench [1] and L-eval are reported below, respectively. All experiments are conducted on one A800-80G GPU:\n\n\n|Method|Single-Doc QA|Multi-Doc QA|Summarization|Few-shot Learning|Synthetic Tasks|Average|\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|\n|Qwen2-7B-Instruct-32k|34.66|35.91|25.77|56.89|33.83|37.41|\n|+ PI|28.28|17.08|24.60|57.51|32.67|32.03|\n|+ NTK|31.35|23.98|24.95|56.64|32.50|33.88|\n|+ Ms-PoE|OOM|OOM|OOM|OOM|OOM|N/A|\n|+ AB|OOM|OOM|OOM|OOM|OOM|N/A|\n|+ MoICE|39.37|37.35|25.81|57.29|34.83|38.93|\n\n\n|Method|Coursera|QuALITY|TOEFL|SFiction|Average|wins|ties|win-rate%|\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|\n|Qwen2-7B-Instruct-32k|78.44|61.88|61.19|69.53|67.76|83|119|40.83|\n|+PI|76.58|61.88|60.32|70.31|67.27|83|107|39.11|\n|+NTK|78.07|62.38|60.32|70.31|67.77|84|111|40.20|\n|+Ms-PoE|75.47|60.89|60.47|71.88|67.18|OOM|OOM|OOM|\n|+AB|78.44|OOM|OOM|OOM|N/A|OOM|OOM|OOM|\n|+MoICE|78.44|62.87|61.77|71.09|68.54|91|105|41.59|\n\n\"OOM\" indicates that due to the extra memory cost required by Ms-PoE and AB, the inference on the long context failed due to out of memory.\n\nOn LLMs with the 32k context window, MoICE still brings  improvements in terms of context awareness. We will add these results in the revision.\n\n---\n### 2. Evaluations on benchmark beyond L-eval (Weakness 2 & Question 2)\n\nThanks for your valuable suggestions on the generalization of MoICE. We have added an additional benchmark Longbench [1], which is a bilingual multitask, and comprehensive assessment of long context understanding capabilities of large language models. We evaluate 16 tasks in 5 scenarios, and we report the average value for each scenario. All experiments are conducted on one A800-80G GPU.\n\n|Method|Single-Doc QA|Multi-Doc QA|Summarization|Few-shot Learning |Synthetic Tasks|Average|\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|\n|Llama2-7B-chat|25.54|18.47|23.37|51.78|3.94|29.85|\n|+ PI|23.42|23.73|25.34|51.63|7.63|31.30|\n|+ NTK|24.73|23.67|25.41|51.97|8.33|31.58|\n|+ Ms-PoE|23.68|24.59|25.33|51.66|8.04|31.75|\n|+ AB|27.06|22.94|25.52|52.84|8.62|32.21|\n|+ MoICE|26.31|23.70|25.60|52.34|9.71|32.25|\n\n|Method|Single-Doc QA|Multi-Doc QA|Summarization|Few-shot Learning |Synthetic Tasks|Average|\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|\n|Mistral-7B-Instruct-8k|27.20|19.89|24.22|52.41|5.06|25.76|\n|+ PI|30.94|24.94|26.24|49.34|9.35|28.16|\n|+ NTK|30.46|21.21|23.89|52.41|8.44|27.28|\n|+ Ms-PoE|27.90|17.89|20.28|48.59|8.95|24.72|\n|+ AB|29.81|21.95|25.58|54.42|7.89|27.93|\n|+ MoICE|31.09|22.98|26.69|55.76|8.02|28.91|\n\n\nOn LLMs with the 4k, 8k, and 32k context windows, MoICE consistently brings improvements in terms of context awareness. We will add these results in the revision.\n\n---\nOnce again, we appreciate your thoughtful review and feedback on our paper. Please let us know if you have any additional questions or suggestions.\n\n### References\n[1] Bai Y, Lv X, Zhang J, et al. Longbench: A bilingual, multitask benchmark for long context understanding[J]. arXiv preprint arXiv:2308.14508, 2023.", "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": "TmwvocpvO7", "reviewer": "Reviewer_jb4B", "summary": "The paper introduces the \"Mixture of In-Context Experts\" (MoICE) method to address uneven context awareness in large language models (LLMs) using Rotary Position Embedding (RoPE). The central element of MoICE is a router that selects different RoPE angles. The authors propose a loss function that learns to select RoPE angles for each head based on context information and encourages diverse RoPE angles among attention heads. MoICE is evaluated on two representative models—one with full attention and the other with sliding window attention—to demonstrate its effectiveness in both open-ended and close-ended long context evaluation tasks.", "questions": "1. How does training data impact MoICE's performance? It appears that MoICE uses additional data to learn its parameters. Even if the base model is frozen, this extra training data could positively affect benchmark performance.\n2. Why does Table 3 show that MoICE performs better with Llama2 than with Mistral?\n3. From Tables 4 and 5, should we always choose larger values for N and K? What are the cost implications of using larger N and K?", "limitations": "I have not observed any red flag in terms of potential negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The concept of mixing multiple RoPE angles within each head is innovative.\n2. MoICE achieves state-of-the-art results on multiple benchmarks.\n3. The paper includes sanity checks and analyses to elucidate the MoICE mechanism.", "weaknesses": "1. The auxiliary loss definition (Equations 8–10) appears to be ad hoc.\n2. The method cannot be adapted to non-RoPE models, such as those using Alibi.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your valuable feedback and suggestions! We hope our response could address your concerns.\n\n---\n### 1. The auxiliary loss appears to be ad hoc (Weakness 1)\n\nThanks for your valuable feedback. The auxiliary loss (Eq.8-10) is a widely adopted practice in MoE systems to address imbalanced routing strategies, which means the router shows too much preference for specific experts [1,2,3,4,5].\n\nIn our method, this loss plays a crucial role in alleviating such issues. Without it, the model's performance will decrease because the router might consistently choose specific RoPE bases without considering alternatives. To demonstrate the impact of this loss term, we have conducted an ablation study on Llama2-7b-Chat and Mistral-7B-Instruct-8k by removing it from Eq. 10. The results, shown below, indicate a decline in performance:\n\n|Llama2-7b-Chat|Coursera|QuALITY|TOEFL|SFiction|Average|\n|:-|:-:|:-:|:-:|:-:|:-:|\n|w/o aux loss|39.83|41.58|56.13|62.5|50.01|\n|w/ aux loss|39.83|42.08|56.13|64.84|50.72|\n\n|Mistral-7B-Instruct-8k|Coursera|QuALITY|TOEFL|SFiction|Average|\n|:-|:-:|:-:|:-:|:-:|:-:|\n|w/o aux loss|47.67|46.04|64.68|58.59|54.25|\n|w/ aux loss|47.82|46.53|64.68|62.50|55.38|\n\nWe will include this discussion and the results in the revision. Thank you for your valuable feedback!\n \n ---\n \n### 2. The adaption to non-RoPE models (Weakness 2)\n\nThank you for pointing this out. We discussed this in detail in Appendix B. Specifically, our method mainly focuses on resolving the issues brought by the wave pattern inherent in RoPE. Compared with non-RoPE, RoPE is more prevalently used in modern LLM, such as Llama, Mistral, Qwen, etc [6]. We believe the study of RoPE shortcomings will also help push the development of advanced LLMs.\n\n---\n\n### 3. The impact of training data on MoICE's performance (Question 1)\n\nThe MoICE router assigns dynamic routing weights for each (predefined and non-trainable) RoPE angle, which are used to calculate a weighted sum of attention scores. As a result, no extra knowledge or ability in additional data is introduced into the base model.\n\nIn Table 6, We observed **very similar** performance improvement when given varied datasets to train MoICE. This ablation study demonstrates that the improvements come from the routing strategies the router learned, not due to any supplementary knowledge derived from the extra training data.\n\n---\n\n### 4. Why does MoICE perform better with Llama2 than with Mistral? (Question 2)\n\nWe implemented MDQA tasks using a 30-doc context for Mistral according to its longer context window. MDQA tasks with a 30-doc context are more challenging than those with a 10-doc context, resulting in Mistral's weaker performance than Llama2.\n\n---\n\n### 5. The effect of N and K Values (Question 3)\n\nWe should not always choose larger values for N and K. \n\nFor N, too large values are unnecessary. As shown in Table 4, performance improves with increasing N up to N=7, after which it plateaus at N=9. Despite testing with larger N values, we observed no significant improvement in performance, while incurring additional costs in both training and inference.\n\nRegarding K, Table 5 demonstrates that the best performance is achieved when K equals N. Although increasing K raises inference costs, our efficiency remains to be the most practical compared to other baselines, even when the cost is at its maximum given the fixed N (K=N=7, as seen in Table 2).\n\n---\nOnce again, we appreciate your thoughtful review and feedback on our paper. Please let us know if you have any additional questions or suggestions.\n\n### References\n[1] Fedus W, Zoph B, Shazeer N. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity[J]. Journal of Machine Learning Research, 2022, 23(120): 1-39.\n\n[2] Zeng Z, Miao Y, Gao H, et al. AdaMoE: Token-Adaptive Routing with Null Experts for Mixture-of-Experts Language Models[J]. arXiv preprint arXiv:2406.13233, 2024.\n\n[3] Xue F, Zheng Z, Fu Y, et al. Openmoe: An early effort on open mixture-of-experts language models[J]. arXiv preprint arXiv:2402.01739, 2024.\n\n[4] Dai D, Deng C, Zhao C, et al. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models[J]. arXiv preprint arXiv:2401.06066, 2024.\n\n[5] Zoph B, Bello I, Kumar S, et al. St-moe: Designing stable and transferable sparse expert models. arXiv 2022[J]. arXiv preprint arXiv:2202.08906.\n\n[6] Zhang Z, Chen R, Liu S, et al. Found in the middle: How language models use long contexts better via plug-and-play positional encoding[J]. arXiv preprint arXiv:2403.04797, 2024", "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": "5ZSa8JHiIH", "reviewer": "Reviewer_uPsd", "summary": "This paper presents an approach, Mixture of In-Context Experts (MoICE) for enhancing the long-context awareness of LLMs with RoPE. Specifically, the authors use a router to dynamically select multiple RoPE angles for each attention head and token. They also use a lightweight router-only training strategy and freeze LLM parameters to only update the routers. Empirical evaluation shows that MoICE outperforms existing methods on long context understanding and generation tasks while maintaining efficiency.", "questions": "Please refer to the section above.", "limitations": "Yes, the authors have discussed limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed MoICE approach deals with the challenge of limited context awareness in LLMs. The idea of dynamically selecting RoPE angles is novel and effectively addresses limitations of the original RoPE technique.\n- The authors conduct extensive experiments across multiple tasks and datasets with LLaMA2-7B and Mistral-7B, demonstrating comparable performance of MoICE with competitive baselines while maintaining inference efficiency.\n- The paper also includes detailed ablation studies and analyses on different hyperparameters: expert total number N, selected expert number K, as well as different training data, showing that the method is robust.", "weaknesses": "- There is a lack of open-ended tasks in the experiments. The authors use a very small open-ended task which contains only 181 questions from 29 long documents. This is far from enough to show that the method could work well on general open-ended tasks. They should do more experiments on open-ended tasks, such as TriviaQA.\n- The proposed approach slightly modifies the language model architecture by adding a router layer and train it for long-context awareness. In fact, it would be more natural to apply this technique to pre-training stage to enhance the model's original ability to understand long contexts. The authors should discuss more about this, and if possible, show whether their method can be generalized to pre-training (even on smaller models, such as GPT-2).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your valuable feedback and suggestions! We hope our response could address your concerns.\n\n---\n### 1. Performance on general open-ended tasks (Weakness 1)\n\nThanks for your valuable suggestions. We have added an additional benchmark Longbench [1], which is a bilingual multitask, and comprehensive assessment of long context understanding capabilities of large language models. We evaluate 16 tasks in 5 scenarios, and we report the average value for each scenario. All experiments are conducted on one A800-80G GPU. **TriviaQA** is included in few-shot learning, and we report the results below.\n\n|Method|Single-Doc QA|Multi-Doc QA|Summarization|Few-shot Learning |Synthetic Tasks|Average|\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|\n|Llama2-7B-chat|25.54|18.47|23.37|51.78|3.94|29.85|\n|+ PI|23.42|23.73|25.34|51.63|7.63|31.30|\n|+ NTK|24.73|23.67|25.41|51.97|8.33|31.58|\n|+ Ms-PoE|23.68|24.59|25.33|51.66|8.04|31.75|\n|+ AB|27.06|22.94|25.52|52.84|8.62|32.21|\n|+ MoICE|26.31|23.70|25.60|52.34|9.71|32.25|\n\n\n|Method|Single-Doc QA|Multi-Doc QA|Summarization|Few-shot Learning |Synthetic Tasks|Average|\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|\n|Mistral-7B-Instruct-8k|27.20|19.89|24.22|52.41|5.06|25.76|\n|+ PI|30.94|24.94|26.24|49.34|9.35|28.16|\n|+ NTK|30.46|21.21|23.89|52.41|8.44|27.28|\n|+ Ms-PoE|27.90|17.89|20.28|48.59|8.95|24.72|\n|+ AB|29.81|21.95|25.58|54.42|7.89|27.93|\n|+ MoICE|31.09|22.98|26.69|55.76|8.02|28.91|\n\n\n|Method|Single-Doc QA|Multi-Doc QA|Summarization|Few-shot Learning|Synthetic Tasks|Average|\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|\n|Qwen2-7B-Instruct-32k|34.66|35.91|25.77|56.89|33.83|37.41|\n|+ PI|28.28|17.08|24.60|57.51|32.67|32.03|\n|+ NTK|31.35|23.98|24.95|56.64|32.50|33.88|\n|+ Ms-PoE|OOM|OOM|OOM|OOM|OOM|N/A|\n|+ AB|OOM|OOM|OOM|OOM|OOM|N/A|\n|+ MoICE|39.37|37.35|25.81|57.29|34.83|38.93|\n\n\n|TriviaQA|Llama2-7B-chat|Mistral-7B-Instruct-8k|Qwen2-7B-Instruct-32k|\n|:-|:-:|:-:|:-:|\n|Origin|84.44|85.61|86.26|\n|+ PI|84.75|84.39|85.26|\n|+ NTK|86.16|85.61|86.56|\n|+ Ms-PoE|85.65|84.39|OOM|\n|+ AB|85.82|85.9|OOM|\n|+ MoICE|86.01|86.44|87.14|\n\n\"OOM\" indicates that due to the extra memory cost required by Ms-PoE and AB, the inference on the long context failed due to out of memory.\n\nOn LLMs with the 4k, 8k, and 32k context windows, MoICE consistently improves the performances on various language tasks including general open-ended tasks. We will add the results in the revision.\n\n---\n\n### 2. Applying MoICE to the pre-training stage (Weakness 2)\n\nThanks for your valuable suggestions. Due to limited time and computing resources, it is not feasible to train such a large model. Therefore, we pre-train a small model from scratch and observe the effectiveness of MoICE. This demonstrates the potential of scaling up our method.\n\nSpecifically, we train a language model with a Llama architecture of 49M parameters, with and without MoICE respectively. The model has 4 layers, 6 heads per layer, and a hidden layer dimension of 512. We train the model with the OpenWebText dataset [2].\n\nWe use 4 GTX A800-80Gs for training for 600k steps, with a context window of 512, which takes 96 hours. (Given the limited time, this is the most extensive scenario we are able to test. We appreciate your understanding regarding these limitations.)\n\nWe measure the model's context awareness on the Key-Value Retrieval task [3]. The prompt for key-value retrieval is shown below:\n\n\t\"eb098018-bdb5\": \"970cbed8-3665\",\n\t\"0a9d957f-2256\": \"be09fd63-4dfa\",\n\t\"e2b49af9-d0e3\": \"c5ed6251-085d\",\n\t\"8ece1451-05e1\": \"2d5932f7-acd8\",\n\t\"eb2f4a8d-e0b7\": \"e0acbc2c-d478\",\n\t\"0c8c0695-dd3c\": \"086d71cb-35c0\",\n\t\"79a1c002-4ba6\": \"e69f5f62-250e\",\n\t\"b0c1c9df-c13f\": \"3ce6b12e-6223\",\n\t\"ee17cc77-6342\": \"41c410e1-776c\",\n\t\"483f6a4d-9aa4\": \"3711356c-6df1\",\n\t\"ee17cc77-6342\": \"41c\n\nWe use 10 key-value pairs as examples in prompt, which includes a query key. We insert the query key-value pair in different positions of examples (In the prompt example above, the query key is inserted in the 9th position). The model's task is to find the value corresponding to the query key and output it, which evaluates its ability of context awareness.\n\nThe performance of a pre-trained Llama model and a pre-trained Llama model with MoICE are shown below, respectively:\n\n||1|3|5|7|9|\n|:-|:-:|:-:|:-:|:-:|:-:|\n|Baseline|0.476|0.324|0.328|0.344|0.502|\n|+ MoICE|0.652|0.762|0.634|0.622|0.814|\n\nFrom the results, we can see that our model can significantly increase the contextual capabilities of the pre-trained language model.\n\n---\nOnce again, we appreciate your thoughtful review and feedback on our paper. Please let us know if you have any additional questions or suggestions.\n\n### References\n[1] Bai Y, Lv X, Zhang J, et al. Longbench: A bilingual, multitask benchmark for long context understanding[J]. arXiv preprint arXiv:2308.14508, 2023.\n\n[2] Peterson J, Meylan S, Bourgin D. Open clone of openai’s unreleased webtext dataset scraper[J]. 2019.\n\n[3] Liu N F, Lin K, Hewitt J, et al. Lost in the middle: How language models use long contexts[J]. Transactions of the Association for Computational Linguistics, 2024, 12: 157-173.", "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": "PjspGLzIHR", "reviewer": "Reviewer_K8VP", "summary": "Large language models (LLMs) have shown promise in various NLP tasks but often fall short in tasks requiring deep contextual understanding, such as coherent long-text generation and Retrieval-Augmented Generation (RAG). Challenges like the \"lost-in-middle\" phenomenon, where LLMs struggle with middle context information, and limitations from the widely-used Rotational Position Encoder (RoPE) significantly impact performance.\nThis work introduces the Mixture of In-Context Experts (MoICE) that dynamically selects optimal RoPE angles within each attention head to direct the attention of a head to specific contextual positions. Experiments are conducted on open-source models such as Mistral by freezing LLM parameters and exclusively updating routers for only a few steps.", "questions": "No specific questions. I would appreciate a response with respect to the weakness stated above.", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is very well written and easy to follow.  The claims are mostly well-substantiated with extensive experimentation supporting them. Background information is provided as needed without overwhelming the reader. The paper provides details on hyperparameters and compute to ensure reproducibility. The ablations, especially the one on visualization of dynamic routing states is very interesting.", "weaknesses": "1. It seems that the main weakness of the paper is in the evaluation section. Firstly in Table 1, the gains in performance by using MoICE are minimal. For instance, the gains on majority on the datasets are not before than 1%. It raises the question of the actual significance and practical implications of this approach. It would be great if authors could report mean and standard deviation of their results. \n\n2. MoICE seems promising for endowing LLMs with the ability to improve context awareness even at pretraining. While all experiments are currently conducted using pretrained LLMs, it will be interesting to see if one could pretrain LLMs with MoICE (maybe 2B size) on datasets such as C4 etc and then test on standard benchmarks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your valuable feedback and suggestions! We hope our response could address your concerns.\n\n### 1. Mean and standard deviation of Table 1\nThanks for your valuable comment. We reported the t-test results of MoICE in Table 1: the p-values ​are both less than 0.02, which illustrates the significant improvement of our method. In addition, we also have set different random seeds and repeated L-eval experiments 5 times, the mean and standard deviation of MoICE are reported below.\n\n||Coursera|QuALITY|TOEFL|SFiction|Average|wins|ties|win-rate%|\n|:-|:-|:-|:-|:-|:-|:-|:-|:-|\n|Llama-2-Chat|36.77 $\\pm$ 0.00|38.12 $\\pm$ 0.00|55.02 $\\pm$ 0.00|60.16 $\\pm$ 0.00|47.52 $\\pm$ 0.00|68.00 $\\pm$ 0.00|117.00 $\\pm$ 0.00|34.94 $\\pm$ 0.00|\n|+ MoICE|39.65 $\\pm$ 0.32|41.88 $\\pm$ 0.27|56.28 $\\pm$ 0.21|64.84 $\\pm$ 0.00|50.66 $\\pm$ 0.05|89.00 $\\pm$ 1.00|117.20 $\\pm$ 1.48|40.77 $\\pm$ 0.20|\n\n\n\n||Coursera|QuALITY|TOEFL|SFiction|Average|wins|ties|win-rate%|\n|:-|:-|:-|:-|:-|:-|:-|:-|:-|\n|Mistral-7B-Ins.|45.20 $\\pm$ 0.00|44.06 $\\pm$ 0.00|62.08 $\\pm$ 0.00|61.72 $\\pm$ 0.00|53.27 $\\pm$ 0.00|71.00 $\\pm$ 0.00|105.00 $\\pm$ 0.00|34.11 $\\pm$ 0.00|\n|+ MoICE|48.08 $\\pm$ 0.24|46.73 $\\pm$ 0.27|65.35 $\\pm$ 0.81|62.18 $\\pm$ 1.19|55.59 $\\pm$ 0.16|85.00 $\\pm$ 1.10|115.20 $\\pm$ 2.05|39.39 $\\pm$ 0.21|\n\nAll methods in our paper use greedy decoding which is determinate, the randomness of MoICE results from the initialization of MoICE router when training.\n\n---\n\n### 2. Applying MoICE to the pre-training stage (Weakness 2)\n\nThanks for your valuable suggestions. Due to limited time and computing resources, it is not feasible to train such a large model. Therefore, we pre-train a small model from scratch and observe the effectiveness of MoICE. This demonstrates the potential of scaling up our method.\n\nSpecifically, we train a language model with a Llama architecture of 49M parameters, with and without MoICE respectively. The model has 4 layers, 6 heads per layer, and a hidden layer dimension of 512. We train the model with the OpenWebText dataset [2].\n\nWe use 4 GTX A800-80Gs for training for 600k steps, with a context window of 512, which takes 96 hours. (Given the limited time, this is the most extensive scenario we are able to test. We appreciate your understanding regarding these limitations.)\n\nWe measure the model's context awareness on the Key-Value Retrieval task [3]. The prompt for key-value retrieval is shown below:\n\n\t\"eb098018-bdb5\": \"970cbed8-3665\",\n\t\"0a9d957f-2256\": \"be09fd63-4dfa\",\n\t\"e2b49af9-d0e3\": \"c5ed6251-085d\",\n\t\"8ece1451-05e1\": \"2d5932f7-acd8\",\n\t\"eb2f4a8d-e0b7\": \"e0acbc2c-d478\",\n\t\"0c8c0695-dd3c\": \"086d71cb-35c0\",\n\t\"79a1c002-4ba6\": \"e69f5f62-250e\",\n\t\"b0c1c9df-c13f\": \"3ce6b12e-6223\",\n\t\"ee17cc77-6342\": \"41c410e1-776c\",\n\t\"483f6a4d-9aa4\": \"3711356c-6df1\",\n\t\"ee17cc77-6342\": \"41c\n\nWe use 10 key-value pairs as examples in prompt, which includes a query key. We insert the query key-value pair in different positions of examples (In the prompt example above, the query key is inserted in the 9th position). The model's task is to find the value corresponding to the query key and output it, which evaluates its context awareness.\n\nThe performance of a pre-trained Llama model and a pre-trained Llama model with MoICE are shown below, respectively:\n\n||1|3|5|7|9|\n|:-|:-:|:-:|:-:|:-:|:-:|\n|Baseline|0.476|0.324|0.328|0.344|0.502|\n|+ MoICE|0.652|0.762|0.634|0.622|0.814|\n\nFrom the results, we can see that our model can significantly increase the contextual capabilities of the pre-trained language model.\n\n---\nOnce again, we appreciate your thoughtful review and feedback on our paper. Please let us know if you have any additional questions or suggestions.\n\n### References\n[1] Peterson J, Meylan S, Bourgin D. Open clone of openai’s unreleased webtext dataset scraper[J]. 2019.\n\n[2] Liu N F, Lin K, Hewitt J, et al. Lost in the middle: How language models use long contexts[J]. Transactions of the Association for Computational Linguistics, 2024, 12: 157-173.", "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:

Meta-review: This paper proposes a new effective method for long context length for LLMs. 
All reviewers agree the contributions of this paper with all positive scores.

AC also agree with the reviewers, so recommends accepting this paper.
Decision: Accept (poster)