Mixture of In-Context Experts Enhance LLMs' Long Context Awareness

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. This 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.

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. 2. 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.

Questions

No specific questions. I would appreciate a response with respect to the weakness stated above.

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

N/A

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.

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. - 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. - 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. - 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).

Questions

Please refer to the section above.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, the authors have discussed limitations.

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.

Strengths

1. The concept of mixing multiple RoPE angles within each head is innovative. 2. MoICE achieves state-of-the-art results on multiple benchmarks. 3. 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. 2. The method cannot be adapted to non-RoPE models, such as those using Alibi.

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. 2. Why does Table 3 show that MoICE performs better with Llama2 than with Mistral? 3. 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?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

I have not observed any red flag in terms of potential negative societal impact.

Summary

After rebuttal: raised score by 1 point after discussion. --- The 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. They 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.

Strengths

* The paper proposes an interesting idea and explains it reasonably well. * 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. 2. 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] [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.

Questions

Similar to weaknesses. 1. Have you tried tuning and evaluating on context length greater than 8k? 2. Have you considered evaluations on other benchmark tasks for long context?

Rating

6

Confidence

2

Soundness

2

Presentation

3

Contribution

3

Limitations

No limitations have been listed. Clear 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.

Satisfied with the additional evaluations

Thanks for adding these additional evaluations. I will let the ACs decide if additional experiments are acceptable at this time. The additional experiments do address the weaknesses I had noted in the submitted paper. Please make sure to include these in the main paper in future versions. If the new experiments are acceptable then I can increase my score by a point. I'll keep my score as is until we get a clarification.

Thanks for your rebuttal

The author's rebuttal has addressed my concerns and I raised my score accordingly.