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

Authors: junjie  fang
likai  tang
hongzhe  bi
yujia  qin
si  sun
zhenyu  li
haolun  li
yongjian  li
xin  cong
yankai  lin
yukun  yan
xiaodong  shi
sen  song
zhiyuan  liu
maosong  sun
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-23 06:35:50

Abstract: Long-context processing is a critical ability that constrains the applicability of large language models (LLMs). Although there exist various methods devoted to enhancing the long-context processing ability of LLMs, they are developed in an isolated manner and lack systematic analysis and integration of their strengths, hindering further developments. In this paper, we introduce UniMem, a Unified framework that reformulates existing long-context methods from the view of Memory augmentation of LLMs.  Distinguished by its four core dimensions—Memory Management, Memory Writing, Memory Reading, and Memory Injection, UniMem empowers researchers to conduct systematic exploration of long-context methods. We re-formulate 16 existing methods based on UniMem and analyze four representative methods: Transformer-XL, Memorizing Transformer, RMT, and Longformer into equivalent UniMem forms to reveal their design principles and strengths. Based on these analyses, we propose 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. The code is publicly available at https://github.com/thunlp/UniMem.

Main paper: 
 introduction: 
 Transformer-based (Vaswani et al., 2017) large language models (LLMs) have ushered in a new era of AI (Brown et al., 2020; Touvron et al., 2023; Bommasani et al., 2021; OpenAI, 2022) , leading to various applications (e.g., natural language processing (Brown et al., 2020) , code generation (Li et al., 2022) ). As the scope of LLM applications expands, the demand for handling longer contexts becomes paramount (e.g., parsing long documents (Gao et al., 2023) and managing intricate dialogues (Zheng et al., 2023) ). However, the inherent computational limitations of traditional Transformer architectures, stemming from the quadratic complexity of self-attention mechanisms, pose significant challenges for scaling to long contexts.Various methods have been developed to address this challenge, including optimizing the quadratic complexity of Transformer Dai et al. (2019) ; Beltagy et al. (2020) ; Wu et al. (2022) , adjusting the position encoding of Transformer to handle longer contexts than during pre-training (Chen et al., 2023a; Peng et al., 2023) , and designing non-Transformer architectures (Gu et al., 2021; Xiong et al., 2023; Gu & Dao, 2023; Sun et al., 2023) . Among them, the first type of method has been extensively explored and utilized, but lacks integration among its various types, prompting our primary focus on it. Upon reviewing this type of Table 1 : Long-context methods are categorized under four dimensions using our UniMem framework (Section 3), introducing UniMix to combine strengths. Excluding FoT, KNN-LM, and TRIME-focused on inference-other methods apply to both training and inference. method, we further divide it into three categories: (1) Context Caching, which stores the intermediate hidden states of context and retrieves relevant ones when processing successive contexts (Wu et al., 2022; Dai et al., 2019; Bertsch et al., 2023; Tworkowski et al., 2023) ; (2) Context Compression, which compresses the context into condensed tokens and prepends them before the successive context (Bulatov et al., 2022; Chevalier et al., 2023; Guo et al., 2021; Zhang et al., 2021) ; (3) Sparse Attention: which elaborately designs sparse attention masks to reduce computational complexity to extend input length (Beltagy et al., 2020; Zaheer et al., 2020; Ding et al., 2023; Kitaev et al., 2020; Roy et al., 2021; Chen et al., 2023b) .While the above approaches effectively reduce computational complexity and enhance longcontext capabilities, they still face the following problems. Firstly, there is a notable absence of equitable evaluation across various approaches. Given that these approaches adhere to distinct design principles and are tested within disparate datasets and settings, assessing their efficacy and adaptability fairly becomes exceedingly complex. Such complexity complicates the comparative analysis to identify the strengths of each approach. Secondly, the diverse implementations of these approaches significantly impede the strength integration of each one into a singular model. This impediment to synergistic integration restricts the potential for substantial progress in the effective processing of long contexts by LLMs. This work introduces UniMem, a unified framework that integrates existing long-context methods for optimizing Transformer computation complexity, viewed via memory augmentation. Specifically, UniMem consists of four essential dimensions: (1) Memory Management determines how much past information is stored and how old memory is replaced, impacting how LLM recalls and uses previous context; (2) Memory Writing describes how the model converts recent information into a memory format, affecting the way that past data is summarized and accessed; (3) Memory Reading focuses on how the model retrieves information from the memory bank, crucial for efficiently using stored information to understand the current context; (4) Memory Injection determines which model layers to augment memory information, influencing overall efficiency and effectiveness in processing long contexts.Based on the UniMem framework, as illustrated in Table 1 , we re-formulate 16 prevailing long-context approaches from the unified view of the four fundamental dimensions. We further select 4 representative approaches from each long-context category: Transformer-XL (Dai et al., 2019 ), Memorizing Transformer (Wu et al., 2022) , RMT (Bulatov et al., 2022) , and Longformer (Beltagy et al., 2020) , and introduce their equivalent formula under the UniMem framework. This unification facilitates a clearer understanding of their subtle interconnections and distinct characteristics. Our analysis indicates a shared foundation in design principles across these methodologies, despite their superficial differences. Additionally, built upon the unified framework, we introduce a new method UniMix which synthesizes the strengths of three categories of long-context algorithms. The experimental results demonstrate that UniMix can achieve superior performance compared to the existing long-context approaches, including position interpolation techniques. Furthermore, our investigation into UniMix reveals two beneficial conclusions: (1) By integrating different memory dimensions, UniMix demonstrates strong robustness and is less sensitive to hyper- (2) Increasing the number of memory layers does not necessarily correlate with performance gains. Rather, the strategic positioning of a single memory layer can achieve comparable effectiveness to integrating additional memory across half of the model's layers. 
vanilla transformer: 
 Given a Transformer-based LLM featuring N stacked layers, and an input sequence X = {x 1 , x 2 , ..., x T } consisting of T tokens, the LLM processes the sequence X layer-wisely, constructing a set of hidden statesH n = {h n 1 , h n 2 , ..., h n T }, where n ∈ [1, N].The n-th LLM layer's hidden states H n are derived from the (n -1)-th layer using feed-forward networks (FFN), multi-head attention (Attn), and softmax (Sftx):H n FFN ←-Attn(Q n , K n , V n ; A) = Sftx Q n (K n ) T √ d + A V a , A ij = 0 if j ∈ [0, i],-∞ else.(1)EQUATION)where A ∈ R T×T is often a causal mask matrix controlling the visible context field and W n q/k/v ∈ R d×d is the projection matrix of multi-head attention in n-th LLM layer.Despite LLMs' strong context modeling, their attention mechanism's T 2 computational complexity challenges tasks with very long contexts. To alleviate this challenge, the community has made many improvements, and we analyze three main types of long-context approaches including context caching, context compression, and sparse attention methods.Context caching methods innovatively tackle long-context modeling by splitting them into segments, and storing the hidden states of historical segments and retrieves relevant ones when processing the current segment. To be specific, a long input sequence X = {x 1 , x 2 , ..., x T } can be split into a succession of segments S τ = {x (τ-1)×L+1 , x (τ-1)×L+2 , ..., x (τ-1)×L+L } with the fixed token length L (L ≪ T), where τ ∈ [1, T/L]. These methods utilize not only the hidden states from the current segment but also those from previous segments during sequential processing. Broadly, there are two categories: one involves storing the hidden states of a previous segment in an external memory cache (Dai et al., 2019) , and the other involves caching key-value pairs of multiple previous segments (Wu et al., 2022; Bertsch et al., 2023; Tworkowski et al., 2023) . This paper highlights Transformer-XL (Dai et al., 2019) and Memorizing Transformers (Wu et al., 2022) as exemplars to demonstrate how external memory integration enhances the capacity for long-context modeling for these two categories.Transformer-XL utilizes both hidden states H n τ of the current segment S τ and the hidden states H n τ-1 processed by the previous segment S τ-1 , as shown in Figure 1 (a). This recurrence way enables Transformer-XL to learn dependencies beyond long contexts of fixed length without breaking temporal coherence. Formally, the n-th hidden states H n τ of segment S τ is constructed layer-wisely as follows:H n τ FFN ←-Attn(Q n τ , K n τ , V n τ ; A), A ij = 0 if j ∈ (i, i + L],-∞ else.(3)EQUATIONwhere mask matrix A ∈ R L×2L is illustrated in Figure 2 Memorizing Transformer designates a specific layer for memory augmentation. When modeling each segment S τ , the modeling process for non-memory extension layers aligns with the vanilla Transformer (Eq. 1).In the memory augmentation layer, as depicted in Figure 1 (b), each query token q n i ∈ Q n τ can access additional k pairs of memory keys and values from all historical processed keysEQUATION)where g is a learnable parameter and ⊙ denotes element-wise multiplication. The mask matrix (Figure 2 (b)) Â ∈ R L×((τ-1)×L) , j ∈ kNN(i) means the dot product of the j-th previous key of K n 1:τ-1 and the i-th current query token of Q n τ ranks in the top k positions among the dot product of the i-th current query token and all previous key tokens. The second mask A ∈ R L×L is the vanilla causal mask.Context compression methods typically compress the hidden states of history tokens into high-level memory representations, which combine multiple tokens into one concise representation. This technique can be implemented through either a pooling-based approach, exemplified by LongT5 (Guo et al., 2021) and PoolingFormer (Zhang et al., 2021) , or a model forward strategy, such as Recurrent Memory Transformer (RMT) (Bulatov et al., 2022) and AutoCompressor (Chevalier et al., 2023) . We choose RMT to demonstrate how compressed memory representations are integrated into the modeling process.RMT incorporates m tunable memory tokens "[mem]" before and after each segment S τ (Figure 1(c) ). The succeeding memory tokens (marked as M n τ ) are employed as a medium to read the historical information from the last segment S τ-1 while the succeeding memory tokens (M n τ ) can be used to write the current segment information for modeling next segment S τ+1 . In this way, the n-th hidden states of segment S τ is expanded toEQUATION)Where the mask matrix A ∈ R (L+2m)×(L+2m) is illustrated in Figure 2 (c). To implement the recurrence mechanism, RMT initializes the first layer representation of read and write memory tokens M 1 τ and M 1 τ of current segment S τ with the last layer representations of write memory tokensM N τ-1 of last segment S τ-1 . 1 H n-1 τ is spliced after H n-1 τ-1 , resulting in A ij = 0 if j ∈ (i, i + L].2 If the memory pool size (m) is limited, only the τm to τ -1 historical segments can be accessed. Roy et al. (2021) . Formally, given an input sequence X = {x 1 , x 2 , ..., x L }, these methods modify the standard attention mechanism by creating diverse sparse attention masks:EQUATIONwhere H n ∈ R T×d is the n-th layer hidden states of X and A ∈ R T×T is the mask matrix. We illustrate sparse attention mask construction using Longformer (Beltagy et al., 2020) as an example. See Appendix A.1 for BigBird details.Longformer utilizes a combination of sliding window and global attention to establish the sparse attention pattern (Figure 1(d) ). Sliding window attention constrains attention for each token within a context of L tokens (We set window length to be equal to segment length and ignore stride here), while global attention allows any token in the sequence to capture global tokens G. Hence, its mask matrix (Figure 2 (d) ) can be formulated as:EQUATION) 
the unified memory framework: 
 Through reviewing the aforementioned methods, we investigate that these methods can be integrated into a unified memory-augmented modeling framework, featuring a segmentlevel streaming input mode. We name this framework as UniMem. In this section, we first reformulate the aforementioned long-context modeling methods within the UniMem framework, distinguishing them across precisely defined memory dimensions, and then integrate multiple memory dimensions to propose new modeling approaches.Formally, UniMem can augment the memory keys K n mem and values V n mem from previously processed segments when modeling the hidden state H n τ for each segment S τ :EQUATIONThe distinctions among existing long-context modeling methods hinge on two critical elements: the construction of memory keys and values (K n mem /V n mem ), and the creation of the attention mask matrix to access memory (A). Next, we present the aforementioned longcontext modeling methods in UniMem, focusing on constructing these essential elements.UniMem: Transformer-XL, as shown in Figure 3 (a), designates the keys and values from the last segment S τ-1 as memory keys and valuesK n mem /V n mem = SG(K n τ-1 )/ SG(V n τ-1), with the mask remaining consistent with that in Eq. 3. Transformer-XL considers the keys K n τ-1and values V n τ-1 from the last segment S τ-1 as the memory keys and values:EQUATION)UniMem: Memorizing Transformer, as illustrated in Figure 3 (b), preserves the keyvalue pairs of all processed segments as memory key and values K n mem /V n mem = SG(K n 1:τ-1 )/ SG(V n 1:τ-1 ) and accesses the memory using kNN's attention mechanism:EQUATION)Given the UniMem formulation, we can deduce the original formulation of Memorizing Transformer in reverse (we omit attention masks for simplicity):EQUATION)where g is constructed as follows:EQUATION)There is also experiment evidence that two implementations can be converted to each other without loss of language modeling performance Tworkowski et al. (2023) .UniMem: RMT constructs the memory keys and values K n mem /V n mem by performing a model forward pass on the read memory token M n τ within the current segment S τ :EQUATIONwhere Q n-1 mem denotes memory token queries for retrieving key-value pairs, K n mem /V n mem , in the model forward pass. The entire modeling process is shown in Figure 3(c) .UniMem: Longformer uses sliding window attention, like other sparse attention methods, to handle long sequences X = x 1 , x 2 , ..., x T , where each token attends to a limited number of previous tokens. This approach facilitates transitioning input into a streaming format, processing fixed-length segments S 1 , S 2 , ..., S τ sequentially. Longformer also applies global attention based on sliding window attention, as illustrated in Figure 3(d) . The memory keys and values K n mem /V n mem are derived from global tokens G and the last segment S τ :EQUATION)where |G| is the number of global tokens and K n glob /V n glob are their key-value pairs.In UniMem, the differences among these methods can be further standardized into four dimensions related to operating the added memory keys and values (Table 1 ).Memory Management: This dimension manages memory cache storage, involving two crucial elements: (1) Memory Size refers to the amount of stored memory key-value pairs from previous segments. Some methods store from just one prior segment ("Single-Sgm"), such as RMT and Transformer-XL as Figure 3 (a) and 3(c) illustrates, whereas other methods implement a large memory cache incorporating multiple past segments ("Multi-Sgm"), (e.g. Memorizing Transformer and Sparse Attention, as depicted in Figure 3 (b) and 3(d)).(2) Overflow Handling denotes the strategies for updating the memory cache when it is full. Some use First-In-First-Out ("FIFO") to discard old memory, while sparse attention methods clear all when max capacity is reached (Figure 3(d) ).Memory Writing: This dimension concerns transforming processed data from past segments into memory keys and values for storage and selective access by the current segment. Memory Reading: This dimension explains fetching keys and values from the memory cache, guided by mask matrix A (Eq. 9). The "Position" method (e.g., Longformer) selects based on relative positions, influenced by Window Length and Global Tokens (Eq.15). The "Similarity" method (e.g., Memorizing Transformer, Routing Transformer) uses similarity to current queries, regulated by Topk (Eq. 11). RMT accesses the whole cache ("All"). Some models (e.g., KNN-LM (Khandelwal et al., 2019) , TRIME (Zhong et al., 2022) ) also map memory cache contents to vocabulary space, integrating them with current tokens.Memory Injection: This aspect defines how LLM layers add an external memory cache, controlled by the Memory Layer parameter for its quantity and placement. While most methods integrate memory uniformly across all layers ("All-Lyr"), selectively adding memory to specific layers can lower CUDA memory usage ("Certain-Lyr").Others: Besides design elements, methods vary in training strategies. Transformer-XL and Memorizing Transformer use stop-gradient to block memory gradients, while RMT and Longformer enable back-propagation through time (BPTT). Additionally, Focused Transformer uses cross-batch training, and TRIME utilizes in-batch contrastive loss.We analyzed existing methods from the UniMem perspective, as shown in Figure 3 . From this analysis, we developed UniMix, which merges the best features in each Memory Dimension. UniMix utilizes "Similarity" and "Position" for memory reading and a "Direct" plus "Model Forward" approach for memory writing to yield empirically robust-but not necessarily optimal-results. This approach serves as a starting point for further exploration. Ongoing research will assess UniMix's performance, the effects of different dimensions on outcomes, and the potential of this new dimensional synergy for enhanced results.We use the language modeling task as the testbed. We select two widely-used datasets: (1) PG-19 (Rae et al., 2019) , from the Project Gutenberg archive, consisting of pre-1919 English books (Sun et al., 2021) ; and (2) GitHub dataset, including a broad collection of code and documentation from RedPajama's GitHub repositories (Computer, 2023) . We apply perplexity as the metric for evaluation.We evaluate two scales of LLaMA: the 22-layer TinyLLaMA-1.1B (Zhang et al., 2024) and 32-layer LLaMA2-7B (Touvron et al., 2023) . Unless specifically stated, our experiments predominantly utilized the TinyLLaMA model. We fine-tune models for one epoch using 0.1B tokens across two datasets (We have found that this amount of training volume is sufficient for achieving solid long-context capabilities and stable method comparisons). The default sequence length was 2, 048 tokens, aligned with TinyLLaMA's positional embeddings. See Appendix D for more implementation details.Baselines We compare our UniMix with vanilla Transformer (Vaswani et al., 2017) , Longformer (Beltagy et al., 2020) , Memorizing Transformers (MemTrans) (Wu et al., 2022) , Transformer-XL (Dai et al., 2019) , RMT (Bulatov et al., 2022) , and positional interpolation (Chen et al., 2023a) . We re-implement all of the existing methods and our UniMix under the UniMem framework. See Table 2 The results are shown in Table 3 , from which we conclude that:(1) We compare the existing classical methods under the same setting and find that the performance trend remains consistent across different datasets. The performance ranking is Transformer-XL > Longformer > MemTrans > RMT. We hope this comparative result can serve as a reference for the community.(2) Our UniMix method outperforms or matches existing methods in both text and code datasets, demonstrating that by effectively organizing each dimension and combining optimization of various methods, we can achieve a better long-context processing method that is applicable across various fields.(3) The UniMix method remains superior on the larger-scale LLaMA2-7B model, emphasizing the method's scalability and robustness. 
effects of different unimem dimensions: 
 The UniMem framework integrates various approaches to long-context language modeling, each characterized by distinct hyperparameters. In this section, our objective is to dissect the impact of each dimension within the UniMem framework. Subsequent analyses have focused on discerning the pivotal role of Memory Read and Memory Injection within the UniMix method by fine-tuning associated hyperparameters. Experiments on Memory Management and Memory Writing adjustments are detailed in the Appendix B.Our study delved into the hyperparameters affecting memory reading-specifically, Topk and Window Length-and their impact on model perplexity (PPL). Key insights are listed as follows:(1) TopK Increment Does Not Guarantee Improved Model Performance: Increasing TopK from 0 to 64 did not uniformly boost performance. While the Memorizing Transformer's PPL steadily dropped with higher TopK, in the UniMix framework (MemTrans settings), optimal performance was noted at TopK=4 and 8, with negligible benefits beyond these points, as illustrated in Figure 4 .(2) Window Length's Minimal Impact on UniMix: Adjusting the window length from 0 to 2048 had little effect on UniMix, unlike Transformer-XL, which showed significant performance changes. This highlights a distinct contrast in how these models respond to window length adjustments.(3) UniMix's Robustness: The observations underscore UniMix's resilience against single-dimensional parameter shifts like TopK and window length, marking a notable strength for long-context processing. Such robustness minimizes the necessity for intricate hyperparameter optimization, ensuring stable performance across diverse settings.To investigate the impact of Memory Injection dimensions on performance, we conducted further analysis on the GitHub dataset. Our findings are summarized as follows:(1) Considerable Impact of Memory Layer Placement: For the TinyLLaMA model, inserting the memory layer at higher levels outperforms its insertion at lower levels, as demonstrated in Figure 4 (c).(2) Significant Influence of Single Memory Layer Position: The positioning of the memory layer has a substantial effect on performance, with an optimal layer significantly outperforming others, as shown in Figure 4 (d).(3) Task-Independent Optimal We evaluate methodologies for processing longer texts, comparing their performance on a GitHub dataset with segments over 64k (Table 4 ). Our experiment involved: 1) Extending vanilla transformers to 64k inputs and fine-tuning; 2) Using Position Interpolation (PI) for increased input to 64k by adjusting position indices; 3) Implementing UniMix, which processes up to 64k, combining 62k memory and 2k local attention. After fine-tuning on a 0.1B token dataset, UniMix showed significant perplexity reduction over PI, benefiting from a memory strategy similar to the LLaMA model's pretraining, without new positional embeddings. Crucially, UniMix demonstrates a computational complexity that scales linearly with text length, thereby surpassing the Vanilla Transformer and PI in terms of inference and training efficiency due to their exponential time complexity.Downstream task experiments were conducted on the TinyLLaMA model with a context length of 512 tokens and a UniMix memory size of 10k tokens. We proportionally adjusted other hyperparameters based on those specified in the original paper. The process begins with fine-tuning the model on a 0.1B token dataset. Subsequently, we apply supervised finetuning (SFT) using 12k instances from the LongAlpaca dataset. Although both fine-tuning and SFT experiments are moderate in scope, potentially affecting the models' performance on downstream tasks, the results still demonstrate the expected performance trends. In addition to our evaluations on LongBench, we conduct experiments using the "Needles in the Haystack" dataset. This dataset includes text lengths ranging from 4k to 32k tokens, as shown in Table 6 . UniMix is the only method to accurately identify the "needles" in texts exceeding 4k tokens in length. 
conclusion: 
 In this paper, we introduce UniMem to unify various long-context approaches under the view of memory augmentation of LLMs. UniMem encompasses four dimensions: memory management, memory writing, memory reading, and memory injection. UniMem provides a structured categorization that enhances understanding, revealing, for instance, that StreamingLLM and Memorizing Transformer differ mainly in their memory reading techniques. This framework allows researchers to optimize long-context modeling by exploring different strategies like Least Recently Used (LRU) for Memory Management and various compression formats for Memory Writing. Our findings show that the memory injection layer significantly impacts performance, with each language model having an optimal layer independent of the dataset. By enabling granular control through tunable hyperparameters, UniMem ensures standardized and fair evaluations. Furthermore, we propose UniMix, which combines the strengths of these methods and significantly outperforms existing approaches across diverse datasets. Future work will integrate UniMem with other techniques, such as diverse position encoding and alternative overflow handling methods, focusing on differential memory strategies across layers.  
 
 Appendix: We explored the impact of varying fine-tuning data volumes on long-context modeling methods. We plotted the Perplexity (PPL) for TinyLLAMA and LLAMA2-7B using different scales of training data, as shown in Figure 14 . Our results indicate that increasing the volume of training tokens significantly reduces perplexity and enhances model performance.UniMix consistently emerged as the most effective method across all data scales. Moreover, methods such as UniMix and Transformer-XL demonstrated excellent scalability, suggesting their potential for even better performance with larger datasets.  


Decision: The paper introduces a unified framework for memory augmentation of LLMs. Long-term memory is an important topic for LLM research. The overall rating turns out to be positive. It would be good to include the submission in the conference proceedings.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: please see the comments in the previous section.
paper_summary: paper_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.
strength_weakness: strength_weakness: **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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: See above.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: 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.

Comments: title: Thanks for the response
Thanks for rebuttal.
comment: Thanks for the response which addresses most of my questions.
I will raise my scores.
Thanks for the clarification. The limited computational resource is a large problem for academia research. I hope that you can find more GPUs to consolidate your work and claims via scaling up the training scale. I will keep the rating unchanged.
Thanks. I will keep my rating unchanged.
Thanks for the clarifications! I think the authors do make a hard effort in justifying the value of this paper, though the additional result on Longbench seems too low to present a fully persuasive conclusion.  I appreciate the novel perspective of this paper and hope the authors continue advancing their research. I will give a positive assessment to encourage this kind of work.
Dear Reviewer, I hope you're doing well. Following up on our recent exchange regarding this paper, I wanted to check if there are any further concerns or feedback from your side. Your insights are invaluable to us, and we're keen to address any remaining issues.