MELODI: Exploring Memory Compression for Long Contexts

## Abstract

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

## 1 Introduction

Figure: Figure 1: Overview of Melodi. Melodi employs a hierarchical memory representation, incorporating both short-term and long-term compression mechanisms, integrated with a transformer-based language model. It utilizes a stack of short-term layers to recurrently compress each context window $x_{k}^{0}$ into short-term memory tokens $\{z^{l}_{k}\}$, and inserts a long-term layer to store compressed key-value pairs within a long-term memory $m_{1: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 1 long-term layer and $N-M-1$ short-term layers following it.
Refer to caption: extracted/5900799/figure/intro.png

Long-context language models, exemplified by Gemini (Gemini-Team, ) and GPT (OpenAI, ), 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., ), chain-of-thought prompting (Wei et al., ), and the incorporation of explicit instructions (Chung et al., ; Wei et al., ). However, the quadratic complexity of attention mechanisms within transformer models necessitates significant computational resources to handle long contexts effectively. This has spurred the development of efficient solutions (Dai et al., ; Wu et al., ; Bulatov et al., ) that process long contexts via short context windows, much like how humans process information by reading a book chapter by chapter. A central question underlying these solutions is: how can we 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, ) tackles this by recurrently compressing historical information into hidden states after processing each token. With the rise of Transformer models (Vaswani et al., ) dominating the language modeling landscape, recent memory designs have shifted towards utilizing Transformers to process a context window, thereby moving the focus of memory management from the token level to the context window level.

Transformer-XL (Dai et al., ) employs a caching mechanism to store multi-layer key-value (KV) pairs from the preceding window as memory. Memorizing Transformer (Wu et al., ) builds upon this foundation by incorporating a dedicated layer to memorize all KV pairs from that layer across all prior windows. Meanwhile, Block Recurrent Transformer (Hutchins et al., ) and Recurrent Memory Transformer (Bulatov et al., ) 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 memory 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 a compression-based approach. The short-term memory, recurrent in nature and possessing low capacity, spans multiple network layers, progressively compressing context tokens and prior memory at each layer. In contrast, the long-term memory, incremental and high-capacity, resides within a single network layer. It maintains a record of the entire history by further compressing each context window and stacking them. Both short-term and long-term memory are seamlessly incorporated into a multi-layer transformer model using a “sandwich” structure (see Figure 1), incurring 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 Transformer (10.62 on PG-19, 2.14 on arXiv) with dense attention (as opposed to top-k attention), while significantly reducing memory usage by a factor of 8. Furthermore, ablation studies confirm the complementary nature of short-term and long-term memory in Melodi, highlighting their synergistic contribution to an efficient and effective memory architecture.

## 2 Our Method: Melodi

Melodi focuses on efficiently comprehending long contexts through the utilization of short context windows, thus circumventing the quadratic complexity associated with attention mechanisms over long sequences. This approach necessitates a memory design that not only ensures smooth transitions between windows but also preserves crucial information from all preceding windows.

### 2.1 Architecture Overview

Design principle:
The core principle behind Melodi is to represent short-term and long-term memory through a hierarchical compression scheme. Specifically, the short-term memory recurrently compresses context tokens across multiple network layers (e.g., condensing a 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 a fixed-size multi-layer long short-term memory (LSTM) mechanism (Hochreiter & Schmidhuber, ). Furthermore, each context window undergoes additional compression within a middle layer and is then stacked into a long-term memory. This long-term memory retains essential information from the entire history, 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 comprehension of long contexts even under the constraint of short context windows.

Sandwich architecture:
Melodi’s network architecture (Figure 1) features a ”sandwich” structure, with a long-term compression layer inserted between multiple recurrent short-term compression layers. 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 $l$ to index network layers. Within the $l^{th}$ layer of the $k^{th}$ context window, the input context tokens are represented by $x_{k}^{l-1}$. The output context tokens, denoted as $x_{k}^{l}$, serve as input for the subsequent layer. The input and output of the short-term memory are $z_{k-1}^{l}$ and $z_{k}^{l}$, respectively. The long-term memory preceding window $k$ is denoted as $m_{1:k-1}$. Note that we omit the subscript $l$ for the long-term memory since it resides within a single long-term layer.

Next, we will discuss both short-term and long-term layers in detail.

### 2.2 Short-Term Memory: Multi-Layer Recurrent Compression

Figure: Figure 2: Short-term layer. The figure illustrates the processing of the $k^{th}$ context window at the $l^{th}$ short-term layer. It takes the memory from the previous window $z_{k-1}^{l}$ and the current context/summary ($x_{k}^{l-1}$, $u_{k}^{l-1}$) from the previous layer as input. The short-term layer adds two linear token mixers (Tolstikhin et al., ) on top of a standard transformer layer (including attention and FFN) to separate the summary for the next layer $u_{k}^{l}$ and the memory for the next window $z_{k}^{l}$. Best viewed in color.
Refer to caption: extracted/5900799/figure/short-term-2.png

The short-term memory is distributed across multiple short-term layers (see Figure 1). This subsection delves into the specifics of the short-term layer, using the $l^{th}$ layer as an illustrative example for processing the $k^{th}$ context window.
The short-term layer serves a dual purpose: (a) transforming context tokens $x_{k}^{l-1}$ via a transformer block (yielding output $x_{k}^{l}$), and (b) recurrently compressing the current context window into the short-term memory $z_{k}^{l}$. 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 tokens traverse vertically across layers (from $x_{k}^{l-1}$ to $x_{k}^{l}$), whereas short-term memory flows horizontally across context windows (from $z_{k-1}^{l}$ to $z_{k}^{l}$). To enable inter-layer communication within the short-term memory, we introduce summary tokens $u_{k}^{l}$ that propagate through the layers. We elaborate on the key components below.

Short-term memory $z_{k}^{l}$:
The short-term memory (illustrated in Figure 2) is implemented as a sequence of vectors with length $S$, 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 memory tokens per window of 512 context tokens). Within each context window, the short-term memory 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.

Transforming context tokens:
The context tokens $x_{k}^{l}$ are generated through causal attention to both (a) the preceding short-term memory $z_{k-1}^{l}$ and (b) preceding tokens within the current context window. This attention mechanism is followed by the application of a feed-forward network (FFN). Relative position embeddings are employed for both the context tokens $x_{k}^{l-1}$ and the preceding short-term memory $z_{k-1}^{l}$. Mathematically, this can be represented as:

$$ $\displaystyle x_{k}^{l}=\mathcal{T}(x_{k}^{l-1}|z_{k-1}^{l}),$ (1) $$

where $\mathcal{T}(x|z)$ indicates applying a transformer block on $x$ for a given preceding context $z$.

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., ) and AutoCompressors (Chevalier et al., ), this compression is achieved by appending summary tokens $u$ 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, expressed as: $\hat{u}_{k}^{l}=\mathcal{T}(u_{k}^{l-1}|z_{k-1}^{l},x_{k}^{l-1})$, where the input summary tokens $u_{k}^{l-1}$ originate from the previous layer (refer to Figure 2). We use $\hat{u}_{k}^{l}$ (instead of $u_{k}^{l}$) to denote an intermediate result that is further processed in the subsequent summary branching step. In practice, both context and summary tokens can be updated simultaneously within a single transformer operation: $x_{k}^{l},\hat{u}_{k}^{l}=\mathcal{T}(x_{k}^{l-1},u_{k}^{l-1}|z_{k-1}^{l})$. Relative position embeddings are applied on the short-term memory $z_{k-1}^{l}$, context $x_{k}^{l-1}$, and summary $u_{k}^{l-1}$, while a causal mask is applied on the combined sequence of $x_{k}^{l-1}$ and $u_{k}^{l-1}$.

Summary tokens (containing $U$ tokens) are initialized from learnable embeddings (prior to the first layer) and set to the same length as the short-term memory ($U=S$). Propagating through all layers, they facilitate inter-layer communication within the short-term memory.
Moreover, branching summary tokens both upwards to the next layer and rightwards to the next window improves performance, a strategy we will discuss in more detail subsequently.

Summary branching:
We employ distinct linear token mixers (Tolstikhin et al., ) on the context and summary tokens to generate separate summary tokens for the subsequent layer and short-term memory tokens for the next window. Unlike channel mixing,
a linear token mixer linearly combines the $M_{i}$ input tokens across each channel to produce $M_{o}$ output tokens with the same dimensionality by using an $M_{i}\times M_{o}$ matrix. This is mathematically represented as follows:

$$ $\displaystyle u_{k}^{l}=\mathcal{M}_{\uparrow}(x_{k}^{l},\hat{u}_{k}^{l}),\;\; z_{k}^{l}=\mathcal{M}_{\rightarrow}(x_{k}^{l},\hat{u}_{k}^{l}),\;\;\text{where }\;\;x_{k}^{l},\hat{u}_{k}^{l}=\mathcal{T}(x_{k}^{l-1},u_{k}^{l-1}|z_{k-1}^{l}),$ (2) $$

where $\mathcal{M}_{\uparrow}$ and $\mathcal{M}_{\rightarrow}$ denote the two linear token mixers towards the subsequent layer and the next window respectively. The resulting summary and memory tokens exhibit distinct linear combination of context $x_{k}^{l}$ and summary tokens $\hat{u}_{k}^{l}$, implying divergent compression flows across layers and windows. Since the summary and short-term memory share the same number of tokens ($S$), each mixer
comprises $(W+S)\times S$ parameters, where $W$ represent the number of context tokens per window. This parameter count is negligible for short context windows. For instance, with a context window of 512 tokens and 128 summary tokens, the two mixers collectively require (512+128)$\times$128$\times$2=164K parameters, constituting a mere 1.3% of a transformer block with 1024 dimensions.

Summary:
The short-term memory layer can be succinctly represented as a function $z_{k}^{l},x_{k}^{l},u_{k}^{l}=h(z_{k-1}^{l},x_{k}^{l-1},u_{k}^{l-1})$. This function transforms context tokens $x_{k}^{l-1}$ and summary tokens $u_{k}^{l-1}$ upward to the next layer (from $l-1$ to $l$) while simultaneously propagating short-term memory $z_{k-1}^{l}$ rightward to the next window (from $k-1$ to $k$). Built upon a standard transformer block, this layer introduces negligible additional parameters through the summary branching mechanism.

Relation to Block Recurrent Transformer (BRT) (Hutchins et al., ):
While BRT represents short-term memory by combining multi-layer KV caching from Transformer XL with a dedicated recurrent memory layer, Melodi models short-term memory as a consistent recurrent compression mechanism across multiple layers. Moreover, Melodi updates the short-term tokens by adding residual connections (cross-attending to the previous memory) over the previous layer, instead of over the previous memory tokens directly as in BRT.

While short-term memory facilitates smooth transitions between context windows, the inherent limitation of its capacity can lead to the inevitable loss of information, particularly for contexts located further back in the sequence. In the following subsection, we will demonstrate how long-term memory can be leveraged to mitigate this forgetting issue.

### 2.3 Long-Term Memory: Single-Layer Memorizing Compressed Key-Value Pairs

Figure: Figure 3: Long-term layer. The long-term layer adds three components on top of the short-term layer (see Figure 2). Firstly, it introduces a long-term memory $m_{1:k-1}$ by caching the compressed key-value pairs and allows the current context/summary ($x_{k}^{l-1}$, $u_{k}^{l-1}$) to cross attend to them. Secondly, the self-attention and cross-attention are integrated via gating. Finally, the linear token mixing output additional compressed tokens and appends their key-value pairs $m_{k}$ into the long-term memory (as $m_{1:k}$) for the next window. Best viewed in color.
Refer to caption: extracted/5900799/figure/long-term.png

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 representations across the entire history.

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 a context window into $L$ long-term tokens at a designated middle layer (refer to Figure 1), where $L$ is less than the number of short-term tokens per layer ($L<S$). Illustratively, we might compress a context window of 512 tokens into $S=128$ short-term tokens at each layer, but further condense it into $L=64$ long-term tokens at a single layer.

Storing long-term memory:
The key-value (KV) pairs of long-term tokens are sequentially stored in a first-in-first-out (FIFO) queue with a maximum capacity of $Q_{max}$ windows. Consequently, the long-term memory can hold up to $L\times Q_{max}$ KV pairs. For contexts shorter than $Q_{max}$ windows, a compressed representation of the entire prior history is preserved. For longer documents exceeding $Q_{max}$ windows, a substantial portion of the recent history ($Q_{max}$ 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, a point we will elaborate on later.

Long-term layer:
Figure 3 illustrates the implementation of a long-term layer, which builds upon a short-term layer but incorporates three key additions. First, it introduces a long-term memory component (denoted as $m_{1:k-1}$ prior to the $k^{th}$ context window) and enables the current context tokens $x_{k}$ and summary tokens $u_{k}$ to cross-attend to it. Second, the cross-attention shares parameters with the self-attention and their results (cross attention: $\mathcal{A}_{x}$, self-attention: $\mathcal{A}_{s}$) are combined through a gating mechanism using a learnable scalar $\alpha$ per attention head, formulated as $\alpha\mathcal{A}_{x}+(1-\alpha)\mathcal{A}_{s}$. Finally, an additional linear token mixer is introduced to generate long-term tokens for the current window, and their key-value (KV) pairs $m_{k}$ are appended to the long-term memory. This token mixer comprises $(W+U)\times L$ parameters, where $W$, $U$ and $L$ 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 2 provides a comparative overview of the long-term and short-term memory mechanisms. In contrast to the short-term memory, which operates recurrently across multiple layers, the long-term memory functions incrementally within a single layer. The short-term memory spans $N$ network layers, with each layer utilizing $S$ tokens, while the long-term memory queue spans $Q_{max}$ context windows, with each window represented by $L$ tokens.
To illustrate, consider a 12-layer transformer model processing context windows of 512 tokens each, with embedding dimension 1024. If we employ $S=128$ short-term tokens per layer and $L=64$ long-term tokens per window, with a maximum capacity of $Q_{max}=128$ windows for the long-term memory, the short-term and long-term memory caches would have capacities of 1.6M (128${}^{tokens}\times$1024 ${}^{dim}\times$12 ^layers) and 16.8M (64${}^{pairs}\times$2${}^{tokens/pair}\times$1024${}^{dim}\times$128^windows) floats, respectively. It’s worth noting that, for computational 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., ):
Similar to Memorizing Transformer, Melodi incorporates key-value pairs from a middle layer into its long-term memory. Experiments 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, as 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, Melodi achieves an 8-fold reduction in long-term memory size.

## 3 Experimental Results

We evaluate Melodi on three long-context datasets: PG19 (Rae et al., ), arXiv Math (Wu et al., ), and C4 (Raffel et al., ), using the standard auto-regressive language modeling task, where the objective is to predict the next token in a sequence. All models are trained from scratch, and we report the average perplexity on the respective test sets as our evaluation metric.

### 3.1 Datasets

PG19:
The PG19 dataset (Rae et al., ) consists of 28,752 English books published before 1919, averaging around 68,972 tokens per book. We utilize three 32k vocabularies: (a) Meena (Thoppilan et al., ), (b) T5 (Raffel et al., ) and (c) a custom SentencePiece vocabulary (Kudo & Richardson, ) trained specifically on PG19.

arXiv Math:
The arXiv dataset (Wu et al., ) comprises technical math papers from arXiv, with token counts comparable to PG19 due to LaTeX’s use of special characters, resulting in smaller subwords. We use a 32k Meena vocabulary (Thoppilan et al., ) and a 32k custom vocabulary.

**Table 1: Long-term vs. short-term memory. In our notation, $L$ and $S$ represent the number of long-term and short-term tokens per network layer, respectively. Note that long-term has fewer tokens ($L<S$). $Q_{max}$ denotes the maximum number of windows encompassed by the long-term memory queue. $N$ signifies the total number of network layers.**
| PROPERTY | LONG | SHORT |
| --- | --- | --- |
| Number of layers | single | multiple |
| Update per window | incremental | recurrent |
| Capacity | $L\times Q_{max}$ | $S\times N$ |

C4(4K+):
The C4 dataset (Raffel et al., ) is a large collection of 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.

### 3.2 Setup

We utilized a decoder-only transformer architecture (with 12 or 13 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 implemented in JAX and Flax and trained from scratch for 500k steps on 32 TPU cores. Further training details are provided in the Appendix B.

During training, each long document was segmented into 4096-token chunks to facilitate batch processing. These chunks were then organized into training batches, each comprising 8 context windows of 512 tokens. In our ablation study, Melodi’s default configuration compressed each context window into $S$=128 short-term memory tokens per layer and $L$=64 long-term memory tokens.

### 3.3 Comparison with Baselines

Baselines:
We benchmark Melodi against three well-established prior works: Transformer XL (Dai et al., ), Block Recurrent Transformer (Hutchins et al., ), and Memorizing Transformer (Wu et al., ). 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 models 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 learning rate schedule (Hoffmann et al., ) instead of the inverse square root decay, (b) using dense cross-attention instead of top-k attention for the Memorizing Transformer. Our re-implementations provide stronger baselines against which to evaluate Melodi’s effectiveness.

**Table 3: Comparisons with baselines on 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: $S_{192}+L_{32}$, $S_{128}+L_{64}$, and $S_{192}+L_{96}$. For instance, $S_{192}+L_{32}$ indicates $S=192$ short-term and $L=32$ long-term tokens per context window. All Melodi models utilized a long-term memory spanning 128 context windows. Melodi $S_{192}+L_{32}$ outperformed Transformer XL and Block Recurrent Transformer while consuming less memory. Notably, Melodi $S_{192}+L_{96}$ clearly surpassed Memorizing Transformer, using only a fifth of its memory.**
| MODEL | MEMORY | PG19 | arXiv | C4(4K+) |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | All | Short | Long | Meena | T5 | Custom | Meena | Custom | Custom |
| Transformer XL | 13.6M | 13.6M | 0M | 8.65 | 11.41 | 12.42 | 2.60 | 3.22 | 18.22 |
| Block Recurrent | 13.1M | 13.1M | 0M | 8.30 | 10.98 | 11.90 | 2.26 | 2.70 | 17.82 |
| Melodi $S_{192}$+$L_{32}$ | 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 $S_{128}$+$L_{64}$ | 18.5M | 1.7M | 16.8M | 8.06 | 10.44 | 11.42 | 2.11 | 2.52 | 17.53 |
| Melodi $S_{192}$+$L_{96}$ | 27.8M | 2.6M | 25.2M | 7.91 | 10.29 | 11.27 | 2.09 | 2.49 | 17.25 |

Comparisons:
Table 3 presents a comparison of Melodi against three baseline models (Transformer XL, Block Recurrent Transformer, and Memorizing Transformer) across three datasets. We evaluate three Melodi configurations: $S_{192}+L_{32}$, $S_{128}+L_{64}$, and $S_{192}+L_{96}$, where $S$ and $L$ denote the number of short-term and long-term tokens per context window, respectively. All models (Melodi and baselines) utilize a 13-layer transformer architecture, except for Block Recurrent Transformer, which inserts a block recurrent layer into a 12-layer transformer, ensuring a similar parameter count for all models.

Melodi $S_{192}$+$L_{32}$ consistently outperforms both Transformer XL and Block Recurrent Transformer across all datasets while utilizing less memory. For instance, it achieves a perplexity of 10.51 on PG-19 (T5 vocabulary), surpassing Transformer XL (11.41) and Block Recurrent Transformer (10.98) while consuming only 11.0M memory compared to their 13.6M and 13.1M, respectively.

In comparison to the Memorizing Transformer, Melodi $S_{128}+L_{64}$ exhibits slightly improved performance while significantly reducing memory consumption by approximately 8 times. Melodi $S_{192}+L_{96}$ further improves perplexity across all datasets, achieving a substantial reduction in memory usage exceeding a factor of 5.
These trends remain consistent across different network depths (12-layer and 13-layer), as shown in Table 5 in Appendix C. These results collectively highlight Melodi’s efficacy and efficiency as a memory architecture for language models.

### 3.4 Ablations

We conduct an ablation study of Melodi using the default configuration ($S_{128}$+$L_{64}$) 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.

Figure: Figure 4: Ablation of memory size on PG-19. The token perplexity is reported for various combinations of short-term and long-term memory sizes. Each curve represents a fixed size of long-term memory, with points along the curve indicating different short-term memory sizes. For example, the blue curve ($L_{96}$), uses 96 long-term tokens per window over 128 windows, totaling 12,288 tokens. Each point on this curve represents a different short-term memory size (e.g. $S_{8}$ denotes 8 short-term tokens per context window). Memory size is measured by the number of floating-point numbers (floats). For instance, $L_{96}$ stores 12,288 long-term key-value (KV) pairs, each with 1024 dimensions, resulting in a total of 12,288$\times$1024$\times$2=25.2 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 to baselines like Transformer XL, Block Recurrrent Transformer and Memorizing Transformer while utilizing fewer memory resources. Best viewed in color.
Refer to caption: extracted/5900799/figure/memory_size.png

Complementary roles of short-term and long-term memory:
Figure 4 illustrates how the sizes of both short-term and long-term memory jointly influence perplexity. Each curve represents a 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 increasing either short-term or long-term memory capacity leads to improved perplexity, highlighting their complementary roles in performance. Notably, by judiciously selecting memory sizes (e.g., $S_{192}$ for short-term and $L_{32}$ for long-term), we can outperform the Memorizing Transformer while utilizing less memory than Transformer XL and Block Recurrent Transformer.

Table: Figure 5: Long-term memory coverage. The coverage metric indicates the number of preceding context windows spanned by the long-term memory. All data points use $L$=64 long-term and $S$=128 short-term tokens per window. However, they vary in long-term memory capacity. For instance, ‘32’ denotes covering 32 context windows, totaling 32$\times$64=2048 long-term tokens. Perplexity improves marginally with long-term memory coverage of 2-4 windows, then accelerates until 32 windows, after which it levels off.

Impact of long-term memory coverage:
In contrast to the previous ablation, we now maintain a fixed number of short-term ($S=128$) and long-term tokens ($L=64$) per context window while varying the number of windows encompassed by the long-term memory. Figure 5 demonstrates that performance improves as the long-term memory covers a wider range of context windows.

Interestingly, incorporating long-term memory for only the preceding 2 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 encompass 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-term memory mechanisms.

Performance with shorter context windows: Figure 6 examines the impact of reducing context window size on model performance. The plot on the left displays perplexity scores for context window sizes of 512, 256, and 128 tokens. The number of short-term and long-term memory tokens in Melodi is proportionally adjusted ($\times\frac{1}{2}$, $\times\frac{1}{4}$, respectively) to ensure consistent long-term memory coverage across different window sizes. Even with shorter context windows, Melodi with appropriate memory allocation (e.g., $S_{128}+L_{64}$) 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-term layers on model performance. For a given layer count (e.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 1 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 a more efficient approach with negligible performance degradation.

Number of long-term layers:
Figure 8 illustrates the effect of varying the number of long-term layers on perplexity. We consider two configurations: $L_{64}$ and $L_{32}$, which compress each context window into 64 and 32 long-term tokens, respectively.
With $L_{64}$, introducing the first long-term layer yields a substantial perplexity improvement, after which performance plateaus. In contrast, with $L_{32}$, while adding a second long-term layer provides a notable gain, the resulting performance remains inferior to that of a single $L_{64}$ layer, even though both configurations utilize the same total number of long-term tokens. These findings validate our design choice of employing a single long-term memory layer with a sufficient number of long-term tokens.

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.

**Table 4: Summary branching. Summary branching provides a consistent gain of approximately 0.3 in perplexity, both with (column ST+LT) and without (column ST) long-term memory.**
| Branching | ST | ST+LT |
| --- | --- | --- |
| No | 11.68 | 11.24 |
| Yes | 11.39 | 10.95 |

## 4 Related Work

Memory in language models:
While LSTMs (Hochreiter & Schmidhuber, ) achieve long-range memory through recurrent compression at the token level, the advent of Transformers (Vaswani et al., ) has shifted the focus to memory mechanisms operating at the context window level. Transformer-XL (Dai et al., ) introduces a caching mechanism to store key-value (KV) pairs from the preceding context window as a form of short-term memory. Block Recurrent Transformer (Hutchins et al., ) and RMT (Bulatov et al., ), integrate recurrent mechanisms inspired by LSTMs into Transformer architectures at the window level. Munkhdalai et al. () explore the use of additional memory like Hopfield Networks (Hopfield, ; Ramsauer et al., ). Memorizing Transformer (Wu et al., ) introduces a dedicated memory layer to store KV pairs for long-term memory, whereas MemoryLLM (Wang et al., ) incorporates long-term memory in every layer, incurring a substantial memory overhead. In contrast to these approaches, Melodi integrates both short-term and long-term memory into a transformer model via compression.

Extending context length:
Recent research has demonstrated promising progress in scaling the context length of language models. To mitigate the cost of attention mechanisms over long contexts, LongLoRA (Chen et al., ), PCW (Ratner et al., ), and Landmark-Attention (Mohtashami & Jaggi, ) employ sparse local attention for efficient fine-tuning. Position-Interpolation (Chen et al., ) and YaRN (Peng et al., ) extend context size by modifying the RoPE embedding scheme (Su et al., ). Xiong et al. () provide a practical recipe for extending LLAMA2 (Touvron et al., ) to handle up to 32,768 tokens, while Fu et al. () further explore data-centric approaches for extending context length through lightweight continual pretraining.

Compression:
Recent work has demonstrated the effectiveness of Transformer models in compressing input sequences by appending summary tokens (Rae et al., ; Bulatov et al., ; Chevalier et al., ; Ge et al., ). For instance, RMT (Bulatov et al., ) utilizes the output of summary tokens recurrently as short-term memory. AutoCompressor (Chevalier et al., ) aggregates summary tokens across segments to generate a summary representation for long documents used in retrieval tasks. GISTING (Mu et al., ) applies this technique to compress long prompts. ICAE (Ge et al., ) further incorporates LoRA fine-tuning (Hu et al., ) for context compression, while TransformerFAM (Hwang et al., ) introduces feedback attention to enhance performance.

## 5 Conclusion

In this work, we have introduced Melodi, a 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 facilitates smooth transitions between context windows through recurrent compression across multiple layers, and long-term memory preserves crucial 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 as 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.

## Appendix A Limitations

A primary limitation of our current method is its focus on training from scratch, without addressing the fine-tuning of pre-trained models with fixed context window sizes. In future work, we plan to explore adapting Melodi (short-term and long-term memory) to enhance the long-context capabilities of pre-trained models through techniques like LoRA fine-tuning.

## Appendix B Training Details

We use Adafactor optimizer (Shazeer & Stern, ) with a 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. (). A dropout rate of 0.05 is applied. All models are trained for 500k steps (200k for ablations) on 32 TPU cores with a batch size of 32 (1 per core).

**Table 5: Comparisons with baselines on 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: $S_{192}+L_{32}$, $S_{128}+L_{64}$, and $S_{192}+L_{96}$. For instance, $S_{192}+L_{32}$ indicates $S=192$ short-term and $L=32$ long-term tokens per context window. All Melodi models utilized a long-term memory spanning 128 context windows. Melodi $S_{192}+L_{32}$ outperformed Transformer XL and Block Recurrent Transformer while consuming less memory. Notably, Melodi $S_{192}+L_{96}$ clearly surpassed Memorizing Transformer, using only a fifth of its memory.**
| MODEL | MEMORY | PG19 | arXiv | C4(4K+) |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | All | Short | Long | Meena | T5 | Custom | Meena | Custom | Custom |
| 12 LAYERS |  |  |  |  |  |  |  |  |  |
| Transformer XL | 12.6M | 12.6M | 0M | 8.76 | 11.54 | 12.63 | 2.61 | 3.23 | 18.61 |
| Block Recurrent | 12.1M | 12.1M | 0M | 8.47 | 11.12 | 12.11 | 2.27 | 2.73 | 18.27 |
| Melodi $S_{192}$+$L_{32}$ | 10.8M | 2.4M | 8.4M | 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 $S_{128}$+$L_{64}$ | 18.4M | 1.6M | 16.8M | 8.16 | 10.61 | 11.66 | 2.13 | 2.55 | 18.01 |
| Melodi $S_{192}$+$L_{96}$ | 27.6M | 2.4M | 25.2M | 8.08 | 10.48 | 11.47 | 2.11 | 2.51 | 17.75 |
| 13 LAYERS |  |  |  |  |  |  |  |  |  |
| Transformer XL | 13.6M | 13.6M | 0M | 8.65 | 11.41 | 12.42 | 2.60 | 3.22 | 18.22 |
| Block Recurrent | 13.1M | 13.1M | 0M | 8.30 | 10.98 | 11.90 | 2.26 | 2.70 | 17.82 |
| Melodi $S_{192}$+$L_{32}$ | 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 $S_{128}$+$L_{64}$ | 18.5M | 1.7M | 16.8M | 8.06 | 10.44 | 11.42 | 2.11 | 2.52 | 17.53 |
| Melodi $S_{192}$+$L_{96}$ | 27.8M | 2.6M | 25.2M | 7.91 | 10.29 | 11.27 | 2.09 | 2.49 | 17.25 |

## Appendix C More Comparison with Baselines

Comparison with baselines using a fixed context window size:
Table 5 presents a comparison of Melodi against three baseline models (Transformer XL, Block Recurrent Transformer, and Memorizing Transformer) across three datasets, using a consistent segment 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 $S_{192}$+$L_{32}$ consistently outperforms both Transformer XL and Block Recurrent Transformer across all datasets while consuming less memory. For instance, the 12-layer variant of Melodi $S_{192}$+$L_{32}$ achieves a perplexity of 10.66 on PG-19 (T5 vocabulary), surpassing the 13-layer variants of Transformer XL (11.41) and Block Recurrent Transformer (10.98).

In comparison to the Memorizing Transformer, Melodi $S_{128}+L_{64}$ exhibits slightly improved performance while dramatically reducing memory consumption by a factor of 8. Furthermore, Melodi $S_{192}+L_{96}$ achieves even better perplexity scores across all datasets with a substantial reduction in memory usage exceeding a factor of 5. The improvement is consistent for using both 12-layer and 13-layer transformer architectures. These results collectively highlight Melodi’s efficacy and efficiency as a memory architecture for language models.

Comparison with baselines that utilize longer context windows:
Table 6 demonstrates that on PG-19 (T5 vocabulary), Melodi (with a context window size of 512 tokens) outperforms three baselines even when they utilize longer context windows (1024 and 2048 tokens), further highlighting its effectiveness and efficiency in processing long sequences. While the baselines generally exhibit improved performance with increasing context window size, Memorizing Transformer shows instability with a 2048-token window. We hypothesize that this instability stems from the limited backpropagation through time (BPTT) depth (two windows) within each segment (4096 tokens).

**Table 6: Comparisons with baselines that utilize longer context windows. Melodi, with a context window size of 512, outperforms three baselines (i.e., Transformer XL, Block Recurrent Transformer and Memorizing Transformer) even when they utilize longer context windows (1024 and 2048 tokens). This highlights Melodi’s effectiveness in processing long sequences despite using shorter context windows.**
| MODEL | Segment | Window | MEMORY | PG19 |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | Length | Length | All | Short | Long | T5 |
| Transformer XL | 4096 | 512 | 13.6M | 13.6M | 0M | 11.41 |
| Transformer XL | 4096 | 1024 | 27.3M | 27.3M | 0M | 11.11 |
| Transformer XL | 4096 | 2048 | 54.5M | 54.5M | 0M | 10.99 |
| Transformer XL | 6144 | 3072 | 81.8M | 81.8M | 0M | 10.96 |
| Block Recurrent | 4096 | 512 | 13.1M | 13.1M | 0M | 10.98 |
| Block Recurrent | 4096 | 1024 | 25.7M | 25.7M | 0M | 10.91 |
| Block Recurrent | 4096 | 2048 | 50.9M | 50.9M | 0M | 10.88 |
| Memorizing Trans. | 4096 | 512 | 147.8M | 13.6M | 134.2M | 10.62 |
| Memorizing Trans. | 4096 | 1024 | 161.5M | 27.3M | 134.2M | 10.53 |
| Memorizing Trans. | 4096 | 2048 | 188.7M | 54.5M | 134.2M | not stable |
| Melodi $S_{192}$+$L_{32}$ | 4096 | 512 | 11.0M | 2.6M | 8.4M | 10.51 |
| Melodi $S_{128}$+$L_{64}$ | 4096 | 512 | 18.5M | 1.7M | 16.8M | 10.44 |
| Melodi $S_{192}$+$L_{96}$ | 4096 | 512 | 27.8M | 2.6M | 25.2M | 10.29 |

**Table 7: Directly copying short-term tokens to long-term memory: Here, we force the short-term and long-term to share the same number of tokens (e.g. $S_{64}$+$L_{64}$), and examine the impact of directly using short-term tokens as long-term tokens, bypassing the linear token mixer. This approach results in performance degradation compared to generating distinct long-term tokens.**
| Copying | $S_{96}$+$L_{96}$ | $S_{64}$+$L_{64}$ | $S_{32}$+$L_{32}$ |
| --- | --- | --- | --- |
| Yes | 11.00 | 11.36 | 11.42 |
| No | 10.92 | 11.08 | 11.25 |

## Appendix D 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 8 indicate that this direct copying method leads to a performance degradation.

Effect of long-term layer position:
We investigate the impact of varying the position of the long-term layer within a 13-layer Melodi model ($S_{128}$+$L_{64}$). Layers are indexed from 0 to 12, with the default long-term layer position at layer 8. Results in Table 8 reveal consistent perplexity scores when the long-term layer is positioned between layers 5 and 11.

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. A preliminary session about the history of recent advances about memory compression would be much appreciated.
soundness: 2
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.
2. The empirical results seem promising.
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.
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.
2. The results only show the efficiency part and do not show the effect of memory network on the performance of LLM.
3. Codes are not shared for checking the implementation.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: Suggestion: Could you provide some more experimental results in the 2k and larger token context regime?
Questions: 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? 
Is 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?

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

The 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.
soundness: 3
strengths: Originality
This 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.

Quality
The 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.

Clarity
The 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. 

Significance
The 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.
The 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.
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.

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

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

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

Compared 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.
weaknesses: (1) The paper has only a limited exploration of longer context window size
The 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.

(2) There appears to be a lack of fine-tuning experiments
MELODI 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.

(3) Limited Analysis of Computational Efficiency
More 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

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

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
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.
* 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.
soundness: 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.
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.
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).
* 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?
* 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?
* 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.

---

confidence: 2
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. Could you add some latest methods, e.g. [1], to related works and baselines?  
2. Could you provide more details about the experiments, e.g. how to choose the M, and how to evaluate the memory footprint?
3. 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?

[1] Wang, Weizhi, et al. "Augmenting language models with long-term memory." Advances in Neural Information Processing Systems 36 (2024).
soundness: 3
strengths: 1. The writing is clear and fluent.
2. Experimental results are promising, demonstrating an improvement in performance while simultaneously reducing memory usage.
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.
weaknesses: 1. The baselines are old. Some related methods are not discussed or compared.
2. Some experimental details are not provided. For example, how to choose the M, and how to evaluate the memory footprint.
3. For Figure 8, with the number of long-term memory layers increasing, the perplexity is reduced. So a single layer may not be sufficient.

Meta-review: This paper proposes a novel and slightly intricate architecture for modelling long context. It is based on two types of layers, short term and long term. The short term modules treat a short span of context and introduces something akin to a recurrence with summary tokens. The long term module cross attends to memory tokens from previous short term modules. The proposed architecture is mostly evaluated on language modeling tasks. The reviewers have raised some practical concerns, and pointed out some lacking baselines (the ones in the paper were old), slightly limited scale of experiments, as well as lack of specialized evaluations (long-context specific, or needle in a haystack). Nonetheless, given the interesting methodological contributions in this work, I recommend accepting this paper for ICLR 2025.

Decision: Accept (Poster)