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: COLM (2024)
Keywords: long-context, large language models
TL;DR: The article introduces UniMem, a unified framework that reformulates existing long-context methods from the view of memory augmentation of large language models (LLMs).

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.

Published as a conference paper at COLM 2024
UniMem: Towards a Unified View of Long-Context Large Language Models
Junjie Fang"; Likai Tang"', Hongzhe Bi', Yujia Qin', Si Sun', Zhenyu Li", Haolun Li², Yongjian Li², Xin Cong', Yankai Lin Yukun Yan', Xiaodong Shitt, Sen Song', Zhiyuan Liu Maosong Sun 122,3
'Department of Computer Science and Technology, Institute for AI, Tsinghua University, China Beijjing National Research Center for Information Science and Technology, China institute for AI Industry Research (AIR), Tsinghua University, Beijing, China
*iaamen University, 'Renmin University of China f fangj j@stu. xmu. wdw..ddu.cmm cn, tang1k200mails. tsinghua edu. cn
Abstract
Long-context processing is 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 inte- gration 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 piveessions-eemory Management, Mem- ory 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 rep- resentative 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, 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 hhtpp:///////t htttts://github.co github.com. http:cculm///////////// UniMem
1 Introduction
Transformer-basee (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
OpenAl
2022), leading to various applications (e.g., natural language processing (Brown et 2020),
code generation eLi 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 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 address this challenge, including optimizing the quadratic complexity of Transformer Dai et al. (2019); [Beltayy et al. (2020); Wu et al.
handle longer contexts than during pre-training (Chen et al.) 2023a |enn et al. (023), and designing non- non-Transformer architec- tures (Gu et al., 2021 Xiong et al. 20023 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 Upon reviewing this type of
20022), adjusting the position encoding of Transtormer
Indicates equal contribution.
Corresponding authors: mandel @xmu, edu. en, liuzy@tainghua..cr en.

---

Published as a conference paper at COLM 2024
UniMem Framework
Methud
Memory Size
Memory Management
Ovettlow Handling
Memory Writing
Memory Reading
Memury Injection
RMT RMT Rulatowet et a
2022
Single-Smm
FIRO
Model forward
All
All-Eyr
Auto
ctal. 20133
Multi
Model forward
All
All-Lyr
Poolinglazmer Zeon
Mulli- Mulli-Sgm
Clear All
Fooling--Die Dinect
Fosilion
All-
Mulli
Clear All
Fooling Direct
Posilion
All-Lyr
2000
Multi 5gm
Position
All-Iyr
Big Bird Zalise MIF-
5gm
Fosition
All-Lyr
LongNer(onng o es14))
5gm
Posttion
All-Lyr
2:19
Single-Sgm
PIPO
Direcl
Posilion
All-Lyr
Streamine111 MM Kenyerel
Multi Sgm
FITO
Direct
Posétion
All-Iyr
Routing Roy 2021
Multi-Sgm
Clear All
Direct
Similarity
ReformerKitzey dr 2000y
Multi Sgm
Clear All
Direct
Similarity
All-Lyz
Monotizzng Tighishrnn or el
2012
Mulli Sgm
HPO
Direct
Similarily
Ceriain-Lyy Lyy
kall Jworkowski el al. 2leas
Mulli Sgm
HPO
Direct
Sunilarily
Ceriain-Lyr
Innhmmtkrrrn/eee Derseer oe
M12
Multti-gii
Direct
Similarity
All-Iyr
KNN- MOIN-LMKKhlleee Khallssseee
Multi 9gm
Direct
Similarity
Certain- Lye
TRIME 2Boogg crah 20124
Multi-5gm
Direct
Similarity
Certain-
UniMix (Ours)
Multi 9gm
PIFO
Dinect-Model forward
Simillcity+Posstion
Certain-
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 peperence-other methods apply to both training and inference.
method, we further divide it into three categories: (1) Context Caching, which stores the in- termediate 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. [022 Chevalier et al. 2023; Guo et al.
2021
Zhang et al. ) (3) Sparse Affention: 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 long- context 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 bilityy fairly becomes exceedingly complex. Such complexity com- plicates 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 singular model. This impediment to synergistic integration restricts the potential for substantial progress in the effective processing of long contexts by LLMs.
their efficacy and adaptability
This work introduces UniMem, unified framework that integrates existing long-context methods for optimizing Transformer computation complexity, viewed via memory augmen- tation. Specifically, UniMem consists of four essential dimensions: (1) Memory Management determines how much past information stored and how old memory replaced, impact- ing 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 sum- marized 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 the UniMem framework, 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 Longtormer (Beltagy 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. Ad- ditionally, built upon the unified framework, we introduce 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

---

Published as a conference paper at COLM 2024
4 Ap
Bpt
(a) Transformer-XL.
(b) Memorizing Transformer
(c) RMT
(d) Longformer
Figure 1: Diagram illustrates long-context methods (segment length = 3). Yellow circles show past segments; blue circles mark the current segment. (a) Transformer-XL caches earlier hidden states.
(b) Memorizing Transformer retrieves past segments with kNN similarity. (c) RMT employs memory tokens for prior segments. (d) Longformer extends segments with global and sliding window attention.
parameters. (2) Increasing the number of memory layers does not necessarily correlate with performance gains. Rather, the strategic positioning of single memory layer can achieve comparable effectiveness to integrating additional memory across half of the model's layers.
Preliminaries
2.1 Vanilla Transformer
Given Transformer-based LLM featuring N stacked layers, and an input sequence X= (2,, consisting of I' tokens, the LLM processes the sequence X layer-wisely, constructing set of hidden states H [h h h4], ) where E [1, The n-th LLM layer's hidden states H" are derived from the (n 1 layer using feed-forward networks (FFN), multi-head attention (Attn), and softmax (Sftx):
H JUN Attn Q", A) Sitx
Q"(K*) + +A v", Ag
Oifje [0,1],
())
else.
/"/K" /V" = H" (WW)) (W") //H" "WW)
(2)
where A RTxT is often causal mask matrix controlling the visible context field and Rdxd is the projection matrix of multi-head attention in n-th LLM layer.
Wa 9/k/o
Despite LLMs' strong context modeling, their attention mechanism's 22 computational complexity challenges tasks with very long contexts. To alleviate this challenge, the com- munity has made many improvements, and we analyze three main types of long-context approaches including context caching, context compression, and sparse attention methods.
2.2 Context Caching
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, long input sequence X = (2, can be split into a succession of segments S, ( (T-1))x=+) with the fixed token length L (L < T), where T [/.]. 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 previous segment an external memory cache (Dai et al. 2019), and the other involves caching key-value pairs of multiple previous segments (uu et al. 2022 Bertsch et al. 2023 Tworkowski et al. 2023). This paper highlights Transformer-X XL. (Dai et al. 2019) and Memorizing Transtormers (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 of the current segment S, and the hidden states H processed by the previous segment ST-1, as shown Figure [(a) This recur- rence way enables Transformer-XT. to learn dependencies beyond long contexts of fixed
3

---

Published as a conference paper at COLM 2024
length without breaking temporal coherence. Formally, the æ-th hidden states H of segment St is constructed layer-wisely as follows:
FEN Att K. VAA) Ag
e (1)+
else.
K = H (W5) /|SG |: (WW""
where mask matrix A E R Lx2L is illustrated in Figure (a). The function SG(-) means to stop gradient back-propagation and [+)]] denotes the concatenation operation
Memorizing Transformer designates a spe- cific layer for memory augmentation. When modeling each segment Sr, the modeling process for non-memory extension layers aligns with the vanilla Transformer (Eq. [ In the memory augmentation layer, de- picted in Figure [(b) each query token q E O can access additional k pairs of mem- ory keys and values from all historical pro- cessed keys k [k] k2 ,,.,, k ((-))x)
(c)
(d)
(b)
Figure 2: Attention patterns for long- methods (L 3).(a) Transformer-Xe (b) Mem-
and values ViT-1 = [vi), v2 .,.,. v (r-))x))
orizing Transformer. (c) RMT. (d) Longformer.
through k-nearesttneighbor (kNN) lookup?
+ (1 Qt, V.,),, ÀG [Oifj e ANN()),
PNN H @
(5)
where g is learnable parameter and @ denotes element-wise multiplication. The mask matrix (Figure (b)) Å E RLx((T-1)XL), e kNN(f) means the dot product of the j-th previous key of K" i:7-1 and the i-th current query token of Q. 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 E R LxL is the vanilla causal mask.
2.3 Context Compression
Context compression methods typically compress the hidden states of history tokens into high-level memory representations, which combine multiple tokens into one concise repre- sentation. This technique can be implemented through either a pooling-basee exemplified by LongT5 [Guo et al. 2021 and PoolingFormer (Zhang et al. 2021), or model forward strategy, such as Recurrent Memory Transformer (RMT) (Bulatov et al. 2022) and wwtpmmeresssrreeallllle et al. 2023). We choose RMT to demonstrate how compressed memory representations are integrated into the modeling process.
approach,
RMT incorporates tunable memory tokens "[mem]" before and after each segment S (Figure I(c). The succeeding memory tokens (marked as M are employed as medium to read the historical information from the last segment ST-1 while the succeeding memory tokens (M") can be used to write the current segment information for modeling next segment St+1- In this way, the n-th hidden states of segment S, is expanded to H = [M;; HP; M"):
Oifj [0, max m)), Delifi e [1. m, I. 2m),
FIN Attml O", Ky A), Ag
(6)
elsc.
Where the mask matrix A E (I+2mm) (L+2mm) illustrated in Figure (c). To implement the recurrence mechanism, RMT initializes the first layer representation of read and write memory tokens M and M of current segment Sr with the last layer representations of write memory tokens M of last segment Sz-1.
HH spliced after H resulting Ay Oif E (i,i L]. 2If the memory pool size (m)) is limited, only the
historical segments be accessed.
4

---

Published as a conference paper at COLM 2024
Upente
Scale Lppdai dea---daaaa
Glaaal
Monery Road
0
Javee
fewad Upko.
Scake L Updnn 280
dla
Sp1
Stit
9x2
3p1
Menory
(a) Transformer-Xi.
(b) Trans
(c) RMT
(d) Longformer
Figure 3: Four typical methods in UniMem. (a) Transformer-XL uses single past hidden
state, shown as memory cache, with sliding window attention. (b) Memorizing Transformers (MemTrans) saves multiple segment states in layer, using kNN similarity for retrieval. (c) RMT utilizes a memory token for prior segment data, stored post-model forwarding. (d)
Longformer applies global and sliding window attention based on token positioning.
2.4
Sparse Attention
Besides segment-level methods, sparse attention techniques tackle long-context tasks by sparse mechanisms to avoid the quadratic computational increase of vanilla attention with long sequences Beltagy et al. (2020) Zaheer et al. [2020] Ding et al. (2023),
x Kitaev et al. (2020); {2,, X2, XL), these methods
Roy al. (202T). Formally, given an input sequence
modify the standard attention mechanism by creating diverse sparse attention masks:
FFN Attn(Q", K", :") H
where H Txd is the n-th layer hidden states of X and E R 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 combination of sliding window and global attention to establish the sparse attention pattern (Figure [(d). Sliding window attention constrains attention for each token within 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:
Oifje e (min(0.--L),,) index(().
(8)
elsse.
The Unified Memory Framework
Through reviewing the aforementioned methods, we investigate that these methods can be integrated into unified memory-augmenter modeling framework, featuring a segment- level 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.
3.1 UniMem Formulation
Formally, UniMem can augment the memory keys K Kmem and values V Vhnem from previously processed segments when modeling the hidden state H for each segment S,:
FEN Autn( [Kinem V...Vemmm V.).A).
(9)
The distinctions among existing long-context modeling methods hinge on two critical elements: the construction of memory keys and values (Kr Vi ) and the creation of the attention mask matrix to access memory (A). Next, we present the aforementioned long- context modeling methods in UniMem, focusing on constructing these essential elements.
UniMem: Transformer-XL. as shown in Figure (aa) designates the keys and values from the
last segment S--1 memory keys and values K K V W SG(K",, SG(V") (V" 1 with
5

---

Published as a conference paper at COLM 2024
the mask remaining consistent with that in Eq. 3 Transformer-XL considers the keys K and values VT+- from the last segment S7-1 as the memory keys and values:
lifje (1.+ 4),
Kine / Vinem SG(K)-)/SG(VS.))
(10)
-oo else.
UniMem: Memorizing Transformer, illustrated in Figure 3(b) preserves the key- value pairs of all processed segments as memory key and values / V (G((K. SG(VT 1r-11) and accesses the memory using kNN's attention mechanism:
0 ENN(() (,T-))))
(1))
else.
Given the UniMem formulation, we can deduce the original formulation of Memorizing Transformer in reverse (we omit attention masks for simplicity):
Qt, Q" Kiamm Ki [Vmami
V!l)
Sitx [Q
K).") [Vesmi V.),
Sitx (Q) (K ))) Vime + (1 g) Sftx (Q")) Vi,
(12)
Atin Q* K V (1 g) AQ" KD V
where g is constructed as follows:
exp ( (eimm)) ))
exp Q% (iim 1 EL, (K))):
There is also experiment evidence that two implementations can be converted to each other without loss of language modeling performance [Iworkossi cefformane/lwooroowkiiii et al. (2023)
UniMem: RMT constructs the memory keys and values K m / /em by performing model forward pass on the read memory token M within the current segment Sy:
M (WT)" /M (WW)) "(W)"" M.
JeN
Attn Attn(Qu Ku-. Vi-.1
(14)
where Qhen denotes memory token queries for retrieving key-value pairs, K the model forward pass. The entire modeling process is shown in Figure 3(c)
V/nemm in
Vnn
Longformer uses sliding window attention, like other sparse attention methods, to handle long sequences X X2,,, XT, where each token attends to a limited number of previous tokens. This approach facilitates transitioning input into a streaming format, processing fixed-length segments S2, -.., Sr sequentially. Longformer also applies global attention based sliding window attention, illustrated in Figure 3(d) The memory keys Vem are derived from global tokens G and the last segment Sy:
UniMem:
and values K!!l
Kmeem/
[Oifje (Li+L+IG),
Ky
V =
KKa // Robi V" Ay
(15)
-ou else,
where G is the number of global tokens and K V glob glab their key-value pairs.
3.2 Memory Dimensions
In UniMem, the differences among these methods can be further standardized into four dimensions related to operating the added memory keys and values (Table
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-X1 as Figure 3(a) and 3(c) illustrates, whereas other methods implement 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 is full. Some use First-In-First-One ("FIFO") to discard old memory, while sparse attention methods clear all when max capacity is reached (Figure (d)).
Memory Writing: This dimension concerns transforming processed data from past seg- ments into memory keys and values for storage and selective access by the current segment.

---

Published as a conference paper at COLM 2024
Memory Management
Memory Reading
Memory Writing
Mem Injection Memory Layer
Model
Memory Size
Topk
Window Length
Global Tokens
Compressed Tokens
Vanilla
Longformer
4,196
2,048
All
Memlians
11,21
Trans-XI.
2,048
2,048
All
UniMix
20,520
2,048
40
12-22
Table 2: Long-context methods decomposed along the defined UniMem dimensions in hyperparameters view for TinyLLaMA (Section 3.2)
Context Caching and Sparse Attention insert keys and values directly during the forward pass ("Direct"), ), shown in Figure 3(b) Conversely, Context Compression employs op- erations like dual forward passes (e.g., RMT, "Model Forward")) and pooling (e.g,, Pool- ingformer (Zhang et al. 2 Long T5 (Guo et al. 2021), "Pooling") to compress the last segment's data into tokens, dictated by the Compressed Tokens count.
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 relative positions, influenced by Window Length and Global Tokens (Eq [5). The "Similarity" method (e.g. Memorizing Transformer, Routing Transformer) uses similarity to current queries, regulated by Topk (Eq.[1) RMT accesses the whole cache (" (""ll". 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 Reading:
This aspect defines how LLM layers add an external memory cache, controlled by the Memory Layer parameter for its quantity and placement. While most meth- ods integrate memory uniformly across all layers (" All-Lyr")) selectively adding memory to specific layers can lower CUDA memory usage ("Certain-Lyrr")
Memory Injection:
Others: Besides design elements, methods vary in training strategies. Transformer-XI. 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.
3.3 UniMix: Synthesizing Different Dimensions
We analyzed existing methods from the UniMem perspective, as shown in Figure 3 From this analysis, we developed UniMix, which merges the best features each Memory Dimension. UniMix utilizes "Similarity" and "Position" for memory reading and "Direct" plus "Model Forward" approach for memory writing to yield empirically robust-but not necessarily optimal-resulls. This approach serves as 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.
synergy
Experiments
4.1 Experiments Settings
Datasets and Evaluation 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, con- sisting of 1919 English books (Sun al. 2021); and (2) GitHub dataset, including broad collection of code and documentation from RedPajamma' GitHub repositories (Computer 2023). We apply perplexity as the metric for evaluation.
Implementations We evaluate two scales of LLaMA: the 22-layer TinyLLaMA-1.1B Zhang et al.. 2024) and 32-layer LLaMA2-7B (Touvron et al.) (023). Unless specifically stated, our experiments predominantly utilized the Tıny model. We fine-tune models for one epoch using O.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
7

---

Published as a conference paper at COLM 2024
comparisons). The default sequence length was 2, 048 tokens, aligned with TinyLLaMA's positional embeddings. See Appendix for more implementation details.
Baselines We compare UniMix with vanilla Transformer (Vaswani et al. 2017 Long-
former [Beltagy et al. 2020), Memorizing Transformers (MemTrang (Wu et al., 2022), ,
Transformer-X XL (Dai et al. 2019), RMT (Bulatov et al. 2022), and positional interpola- tion (Chen et al. 2023a J. We re-implement all of the existing methods and our UniMix under the UniMem framework. See Table for detailed hyperparameters.
4.2 Comparison of Existing Methods Under UniMem
The results are shown in Table
3 from which we conclude that:
TinyLLaMA
LLaMA2-7B
Model
(1) We compare the existing clas- sical methods under the same set- ting and find that the performance trend remains consistent across different datasets. The perfor- mance ranking is Transformer-XL Longformer Memlrans > RMT. We hope this comparative result can serve as reference for Table eetpertorns the community. (2) Our UniMix method outperforms ör 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 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.
PG19
Github
PG19
Github
Vanilla
14.53
2.66
11.01
2.53
Longformer
13.92
2.42
10.34
2.27
Mem
14.34
2.57
10.66
2.34
Transformer-XL
13.78
2.36
10.28
2.23
RMT
14.68
2.70
10.64
2.55
UniMix 13.78
2.32
10,28
2.20
Perplexity of different methods in UniMem,
4.3 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
Memory Reading Our study delved into the hyperparameters affecting memory read- ing ing-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 's PPL steadily dropped with higher TopK, in the UniMix frame- work (MemTrane 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 Robust- ness: The observations underscore UniMix's resilience against single-dimensinnl parameter shifts like TopK and window length, marking notable strength for long-context process- ing. Such robustness minimizes the necessity for intricate hyperparameter optimization, ensuring stable performance across diverse settings.
Memory Injection To investigate the impact of Memory Injection dimensions on perfor- mance, we conducted further analysis on the GitHub dataset. Our findings are summarized as follows: (1) Considerable Impaci of Memory Layer Placement: For the Liny model, inserting the memory layer at higher levels outperforms its insertion at lower levels, as demonstrated in Figure ((c) (2) Significant Influence of Single Memory Layer Position: The po- sitioning 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 fisk-Inndpeedeett Optimal
8

---

Published as a conference paper at COLM 2024
Cablic Nextlons
Ceptic Mere Frems
X2
Dafth Banwkitt
2.4
Ualtia Baneth 12:22
Perplestt
Perphestyy
perptestty
1,0
Parplesity
2.
1.5
24
Gattas P01e
24
9.4
35
04
512
1824 Window Length
2048
1933
2009
3000
10 13 12
Topk
Step
Layer
(a)
(b)
(c)
(d)
Figure 4: Effects of different UniMem dimensions on perplexity across datasets. (a) Topk's role for MemTrans and UniMix; (b) Combined effects with Window Length; (c) Memory Layer Distribution's s impact; (d) Memory Layer Position's influence (Single Layer Injection).
Memory Layer: The optimal layer remains consistent across different datasets HGitHub and PG19) for the same model configuration, according to Figure 4[d) Model-Depentee Optimal Memory Layer: On the GitHub dataset, the optimal layer for the 7b model is the 11th layer, whereas, for the 1b model, it's the 16th layer. Notably, the 1b model has a singular optimal layer, whereas multiple optimal layers exist for the 7b model, as illustrated in Figure and Figure 10 (4) The superiority of Optimal Memory Layer in Practical Applications: When inserting a single memory layer across different levels of the model, the performance at the 16th layer significantly surpasses that of other levels. Remarkably, the enhancement achieved by single layer at the 16th level is comparable to the improvements observed when memory layers are inserted across layers 12 to 22, far exceeding the performance improvements seen when inserted in layers 1 to 11. This indicates that a single memory layer, optimally positioned, can reach more than 98% of full performance using just 1/11 the computational time and resources (Figure
4.4 Comparison with Position Interpolation
We evaluate methodologies for processing longer texts, comparing their performance on GitHub dataset with segments over 64k (Table 4 Our ex- periment involved: 1) Extending vanilla transform- ers to 64k inputs and fine-tuning; 2) Using Position Interpolation (PT) for increased input to 64k by ad- justing 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 memory strategy similar to the LLaMA model's pretrain- ing, 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.
Model
Perplexity
Vanilla-64K Position Interpolation UniMix
20.57
5.87
2.75
Table 4: Perplexity comparison among Vanilla, Position Interpola- tion (PI), and UniMix models.
4.5 Downstream Task Experiments
Downstream task experiments were conducted on the TinyL.LaMA model with a context length of 512 tokens and 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 fine- tuning (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.
Results on LongBench As shown in Table The results of the downstream tasks indicate that the performance ranking almost aligns with the perplexity (PPI.) evaluation in our paper: UniMix > Transformer-XL >> >Longformer >RMT --Mmirrnn, This consistency across different evaluation metrics highlights the robustness of our proposed UniMix method.

---

Published as a conference paper at COLM 2024
Model
NQA
Qasper
MFQA
HQA
2WikiMQA
Musique
Vanilla
5.2
8.1
19.45
6.19
16.6
3.58
Longformer
4.5
8.01
19.52
6.11
15,74
3.8
Mem'Irans
5,88
8.69
19.18
5.55
14.12
3.49
Transformer-XT.
4.54
8.11
18.98
5.15
14.33
3.23
RMT
4.36
7.58
18.59
5.9
13.75
3.78
UniMix (Ours)
6.19
8.26
19.27
6.82
14.45
3.77
Model
GovReport
QMmum
MultiNews
TREC
TQA
SAMSum
Vanilla
12.47
19
10.06
11
11.9
0.42
Longformer
19.45
11.28
16.5
14.54
TemnTrans
18.72
13.14
11.82
Transformer-XL.
13.33
19.6
13.57
13
17.79
0.0
RMT
11.6
19.54
11.08
13
12.33
0.0
UniMix (Ours)
16
20.02
12.09
20
12.24
3.01
Model
PsgCount
PsgRetrieval
LCC
RB-P
Avg
Vanilla
0.0
4.0
25.34
32.53
11.62
Longformer
1.0
4.0
24.56
31.55
11.90
MemTrans
2.75
0.25
26.48
30.78
11.48
Transformer-XL
1.12
0.92
24.75
32.38
11.93
RMT
1.55
1.0
28.4
32.02
11.53
UniMix (Ours)
1.62
2.5
28.37
33.65
13.01
Table 5: Performance of Different Methods on LongBench.
Model
4k
8k
16k
32k
Avg
Vanilla
0.05
0.04
0.03
0
0.03
Longformer
0.07
0.03
0.03
0.03
MemTrans
0.05
0
0
0
0,01
Transformer-XL
0.20
0.04
0.03
0.01
0.07
RMT
0.04
0.03
0.03
0.01
0.03
UniMix (Ours)
0.37
0.39
0.09
0.04
0.22
Table 6: Accuracy of different methods Needles in the Haystack Dataset.
Needles the Haystack Dataset 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 Table 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-cone context approaches under the view of memory augmentation of LLMs. UniMem encompasses four dimensions: memory management, memory writing, memory reading, and memory injection. UniMem pro- vides structured categorization that enhances understanding, revealing, for instance, that StreamingLLM and Memorizing Transformer differ mainly in their memory reading tech- niques. This framework allows researchers 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 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 ap- proaches 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.
10

---

Published as a conference paper at COLM 2024
Acknowledgement
This work is supported by the National Key R&D Program of China (No.20222D0111332)) National Natural Science Foundation of China (No. 62236004), the Postdoctoral Fellow- ship Program of CPSE (Grant No. GZB20230343), China Postdoctoral Science Foundation (Grant No. 2023M741945), National Science and Technology Major Project (Grant No. 20222D0116101), the Key Support Project of NSFC-Liaoning Joint Foundation (Grant No. U1908216), the Major Scientific Research Project of the State Language Commission in the 13th Five-Year Plan (Grant No. WT135-38) and the public technology service platform project of Xiamen City (Grant No. 3502220231049). Yujia Qin is sponsored by the Baidu Scholarship.
References
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long- -ocummecument trans- former. ar Xiv preprint ar Xix:2004.005150, 2020.
Amanda Bertsch, Uri Alon, Graham Neubig, and Matthew R Gormley. Unlimiformer:
Long-range transformers with unlimited length input. arXiv preprint arXiv:2305.01625, 2023.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. ArXiv preprint, abs/ abs/2108.07258, 2021. URL htttss:/:///ass wwtpxiiv.rrg//ae hogg/bbs/1100 0008.0725888888
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Adoances in neural information processing systems, 33:1877-1901, 2020.
Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. Recurrent memory transformer. Advances in Neural Information Processing Systems, 35:11079-1109,, 2022.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending con- text window of large language models via positional interpolation. arXiv preprint arXiv::306. 15595, 2023a.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint drXiv:2309.12307, 2023b.
Alexis Chevalier, Alexander Wettig, Anirudh Ajith, and Danqi Chen. Adapting language models to compress contexts. ar arXio preprint ar 2tt::2305.1488. 2023.
Together Computer. Redpajama: an open dataset for training large language models, 2023. (Reettaeeeeeeeaaaeeaaeaaaa Data
https:/t::/ gighhu. com/
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhut- dinov. Transformer-xl. Attentive language models beyond a fixed-length context. ar Xiv preprint ar Xiv: 100:00900.00860, 2019.
Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, Nanning Zheng, and Furu Wei. Longnet: Scaling transformers to 1,000,000,000 tokens. ar arviv preprint XrXiv=2307.02466 2023.
Yunfan Gao, Yun Xiong, Xinyu Gao, Kangxiang Jia, Jinliu Pan, Yuxi Bi, Yi Dai, Jiawei Sun, and Haofen Wang. Retrieval-augmentee generation for large language models: A survey. Xrviv preprint ar Xxio:2312.110997 2023.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. ar Xiv preprint ar 20X002231..00752, 2023.
11

---

Published as a conference paper at COLM 2024
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, arx preprint urXio:2111.00396, 2021.
Mandy Guo, Joshua Ainslie, David Uthus, Santiago Ontanon, Jianmo Ni, Yun-Hsuan Sung, and Yinfei Yang. Longt5: Efficient text-to-text transformer for long sequences. arXiv preprint ar Xit:2112.07916, 2021.
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Gen- eralization through memorization: Nearest neighbor language models. arXiv preprint (rXiv:1911.00172, 2019.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. ar arviv preprint arXio:2001104451, 2020.
Yujia Li, David Choi, Junyoung, Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, et al. Competition-level code
generation with alphacode. Science, 378(6624; 2022.
OpenAL. OpenAl: Introducing ChatGPT, 2022. URL httppt:s:///ppntaa http:////ppaaaccoooooooo bltg//lloo///////ooo hamm//lgg/chhhtttttt
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.00071 2023.
Jack Rae, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for rangerange sequence modelling. arv preprint ar xrXo:1991.0050)/ 2019.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, 2021.
Simeng Sun, Kalpesh Krishna, Andrew Mattarella-Micke, and Mohit lyyer. Do long-range language models actually use long-range context? arXiv preprint Xiix:2109.09115, 2021.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor transformer for large language models. arx preprint arXto:2307.0862) 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint ar ar3i1:2302.13971, 2023.
Szymon Tworkowski, Konrad Staniszewski, Mikolaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Milos. Focused transformer: Contrastive training for context scaling. arXiv preprint ar arXiv:2307.03170 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances neural informa- fion processing systems, 30, 2017.
Yuhuai Wu, Markus N Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. Xiv preprint ar arXiv:2203.08913, 2022.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient stream- ing language models with attention sinks. arv preprint ar wxXi:23099.1455. 2023.
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Ábinav Sankararaman, Barlas Oguz, et al. Effective long-context scaling of foundation models. arv preprint axXt::239911603 2023.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33: 17283-17297, 2020.
12

---

Published as a conference paper at COLM 2024
Hang Zhang, Yeyun Gong, Yelong Shen, Weisheng Li, Jiancheng Lv, Nan Duan, and Weizhu Chen. Poolingformer: Long document modeling with pooling attention. In International Conference on Machine Learning, pp. 12437-12446. PMLR, 2021.
Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model, 2024.
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric. P Xing, Hao Zhang, Joseph E. Gonzalez, and lon Stoica. Judging llm-as-a-judge with mt-bench and chatbot arena, 2023.
Zexuan Zhong, Tao Lei, and Dangi Chen. Training language models with memory augmen- tation. Xiv preprint arxiv:2055.12677 2022.
13

---

Published as a conference paper at COLM 2024
A Existing Method
A.1 BigBird
Big Bird Zaheer et al. (2020) introduces random attention in addition to sliding window attention and global token attention, i.e,, it randomly selects some key tokens for each query token where the query can attend over these keys:
O j E (min(0, UGUR(I),
Aij
(16)
oo else.
where Rii) denotes the position set of randomly selected key tokens for the i-th query token.
UniMem: Big Bird can also be reformulated as the following memory-augmented modeling process:
K
Vil
K glob? K rand /
globi V
rand V T-1
mem
0 if E [i, + |G] + [R]],
(17)
A where || is the number of random attention tokens and K rand V values.
0o else,
/and their keys and
B Effects of Different UniMem Dimensions
B.1 Memory Management
We focus on Overflow Handling in the Memory Management dimension. Specifically, tune the Overflow Handling dimension from default First-In-First-Our ("FIFO") to "Clear all" for Longformer, Memtrans, and Mix. We find that Longformer and UniMix achieve worse perplexity with "Clear all",, which can be the result of the reduced horizon of attention of "Clear all" (Figure On the other hand, Memtrans gets slightly better perplexity with "Clear all". We hypothesize that this could have originated from the "Memory staleness problem" as in Wu al. (2022).
2.7
FIFO
Clear All
leplexityy
2.5
Longformer
MemTrans Methods
UniMix
Figure 5: Impact of Overflow Handling on perplexity for Longformer, Mem Trans and UniMix.
B.2 Memory Write
We adjust the "Memory Tokens" for UniMix. to alter the way Memory writes are conducted. The larger the "Memory Tokens", the more memory is written in the form of Model forward. When the "Memory Tokens" is set to 0, all memory writes are direct. It can be observed in Figure 6 that increasing the "Memory Tokens" does not demonstrate positive effect. This is consistent with the experimental results of existing methods discussed earlier.
14

---

Published as a conference paper at COLM 2024
2.7
UniMix (Layer1-7))
(1ayyer8-14)
3.0
UniMis (Layeri5-21)
Paplexity
Perplexity
2.5
UniMix
2.4
2,4
2.3
20
40
80
160
1060
2000
3000
Compressed Tokens
Step
Figure 6: Impact of Compressed Tokens on perplexity for UniMix.
Figure 7: Perplexity of Layer(1-7), Layer(8-14) and Layer(15-21) on UniMix
B.3
Memory Injection
14.6
2.70
2.65
2.50
Popletity
14.4
2.60
Peplexity
2.45
2.55
2,50
14.2
2.15
2.35
14.
2.40
2.30
UniMix
Men:Tran
2.35
2.35
10
13
16
19
22
11
16
21
26
31
Layer
Layer
Figure Perplexity Sin- Figure 9: Perplexity of Single gle Layer UniMix(PG19, Tiny LLaMA)
Figure 10: Perplexity of Sin- gle Layer on UniMix(Github, LLaMA2-7B).
Layer Memtrans(Github LLMMA)
Tram-XL, (Laye16) Trans (Layerl II)
Longiormer (Layer16)
Mem Trans (Laye16)
3.0
Longlormer (Layer1-11)
Memiraas (Layer) II)
3.0
2.9
peplesitt
2.8
Perpessit
2.8
Perplesity
2.88
2.7
2.7
2.6
2.5
2.4
1060
2000
3000
1600
2000
3000
1000
2000
3000
Step
Stop
Stop
Perplexity of Figure 12: Perplexity of Figure 13: Perplexity of Layer(1-11) and Layer(16) on Layer(1-11) and Layer(16) on Layer(1-11) and Layer(16) on Transformer-Xi (Github).
Figure 11:
Longformer (Github).
Memtrans (Github).
15

---

Published as a conference paper at COLM 2024
3.2
Longfommer
Transformer
Longformer
Transformer
3.0
MemTrans
UniMix
UniMix
3.0
Vanilla
2.8
Perplexity
Perplext plexitty
2.8
2.6
2.6
2,4
2.4
2.2
SM
24M
40M Training Tokens
56M
72M
8SM
104M
24M
40M
56M 72M
88M
104M
Training Tokens
(a)
(b)
Figure 14: Perplexity Curves of Different Methods with Varying Training Token Amounts on the GitHub Dataset. (a) TinyLLaMA model. (b) LLaMA2-7B model.
C Impact of different training data scales
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-Xe demonstrated excellent scalability, suggesting their potential for even better performance with larger datasets.
D Hyperparameters
Hyperparameter
Value
Common
Learning rate Learning rate schedule Optimizer DeepSpeed Config
5B-5
Linear
AdamW
ZeRO-3 CPU Offload
Eval Interval
256
Model-specific
TinyLLaMA-1.1B
LLaMA2-7B
#Layers #Heads Embedding dim
22
32
4
32
2048
4096
Intermediate dim
5632
11008
Local context length
2048
2048
16


Reviews:
[{"review_id": "CGNuqqvGCe", "reviewer": "Reviewer_95Yk", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": true, "final_justification": "", "rebuttal": "Thank you very much for your comments! Regarding your valuable suggestions, we here commit to the following: \n\n## 1. About discussing more recent work about position interpolation\nFollowing your suggestion, we will add a \"Related Work\" section to discuss long-context methods parallel to UniMem, especially Position Interpolation (PI). UniMem and PI use different strategies for longer contexts: UniMem optimizes attention mechanism complexity via memory mechanisms, while PI adjusts the Transformer's position encoding for longer contexts than those used in pre-training.\n\n## 2. About adding more downstream evaluation tasks\nBased on your suggestions, we have supplemented downstream experiments on LongBench. The results are presented in the table below. \n\n[Table 1: Comparative Performance on Longbench](https://anonymous.4open.science/r/Downstream_tasks_results/README.md)\n\nDue to time limitations, our experiments were performed on tinyLlama, setting the context length to 512 and the memory size of Unimix to 10k tokens, while other hyperparameters were reduced proportionally compared to those specified in the paper. \n\nThe above results show that the performance ranking of downstream tasks aligns with the PPL evaluation in our paper: **UniMix > Transformer-XL > Longformer > MemTrans > RMT**. \n\nWe also conducted experiments using the \"Needles in the Haystack\" task, which comprises text lengths ranging from 4k to 32k, as depicted in the table below. UniMix is the sole method capable of identifying the \"needles\" in texts exceeding 4k in length.\n\n[Table 2: Performance on Needles in the Haystack across different text lengths](https://anonymous.4open.science/r/Downstream_tasks_results/README.md)\n\n## 3. About supplementary analysis of the impact of different training data scales.\n\nBased on your suggestion, 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 the figure below.\n\n[PPL trend of TinyLLaMA and  LLaMA2-7B on the GitHub Dataset](https://anonymous.4open.science/r/exp_mem-2EF0/README.md)\n\nWe will include the above experimental results in the paper.\n\n## 4. About the concern of Section 4.4\nCould you provide more detailed suggestions or guidelines? Specifically, how should we design our experiments to address your concerns?\n\nThank you once again. Please don't hesitate to reach out if you have any further questions!", "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": "1, To consider the long-context modeling works from a unified perspective is new and seems to be inspiring for the community.\n2, 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.\n2, 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.\n3, 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. \n4, The section 4.4 may not be a fair comparison.", "questions_to_authors": "please see the comments in the previous section.", "extra_scores": {}}, {"review_id": "PiIfV4MDoW", "reviewer": "Reviewer_HrNV", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your time and efforts. \n\n## 1. About the value and novelty of UniMem\nThe variety of extensive long-context methods can be challenging for researchers and industry users of this important area. UniMem provides them with a unified memory perspective to systematically understand these approaches. Thus, we think UniMem's value lies in the three aspects:\n### (1) Better understanding\nUniMem categorizes diverse long-context approaches into four clear memory dimensions. For instance, StreamingLLM and Memorizing Transformer were initially seen as entirely different methods, but from the UniMem perspective, their main difference lies in their use of different memory reading methods (position reading vs. similarity reading). UniMem's four memory aspects give a broad, clear adjustment range for long-context modeling. For instance, from the perspective of Memory Management, researchers can explore various overflow handling strategies, such as Least Recently Used (LRU), to maximize memory capacity utilization. Moreover, Memory Writing can be further optimized by experimenting with different compression and storage formats to achieve various levels of granularity and effectiveness in storage and retrieval.\n### (2) Insightful findings\nFor instance, we find varying the memory injection layer notably affects performance. Each language model has a distinct best injection layer, regardless of the downstream dataset.\n### (3) Fair evaluation\nUniMem categorizes long-context techniques into four memory dimensions, enabling granular control through tunable hyperparameters (Table 2). This standardization ensures equitable assessments, enhancing the reliability of comparative analyses.\n\nIn summary, UniMem is more of an analytical paper rather than solely proposing a stronger combined long-context method like UniMix.\n\n## 2. About adding more downstream evaluation tasks\nBased on your valuable suggestions, we have supplemented downstream experiments on LongBench. The results are presented in the table below. \n\n[Table 1: Comparative Performance on Longbench](https://anonymous.4open.science/r/Downstream_tasks_results/README.md)\n\nThe results show that the performance ranking aligns with the PPL results: **UniMix > Transformer-XL > Longformer > MemTrans > RMT**. \n\nThank you once again. We look forward to discussing this with you. Please let us know if you have any other questions!", "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": "**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. \n\n- 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. \n\n-  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. \n\n- 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.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VyCfXdoGcv", "reviewer": "Reviewer_Hfb4", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your comments! Regarding your valuable suggestions, we here commit to the following:\n\n## 1. About the supplemented baseline\n\nWe agree that Giraffe, YaRN, and Focused Transformer represent the latest advancements in this field, and including them in the baseline comparisons is essential for comprehensively evaluating the performance of UniMix. We will now discuss each of these three models individually:\n- Giraffe [1] and YaRN [2]: Giraffe and YaRN offer insights into positional interpolation and linear scaling methods, which can provide valuable comparisons for UniMix. In future work, We will include a detailed analysis of these position interpolation methodologies and their relationship with the UniMem framework.\n- Focused Transformer [3]: The model can access an additional context of (key, value) pairs through the k-nearest neighbors (kNN) algorithm. Focused Transformer's approach aligns with the Memory Reading and Memory Injection dimensions of the UniMem framework, as it utilizes similarity-based retrieval and augments specific layers with additional context, as shown in Table 1.\nDue to computation resource limitations, we plan to supplement these three models as new baselines in our future work. These works will further refine the UniMem framework, guiding our ongoing efforts to enhance long-context processing capabilities.\n\n[1] Pal et al., Giraffe: Adventures in Expanding Context Lengths in LLMs.\n\n[2] Peng et al., YaRN: Efficient Context Window Extension of Large Language Models.\n\n[3] Tworkowski et al., Focused Transformer: Contrastive Training for Context Scaling.\n\n## 2. About adding more downstream evaluation tasks\nBased on your valuable suggestions, we have supplemented downstream experiments on LongBench (including a long document QA task) and \"Needles in the Haystack\". The results are presented in the tables below. \n\n[Table 1: Comparative Performance on Longbench and Table 2: Performance on Needles in the Haystack across different text lengths.](https://anonymous.4open.science/r/Downstream_tasks_results/README.md)\n\nTable 1 shows that the performance ranking aligns with the PPL results: **UniMix > Transformer-XL > Longformer > MemTrans > RMT**.  In Table 2, UniMix is the sole method to identify the \"needles\" in texts exceeding 4k in length.\n\n\nThank you once again. Please don't hesitate to reach out if you have any further questions!", "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": "1. The unified framework for existing long-context methods is elegant in math.\n2. 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.\n\n[1] Pal et al, Giraffe: Adventures in Expanding Context Lengths in LLMs, https://arxiv.org/pdf/2308.10882\n\n[2] Peng et al, YaRN: Efficient Context Window Extension of Large Language Models, https://arxiv.org/pdf/2309.00071\n\n[3] Tworkowski et al, Focused Transformer: Contrastive Training for Context Scaling, https://arxiv.org/pdf/2307.03170\n\n2. 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.\n\n[4] Kuratov et al, In Search of Needles in a 11M Haystack: Recurrent Memory Finds What LLMs Miss, https://arxiv.org/pdf/2402.10790\n\n[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.", "extra_scores": {}}, {"review_id": "1Ct4NvnjpN", "reviewer": "Reviewer_pBGi", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your time and efforts. Regarding the weaknesses you mentioned, we are eager to engage in a deeper discussion with you.\n\n## 1. About adding more downstream evaluation tasks\nBased on your valuable suggestions, we have supplemented downstream experiments on LongBench and \"Needles in the Haystack\". The results are presented in the table below. \n\n[Table 1: Comparative Performance on Longbench and Table 2: Performance on Needles in the Haystack across different text lengths.](https://anonymous.4open.science/r/Downstream_tasks_results/README.md)\n\nTable 1 shows that the performance ranking aligns with the PPL results: **UniMix > Transformer-XL > Longformer > MemTrans > RMT**.  In Table 2, UniMix is the sole method to identify the \"needles\" in texts exceeding 4k in length.\n\n## 2. About the data size of training tokens\nWe acknowledge that 0.1B tokens is a small data size. Our decision to use 0.1B tokens in the current experiments was primarily based on the following considerations:\n\n### (1) Limited Computational Resources\nOur current equipment cannot support multiple model training and tuning on large-scale datasets. The 0.1B tokens data size was a compromise considering training time and cost.\n### (2) Reference to Existing Work\nWe noted that some representative works in long-context modeling, such as Longformer [1], also conducted fine-tuning experiments on datasets of a similar scale. \n\n[1] Longformer: The Long-Document Transformer. \n\n## 3. About Questions\n\n### (1) Data size selection\nWe selected and filtered 0.1 billion tokens from two primary datasets, the PG-19, and the GitHub datasets, and we focused on instances longer than 20k characters(as illustrated in Section 4.1).\n\n### (2) Number of trainable parameters and learning rate\nThe number of parameters is consistent across different methods, except that RMT introduces a small number of additional parameters. For the learning rate, we found that learning rate does not affect performance a lot.\n\n### (3) Dataset choice\nGitHub has been used by significant related works like the Memorizing Transformer[2][3]. We plan to include more datasets, such as Arxiv, in the future to enhance the robustness and applicability of our research findings.\n\n[2] Wu et al., Memorizing Transformers.\n\n[3] Tworkowski et al., Focused Transformer: Contrastive Training for Context Scaling.\n\n\nThank you once again. We look forward to discussing this with you. Please let us know if you have any other questions!", "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": "1. The research problem is significant to the community. The long-context LLM suffered from the unfair comparison between different methods for quite long.\n\n2. 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.\n\n2. 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.\n\n[Wang et al., 2023] Augmenting Language Models with Long-Term Memory\n[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?\n\n2. The number of trainable parameters for different methods vary a lot. Is the same lr 5e-5 fair for the implementation of each method?\n\n3. Why Github is considered as a long-context corpus for evaluation? The more common practice is using ArXiv paper dataset.", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept