UniMem: Towards a Unified View of Long-Context Large Language Models

Summary

This paper summarized the previous representative works of memory-augmented long LLMs and proposed a unified framework to formulate the methodology of these methods. This paper well resolved the previous issue of incompatible evaluation suites between different method and instead evaluated all methods under the same setting. Additionally, the comprehensive ablation studies are conducted to find the best design choices for each component in the memory module, and the authors used this best hyperparameter groups to train a best model, UniMix.

Rating

7

Confidence

4

Ethics flag

1

Reasons to accept

1. The research problem is significant to the community. The long-context LLM suffered from the unfair comparison between different methods for quite long. 2. The architecture design lessons and conclusions drawn from the experiments are very useful for other researchers. It helps everyone working on memory-augmented LLMs to have an optimal architecture design.

Reasons to reject

1. The evaluations are limited to ppl metrics on language modeling dataset. However, the recent research trend in long-context llm community do not regard the lower ppl as the sign for strong long-context understanding capability. I strongly suggested that the authors could consider the e Needle-ina-Haystack tests, the long-document qa experiments, and the long-context in-context learning experiments. 2. The baselines are only trained for 0.1B tokens, which is a small data size for long-context adaptation. Even if the recent consensus on the optimal long-context extension data size has decreased from 100B (Wang et al., 2023) to 5B (Fu et al., 2024), the 0.1B is too small to train the model to understand the long-context input well. [Wang et al., 2023] Augmenting Language Models with Long-Term Memory [Fu et al. 2024] Data Engineering for Scaling Language Models to 128K Context

Questions to authors

1. How you sample and filter the 0.1B tokens? What are the data resources for that? 2. The number of trainable parameters for different methods vary a lot. Is the same lr 5e-5 fair for the implementation of each method? 3. Why Github is considered as a long-context corpus for evaluation? The more common practice is using ArXiv paper dataset.

Summary

This paper focused on extending the context length of large language models. It firstly reviewed several common existing context modeling approaches in Transformer-based LLMs, including Memory Management, Memory Writing, Memory Reading, and Memory Injection. Then it proposed an unified framework that reformulates existing long-context methods from the view of Memory augmentation of LLMs. It also proposed UniMix, an innovative approach that integrates the strengths of these algorithms. Experimental results show that UniMix achieves superior performance in handling long contexts with significantly lower perplexity than baselines.

Rating

6

Confidence

3

Ethics flag

1

Reasons to accept

1. The unified framework for existing long-context methods is elegant in math. 2. The proposed approach has some advantage on metric of perplexity.

Reasons to reject

1. The baselines can be supplemented. There are quite a lot of recent works on extending context of LLMs [1][2][3], I would suggest this paper should compare with them. [1] Pal et al, Giraffe: Adventures in Expanding Context Lengths in LLMs, https://arxiv.org/pdf/2308.10882 [2] Peng et al, YaRN: Efficient Context Window Extension of Large Language Models, https://arxiv.org/pdf/2309.00071 [3] Tworkowski et al, Focused Transformer: Contrastive Training for Context Scaling, https://arxiv.org/pdf/2307.03170 2. The evaluation method can be enhanced. The paper only considered evaluating perplexity of differenct models, however, it ignored some popular downstream tasks tailored for long-context LLMs, such as Needles in a Haystack [4], Document-based Q&A task [5]. I would suggest the author can evaluate their approach in those benchmarks. [4] Kuratov et al, In Search of Needles in a 11M Haystack: Recurrent Memory Finds What LLMs Miss, https://arxiv.org/pdf/2402.10790 [5] Ma et al, MEGALODON: Efficient LLM Pretraining and Inference with Unlimited Context Length, https://arxiv.org/pdf/2404.08801

Questions to authors

See above.

Summary

**Summary.** In this paper, the author aim to make a systematic review for the existing techniques on long-context LM. The paper proposes a new framework referred as UniMem, in which four dimensions are introduced for analysis: memory management, memory reading, memory writing, memory injection. The paper also presents a UniMix where the existing memory reading and writing techniques are combined to handle the long-context. Finally the paper conducts empirical study based on several long-context language modeling datasets.

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

**Strength.** Long-context language modeling is a popular and critical research topic today. Given the quick emergence of new models and algorithms, systematic discussions in this field will be beneficial to the technical community.

Reasons to reject

**Weakness.** However, there are concerns in the following perspectives regarding the paper's quality. - The value of the proposed framework is unclear. Although the existing techniques are discussed based on the four dimensions introduced in this paper, very few insights can be drawn from the discussions. - The technical novelty of this paper is limited. This paper seems like a research-oriented paper. However, the proposed method is merely a simple combination of the existing approaches. - The experimental study needs to be improved. Despite that the experiment presents preliminary results on language modeling, none of other important evaluation methods and benchmarks are included, like NIHS, topic retrieval, longbench, et. al.

Summary

This paper presents a unified framework, UniMem, to reformulate most existing long-context modeling works from a perspective of memory augmentation for LLMs. The paper briefly reviews 4 representative efforts from four memory related dimensions: memory management, writing, reading and injection, in each of which the authors select one model to explain. After analysis regarding the 4 dimensions, the authors develop a new method, UniMix, to benefit from the best features/advantages from the discussed four dimensions. Experiments on two different scales of models, tinyLLaMA-1.1b and LLaMA2-7, to discuss the effect of different long-context solutions as well as position interpolation.

Rating

7

Confidence

4

Ethics flag

1

Reasons to accept

1, To consider the long-context modeling works from a unified perspective is new and seems to be inspiring for the community. 2, a new method is proposed based on the analysis using the proposed framework UniMem, which seems to be effective.

Reasons to reject

1, The existing works covered in the proposed analyzing framework are still limited. For example, many recent works addressing modified position interpolations are not mentioned or discussed. 2, Using perplexity to evaluate long-context modeling ability solely may not be a choice. The authors should consider using different down streaming tasks, even using \textit{needle in a haystack style} evaluation, which could be a better choice for illustration. 3, Different long-context modeling methods may need different amounts of fine-tuning data, it would be more convincing to plot the trends of using different amounts of fine-tuning data. And it would be better give more about the statistics of the data (there have been several data engineering works to address the importance of fine-tuning data to improve the long-context modeling abilities. 4, The section 4.4 may not be a fair comparison.

Questions to authors

please see the comments in the previous section.

Thanks for the response

Thanks for the response which addresses most of my questions. I will raise my scores.