MELODI: Exploring Memory Compression for Long Contexts

$\color{blue}{**[Weakness 4]:**}$

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.

Thank you for the feedback. We acknowledge that the models evaluated in the main paper (12-layer Transformers with an embedding dimension of 1024, approximately 150M parameters) were relatively small. To address concerns about scalability, we conducted additional experiments on the PG-19 dataset (using the T5 vocabulary) with larger architectures. We explored two scaling approaches:

1. Increased Depth: We trained models with 24 and 36 layers, effectively scaling the number of parameters by 2 and 3 times, respectively.

2. Increased Width and Depth: We trained a model with 16 layers and an embedding dimension of 1536, which also scales the parameter count by 3 times.

The results, presented in the table below, demonstrate that MELODI's advantages hold even with larger models. It consistently outperforms the Memorizing Transformer while using significantly less memory (approximately 8 times less).

---

$\color{blue}{**[Question 1]:**}$

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.

Thank you for the suggestion. Due to limited computational resources, we are unable to train a conventional Transformer model with a context window of 256K tokens or more. The maximum context window length we can achieve for a 12-layer transformer with an embedding size of 1024 is 4K tokens.

However, we do provide a comparison against Memorizing Transformer, which does have a long context of 64k, and uses full dense attention. Memorizing transformer only has one long-context layer, but prior work has demonstrated that there is limited benefit to adding multiple long-context layers. At this scale, a single long-range layer seems to be sufficient. (Our experiments confirm this finding, see response to reviewer i9z9 below.) Going beyond 64k on PG19 also has a very limited benefit, since the average book length is only 69k. Thus, the Memorizing Transformer baseline is a reasonable approximation of a conventional long-context model.

To provide a further reference point for longer context windows using a conventional architecture, in which every layer has the same context length, we compare MELODI (with a context window of 512 tokens) to Transformer XL with context windows ranging from 512 to 4K tokens. 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.

As shown in the table below, Transformer XL's performance improves with increasing context length. Notably, MELODI (512 tokens) outperforms Transformer XL even with an 8 times longer context window (4K tokens).

Slower Perplexity Increase for MELODI with Increasing Mask Ratio

Interestingly, even with a high masking ratio of 0.5, the 12-layer MELODI variant (S192+L96) achieves a perplexity of 52.59, which is remarkably close to the 36-layer TransformerXL's perplexity of 51.23. This suggests that MELODI can achieve comparable performance with significantly fewer parameters.

To further analyze this trend, we examined the change in perplexity as the masking ratio increases: 0 → 0.125, 0.125 → 0.25, and 0.25 → 0.5. As shown in the table below, the perplexity increase for MELODI is consistently 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 dependencies.

---

$\color{blue}{**[Weakness 2]:**}$

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?

Thank you for the suggestion. We will add the following details about efficient parallel training, addressing the challenges posed by the model's recurrent nature.

Batch Processing Strategy

To enable efficient parallel training despite the recurrent dependencies in the short-term and long-term memory, we employ a combined layer-wise and chunk-wise processing strategy.

1. Layer-wise Processing: As is standard practice, MELODI uses a Transformer architecture with multiple layers, and the layers are processed in sequence, with the output of each layer becoming the input to the next. Both short-term and long-term memory operations are confined within their respective layers.

2. Unrolling Multiple Windows: Within a layer, each training batch covers 8 context windows (512 tokens per window) with a total of 4096 tokens, allowing for backpropagation through time (BPTT) across multiple windows. Although the windows must be processed sequentially due to window (i.e. block) level recurrence, self-attention within each window can be done in parallel. We achieve good device utilization so long as the window size is large enough.

3. Parallel and Sequential Operations: The computation of query, key, and value matrices (of context and summary tokens) for attention are computed in parallel across all chunks, as they do not depend on the recurrent memory updates. Therefore, each layer 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 strategy differs from a chunk-wise then layer-wise approach, where each chunk must be processed by all layers for the entire network before the next chunk can be processed.

---

$\color{blue}{**[Weakness 3]:**}$

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?

Excellent point! MELODI's chunk-wise approach operates in two steps during generation:

• Token-wise Generation: Within each chunk, MELODI generates tokens sequentially, similar to GPT-like LLMs, utilizing both short-term and long-term memory.

• Summary Generation: After completing a chunk, MELODI generates all summary tokens in parallel, and updates short-term and long-term memory accordingly.

This ensures that token-wise generation within a chunk remains unaffected by the chunk boundaries, preserving the key mechanism of GPT-like LLMs.

We sincerely thank the reviewer for the thoughtful feedback and valuable suggestions, which have significantly improved the quality of our paper.

$\color{blue}{**[Weakness 1]:**}$

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

We acknowledge the limitation of not evaluating MELODI on benchmarks such as LongBench or Needle-in-a-Haystack. This will involve scaling up the model and incorporating MELODI into a pre-train LLMs using fine-tuning or LoRA, which we plan to explore in future work. (Please see our response to other reviewers.)

Although a comprehensive evaluation on downstream tasks is impractical within the rebuttal period, we introduce a novel task designed to further evaluate MELODI's capacity for handling long-range dependencies.

New Task: Masked Next Token Prediciton

Similar to our primary evaluation, we use next token prediction and measure performance with perplexity. However, we introduce a key challenge: masking out a portion of the input tokens. Unlike BERT's bidirectional masking approach, this new task employs a unidirectional, left-to-right prediction scheme.

Why evaluate on this task?

This masking strategy presents a 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 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 experiments 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 systematically vary the degree of contextual disruption.

As expected, perplexity increased across all models as the masking ratio increased, indicating the growing difficulty of the task. However, MELODI consistently outperformed both TransformerXL and the Memorizing Transformer across all masking ratios and architectural configurations. These results, presented in the table below, strongly suggest that MELODI exhibits superior capabilities in capturing and utilizing long-range dependencies compared to these baseline models.

Slower Perplexity Increase for MELODI with Increasing Mask Ratio

Interestingly, even with a high masking ratio of 0.5, the 12-layer MELODI variant (S192+L96) achieves a perplexity of 52.59, which is remarkably close to the 36-layer TransformerXL's perplexity of 51.23. This suggests that MELODI can achieve comparable performance with significantly fewer parameters.

To further analyze this trend, we examined the change in perplexity as the masking ratio increases: 0 → 0.125, 0.125 → 0.25, and 0.25 → 0.5. As shown in the table below, the perplexity increase for MELODI is consistently 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 dependencies.

$\color{blue}{**[Weakness 1.2]:**}$

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.

Figure 5 provides insights into MELODI's ability to utilize information from a much larger effective context length than size of the context window. In this experiment, the context window size is fixed at 512 tokens, and the x-axis in Figure 5 represents the number of previous context windows covered by the long-term memory. Each window is compressed to 64 long-term tokens, e.g. an 8x compression ratio. For instance, the point labeled '192' indicates that the long-term memory stores information from the previous 192 context windows, so the effective context length is 192 windows × 512 tokens/window = 96K tokens. The results in Figure 5 show a clear performance improvement as the long-term memory coverage increases, especially from 4 to 32 previous windows.

---

$\color{blue}{**[Weakness 2]:**}$

lack of fine-tuning experiments MELODI and the baselines. Evaluating the performance of fine-tuned MELODI models on downstream tasks would provide insights into its effectiveness in practical applications.

Thank you for the insightful suggestion. We totally agree that it is important to test MELODI on downstream tasks. However, testing on downstream tasks will require larger models and more extensive pre-training than we currently have been able to do, so we are leaving those studies for future work.

The goal of this paper was simply to determine whether the MELODI memory mechanism works, by using small models trained from scratch, and comparing those models against strong baselines which were also trained from scratch.

In future work, we plan to explore techniques like LoRA (Low-Rank Adaptation) and/or fine-tuning to efficiently integrate MELODI's short-term and long-term memory mechanisms into pre-trained models. This will allow us to assess MELODI's effectiveness on various downstream tasks, such as summarization and question answering, and compare its performance against SOTA baselines.

We sincerely thank the reviewer for the insightful comments and detailed explanations, which have significantly helped us improve the quality of this paper.

$\color{blue}{**[Weakness 1.1 + Question 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. ... ... Could you provide more experimental results in the 2k and larger token context regime? 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? 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?

Experiments with Longer Context Window

Thank you for the suggestion. To explore the performance with larger context windows, we conducted additional experiments with batch sequence lengths of 4K and 8K tokens and window lengths up to 2K tokens. All models were trained on PG-19 (T5 vocabulary) with 12-layer architectures.

As shown in the table below, MELODI continues to outperform the Memorizing Transformer with significantly less memory (approximately 8 times), even with larger context windows. Due to limited computational resources, we were unable to train MELODI with context windows exceeding 2K tokens.

Challenges for Scaling MELODI Up

Scaling MELODI to larger context windows presents two primary challenges:

1. Memory Capacity: The memory requirements of MELODI grow with both the context window size and the batch sequence length. With our current hardware and codebase, we are unable to train MELODI with context windows exceeding 2K tokens. But this can be addressed by using TPUs with more memory and improving model parallelism in our codebase. We plan to investigate it in the future work.

2. Backpropagation Through Time (BPTT): MELODI's recurrent connections necessitate BPTT, which further increases memory consumption. Each batch spans 8 context windows by default, allowing for BPTT across multiple windows. While crucial for the recurrent short-term memory, BPTT contributes to memory usage. As shown in the table above, using longer batch sequences (8K tokens) with 1K token windows yields better performance than shorter sequences (4K tokens) due to the increased BPTT depth.

We sincerely thank the reviewer for the thoughtful feedback and valuable suggestions, which have significantly improved the quality of our paper.

$\color{blue}{**[Weakness 1 + Question 1]:**}$

The baselines are old. Some related methods are not discussed or compared. Could you add some latest methods, e.g. [1], to related works and baselines?

[1] Wang, Weizhi, et al. "Augmenting language models with long-term memory." Advances in Neural Information Processing Systems 36 (2024).

Thank you for the suggestion. We agree that comparing MELODI with LONGMEM [1] would provide valuable insights. We will add LONGMEM to the related work section and discussed its relationship to MELODI.

While both LONGMEM and MELODI improve upon the Memorizing Transformer, they do so with different focuses and approaches:

• LONGMEM: LONGMEM focuses on improving prediction accuracy in a fine-tuning setting by introducing a SideNet 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-scratch setting by compressing both long-term and short-term memory. It introduces a negligible number of extra parameters and is compared against the Memorizing Transformer variant that uses all KV pairs in memory.

A direct experimental comparison with LONGMEM in the fine-tuning (of LLMs) setting is currently out of the scope of this paper, as our primary focus is on evaluating MELODI's performance in the train-from-scratch setting. However, we plan to explore extending MELODI to the fine-tuning setting in our future work, where we will include a comprehensive comparison with LONGMEM and other relevant baselines. This will provide a more complete picture of MELODI's capabilities and its relative strengths compared to other memory-augmented language models.

---

$\color{blue}{**[Weakness 2 + Question 2]:**}$

more details about the experiments, e.g. how to choose the $M$, and how to evaluate the memory footprint?

Hyperparameter $M$:

$M$ in Figure 1 denotes the index of the long-term layer within the transformer architecture. For a 13-layer network, the default value is $M=8$. Ablation in Appendix D (Table 8) demonstrates that model performance is robust to variations in $M$. For convenience, we attach the table below.

Memory Footprint:

Memory size is measured as the total number of floating-point numbers (floats) used to store model activations. For example, MELODI S128+L64 utilizes 128 short-term memory tokens per layer. In a 13-layer transformer with a token dimension of 1024, this results in 128 × 1024 × 13 = 1.7M floats for short-term memory. Additionally, the long-term memory stores 64 key-value pairs (128 tokens) per context window. To cover 128 context windows (64k context tokens if each window has 512 tokens), the long-term memory requires 64 × 2 × 1024 × 128 = 16.8M floats.

---

$\color{blue}{**[Weakness 3 + Question 3]:**}$

Only one long-term layer in the proposed method may not be sufficient to model long-term memory. Could you conduct experiments with more long-term memory layers in Figure 8?

Excellent point. We conducted additional ablations to extend the number of long-term layers to 6 and report the results in the Table below. The perplexity plateaus when using more than 3 long-term layers. Compared to the maximum improvement of 0.60 (from 11.39 with no long-term layer to 10.79 with 6 layers), the first long-term layer alone contributes 73% (0.44/0.60) of this gain. Since additional long-term layers increase both memory footprint and computational cost due to cross-attention, a single long-term layer offers a good balance between performance and efficiency.

Thanks for the follow-up question. We agree that it's important to compare the memory footprint of the memory tokens to the overall model size, especially when considering larger models. To facilitate this analysis, we first define the following notation:

$N$: Number of layers in the model

$W$: Number of context tokens per window

$D$: Dimension of the token embeddings

$S$: Number of short-term memory tokens per window

$Q$: Number of windows covered by long-term memory in Memorizing Transformer

$R$: MELODI's Long-term memory redunction rate compared to Memorizing Transformer

The table below compares the short-term and long-term memory size of MELODI and the baseline models.

For example, the MELODI S128+L64 variant uses a quarter of the window length for short-term memory ($S=W/4$) and achieves a long-term memory reduction rate of $R=8$. This results in an 8x reduction in memory usage compared to the Memorizing Transformer.

To further illustrate the relationship between memory footprint and model size, we present a comparison for three different model configurations in the table below. The first two are the smaller Transformer models, while the third is the Llama-2 7B model with 32 layers, an embedding dimension of 4096, and a context window of 4096. For this comparison, we set the long-term memory to cover 128 context windows.

As you can see, the memory size for the Memorizing Transformer is comparable to the model size, while MELODI's memory size remains significantly smaller—by an order of magnitude—even for a large model like Llama-2. We will add this analysis to the revised manuscript to provide a clearer comparison of memory footprint relative to model size.

Compression:

Recent work has explored using summary tokens for compression in Transformers [5, 13, 14, 15]. Recurrent Memory Transformer (RMT) [5] utilizes the output of summary tokens recurrently as short-term memory. AutoCompressor [14] aggregates summary tokens across segments to generate a summary representation for long documents used in retrieval tasks. Gisting [16] applies this technique to compress long prompts. The In-context Autoencoder (ICAE) [15] further incorporates LoRA fine-tuning [17] for context compression, while TransformerFAM [18] introduces feedback attention to enhance performance.

Unlike these methods that compress input tokens, MELODI compresses network activations over multiple layers.

[13] Jack W Rae, et al, "Compressive transformers for long-range sequence modelling", arXiv preprint, 2019. https://arxiv.org/abs/1911.05507.

[14] Alexis Chevalier, et al, "Adapting language models to compress contexts", Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, 2023.

[15] Tao Ge, et al, "In-context autoencoder for context compression in a large language model", 2024. https://arxiv.org/abs/2307.06945.

[16] Jesse Mu, et al, "Learning to compress prompts with gist tokens", 2024. https://arxiv.org/abs/2304.08467.

[17] Edward J. Hu, et al. "Lora: Low-rank adaptation of large language models", 2021. https://arxiv.org/abs/2106.09685.

[18] Dongseong Hwang, et al. "Transformerfam: Feedback attention is working memory", 2024. URL https://arxiv.org/abs/2404.09173.

---

$\color{blue}{**[Weakness 2]:**}$

The results only show the effiency part but not the effect of memory on the performance of LLM.

We apologize for the lack of clarity. Our results demonstrate both the efficiency and impact on LLM performance of MELODI's memory mechanisms.

• Impact of memory on LLM performance: To showcase the impact of memory on LLM performance, we compare MELODI to Transformer XL, a transformer-based language model with a KV cache of the previous context window. MELODI uses about the same amount of memory as TransformerXL, but significantly outperforms Transformer XL in terms of perplexity. For instance in Table 3 (of the paper), on the PG-19 dataset with the T5 vocabulary, MELODI S192+L32 achieves a perplexity of 10.51, compared to 11.41 for Transformer XL, demonstrating a clear improvement. This highlights the effectiveness of MELODI's hierarchical memory mechanism in capturing and utilizing long-range context.

• Efficiency: We also compare MELODI to the Memorizing Transformer. MELODI achieves comparable or slightly better perplexity as memorizing transformer, but uses approximately 8 times less memory.

---

$\color{blue}{**[Weakness 3]:**}$

Codes are not shared for checking the implementation.

Thank you for your feedback. We understand the importance of code sharing for reproducibility. Our organization's policy requires an internal review before code release, which can take some time. We are committed to making our code publicly available upon the paper's acceptance. In the meantime, we are happy to answer any specific questions you may have about the implementation.

We sincerely thank the reviewer for the thoughtful feedback and valuable suggestions, which have significantly improved the quality of our paper.

$\color{blue}{**[Weakness 1 + Question 1]:**}$

Prior knowledge about how the memory network works with the LLMs. Adding a preliminary session about the history of recent advances about memory compression.

Thank you for the suggestion. Although the prior works are discussed in the related work section (i.e. Section 4 after the experiment section), we agree that providing the context on memory networks and compression earlier in the paper would improve readability. We will move the related work section before the method section in the revised manuscript.

For your convenience, we have included an overview of relevant work on memory networks and compression below:

Memory in language models:

Long Short-Term Memory (LSTMs) [1] use token-level recurrence to compress prior context into a state vector, which is a limited form of memory. With the advent of Transformers [2], 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) to be processed in parallel. The Block Recurrent Transformer [4] and Recurrent Memory Transformer (RMT) [5] integrate recurrent mechanisms inspired by LSTMs into the Transformer architecture, but the recurrence is over blocks, rather than individual tokens.

Transformer-XL [3] introduces a caching mechanism to store key-value (KV) pairs from the preceding context window as a form of short-term memory. Memorizing Transformer [9] dramatically expands the size of the cache to store KV pairs for long-term memory, but only uses the long-range cache in a single layer, for efficiency reasons. MemoryLLM [10] incorporates long-term memory in every layer, incurring a substantial memory overhead. Infini-Transformer [6] explore the use of additional memory like Hopfield Networks [7, 8]. LONGMEM [11] improves over Memorizing Transformer by introducing a SideNet for memory retrieval and fusion. You Only Cache Once (YOCO) [12] 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 a transformer model via compression.

[1] Sepp Hochreiter and Jurgen Schmidhuber, "Long short-term memory", Neural Computation, 9(8): 1735–1780, 1997.

[2] Ashish Vaswani, et al, "Attention is all you need", Advances in Neural Information Processing Systems, 2017.

[3] Zihang Dai, et al, "Transformer-XL: Attentive language models beyond a fixed-length context", Annual Meeting of the Association for Computational Linguistics, 2019.

[4] DeLesley Hutchins, et al, "Block-recurrent transformers", Advances in Neural Information Processing Systems, 2022.

[5] Aydar Bulatov, et al, "Recurrent memory transformer", Advances in Neural Information Processing Systems, 2022.

[6] Tsendsuren Munkhdalai, et al, "Leave no context behind: Efficient infinite context transformers with infini-attention", 2024. https://arxiv.org/abs/2404.07143.

[7] J. J. Hopfield. "Neural networks and physical systems with emergent collective computational abilities", Proceedings of the National Academy of Sciences of the United States of America, 79(8): 2554–2558, April 1982.

[8] Hubert Ramsauer, et al, "Hopfield networks is all you need", 2021. https://arxiv.org/abs/2008.02217.

[9] Yuhuai Wu, et al, "Memorizing transformers", International Conference on Learning Representations, 2022.

[10] Yu Wang, et al, "Memoryllm: Towards self-updatable large language models", 2024. https://arxiv.org/abs/2402.04624.

[11] Wang, Weizhi, et al. "Augmenting language models with long-term memory." Advances in Neural Information Processing Systems 36 (2024).

[12] You Only Cache Once: Decoder-Decoder Architectures for Language Models, 2024.