MELODI: Exploring Memory Compression for Long Contexts
Authors: Yinpeng Chen
DeLesley Hutchins
Aren Jansen
Andrey Zhmoginov
David Racz
Jesper Sparre Andersen
Venue: ICLR 2025 Poster (2025)
Keywords: Memory
Compression
Long Context
Primary area: applications to computer vision, audio, language, and other modalities

Abstract: We present MELODI, a novel memory architecture designed to efficiently process long documents using short context windows. The key principle behind MELODI is to represent short-term and long-term memory as a hierarchical compression scheme across both transformer layers and context windows. Specifically, the short-term memory is achieved through recurrent compression of context windows across multiple layers, ensuring smooth transitions between windows. In contrast, the long-term memory performs further compression within a single middle layer and aggregates information across context windows, effectively consolidating crucial information from the entire history. Compared to a strong baseline - the Memorizing Transformer employing dense attention over a large long-term memory (64K key-value pairs) - our method demonstrates superior performance on various long-context datasets while remarkably reducing the memory footprint by a factor of 8.

Published as a conference paper at ICLR 2025
MELODI: EXPLORING MEMORY COMPRESSION FOR LONG CONTEXTS
Yinpeng Chen Andrey Zhmoginov
DeLesley Hutchins
Aren Jansen Jesper Andersen
David Racz
Google DeepMind
(yinpenge, delesley, jannjansen, azhmogin, dracz, jespera) agoogle.com
ABSTRACT
We present MELODI, a novel memory architecture designed to efficiently process long documents using short context windows. The key principle behind MELODI is to represent short-term and long-term memory as hierarchical compression scheme across both transformer layers and context windows. Specifically, the short-term memory is achieved through recurrent compression of context windows across multiple layers, ensuring smooth transitions between windows. con- trast, the long-term memory performs further compression within single middle layer and aggregates information across context windows, effectively consolidat- ing crucial information from the entire history. Compared to strong baseline the Memorizing Transformer employing dense attention over a large long-term memory (64K key-value pairs) our method demonstrates superior performance on various long-context datasets while remarkably reducing the memory footprint by factor of 8.
INTRODUCTION
output #
autput ar kel
*(N-M-11
x/NM-11
Short-term Layer
Short-term Layer
Long-term Layer
Long-term Layer
m auv
*+1
*M
*M
Short-term Layer
Short-term Layer
short-term
short-term memory
short-term memory
memory
window K
window k+1 30
k++1
Figure 1: Overview of MELODI. MELODI employs a hierarchical memory representation, incorpo- rating both short-term and long-term compression mechanisms, integrated with transformer-based language model. It utilizes stack of short-term layers to recurrently compress each context win- dow z into short-term memory tokens {zf{), and inserts long-term layer to store compressed key-value pairs within long-term memory mi:k- Both short-term and long-term layers leverage modified transformer blocks. In this illustration, we assume a total of N layers, with M short-term layers preceding long-term layer and N M short-term layers following it.
Long-context language models, exemplified by Gemini (Gemini- Team 2024) and GPT (OpenAl 2024). showcase remarkable capabilities across diverse modalities (e.g., text, images, audio, code, video) and seamlessly integrate various machine learning techniques, including many-shot in- context learning (Agarwal et al. 2024), thain-of-thought prompting (Wei et al.) 2022b), and the incorporation of explicit instructions (Chung et al. 2024 Wei et al. 2022a). However, the quadratic complexity of attention mechanisms within transformer models necessitates significant computa- tional resources to handle long contexts effectively. This has spurred the development of efficient solutions (Dai et al,) 2019: Wu et al. 2022 Bulatov et al. (022) that process long contexts via short context windows, much like how humans process information by reading book chapter by chapter.

---

Published as a conference paper at ICLR 2025
A central question underlying these solutions is: how can effectively model and manage memory to bridge the gaps between these short context windows over long context?
Memory fundamentally revolves around compressing and storing information for future utilization, all within the constraints of limited capacity. The Long Short-Term Memory (LSTM) architecture (Hochreiter & Schmidhuber, 1997) tackles this by recurrently compressing historical information into hidden states after processing each token. With the rise of Transformer models (Vaswani et al.) 2017) dominating the language modeling landscape, recent memory designs have shifted towards utilizing Transformers to process a context window, thereby moving the focus of memory manage- ment from the token level the context window level.
Transformer-X0 XL Dai et al. 2019) employs a caching mechanism to store multi-layer key-value (KV) pairs from the preceding window as memory. Memorizing Transformer (Wu al., 2022) builds upon this foundation by incorporating a dedicated layer to memorize all K pairs from that layer across all prior windows. Meanwhile, Block Recurrent Transformer (Hutchins et al. 2022) and Recurrent Memory Transformer (Bulatov et al.] 2022) introduce distinct recurrent compression mechanisms, implemented in a middle layer and at the output, respectively.
In this paper, we introduce MELODI (short for "MEmory with LOw DImension"), an efficient mem- ory architecture designed to handle long contexts despite operating on short context windows (e.g., 512 tokens per window). MELODI integrates both short-term and long-term memory through compression-basee approach. The short-term memory, recurrent in nature and possessing low ca- pacity, spans multiple transformer layers, progressively compressing context tokens and prior mem- ory at each layer. In contrast, the long-term memory, incremental and high-capacity, resides within single transformer layer. It maintains record of the entire history by further compressing each context window and concatenating them. Both short-term and long-term memory are seamlessly incorporated into multi-layer transformer model using a "sandwich" structure (see Figure curring negligible additional parameters.
MELODI demonstrates strong performance on various long-context datasets. For instance, utilizing a 13-layer transformer network with 1024 embedding dimensions and 512-token context windows, MELODI achieves perplexity scores of 10.44 and 2.11 on PG-19 (T5 vocabulary) and arXiv Math (Meena vocabulary), respectively. This represents a clear improvement over the Memorizing Trans- former (10.62 on PG-19, 2. 14 on arXiv) with dense attention (as opposed to top-k attention), while significantly reducing memory usage by factor of 8. Furthermore, ablation studies confirm the complementary nature of short-term and long-term memory in MELODI, highlighting their syner- gistic contribution to efficient and effective memory architecture
2 RELATED WORK
Memory in language models: Long Short-Term Memory (LSTMs) (Hochreiter & Schmidhuber 1997] use token-level recurrence to compress prior context into a state vector, which limited form of memory. With the advent of Transformers (Vaswani et al. 2017). the focus has shifted to memory mechanisms operating at the level of the context window; this shift allows blocks of tokens (i.e. all tokens within the window) be processed in parallel. The Block Recurrent Transformer (Hutchins et al. (022) and Recurrent Memory Transformer (Bulatov et al., (002)) integrate recurrent mechanisms inspired by LSTMs into the Transformer architecture. but the recurrence is over blocks,
rather than individual tokens. Transformer-Xl (Dai et al. 2019) introduces a caching mechanism to store key-value (KV) pairs from the preceding context window form of short-term memory. Memorizing Transformer (Wu et al. 2022) uses large cache to store KV pairs for long-term mem- ory, but only in one layer. Memory LLM (Wang et al. 2024) incorporates long long-term memory in every layer, incurring substantial memory overhead. Infini- Transfrrrafo (Munkhdalai et al. (202]) ex- plore the use of additional memory like Hopfield Networks (Hopfield) 1982] Ramsauer al.. 202T). LongMem (Wang et al., 2023) improves over Memorizing Transformer by introducing SideNet for memory retrieval and fusion. You Only Cache Once (Sun et al., 2024 shows that long-term KV cache is reuable for the latter half of the network, significantly improving pre-filling efficiency by enabling early exit. MELODI, in contrast, integrates integrates both short-term and long-term memory into transformer model via compression.
Compression: Recent work has explored using summary tokens for compression in Transformers (Rae et al. 2019 Bulatov et al. 2022 Chevalier et al. 2023 Ge et al. 2024). Recurrent Memory
2

---

Published as a conference paper at ICLR 2025
Transformer (Bulatov et al. 2022) utilizes the output of summary tokens recurrently short-term memory. AutoCompressor (Chevalier et al. 2023) aggregates summary tokens across segments to generate summary representation for long documents used in retrieval tasks. Gisting (Mu et al. applies this technique to compress long prompts. The In-context Autoencoder eG 2024) further incorporates LoRA fine-tuning (Hu al. 202T for context compression, while Transformer- FAM (Hwang et al., 2024) introduces feedback attention to enhance performance. Unlike these methods that compress input tokens, MELODI compresses network activations over multiple layers.
Extending context length: Recent research has demonstrated promising progress in scaling the context length of language models. To mitigate the cost attention mechanisms over long contexts, LongLoRA (Chen et al.) (2024), PCW (Ratner et al. (023). and Lundmark-Atteninn (Mohtashami (Chen
Iosition-Interpllane et al. 2023) and YaRN (Peng et al., 2023) extend context size by modifying the RoPE embedding scheme (Su et al., 2023). Xiong et al. (2023) provide a practical recipe for extending LLAMA2
Jaggi 2023) employ sparse local attention for efficient fine-tuning.
(Touvron al. 2023) to handle up to 32,768 tokens, while Fu et al. (2024) further explore data-
centric approaches for extending context length through lightweight continual pretraining.
3
OUR METHOD: MELODI
MELODI focuses efficiently comprchending long contexts while still using short context win- dows, thus circumventing the quadratic complexity associated with attention mechanisms over long sequences. This approach necessitates memory design that not only ensures smooth transitions between windows but also preserves crucial information from all preceding windows.
3.1 ARCHITECTURE OVERVIEW
Design principle: The core principle of MELODI is to represent short-term and long-term memory through a hierarchical compression scheme. Specifically, the short-term memory recurrently com- presses context across multiple transformer layers (e.g., condensing 512-token context window into 128 memory tokens). This process not only facilitates seamless transitions between context windows but also aggregates information across them, effectively functioning as fixed-size multi-layer long short-term memory (LSTTM) mechanism (Hochreiter & Schmidhuber 1997). Furthermore, each context window undergoes additional compression within middle layer and is then concatenated into a long-tern memory. This long-term memory retains essential information from the entire his- tory, thus compensating for any potential forgetting in the short-term memory. Both short-term and long-term memory are seamlessly integrated into a transformer- based language model, enabling the under the constraint of short context windows.
comprehension of long contexts
Sandwich architecture: MELODI'S network architecture (Figure) T features a "sandwich" structure, with long-term compression layer inserted between multiple recurrent short-term compression lay- ers. Both layer types utilize a standard transformer block (attention and feed-forward network) with tailored compression modifications. The short-term layers recurrently compress the current context window and update short-term memory, while the long-term layer further compresses information and appends it to long-term memory.
Terminology: In the remainder of this paper, we adopt the following notation. We use k to index context windows and to index transformer layers. Within the 1th layer of the t context window, the input context tokens are represented by x xf-1 The output context tokens, denoted as 2 serve as input for the subsequent layer. The input and output of the short-term memory are 2/__ and e/., respectively. The long-term memory preceding window & is denoted as mi1kk11 Note that we omit the subscript for the long-term memory since it resides within single long-term layer. Next, we will discuss both short-term and long-term layers in detail.
3.2 SHORT-TERY MEMORY: MULTI-LAYER RECURRENT COMPRESSION
The short-term memory is distributed across multiple short-term layers (see Figure D. This subsec- tion delves into the specifics of the short-term layer, using the 7th layer as an illustrative example for processing the tth context window. The short-term layer serves dual purpose: (a) transforming 1-1 via transformer block (yielding output z.), and (b) recurrently compressing context tokens z4
3

---

Published as a conference paper at ICLR 2025
the current context window into the short-term memory z It accomplishes this by updating both context tokens and short-term memory through a shared transformer block, albeit along separate pathways. As visualized in Figure 2 context to- kens traverse vertically across layers (from = to .rk), whereas short-term memory flows horizon- tally across context windows (from 2|__ to zg). To enable inter-layer communication within the short- term memory, we introduce summary tokens u' that propagate through the layers. We elaborate on the key components below.
*/ S ui
h( 1 cl
layar
short-term layer!
lineur mixing
FFN
Actention
Short-ferm memory z The short-term mem- ory (illustrated in Figure ( is implemented as a sequence of length S of vectors, each having the same dimensionality as context tokens (e.g.. 1024 channels). Notably, the number of short- term memory vectors is substantially smaller than the length of the context window (e.g., 128 mem- ory tokens per window of 512 context tokens). Within each context window, the short-term mem- ory serves initially as a previous context for auto- regressive prediction of subsequent context tokens (except for the first window in which the short- term memory is empty). It is then updated by compressing and incorporating information from the context tokens within the current window.
memory # from prewious window 2-1
memory next
window E+1
context
summary
previous tayer
Figure 2: Short-term layer. The figure illus- trates the processing of the 4thh context window the 1th short-term layer. It takes the mem- ory from the previous window 2 and the current context/summary (x) uu (-11 from the previous layer as input. The short-term layer adds two linear token mixers (Tolstikhin et al. 2021) top of standard transformer layer (including attention and FFN) to separate the summary for the next layer ut, and the memory for the next window 2/. Best viewed in color.
The context to- kens o are generated through causal attention to both (a) the preceding short-term memory = and (b) preceding tokens within the current context window. This attention mechanism is followed feed-forward network (FFN). Relative position embeddings are employed for
Transforming context tokens:
by the application both the context tokens of and the preceding short-term memory z Mathematically, this can be represented as:
a T 12(--))
(1)
where T(x22) indicates applying transformer block on for given preceding context x.
Recurrent compression: Beyond transforming context tokens, the short-term layer also recurrently compresses the current context window into short-term memory. Similar to the approach in RMT (Bulatov et al., 2022) and AutoCompressors (Chevalier al.] (023), this compression is achieved by appending summary tokens after context tokens x and passing the combined sequence through the transformer block. Consequently, the resulting summary tokens compresses both the preceding short-term memory and the current context window via attentional pooling (Lee et al.) (019)) ex- pressed as: i/. (uu-1122h k= ==/_1-=(/") 11, where the input summary tokens ux 1-1 originate from the previous layer (refer to Figure 2 We use it (instead of ut denote intermediate result that is further processed in the subsequent summary branching step. Both context and summary tokens can be updated simultaneously within a single transformer operation: z/e- at = T
, u/ 14-11=1/-) (=(-1)- Relative position embeddings applied on the short-term memory 2 context - and sum- mary d/²" while a causal mask is applied on the combined sequence of 2-11 and u/
Summary tokens (containing U tokens) are initialized from learnable embeddings (prior the first layer) and set to the same length as the short-term memory (U = 5). Propagating through all layers, they facilitate inter-layer communication within the short-term memory. Moreover, branch- ing summary tokens both upwards to the next layer and rightwards to the next window improves performance, a strategy will discuss in more detail subsequently.
Summary branching: We employ distinct linear token mixers Tolstiikhin et al., 2021 on the con- text and summary tokens generate separate summary tokens for the subsequent layer and short-
4

---

Published as a conference paper at ICLR 2025
next layer
Difference short-term layer
Cross attention long-term memory.
Gate over and cross attention
Appending long-term memory.
linear mixing
FFN
long-term layer
gate
from previous window k-T
Self-Attention
Cross-Attention
short-term memory
m=s-11
17
long-temn memory
shart-term memory
TR1--1
long-term
context summary from previous layer
next window k+
memory
Figure 3: Long-term layer. The long-tern ter-ttemm layer adds three components top of the short-term layer (see Figure Firstly, it introduces a long-term memory m1-k-1 by caching the compressed key-value pairs and allows the current context/summary (x1-1 u 1-11 to cross attend to them. Sec- ondly, the self-attention and cross- coss-attention are integrated via gating, Finally, the linear token mixing output additional compressed tokens and appends their key-value pairs my into the long-term mem- ory (as m1sk) for the next window. Best viewed in color.
term memory tokens for the next window. Unlike channel mixing, linear token mixer linearly combines the M, input tokens across each channel produce M, output tokens with the same di- mensionality by using M, M. matrix. The resulting summary and memory tokens exhibit distinct linear combination of context of. and summary tokens d/g implying divergent compression flows across layers and windows. Since the summary and short-term memory share the same num- ber of tokens (S), each mixer comprises (W + +) x S parameters, where represent the number of context tokens per window. This parameter count negligible for short context windows. For instance, with a context window of 512 tokens and 128 summary tokens, the two mixers collectively require (5 (512+12 x 12642=164K parameters, constituting mere 1.3% of transformer block with 1024 dimensions.
Summary: The short-term memory layer can be succinetly represented as function 2 2/. u/. h= #-1 1-1 u This function transforms context tokens =/-1 and summary tokens / up- ward to the next layer (from - to 1) while simultaneously propagating short-term memory =|_11 rightward to the next window (from i: 1 to k). Built upon standard transformer block, this layer introduces negligible additional parameters through the summary branching mechanism.
Relation Block Recurrent Transformer (BRT) (Hutchins al., 2022): Like MELODI, BRT also combines a block-wise recurrence mechanism with a Transformer- XL style KV-cache. How- ever, BRT only uses recurrence single "memory" layer, while MELODI adds recurrence to all layers. The MELODI recurrence mechanism relies on summary tokens, which do not require any new parameters, while BRT uses a separate "horizontal" transformer block for recurrence instead, which does introduce new parameters. Finally, MELODI builds summaries incrementally over mul- tiple layers (by passing summary tokens), while BRT recurrent state visible only to single layer.
3.3 LONG-TERM MEMORY: SINGLE-LAYER MEMORIZING COMPRESSED KEY- VALUE PAIRS
While short-term memory facilitates smooth transitions between context windows, the inherent lim- itation of its capacity can lead to the inevitable loss of information, particularly for contexts located further back in the sequence. In this subsection, we utilize long-term memory to retain information from all previous windows. thereby alleviating the forgetting inherent in short-term memory. The key idea involves further compressing the context window and storing the compressed representa- tions across the entire history.
5

---

Published as a conference paper at ICLR 2025
Table 1: Long-term vs. short-term memory. In our notation, L and S represent the num- ber of long-term and short-term tokens per net- work layer, respectively. Note that long-term has fewer tokens (L < S). Qmaz denotes the maximum number of windows encompassed by the long-term memory queue. .V signifies the total number of network layers.
Table 2: Baseline implementations. Our re- implementations utilize a cosine decay learning rate schedule (replacing inverse square root de- cay) and dense cross-attention for the Memorizing Transformer (replacing top-k attention). This re- sults in improved performance compared to prior reported results.
Baseline Method
PG-19 (T5) prior
arXiv (Meena)
PROPERTY
LONG SHORT
prior
our
our
Number of layers Update per window Capacity
single incremental recurrent Lx x Qmax S X N
multiple
Transformer XL
11.96 11.54 2.67
2.61
Block Recurrent
11.55 10.98
2.36
2.26
Memorizing Trans.
11.62 10.74 2.31
2.15
Further compression: In contrast to the short-term memory, which compresses information across multiple layers, the long-term memory achieves a higher compression rate within a single layer. Specifically, it compresses context window into L long-term tokens at designated middle layer (refer to Figure D where L is less than the number of short-term tokens per layer (L < S). Illustra- tively, we might compress a context window of 512 tokens into S 128 short-term tokens each layer, but further condense it into I. = 64 long-term tokens at a single layer.
Storing long-term memory: The key-value (KV) pairs of long-term tokens sequentially stored in a first-inffrrst-nn (FIFO) queue with a maximum capacity of Quaz windows. Consequently, the long-term memory can hold up L Qmaz KV pairs. For contexts shorter than Qmaa windows, compressed representation of the entire prior history is preserved. For longer documents exceeding Qmax windows, substantial portion of the recent history (Qmax windows) is still retained. We opt to store KV pairs (rather than the tokens themselves) because they are repeatedly utilized in cross-attention mechanisms for subsequent context windows, point we will elaborate on later.
Long-term layer: Figure 3 illustrates a long-term layer, which builds upon a short-term layer but incorporates three key additions. First, introduces a long-term memory (denoted as m1:k-1 prior to the ttth context window) and enables the current context tokens and summary tokens uy to cross-attend to it. Second, the cross-attention shares parameters with the self-attention, and their results (cross attention: Az, self-attention: A.) are combined through a gating mechanism using learnable scalar a per attention head, formulated as a.Ag + (1 a Ag. Finally, an additional linear token mixer is introduced to generate long term tokens for the current window, and their key-value (KV) pairs my are appended to the long-term memory. This token mixer comprises (W + U) x T, parameters, where W. U and I. represent the number of context, summary and long-term memory tokens, respectively. It is notably smaller than the mixers in the short-term layer because: (a) L is less than S (the number of short-term memory tokens), and (b) it is present in only one layer.
Long-term vs short-term memory: Table D provides a comparative overview of the long-term and short-term memory mechanisms. While short-term memory operates recurrently across N layers with S tokens per layer, long-term memory functions incrementally in one layer, spanning Qmax windows with L tokens per window. To illustrate, consider 13-layer transformer model processing context windows of 512 tokens each, with embedding dimension 1024. If we employ 128 short-term tokens per layer and L. 64 long-term tokens per window, with a max- imum capacity of Qmaz 128 windows for the long-term memory, the short-term and long- term memory caches would have capacities 1.7M (128tokens x 1024 dim x13 layers) and 16.8M (64peirs 22tokess/pan
1024dim x 128 floats, respectively. It's worth noting that, for com- putational efficiency, we store key-value (KV) pairs in the long-term memory, whereas tokens are stored directly in the short-term memory.
Relation to Memorizing Transformer (MT) (Wu et al., 2022): Similar to Memorizing Trans- former, MELODI incorporates key-value pairs from a middle layer into its long-term memory. Ex- periments corroborate the finding in MT that employing additional long-term layers yields only incremental gains. However, a key distinction lies in the fact that MELODI stores compressed KV pairs, rather than the KV pairs of context tokens directly, in MT. This modification substantially reduces the size of the long-term memory. For instance, when compressing a context window of 512 tokens into 64 long-term tokens, MELOODI achieves 8-fold reduction in long-term memory size.

---

Published as a conference paper at ICLR 2025
Table 3: Comparisons with baselines three datasets, The table reports average token-level perplexities for various models on three datasets. All models were trained under the same settings (e.g. segment length 4096, context window 5 512, 500k training steps). Three MELODI configurations were used: S192 L32, S128 + L64. and S192 + Les. For instance, S192 + La2 indicates 5 = 192 short-term and = 32 long-term tokens per context window. All MELODI models utilized a long- term memory spanning 128 context windows. MELODI S192 + Lyz outperformed Transformer XL and Block Recurrent Transformer while consuming less memory. Notably, MELODI S192 + L9
clearly surpassed Memorizing Transformer, using only fifth of its memory.
MODEL
MEMORY
PG19
arXiv Meena
C4(4K+)
All
Short
Long
Meena
T5
Custom
Custom
Custom
Transformer XL
13.6M
13.6
0M
8.65
11.41
12.42
2.60
3.22
18.22
Block Recurrent
13.1M
8.30
10.98
11.90
2.26
2.70
17.82
MELODI S192+ Lg32
11.0M
2.6M
8.4M
8.08
10.51
11.47
2.12
2.54
17,55
Memorizing Trans.
147,8M
13.6M
134.2M
8.07
10.62
11.53
2.14
2.56
17.37
MELODI $128 L64
18.5M
1.7M
16.8M
8.06
10.44
11.42
2.11
2.52
17.53
MELODI +192+
27.8M
25.2M
7.91
10.29
11.27
2.09
2.49
17.25
EXPERIMENTAL RESULTS
We evaluate MELODI on three long-context datasets: PG19 (Rae al.) 2019), arXiv Math (Wu et al. (022), and C4 (Raffel et al., 2020a), using the standard auto-regressive language modeling task, where objective to predict the next token in a sequence. All models are trained from scratch, and report the average perplexity the respective test sets our evaluation metric.
4.1 DATASETS
2019) consists of 28,752 English books published before 1919, averaging around 68,972 tokens per book. We utilize three 32k vocahularies: (a) Meena (Thoppilan ct al., 2022). (b) T5 (Raffel et al.
PG19: The PG19 dataset (Rae et al.
2020b) and (c) a custom SentencePiece vocabulary (Kudo & Richardson 2018) trained specifically on PG19.
arXiv Math: The arXiv dataset (Wu al. 2022 comprises technical math papers from arXiv, with token counts comparable to PG19 due to subwords. We use 32k Meena vocabulary
LaTeX's use of special characters, resulting smaller
TToppilann et al. 2022 and 32k custom vocabulary.
C4(4K+): The C4 dataset (Raffel et al. (020aa is a large collection internet documents. To emphasize long documents where memory is crucial, we filter out those with fewer than 4,096 tokens and utilize a 32k custom vocabulary.
4.2 SETUP
utilized decoder-only transformer architecture (with or layers), incorporating both short- term and long-term memory caches. The model had an embedding size of 1024, 8 attention heads with a dimensionality of 128 each, and an FFN hidden layer of size 4096. All models were imple- mented in JAX and Flax and trained from scratch for 500k steps 32 TPU cores. Further training details are provided in the Appendix
During training, each long document segmented into 4096-token chunks to facilitate batch pro- cessing. These chunks were then organized into training batches, each comprising 8 context win- dows of 512 tokens. In ablation study, MELODI's default configuration compressed each context window into 5=128 short-term memory tokens per layer and L=64 long-term memory tokens.
4.3 COMPARISON WITH BASELINES
Baselines: We benchmark MELODI against three well-established prior works: Transformer XL (Dai al. 2019), Block Recurrent Transformer (Hutchins al. 2022), and Memorizing Trans- former (Wu et al. 2022). To ensure a fair comparison, we re-implement these baselines within our framework and evaluate them under identical settings. Our re-implementations of the baseline mod-

---

Published as a conference paper at ICLR 2025
12.0
Transfomer
Lo
La
Le4
Los
0.0M
Token Perplexisy in PG-19
11.8
516
Block Recurrent Transformer
Ss
11.81
11.69
11.58
11.42
11.34
11.15
11.6
512
S16
11.75
11.83
11.47
11.35
11.24
11.10
0.2M
11.4
5125
Memoniinn Transformer
Saz 0.M
11.80
11.57
11.40
11.25
11.11
11.07
Si92 S25t
11.2
-
Log: 96 tokens x 128 windows
S64
11.52
11.43
11.36
11.23
11.08
10.98
Los: 64 tokens 128 windows
0:9M
11.0
32 tokens 128 windows
S64
S1zn
S120
11.39
11.41
11.27
11.08
10.95
10.90
L
S
1.7M
10.8
Lei tokens 128 windows
Lo: Short term alone
S192
2dM
11.25
10.83
10.6
10a
10
10"
10
S256
Memory Size (Long-term + Short-term)
3.4M
Figure 4: Ablation of memory size PG-19. The token perplexity is reported for various combi- nations of short-term and long-term memory sizes. Each curve represents a fixed size long-term memory, with points along the curve indicating different short-term memory sizes. For example, the blue curve (Lg6), uses 96 long-term tokens per window over 128 windows, totaling 12,288 tokens. Each point this curve represents different short-term memory size (e.g. Sg denotes short-term tokens per context window). Memory size is measured by the number of floating-point numbers (floats). For instance, Lye stores 12,288 long-term key-value (KV) pairs, each with 1024 dimen- sions, resulting in total of 12,288x | 1024x x=25,22 million floats. The table on the right provides the perplexity results for each point on the left plot, using matching colors. These results highlight that long-term and short-term memories play complementary roles, and increasing either type's ' capacity improves performance. Notably, MELODI achieves superior performance compared baselines like Transformer XL, Block Recurrrent Transformer and Memorizing Transformer while utilizing fewer memory resources. Best viewed in color.
els achieve superior performance compared to the results reported in their original papers (see Table 2). This improvement can be primarily attributed to two key factors: (a) using cosine decay learn- ing rate schedule (Hoffmann et al. 2022) instead of the inverse square root decay, (b) using dense cross-attention instead top-k attention for the Memorizing Transformer. Our re-implementations provide stronger baselines against which to evaluate MELODI's effectiveness.
Comparisons: Table presents a comparison of MELODI against three baseline models (Trans- former XL., Block Recurrent Transformer, and Memorizing Transformer) across three datasets. We evaluate three MELODI configurations: S192 + L32, S128 + L64, and S192 + Lys. where and L denote the number of short-term and long-term tokens per context window, respectively. All mod- els (MELODI and baselines) utilize a 13-layer transformer architecture, except for Block Recurrent Transformer, which inserts block recurrent layer into 12-layer transformer, ensuring similar parameter count for all models. Results demonstrate both the impact on LLM performance and efficiency of MELODI'S memory mechanisms.
Impact on LLM performance: Compared to Transformer XL, a transformer-basee language model with KV cache of the previous context window, MELOODI S192+ uses about the same amount of memory but significantly outperforms Transformer XL in terms of perplexity. For instance, on the PG-19 dataset with the T5 vocabulary, MELODI Laz achieves a perplexity of 10.51, compared to 11,41 for Transformer XL and 10.98 for Block Recurrent Transformer (enhanced Transformer XL with recurrent memory), demonstrating clear improvement. This highlights the effectiveness of MELODI'S 's hierarchical memory mechanism in capturing and utilizing long- range context.
Compared to Memorizing Transformer, MELODI variants (S128 - Les and S192 + Lgo) consistently achieves comparable or slightly better perplexity across all datasets, but approxi- mately 8 and 5 times less memory respectively.
Efficiency:
These trends remain consistent across different network depths (12-layer and 13-layer), shown in Ta Appendix(D) [.II] These results collectively highlight MELODI's efficacy and efficiency as a memory architecture for language models.
8

---

Published as a conference paper at ICLR 2025
11.4
MELOOI Sus
MELCOI Sua alone
Block Recument Nemorizing Trans.
Transfomar
MELO01
11.3
MELOOI Sua Aag
Parplaaity increment (512 -128 tokensouiee
0.8
12.75
Block Rocurrett
Tunnformer
Transormer
Perplesity
11.2
0.7
26
12.25
MELODI
11.1
02
Tolen Perpieiit
12.00
Size alone
11.0
64
11.75
0.5
Memorzinn
11.50
0.4
Tanstomer
NELODI
Sus
0
50
100
150
200
MELODI
Long-term Memoryy' Coverage of Windows)
11.00
Sux-Lac
MELODI
Figure 5: Long-term memory coverage. The coverage met- ric indicates the number of pre- ceding context windows (512 to- kens per window) spanned by the long-term memory. All data points use L=64 long-term to- kens per window, but vary in long-term memory capacity. For instance. '192' denotes storing information from previous 192 context windows, with effective context length 192x512=96k tokens. Perplexity improves as the long-term memory coverage increases, especially from 4 to 32 previous windows.
128
512
0.0
0.1
0.2
0.3
0.4
Window Size
Perplexity incremment (512 - 256 takersivin)
Figure 6: Shorter context windows. (Left): Perplexity with varying context window sizes (512, 256, 128 tokens). The num- ber of short-term and long-term memory tokens in MELODI is proportionally adjusted (x 2- x respectively) to ensure consis- tent long-term memory coverage. Even with smaller windows, MELODI with appropriate memory allocation (e.g. S12)+ consistently outperforms baselines. (Right): Perplexity increase due to window size reduction. The x: and y axes represent the perplexity increment when reducing the window size from 5 512 to 256 and 128 tokens, respectively. Models with long-term mem- ory (i.e. MELODI variants with long-term memory and Memo- rizing Transformer) exhibit significantly less performance degra- dation (smaller perplexity increments) compared to those rely- ing solely on short-term memory (MELODI with only short-term memory, Transformer XL, and Block Recurrent Transformer).
4.4 ABLATIONS
We conduct an ablation study of MELODI using the default configuration (612 with 128 short-term and 64 long-term tokens per context window. The long-term memory spans 128 context windows, and the transformer architecture consists of 13 layers. All models are trained on the PG-19 dataset with the T5 vocabulary for 200k steps.
Complementary roles of short-term and long-term memory: Figure illustrates how the sizes of both short-term and long-term memory jointly influence perplexity. Each curve represents fixed long-term memory size, with varying short-term memory sizes depicted by points along the curve. With the exception of the black curve, which solely utilizes short-term memory, all other curves incorporate long-term memory spanning 128 context windows. The figure demonstrates that in- creasing either short-term long-term memory capacity leads to improved perplexity, highlighting their complementary roles in performance. Notably, by judiciously selecting memory sizes (e.g.. outperform the Memorizing Transformer while
S192 for short-term and La2 for long-term), utilizing less memory than Transformer XL and Block Recurrent Transformer.
Impact of long-term memory coverage: In contrast the previous ablation, we now maintain fixed number of short-term ($ = 128) and long-term tokens (L 64) per context window while varying the number of windows encompassed by the long-term memory. Figure performance improves as the long-term memory covers a wider range of context windows. Inter- estingly. incorporating long-term memory for only the preceding or 4 windows yields marginal perplexity improvements, suggesting that recent context is already effectively captured by the short- term memory. However, performance gains accelerate as the long-term memory expands to encom- pass up to 32 windows, after which the improvements level off. This observation indicates that while the middle and distant history are beneficial for language modeling, they are not adequately retained in the short-term memory. These findings further underscore the complementary nature of short-term and long-ted term memory mechanisms.
that
Figure о examines the impact of reducing context window size on model performance. The plot on the Teft displays perplexity scores for context window sizes of 512, 256, and 128 tokens. The number of short-term and long-term memory
Performance with shorter context windows:
9

---

Published as a conference paper at ICLR 2025
Table 4: Summary branch- ing. Summary branching provides a consistent gain of approximately 0.3 in per- plexity, both with (column ST+LT) and without (col- umn ST) long-term memory.
11.8
Without Long term Memory With Long term Memory
11.4
Lagt Long-term
11.6
Leet 64 Long term tokens
Zoken Perpiesty
Nerplesity in
11.2
11.4
11.2
11.0
10.9
10.8
10
12
Number of Shott-term Layerr
Number of Long term Layers
Branching
ST
ST+LT
Number of short- term layers. Perplexity im- proves as the number of short- term layers increases.
Figure 7:
Figure 8: Number of long- term layers. A single layer with sufficient long-term tokens is adequate.
No
11.68
11.24
Yes
11.39
10.95
kens in MELODI is proportionally adjusted (x 2. x1, respectively) to ensure consistent long-term memory coverage across different window sizes. Even with shorter context windows, MELODI with appropriate memory allocation (e.g., S128 + Las) consistently outperforms the baseline models.
The plot on the right illustrates the increase in perplexity resulting from reducing the context window size. Notably, models incorporating long-term memory (MELODI variants with long-term memory and Memorizing Transformer) exhibit significantly less performance degradation (smaller perplexity increments) compared to models relying solely on short-term memory (MELODI variant with only short-term memory, Transformer XL, and Block Recurrent Transformer). This highlights the greater robustness of long-term memory mechanisms to reductions in context window size.
Number of short-term layers: We investigate the effect of varying the number of short-tern layers on model performance. For a given layer count (c.g., 4 layers), the short-term layers are uniformly distributed throughout the network (e.g., layers 1, 5, 9, and 13). To disable short-term memory within a layer, we remove (a) the attention mechanism to the preceding short-term memory and (b) the linear token mixer responsible for updating the short-term memory. Figure 7 shows that perplexity improves rapidly as the number of short-term layers increases from to 4, after which the gains diminish. This observation supports our utilization of multiple layers for effective short-term memory modeling. However, it also suggests that disabling short-term memory in half of the layers offers more efficient approach with negligible performance degradation.
Number of long-term layers: Figure B shows the impact of the number of long-term layers on perplexity for two configurations: Lea and Lz2, which compress each window into 64 and 32 tokens, respectively. Fewer tokens (I32) require more long-term layers for stable perplexity: L32 needs 4 layers to reach a plateau, while Les only needs 3. Notably, with additional tokens (La4), the first long-term layer contributes most significantly to performance improvement. It provides 73% of the total gain achieved with 6 layers. This suggests that a single long-term layer with sufficient tokens offers good balance between performance and efficiency. adding more layers increase memory and computational costs with diminishing returns.
Summary branching: Table 4 examines the effect of summary branching on perplexity, both with (ST+LT) and without (ST) long-term memory. Summary branching consistently yields a perplexity improvement of approximately 0.3, indicating distinct summary information flow across network layers and context windows.
5 CONCLUSION
In this work, we have introduced MELODI, novel memory architecture designed to address the challenges of long document processing within the constraints of short context windows. The core innovation of MELODI lies in its hierarchical compression approach, wherein short-term memory fa- cilitates smooth transitions berween context windows through recurrent compression across multiple layers, and long-term memory preserves erucial information from the entire history by performing further compression and aggregation within a single middle layer. Our empirical evaluations on multiple long-context datasets have validated MELODI an efficient and effective solution. The success of MELODI underscores the potential of hierarchical memory compression for tackling the complexities of long document processing. We anticipate that further research in this direction will enhance long context understanding and generation over multiple modalities.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Rishabh Agarwal, Avi Singh, Lei M. Zhang, Bernd Bohnet, Luis Rosias, Stephanie Chan, Biao Zhang, Ankesh Anand, Zaheer Abbas, Azade Nova, John D. Co-Reyes, Eric Chu, Feryal Be- hbahani, Aleksandra Faust, and Hugo Larochelle. Many-shot in-context learning, 2024. URL
https://arxiv. http:///aaxii...oaa///////////// wwt./bog////oo/./oo 22/.//444.......
Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. Recurrent memory transformer. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 11079- 11091. Curran Associates, Inc., 2022. URL bttps::/proceeediiieeee hetesssneerspppppp cc/pper wwwp://aaeeeeeeeeeeeaeeaaaaee fi 1 47e 2elee288622 9 69 0ateeeeae1eee0000000000000000000 Coneeeeaeaeenooeeeeeeeeeeee
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation, 2023. URL https://arxiv.om org/
abs/2306.155 15595
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuming tuning of long-context large language models, 2024. URL https://arxiy:
Stg/abs/2300 2309. 12307
Alexis Chevalier, Alexander Wettig, Anirudh Ajith, and Danqi Chen. Adapting language models to compress contexts. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 3829-3846, Singapore, December 2023, Association for Computational Linguistics. doi: 10.18653/v http//22233.mmmmm main. 232. URL https:t//////.. //a::///ccc /acl antho logy. org www.orgg2000 222.3 emn p- ma i 2332
Hyung Won Chung, Hou, Shayne Longpre, Barret Zoph, Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Alex Castro-Ros, Marie Pellat, Kevin Robinson, Dasha Valter, Sharan Narang, Gaurav Mishra, Adams Yu. Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Ja- cob Devlin, Adam Roberts, Denny Zhou, Quoc V. Le, and Jason Wei. Scaling instruction- finetuned language models. Journal of Machine Learning Research, 25(70):1-53, 2024. URL http:// http://mmmrr.ooo waammr.org/eaaae www.oo////ppp/.////e.rr
0870.html
Zihang Dai, Zhilin Yang, Yiming Yang. Jaime Carbonell, Quoc Le, and Ruslan Salakhutdi- nov. Transformer-LL: Attentive language models beyond fixed-length context. In Anna Korhonen, David Traum, and Lluís Marquez (eds.), Proceedings of the 57th Annual Meeting of the Association for Computational Linguisties, pp. 2978-2983. Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 10.18663/11/P19-12855 URL https::/ ac anthooooo........ www//oggg///.... 1285
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, year=.
Lizhe Fang, Yifei Wang, Zhaoyang Liu, Chenheng Zhang, Stefanie Jegelka, Jinyang Gao, Bolin Ding, and Yisen Wang. What is wrong with perplexity for long-context language modeling?, 2024. URL URL/https: //tp://aaivvvv...v wwwp:i.ooo//b////1.11..1 2211//0222/2111771111 23771
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannanch Hajishirzi, Yoon Kim, and Hao Peng Data engineering for scaling language models to 128k context, 2024. URL httphtttss//arxi
or org abs/ 2402. 10171
Tao Ge, Jing Hu, Lei Wang, Xun Wang, Si-Qing Chen, and Furn Wei. In-context autoencoder for context compression in a large language model, 2024. URL https://arxiv..rrt www.xxiwv.ooggg..oo 2307.0694
Gemini-Team. Gemini 1.5: Unlocking multimodal understanding across millions of tokens con- 05530
text, 2024. htt.st/:/// U

---

Published as a conference paper at ICLR 2025
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9(8): 1735-1780, 1997.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models, 2022. URL attps::/axxiv..rr//eee 1223.15556
J. J. Hopfield. Neural networks and physical systems with emergent collective computational abil- ities. Proceedings of the National Academy of Sciences of the United States of America, 79(8): 2554-2558, April 1982.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models, 2021. URL https:
//arxiv.org/absst /200..11666000 09685
Yutong Hu, Quzhe Huang, Mingxu Tao, Chen Zhang, and Yansong Feng. Can perplexity reflect large language model's ability in long text understanding?, 2024. URL https://arxiv:
(/55/bb//200 org/abs/2405. 2405./245.0000 06105
DeLesley Hutchins, Imanol Schlag, Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur. Block-recurrent transformers. In S. Koyejo. S. Mohamed, A. Agarwal, D. Bel- grave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 33248-33261. Curran Associates. Inc., 2022.
URL
cc/paper f.les/paper/202//FF fflee/////////////
dbe00bbbfcc444100002222233333333 Pappr//s-0sseeieeoo----r-rreeerr Paper- Confereeecc
Dongseong Hwang. Weiran Wang, Zhuoyuan Huo, Khe Chai Sim, and Pedro Moreno Mengibar. Transformerfam: Feedback attention working memory, 2024. UtL.htttss:/arzxiv.oo
09173
Taku Kudo and John Richardson. SentencePiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In Eduardo Blanco and Wei Lu (eds.), Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 66-71, Brussels, Belgium, November 2018. Association for Com- putational Linguistics. doi: 10. 101118/////D1--//// 2012. URL
2012
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set former: framework for attention-bant permutation-invaare neural networks. Proceedings of the 36th International Conference on Machine Learning, pp. 3744-3753, 2019.
Amirkeivan Mohtashami and Martin Jaggi. Random-access infinite context length for transformers. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 54567-54585. 2023. URL ht httt::::///pooeeengggneeeeeeeeee file/ otl./ wwt.--aaapeee-eeeeeeeeeeee pdf
Curran Associates, Inc.. fileesppper/2023/
Jesse Mu, Xiang Lisa Li, and Noah Goodman. Learning to compress prompts with gist tokens, 2024.
URL /t/htttte///aaoooooooo//a///////
08467
Dor Muhlgay, Ori Ram, Inbal Magar, Yoav Levine, Nir Ratner, Yonatan Belinkov, Omri Abend, Kevin Leyton-Brown, Amnon Shashua, and Yoav Shoham. Generating benchmarks for factuality evaluation of language models. In Yvette Graham and Matthew Purver (eds.), Proceedings of the 18th Conference of the European Chapter of the Association for Computational Linguistics, 2024.
Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: Efficient infinite context transformers with antinintttention, 2024, URL https: // /arxivv.... org. abs/ 2404. 07143
12

---

Published as a conference paper at ICLR 2025
OpenAl. Gpt-4 technical report, 2024. /BL.httpps////xxx///////////////
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yam: Efficient context window extension of large language models, 2023. URL U.L.tttps.//axxx...oroobbb/////p
Jack W Rae, Anna Potapenko. Siddhant M Jayakumar, Chloe Hillier, and Timothy Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint, 2019. URL
https://axxiv http:////aaa...ooo//////.......
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to- text transformer. Journal Machine Learning Research, 21(140):1-66 2020a. URL http:
v22 /20- 074 html
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to- text transformer. Journal of Machine Learning Research, 21(140):1667, 2020b. URL http:
7/jmlr. wttp://ppps/////////// 074 html
Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, Victor Greiff, David Kreil, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. Hopfield networks is all you need, 2021. URL https: http//axivv..o
org.abs//rr//////./.aa http::/:///.axi..oo//aaab.///// 2008.0000000000000000 02217
Nir Ratner, Yoav Levine, Yonatan Belinkov, Ori Ram, Inbal Magar, Omri Abend, Ehud Karpas, Amnon Shashua, Kevin Leyton-Brown, and Yoav Shoham. Parallel context windows for large language models, 2023.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive Learning Rates with Sublinear Memory Cost. In Proceedings of the 35th International Conference Machine Learning, pp. 4603-4611. PMLR, 2018.
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding, 2023. htw.htttpss//exxi..rrg////// 2104. 09864
Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma. Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models, 2024. URL pattpss//aaxxt
ww.. abs /200/ 05254)
Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kuishreshtha, Heng-Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, and et al. Lamda: Language models for dialog applications, 2022. URL http:httss///e/ /ttp:::///axvvv.// wttp://www.rrr/b/////0022222222
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung. Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, and Alexey Dosovitskiy. Mip-mixer: An all-mlp architecture for vision, 2021.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami. Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunoy, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023. URL (RI.htts:: a/arx//rxii org. abs aataaba/2302 /237 09288
13

---

Published as a conference paper at ICLR 2025
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser. and Illia Polosukhin. Attention is all you need. In Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Weizhi Wang, Li Dong, Hao Cheng, Xiaodong Liu, Xifeng Yan, Jianfeng Gao, and Furu Wei. Aug- menting language models with long-term memory. arXiv preprint arXiv:2306.07174 2023.
Yu Wang, Yifan Gao, Xiusi Chen, Haoming Jiang, Shiyang Li, Jingfeng Yang. Qingyu Yin, Zheng Li, Xian Li, Bing Yin, Jingbo Shang, and Julian McAuley. Memoryllm: Towards self-updatable large language models, 2024. URL https:ttttttt..tt.tttt
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoe V Le. Finetuned language models zero-shot learners. In International Conference on Learning Representations, 2022a. URL WRRLhttps:///pereeee
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed Chi, Quoc Le, and Denny Zhou. Chain-of ooain-of-thought prompting elicits reasoning in large language models. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 24824-24837. Curran Associates, Inc., 2022b. URL http://sppppooeeeii...eeeeooooes
files/paper/2022/ File/ Eatp://cc0SE3B22aa1aaaaaaaaaaaa0
.df
Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. In International Conference on Learning Representations, 2022. URL https:
//openzeview. wwtp://eteeettmmmeeemem http://ww...ttttommmmaaaaaaaaaaa jbxzeeff-
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang. and Hao Ma. Effective long-context scaling of foundation models, 2023. URL https: /awiiv..vv.. org/ abs /tt.//////2/ 2309.
14

---

Published as a conference paper at ICLR 2025
A LIMITATIONS AND FUTURE WORK
A.1 LIMITATIONS
primary limitation of our current study is the focus on training models from scratch. We acknowl- edge the importance of evaluating MELODI on downstream tasks and with pre-trained language models, which often have fixed context window sizes. However, adapting MELODI to such settings requires larger models and more extensive vre-training.
Furthermore, our evaluation primarily relies on perplexity as a performance metric. While perplexity provides useful measure of language modeling capability, it may not fully capture real-workd performance (Hu al. 2024 Fang et al. 2024). the gap somewhat emphasized for smaller models Muuhlgay et al. 2024).
While MELODI does not introduce additional parameters, the summary tokens within its short- term and long-term layers, along with the cross-attention mechanism in the long-term memory, do increase FLOPs (see Appendix E.3). Consequently, training time is extended when compared to Transformer XL (Dai et al. 2019). However, it only slightly increases test time (see Appendix E.4
A.2 FUTURE WORK
In future work, we plan to address these limitations by:
Integrating MELODI with pre-trained models: We will explore techniques like LoRA (Low-Rank Adaptation) to efficiently integrate MEL.ODI's short-term and long-term memory mechanisms into pre-trained language models. This will enable us leverage the benefits of pre-training and evaluate MELODI's effectiveness on downstream tasks like summarization and question answering.
Evaluating on downstream tasks: By adapting MELODI to pre-trained models, we can conduct comprehensive evaluation of its effectiveness on broader range of long-context benchmarks and downstream tasks. This will facilitate direct comparisons with sate-of-the-as baselines and provide a more complete picture of MELODI'S 's capabilities.
By addressing these limitations, we aim to provide a more comprehensive evaluation of MELODI and its potential for advancing long-context language modeling.
B CONNECTIONS TO RECENT WORK ON LONG-TERM AND SHORT-TERM MEMORY
This section compares MEL.ODI with recent works that explore different strategies for long-term and short-term memory in language models: LONGMEM (Wang et al. 2023) and YOCO (Sun et al.) 2024).
Comparison with LONGMEM (Wang et al. 2023): Both LONGMEM and MELODI improve upon the Memorizing Transformer (Wu et al. 2022) but with distinct focuses and methodologies:
LoNGMEM: LONGMEM focuses improving prediction accuracy in a fine-luning set- ting by introducing a SideNel for memory retrieval and fusion. It is compared against the Memorizing Transformer variant that uses top-k KV pairs from memory.
MELODI: MELODI prioritizes efficiency in a train-from-scratel setting by compressing both long-term and short-term memory. It introduces a negligible number extra parame- ters and is compared against the Memorizing Transformer variant that uses all KV pairs in memory.
direct experimental comparison with LONGMEM in the fine-tuning setting is beyond the scope of this paper, our primary focus is evaluating MELODI's performance when training from seratch. However, we plan to explore extending MELODI to the fine-tuning of LLMs future work, which will allow for comprehensive comparison with LoNGMEM and other relevant baselines.
Comparison with YOCO (Sun et al. 2024): Both YOCO and MELODI utilize a long-short mem- ory strategy, with short-term memory components spanning multiple layers and long-term informa-
15

---

Published as a conference paper at ICLR 2025
tion cached in middle layer. However, their core strategies differ, detailed below. Furthermore, they exhibit distinct properties (reusability vs. compressibility) in their long-term key-value (KV) memory:
YOCO: YOCO employs Bidding-windon attention across multiple self-decoder layers for short-term memory and KV cache in the middle layer for long-term memory. Importantly YOCO demonstrates that this long-term KV cache is reusable for the latter half of the network, significantly improving pre-filling efficiency by enabling early exit.
MELODI: MELODI implements a different long-short strategy based on compression. Short-term memory achieved through recurrent compression context windows across multiple layers, while long-term memory leverages further compression within a single middle layer. MELODI demonstrates that long-term KV memory compressible, signifi- cantly reducing its size.
YOCO and MELODI enrich the design space for long-short memory models and offer complemen- tary approaches. Combining their strengths could lead to even more effective models. For instance, exploring whether long-term memory can be both reusable (as in YOCO) and compressible (as in MELODI) is a promising avenue for future research. This could enable efficient caching with a compact KV memory footprint.
C IMPLEMENTATION DETAILS
C.I EFFICIENT PARALLEI. TRAINING
To enable efficient parallel training despite recurrent dependencies in short-term and long-term memory, MELODI employs a combined layer-wise and chunk-wise processing strategy. This ap- proach allows for parallel computation within each layer while maintaining the sequential process- ing required for recurrent memory updates.
Layer-wise processing and window unrolling: MELODI utilizes standard Transformer archi- tecture with multiple layers processed sequentially. Both short-term and long-term memory op- erations are confined within their respective layers. Within each layer, a training batch covers context windows (512 tokens window), totaling 4096 tokens. This allows for backpropagation through time (BPTT) across multiple windows. While the windows are processed sequentially due to window-leve level recurrence, self-attention within each window can be parallelized, ensuring good device utilization with sufficiently large window sizes.
Parallel and sequential operations: Within each layer, the computation of query, key. and value matrices for attention is parallelized across all chunks, as these computations are independent of recurrent memory updates. Each layer thus involves two steps: (a) parallel computation of query/key/value matrices, and (b) sequential computation of self-attention, feed-forward networks, and token mixing for memory updates.
This approach contrasts with a chunk wise then layer-wise strategy, where each chunk must be fully processed by the entire network before the next chunk can begin processing.
C.2 WINDOW-WISE PROCESSING IN GENERATION
MELLODI employs a two-step window-wist (or chunk-wise) processing strategy during generation:
Token-wise generation: Within each context window, MELODI generates tokens sequen- tially, similar GPT-like LLMs, utilizing both short-term and long-term memory. This ensures that token generation within a window (or chunk) remains unaffected by window boundaries, preserving the key mechanisms of autoregressive language models.
Summary generation: Upon completing window (or chunk), MELODI generates all sum- mary tokens in parallel and updates both short-term and long-term term memory accordingly.
This approach effectively balances the sequential nature of token generation with the parallel pro- cessing of summary information, optimizing generation efficiency.
16

---

Published as a conference paper at ICLR 2025
Table 5: Comparisons with baselines three datasets, The table reports average token-level perplexities for various models on three datasets. All models were trained under the same settings (e.g. segment length 4096, context window 512, 500k training steps). Three MELODI configurations were used: S192 L32, S128 + L64. and S192 + Les. For instance, S192 + L32 indicates 5 = 192 short-term and = 32 long-term tokens per context window. All MELODI models utilized a long- term memory spanning 128 context windows. MELODI S192 + Laz outperformed Transformer XL and Block Recurrent Transformer while consuming less memory. Notably, MELODI S192 + L9
clearly surpassed Memorizing Transformer, using only fifth of its memory.
MODEL
MEMORY Short
PG19
arXiv Meena
C4(4K+)
All
Long
Meena
T5
Custom
Custom
Custom
12 LAYERS
Transformer XI. Block Recurrent
12.6M
8.76
11.54
12.63
2.61
3.23
18.61
12.1M
12.1M
0M
8.47
11.12
12.11
2.27
2.73
18.27
MELODI +
10.8M
8.22
10.66
11.66
2.13
2.55
18.03
Memorizing Trans.
146.8M
12.6M
134.2M
8.15
10.74
11.68
2.15
2.57
17.88
MELODI S128 Lea
18.4M
8.16
10.61
11.66
2.13
2.55
18.01
MELODI +192+Loo
27.6M
2.4M
25.2M
8.08
10.48
11.47
2.11
2.51
17.75
13 LAYERS
Transformer XI.
13.6M
13.6M
0M
8.65
11.41
12.42
2.60
3.22
18.22
Block Recurrent
13.1M
13.1M
8.30
10.98
11.90
2.26
2.70
17.82
MELODI La22
11.0M
2.6M
8.4M
8.08
10.51
11.47
2.12
2.54
17.55
Memorizing Trans.
147.8M
8.07
10.62
11.53
2.14
2.56
17.37
MELODI S128 + Los
18.5M
1.7M
16.8M
8.06
10.44
11.42
2.11
2.52
17.53
MELODI 5
27.8M
2.6M
25.2M
7.91
10.29
11.27
2.09
2.49
17.25
C.3 TRAINING SETUP
We use Adafactor optimizer (Shazeer & Stern 2018) with learning rate schedule that employs a linear warmup for the first 1000 steps, followed by cosine decay. The maximum and minimum learning rates are set to 0.01 and 0.001, respectively, as recommended in [Hoffmann et al. (0022). dropout rate of 0.05 applied. All models are trained for 500k steps (200k for ablations) on 32 TPU cores with batch size of 32 (1 per core).
D ADDITIONAL EXPERIMENTAL RESULTS
D.1 MORE COMPARISON WITH BASELINES OVER DATASETS
Table 5 presents comparison of MELODI against three baseline models (Transformer XL, Block Recurrent Transformer, and Memorizing Transformer) across three datasets, using a consistent seg- ment length of 4096 tokens and a context window size of 512. The evaluation includes both 12-layer and 13-layer transformer architectures to assess the impact of model depth on performance.
Notably, even with fewer layers, MELODI S192+ consistently outperforms both Transformer XL and Block Recurrent Transformer across all datasets while consuming less memory. For instance, the 12-layer variant of MELODT S192+Lg2 achieves perplexity of 10.66 on PG-19 (T5 vocabu- lary), surpassing the 13-layer variants of Transformer XI. (11.41) and Block Recurrent Transformer (10.98).
In comparison to the Memorizing Transformer, MELODI S128 + Le4 exhibits slightly improved per- formance while dramatically reducing memory consumption by a factor of 8. Furthermore, MELODI S192 + Lys achieves even better perplexity scores across all datasets with a substantial reduction in memory usage exceeding factor of 5. The improvement is consistent for using both 12-layer and 13-layer transformer architectures. These results collectively highlight MELODI'S 's efficacy and effi- ciency memory architecture for language models.
17

---

Published as a conference paper at ICLR 2025
Table 6: Comparison of models with larger sizes evaluated on the PG-19 dataset using the T5 vocabulary. Models are evaluated based on perplexity (lower is better) and memory size. The batch sequence length is 4096, comprising 8 context windows of 512 tokens.
Embedding Dim
Model Size
Memory Size
MODEL
#Layers
Perplexity
Transformer XL
12
1024
151M
12.6M
11.54
Memorizing Transformer
12
1024
151M
146.8M
10.74
MELODI S128 + Lo4
12
1024
153M
18.4M
10.61
MELODI S992 LoG
12
1024
153M
27.6M
10.48
Transformer XL. Memorizing Transformer
1024
302M
25.2M
10.15
24
1024
302M
159.4M
9,45
MELODI S + Lo4
24
1024
306M
20.0M
9.30
MELODI S192 + Lye
24
1024
309M
30.0M
9.16
Transformer XI. Memorizing Transformer
36
1024
453M
37.8M
9.61
36
1024
453M
172.0M
8.92
MELODI S128 + Lo4
36
1024
459M
21.6M
8.81
MELODI S192 LoG
1024
463M
32.4M
8.70
Transformer XL Memorizing Transformer
1536
453M
25.2M
9.69
16
1536
453M
226.5M
9,01
MELODI S128 + Los
16
1536
456M
28.3M
8.89
MELODI S192 + Lgo
1536
459M
42.5M
8.79
D.2 SCALING UP MODEL SIZE
Table 6 presents the results of scaling up MELODI and comparing its performance with baselines on the PG-19 19 dataset (using the T5 vocabulary). We explored two scaling approaches:
Increased depth: We trained models with 24 and 36 layers, effectively scaling the number parameters by and 3 times, respectively, compared to the 12-layer base model.
Increased width and depth: We trained model with layers and an embedding dimension of 1536. which also scales the parameter count by 3 times.
The table shows that MELODI'S advantages hold even with larger models, achieving consistently lower perplexity than the Memorizing Transformer while using significantly less memory (approxi- mately 8 and 5 times less). This demonstrates MELODDI'S effectiveness in scaling to larger architec- tures while maintaining memory efficiency.
D.3 SCALING UP CONTEXT LENGTH
To investigate the performance of MELODI with larger context windows, we conducted experiments with batch sequence lengths of 4096 and 8192 tokens and window lengths of 512. 1024, and 2048 tokens. Increasing the context window allows the model to capture longer-range dependencies in the input sequence, which can be crucial for understanding complex language structures and tasks. All models were trained on PG-19 (T5 vocabulary) with 12-layer architectures.
Table D.2 shows that MELODD continues outperform the Memorizing Transformer with signifi- cantly less memory (approximately and times less), even with larger context windows.
Due to limited computational resources, exceeding 2048 tokens. Scaling MELODI to even larger context windows presents two primary challenges:
unable train MELODI with context windows
Memory capacity: The memory requirements of MELODI grow with both the context win- dow size and the batch sequence length. With our current hardware and codebase, we are unable to train MELODI with context windows exceeding 2048 tokens. However, this can
18

---

Published as a conference paper at ICLR 2025
Table 7: Impact of sequence and window length. This table compares the perplexity and mem- ory usage of Transformer XL, Memorizing Transformer, and MELODI with varying sequence and window lengths on the PG-19 dataset using the T5 vocabulary. All models have 12 layers.
Sequence Length
Window Length
Memory Size
MODEL
Perplexity 4
Transformer XL
4096
512
12.6M
11.54
Memorizing Transformer
4096
512
146.8M
10.74
MELODI S128 + L64
4096
512
18.4M
10.61
MELODI S192 + Lgs
4096
512
27.6M
10.48
Transformer XL
4096
1024
25.2M
11.26
Memorizing Transformer
4096
1024
159.4M
10.64
MELODI S250 + 128
4096
1024
20.0M
10.47
MELODI S384 L192
4096
1024
30.0M
10.36
Transformer XL Memorizing Transformer
8192
1024
25.2M
11.18
8192
1024
159.4M
10.42
MELODI S256 + L128 MELODI S384 L192
8192
1024
20.0M
10.27
8192
1024
30.0M
10.19
Transformer XL Memorizing Transformer
8192
2048
50.3M
10.94
8192
2048
184.5M
10.38
MELODI S512 + T256 MELODI S768 T.84
8192
2048
2
10.20
8192
2048
34.6M
10.10
Table Comparison of MELODI with Transformer XL using varying context window lengths. The perplexity on the PG-19 dataset using the T5 vocabulary is reported. The batch sequence length 8192 for all models. Transformer XL is evaluated with context window lengths of 512, 1024, 2048, and 4096, while MELODI uses context window length of 512.
MODEL
Sequence Length
Context Window Length
Perplexity
Transformer XI.
8192
512
11.40
Transformer XI.
8192
1024
11.18
Transformer XL
8192
2048
10,94
Transformer XL
8192
4096
10.65
MELODI S128 + Les
8192
512
10.34
MELODI Si92 Les
8192
512
10.22
be addressed by using TPUs with more memory and improving model parallelism, which we plan to investigate in future work.
Backpropagation through time (BPTT): MELODI'S recurrent connections necessitate BPTT, which involves unrolling the network over multiple time steps and further increases mem- ory consumption. Each batch spans 8 context windows by default, allowing for BPTT across multiple windows. While crucial for the recurrent short-term memory, BPTT con- tributes significantly memory usage. As shown in Table D.2] using longer batch quences (8K tokens) with IK token windows yields better performance than shorter se- quences (4K tokens) due to the increased BPTT depth.
D.4 COMPARISON WITH TRANSFORMER MODELS WITH EXTENDED CONTEXT
To further explore the impact of context window length, we compare MELODI (with a context win- dow of 512 tokens) conventional Transformer XL architecture with extended context windows ranging from 512 to 4096 tokens. This comparison allows us to assess the efficiency MELODI'S long-term memory mechanism against baseline model that relies solely increasing the context
19

---

Published as a conference paper at ICLR 2025
window within each layer. All models are trained on the PG-19 dataset with the T5 vocabulary for 500k steps, with each batch including 8K tokens to allow backpropagation through time (BPTT) across multiple windows.
Table g shows that Transformer XL's performance improves with increasing context length, as ex- pected. Notably, MELODI (512 tokens) outperforms Transformer XL even with an 8 times longer context window (4096 tokens), demonstrating the effectiveness of MELODI'S approach in capturing long-range dependencies.
D.5 NEW TASK: MASKED NEXT TOKEN PREDICTION
We introduce a novel task designed to further evaluate MELODI'S capacity for handling long-range dependencies. Similar our primary evaluation, formance with perplexity. However, we introduce a key challenge: masking out a portion of the input tokens. Unlike BERT's (Devlin al.) bidirectional masking approach, this new task employs a unidirectional, left-to-right prediction scheme.
token prediction and measure per-
Why evaluate on this task? This masking strategy presents unique challenge to the model. By disrupting the local context with masked tokens, we force the model to rely on information from more distant, unmasked tokens to accurately predict the next word. This effectively evaluates the model's 'ss ability to capture and utilize longer range dependencies. A model with strong long-range capabilities will demonstrate lower perplexity on this task, indicating its proficiency in integrating information across extended sequences.
Experimental results: To evaluate MELODI'S performance on this new task, we conducted exper- iments using the PG-19 dataset (with the T5 vocabulary) and two different architectures: 12-layer and 36-layer. We applied random masking with three masking ratios (0.125, 0.25, and 0.5) to sys- tematically vary the degree of contextual disruption. The batch sequence length is 4096, comprising 8 context windows of 512 tokens.
As expected, perplexity increased across all models the masking ratio increased, indicating the growing difficulty of the task. However, MELODI consistently outperformed both Transformer XL. and the Memorizing Transformer across all masking ratios and architectural configurations. These results, presented in Table strongly suggest that MELODI exhibits superior capabilities in cap- turing and utilizing long-range dependencies compared these baseline models.
Interestingly, even with high masking ratio of 0.5, the 12-layer MELODI variant (S192 + L9a) achieves a perplexity of 52.59, which is remarkably close to the 36-layer Transformer XL's perplex- ity of 51.23. This suggests that MELODI can achieve comparable performance with significantly fewer parameters.
Slower perplexity increase for MELODI with increasing mask ratio: We further examined the change in perplexity as the masking ratio increases: 0-5,.125, 0.125- and 0.25-s0.5. As shown in smaller than that observed for the baseline models. This indicates that MELODI is less affected by the increasing disruption of local context, further highlighting its superior ability to capture and leverage long-range dependencics.
numbers in the bracket), the perplexity increase for MELODI is consistently
D.6 MORE ABLATIONS
In this section, we present additional ablation studies.
Directly copying short-term memory to long-term memory: This ablation experiment, conducted at the long-term layer, explores directly copying short-term memory tokens as long-term tokens. Instead of generating long-term tokens using the linear token mixer, this approach utilizes the short- term tokens present at the long-term layer and stores them directly in the long-term memory. Here, we force the short-term and long-term to share the same number of tokens (S = L). The results in Table 10 indicate that this direct copying method leads to performance degradation.
Effect of long-term layer position: We investigate the impact of varying the position of the long- term layer (M in Figure within 13-layer MEL.DDI model (S128+ T-44). Layers indexed from
20

---

Published as a conference paper at ICLR 2025
Table Performance comparison on masked next token prediction with varying mask ratios. This table presents the perplexity of Transformer XL, Memorizing Transformer, and MELODI on the PG-19 dataset (T5 vocabulary) across different mask ratios. Each column represents a different masking level, and each model is evaluated with both 12 and 36 layers. The number in the bracket indicates the perplexity increase from the previous column due to the increase of mask ratio.
Mask Ratio
MODEL
#Layers
mask
0.125
0.25
0.5
Transformer XI
12
11.54
(+..5((+4.5))
2
63.51 (+43.01)
Memorizing Transformer
12
10.74
14.91 (+4.17)
20. (5..55))
56.28; (6.28(+35.52) (+35..2)
MELODI S128 + Lea
12
10.61
14.58 (+3.97)
20.09 (45.511)
53.63 (+33.54)
MELODI S192 + Las
12
10.48
14.39 (+3.91)
20.00 (+5.61)
52.59
(+32.59)
Transformer XI. Memorizing Transformer
36
9.61
13.10
18. 18.18 (+4.008)
51.23 51.23 (+33.05)
36
8.92
(
16. 16.64(44.49)
46.80; (+30.11)
MELODI S128 + Lo4
36
8.81
11.97 (+3.166
16. (+4.244
45.5 (+229.36)
MELODI S192 + Lge
36
8.70
11.74 (+3.04)
15.98 (44.24)
44.05 (+28.00))
Directly copying short-term to- kens to long-term memory: Here, we force the short-term and long-term same number of tokens (e.g. L64+ and examine the impact of directly using short- term tokens long-term tokens, bypassing the linear token mixer. This approach results in performance degradation compared to gen- erating distinct long-term tokens.
Table 10:
Table 11: Position of long-term layer. Here, we use a 13-layer MELODI model, where layers are indexed from 0 to 12. The model uses 128 short- term and 64 long-term tokens per context window. While the default position of the long-term layer is at layer 8, placing different layers between 5 and yields consistent perplexity scores.
share the
Layer
5
6
7
8
Copying S96+ Tgg Se4+T44
S32+ T-32
Perplexity
11.00
11.01
11.03
10.95
Yes No
11.00 10.92
11.36
11,42 11.25
Layer
9
10
11
11.08
Perplexity
10.96
10.95
10.94
0 to 12, with the default long-term layer position at layer 8. Results in Table 11
reveal consistent
perplexity scores when the long-term layer is positioned between layers 5 and IT.
E ANALYSIS OF MEMORY USAGE AND COMPUTATIONAL COMPLEXITY
This section analyzes the memory usage and computational complexity of MELODI, including em- pirical measurements of training and test times compared to the baselines.
E. NOTATIONS
To facilitate our analysis, we first define the following notations:
N: Number of layers in the model.
W: Number of context tokens per window.
D: Dimension of the token embeddings.
Number of short-term memory tokens per window.
Q: Number of windows covered by long-term memory in the Memorizing Transformer.
R: MELODI's long-term memory reduction rate compared to the Memorizing Transformer.
21

---

Published as a conference paper at ICLR 2025
Table 12: Comparison memory usage between MELODI and the baseline models. Notations introduced in Appendix E.T
MODEL
Short-Term Memory
Long- Tem Memory
Transformer XL
2NWD 2NW D NSD
0 2QW D 2QWD
Memorizing Transformer
MELODI
Table 13: Memory usage and model size for MELODI and baseline models. Memory usage is measured in number floats and includes both short-term and long-term memory. Model size is reported in number of parameters.
MODEL
#Layers
Embedding Dim
Model Size
Window Length
Memory Size
Memory Size Model Size
Transformer XL Memorizing Transformer
12
1024
151M
512
12.6M
0.08
12
1024
151M
512
146.8M
0.97
MELODI S128 + L61
12
1024
153M
512
18.4M
0.12
Transformer XL Memorizing Transformer
16
1536
453M
512
25.2M
0.05
16
1536
453M
512
226.5M
0.50
MELODI S128 + L64
16
1536
456M
512
28.3M
0.06
Transformer XL Memorizing Transformer
32
4096
6.7B
4096
1.07B
0.16
32
4096
6.7B
4096
5.37B
0.80
MEL.ODI S1024 + T.512
32
4096
6.8B
4096
0.67B
0.10
E.2 MEMORY USAGE
Table F. compares the short-term and long-term memory size of MELODD and the baseline models. For example, the MELLODI S128 + T-4 variant uses quarter of the window length for short-term memory (S 4/4) and achieves a long-term memory reduction rate R 8, resulting in an 8x reduction in memory usage compared to the Memorizing Transformer.
To further illustrate the relationship between memory footprint and model size, we compare three different model configurations in Table 13] This includes the two smaller Transformer models used in our experiments and the Llama-2 (Touvron et al. 2023) 7B model with 32 layers, an embedding dimension of 4096, and context window of 4096. For this comparison, ory to cover 128 context windows. The table demonstrates that the memory size for the Memorizing Transformer is comparable to the model size itself. In contrast, MELODD'S memory size remains sig- nificantly smaller-by an order of magnitude- for large model like Llama-2. This highlights the memory efficiency MELODI'S approach, especially for large-scale language models.
the long-term
E.3 COMPUTATIONAL COMPLEXITY
As shown in Table 14 the Memorizing Transformer introduces additional complexity compared to Transformer XL due to the cross-attention computation over the long-term memory its middle layer. MELODI reduces this cross-altention complexity by factor of /t compared to the Memorizing Transformer. However, it introduces additional computations in the query-key-value, feedforward network, and linear token mixing components. The extent of these additional computations is trolled by the number of short-term memory tokens, S, which is about a quarter of window length W.
22

---

Published as a conference paper at ICLR 2025
Table 14: Component-wise comparison of computational complexity between MELODI and the
baseline models. Notations are introduced in Appendix E.
MODEL
Query- Key-Value
Self- Attention
Cross- Attention
Feedforward Network
Linear Token Mixing
Trans. XL Mem Trans.
O(NWD*) ONNW O(N(W+S)D*) O(N(W+S)SD) (N((W)))D O(NW2D) O ((WD)
O(NWD2) O(NWD2)
O(NW2D) O(NW2D)
D(QWD)
MELODI
Table 15: Comparison of training and test times for MELODI and baseline models.. We mea- sured the training and test times of 12-layer networks with an embedding dimension of 1024 on TPU v6e. Experiments were conducted with sequence lengths of 4K and 8K tokens and window lengths of 512 and IK tokens.
MODEL
Segment Length (num of tokens)
Window Length (num of tokens)
Training Time (sec)
Test Time (sec)
Transformer XL
4096
512
0.052
0,022
Memorizing Trans.
4096
512
0.094
0.033
MELODI L628+L64
4096
512
0.076
0.023
MELODI L99
4096
512
0.097
0,027
Transformer XL.
4096
1024
0.068
0.027
Memorizing Trans.
4096
1024
0.108
0.036
MELODI S158+L128
4096
1024
0,094
0.028
MELODI S184+L192
4096
1024
0.111
0.031
Transformer XL
8192
1024
0.108
0.045
Memorizing Trans.
8192
1024
0.196
0.067
MELODI S258+I118
8192
1024
0.164
0.047
MELODI S181+L192
8192
1024
0.204
0.054
E.4 TRAINING AND TEST TIMES
To evaluate the computational efficiency of MELODI, we measured its training and test time and compared it against the baselines. All experiments were conducted using 12-layer networks with an embedding dimension of 1024 on TPU v6e. We evaluated performance with batch sequence lengths of 4K and 8K tokens and window lengths of 512 and 1K tokens.
As shown in Table 15 MELODI consistently demonstrates faster test times than the Memorizing Transformer. Impressively, the MELODI + variant approaches the test speed of Transformer XL (e.g., 0.022 vs. 0.023 seconds with 4K sequence length and 512 window length).
However, when considering training time, both MELODI variants are comparable to the Memoriz- ing Transformer but slower than Transformer XI.. This suggests that the additional computations introduced by MELODI'S short-term memory component (in the query-key-valus feedforward net- work, and linear token mixing) are relatively inexpensive during (forward pass) but become more costly during backpropagation.
23


Reviews:
[{"review_id": "7FSHlSSZQT", "reviewer": "Reviewer_9LJX", "summary": "This paper proposes a new context compression neural network for the long context in today's LLMs. With a hierarchical structure, it achieved a decent performance by reducing the memory footprint by a factor of 8.", "questions": "1. A preliminary session about the history of recent advances about memory compression would be much appreciated.", "limitations": "", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. The methodology seems novel and makes a lot of sense in dividing the tasks into long term and short term memory in the compression.\n2. The empirical results seem promising.", "weaknesses": "1. Without prior knowledge, it's not clear that how the memory network works with the LLMs and this is also not reported as preliminary knowledge.\n2. The results only show the efficiency part and do not show the effect of memory network on the performance of LLM.\n3. Codes are not shared for checking the implementation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jXx9R2beLe", "reviewer": "Reviewer_XNtf", "summary": "MELODI is a novel memory architecture designed to efficiently process long documents using short context windows. MELODI is a hierarchical compression scheme for both short-term and long-term memory.\n\nThe Short-Term Memory uses a Multi-Layer Recurrent Compression Scheme which is distributed across multiple short-term layers, this mechanism progressively compresses context tokens and prior memory at each layer. It is asserted that this ensures smooth transitions between context windows and aggregates information across them. Each short-term layer: (1) Transforms context tokens via a transformer block and (2) Recurrently compresses the current context window into short-term memory *tokens*. Summary tokens are introduced to facilitate inter-layer communication within the short-term memory. Distinct linear token mixers are used to generate separate summary tokens for the next layer and short-term memory tokens for the next window, enabling divergent compression flows.\n\nThe Long-Term Memory is a Single-Layer Module which Memorizes Compressed Key-Value Pairs. It uses a single network layer and maintains a record of the entire history by further compressing each context window and stacking them. The claim is that this addresses the forgetting issue inherent in short-term memory due to limited capacity. Key-value (KV) pairs of long-term tokens are stored in a FIFO queue, ensuring the retention of a substantial portion of the recent history. The long-term layer introduces a long-term memory component and enables cross-attention between current context/summary tokens and the long-term memory. It employs a gating mechanism to combine self-attention and cross-attention results. It uses a linear token mixer to generate compressed KV pairs for the current window, which are then appended to the long-term memory.\n\nThe overall MELODI Architecture utilizes a \"sandwich\" structure with a long-term compression layer inserted between multiple recurrent short-term compression layers. It appears to demonstrate strong performance on various long-context datasets while significantly reducing memory usage compared to baselines like the Memorizing Transformer. Ablation studies appear to confirm the complementary roles of short-term and long-term memory, highlighting their synergistic contribution to an efficient and effective memory architecture.\n\nCompared to LSTMs, which operate at the token level, MELODI focuses on memory management at the context window level. Unlike Transformer XL, which relies on caching KV pairs from the preceding window, MELODI employs a multi-layer recurrent compression for short-term memory. While similar to the Memorizing Transformer in using a dedicated memory layer, MELODI stores compressed KV pairs, leading to a substantial reduction in memory size.", "questions": "Suggestion: Could you provide some more experimental results in the 2k and larger token context regime?\nQuestions: Could you elaborate on any particular challenges and potential solutions for scaling MELODI to handle even larger context windows, such as 4096 or 8192 tokens, which are becoming increasingly common in language modelling? \nIs the key issue \"just\" the size of the GPUs / TPUs that you/one has access to, or could you provide more (concise) insights on this point?\n\nCould you provide a discussion and/or analysis of MELODI's computational efficiency, including complexity analysis or measurements and comparisons of training and inference times compared to baseline models? Latency is an issue in real world situations and it would be useful for the reader to understand the inference time implications here.\n\nThe paper highlights the potential of MELODI for various natural language processing applications, such as document summarization, question answering, and machine translation. Have you considered conducting preliminary experiments on any of these downstream tasks to demonstrate the practical benefits of MELODI in specific applications?  With those kinds of experiments and answers to the questions above I could potentially score this paper higher.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality\nThis paper combines and extends ideas from previous works on memory mechanisms for language models. It integrates both short-term and long-term memory into a unified architecture, leveraging compression techniques inspired by LSTMs, Transformer XL, and Memorizing Transformer. The paper introduces a new hierarchical compression scheme for representing short-term and long-term memory. This scheme differentiates MELODI from previous approaches by employing recurrent compression across multiple layers for short-term memory and further compression within a single middle layer for long-term memory. The paper introduces a summary branching mechanism that uses linear token mixers to generate separate summary tokens for different purposes (next layer and next window). This contributes to the originality of the approach by enabling distinct compression flows within the short-term memory.\n\nQuality\nThe paper presents comprehensive experimental results on three long-context datasets, demonstrating that MELODI outperforms several strong baselines (Transformer XL, Block Recurrent Transformer, and Memorizing Transformer) in terms of perplexity while significantly reducing memory usage. The paper conducts detailed ablation studies to analyze the impact of various design choices, including the sizes of short-term and long-term memory, long-term memory coverage, context window size, the number of short-term and long-term layers, and summary branching. These studies provide good evidence supporting the effectiveness and efficiency of MELODI's architectural choices.\n\nClarity\nThe paper is well-written and easy to follow. The authors present the motivation, architecture, and experimental results in a clear and concise manner. The paper uses several figures to effectively illustrate the key concepts and architectural components of MELODI. These figures significantly aid in understanding the hierarchical compression scheme and the flow of information within the model.  The authors provide detailed explanations of the various components of MELODI, such as the short-term and long-term layers, summary branching, and the gating mechanism used in the long-term layer. \n\nSignificance\nThe paper tackles the important problem of efficiently processing long documents with transformer language models. This is a significant challenge because the quadratic complexity of attention mechanisms makes it computationally expensive to handle long sequences directly.\nThe proposed method has practical implications for various real-world applications that involve processing long texts, such as document summarization, question answering, and machine translation. MELODI's efficiency and effectiveness in handling long contexts could lead to improved performance and reduced computational costs in these areas. The paper does point to some new avenues for future research in memory architectures for language models. The hierarchical compression scheme and the use of summary branching could inspire further investigations into efficient and effective ways to manage information over extended sequences. Additionally, the authors acknowledge the limitation of their current work in addressing the fine-tuning of pre-trained models and suggest exploring adaptation techniques like LoRA fine-tuning in future work.", "weaknesses": "(1) The paper has only a limited exploration of longer context window size\nThe primary focus and the majority of the experiments revolve around a context window size of 512 tokens. This emphasis on relatively short context windows might limit the applicability of the method to tasks requiring significantly longer contexts, such as those involving very long documents or those benefiting from processing larger chunks of information simultaneously. The paper could be strengthened by conducting more extensive experiments with larger context window sizes, particularly those exceeding 2048 tokens or, ideally experiments in the 10k token regime. This would provide a better understanding of MELODI's performance and scalability in handling truly long contexts. Exploring how the model's memory capacity and compression mechanisms need to be adjusted for optimal performance with varying context lengths would be valuable. Figure 5 could be highlighted, discussed and presented more clearly and explicitly in terms of the long context issue. It seems this experiment does give perplexity numbers for 192x64=12k token context windows, but it is a bit hidden amongst the rest of the experiments given the importance of these longer regimes.\n\n(2) There appears to be a lack of fine-tuning experiments\nMELODI and the baselines, seem to have been trained from scratch. The paper doesn't explore the performance of MELODI when fine-tuned on a pre-trained language model. While most LLMs are derived from a larger pre-trained model, this is only a minor limitation given the amount of resources needed to train a SOTA grade LLM.  Investigating the adaptability of MELODI to be applied in fine-tuning scenarios would be beneficial. This could involve exploring techniques like LoRA (Low-Rank Adaptation) to integrate the short-term and long-term memory mechanisms into pre-trained models without significantly increasing the number of trainable parameters. Evaluating the performance of fine-tuned MELODI models on downstream tasks would provide insights into its effectiveness in practical applications. This is mentioned as possible future work.\n\n(3) Limited Analysis of Computational Efficiency\nMore analysis and discussion of Memory Usage and Computational Time or computational complexity would be helpful. The paper could potentially benefit from discussions and comparisons of training and inference times for MELODI and the baselines\n\n(4) Experiments are almost exclusively performed providing test set perplexity results. While there are many downstream tasks that could benefit from models yielding better perplexity, it would be more compelling to see the impact on other kinds of tasks and metrics where long range information is important.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IKYhIV6dGb", "reviewer": "Reviewer_hBrU", "summary": "The paper introduces a novel Transformer-based architecture, MELODI, which combines multiple short-term context window (recurrent + chunk-wise) Transformer layers with a single long-term compression layer. Experimental results show that MELODI achieves the lowest perplexity compared to similar memory-based Transformer variants. In addition, the paper presents ablation studies exploring various combinations of architectural components.", "questions": "* It would be beneficial to provide the upper bound, such as a very long-context Transformer model (ex: 256K or more) for readers to get a sense of the performance gap.\n* Suggestion: It would be interesting to compare MELODI with YOCO [1], which employs a similar long-short strategy. I understand that [1] is a very recent paper and you do not need to compare the work; it would be valuable to share your thoughts on the similarities and differences. [1] You Only Cache Once: Decoder-Decoder Architectures for Language Models, 2024.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The integration of long and short memories is a straightforward yet effective strategy. The paper clearly illustrates the concept through well-written explanations and clear visualizations.", "weaknesses": "* The experiments mainly focus on perplexity comparisons, rather than evaluating the performance on long-context benchmarks such as LongBench or Needle-in-a-Haystack (NIH).\n* The paper lacks discussion regarding the challenges in training the model; the recurrent nature of the model may impede efficient parallel training. Can you provide more detailed explanations?\n* Although the model is Decoder-only, it is unclear how the chunk-wise approach affects token-wise generation behavior (a key mechanism in GPT-like LLMs). Does summary-generating action occur only after completing each chunk during the generation phase?\n* The model appears to be relatively small, raising questions about the scalability of the architecture. Details regarding the model size (e.g., 1B, 2B?) and the model’s potential for scaling up would enhance the paper’s value.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4tyc9PObcJ", "reviewer": "Reviewer_i9z9", "summary": "This paper presents MELODI, a memory architecture for efficient long document processing. It utilizes a hierarchical compression scheme, where the short-term memory performs recurrent compression of context windows, and long-term memory aggregates information from the entire history. Experimental results indicate superior performance along with a reduction in memory footprint.", "questions": "1. Could you add some latest methods, e.g. [1], to related works and baselines?  \n2. Could you provide more details about the experiments, e.g. how to choose the M, and how to evaluate the memory footprint?\n3. There is only one long-term layer in the proposed method, which may not be sufficient to model long-term memory. As shown in Figure 8, could you conduct experiments with more long-term memory layers?\n\n[1] Wang, Weizhi, et al. \"Augmenting language models with long-term memory.\" Advances in Neural Information Processing Systems 36 (2024).", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The writing is clear and fluent.\n2. Experimental results are promising, demonstrating an improvement in performance while simultaneously reducing memory usage.", "weaknesses": "1. The baselines are old. Some related methods are not discussed or compared.\n2. Some experimental details are not provided. For example, how to choose the M, and how to evaluate the memory footprint.\n3. For Figure 8, with the number of long-term memory layers increasing, the perplexity is reduced. So a single layer may not be sufficient.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)