Thank you for your thoughtful reviews and insightful suggestions. We appreciate the opportunity to address your concerns. Most of the issues stem from insufficient clarification and misunderstanding regarding the primary goal of this paper. As suggested by reviewer vHo6, we would like to highlight footnote 5 in our paper, emphasizing that "this work is neither an attempt at linearizing attention mechanisms nor focused on model efficiency or model serving".

The goal of this paper is to develop new architectures capable of effectively learning novel tasks. We do not focus on the performance metrics typically associated with persistent-knowledge storage & retrieval, nor model efficiency, model serving, or throughput. While model efficiency could be a future research direction once one wants to establish model serving at a product level, this paper is currently focused on the early stages of architecture exploration and new-task learning ability.

Therefore, most concerns can be addressed by providing more clarity regarding our goals. We will emphasize the goal of this work in the final version!

---

Q1: Additional ruler benchmarks

We chose two question-answering tasks in RULER, because they are notably more challenging than other "needle in a haystack" tasks like multi/single-key retrieval, where many models can easily achieve >99% performance (see Appendix E and F in RULER [4]). Note that the information entropy of multi/single-key retrieval tasks is extremely low, making these tasks easy to compress.

On the other hand, we are happy to share experimental results of Memory Mosaics v2 on these key retrieval tasks. Table A and B show that Memory Mosaics v2 outperforms Transformers.

Thank you for suggesting additional benchmarks such as HELMET, BabiLong, PG-19, SWDE, SQA, and SNQA. However, these benchmarks overlap with each other, RULER, and the three in-context learning tasks presented in this paper.

For instance:

• HELMET includes HotpotQA, banking77, and some ruler tasks.
• The HotpotQA is used in one RULER question-answering task, which is used in this paper.
• Banking77 is one of 3 in-context learning tasks used in this work.
• BabiLong consists of "needle-in-a-haystack" tasks with low information entropy, similar to many RULER tasks.

Therefore, this paper doesn't evaluate similar or redundant tasks repeatedly.

[4] Hsieh, Cheng-Ping, et al. "RULER: What's the Real Context Size of Your Long-Context Language Models?." arXiv (2024).

---

Q2: The goal of this work

As mentioned previously, we would like to emphasize that Footnote 5 states: This work is neither a linearization of attention nor attention efficiency.

Thanks for the valuable suggestion! We will emphasize this goal more clearly in the final version.

---

Q3: The number of in-context learning benchmarks

This paper includes 3 in-context learning datasets, each paired with target label strategies, 2 few-shot prompt strategies, 2 few-shot examples ordering strategies, and 10~16 possible number of shots. This results in over 300 different tasks. Despite that the more benchmarks the better understanding, we find that the 300 tasks in this paper are sufficient to support the in-context learning dimension within a limited space. Note that there are two other evaluation dimensions to explore within the limited pages.

Thanks for the suggestions. We will leave more in-context learning benchmarks as future work.

---

Q4: training + inference throughput profiling.

The primary concern regarding model efficiency and throughput appears to be based on the assumption that this paper targets model efficiency or serving, which is not the case. As mentioned earlier, this work is neither a linearization of attention nor attention efficiency (Page 3, Footnote 5). The goal of this paper is to develop new architectures capable of better learning new tasks.

Despite these clarifications, we are also happy to share the training and inference costs for our current Memory Mosaics v2 implementation.

Currently, Memory Mosaics v2 is implemented based on FlexAttention [5] due to its flexibility. Compared to highly optimized transformer implementations over many years, like flash-attention or fused kernels, FlexAttention [5] achieves approximately 85%–90% of their speed in practice. Therefore, the training cost of Memory Mosaics v2 and baseline Transformer models remains similar. We do not observe any significant training or inference overhead.

Thanks for the suggestions on model efficiency. We leave it as a future direction when our model is ready for model serving.

[5] Dong, Juechu, et al. "Flex attention: A programming model for generating optimized attention kernels." arXiv (2024).

---

Q5: What does the design of long/short-term memory target on, the separation of functionality or computational efficiency?

As emphasized by the goal of this paper, the design of long/short-term memory targets the separation of functionality, rather than computational efficiency.

---

Q6: Re Figure 1: Do the different lines correspond to different layers?

The figure is borrowed from [1], where different lines correspond to different attention/memory heads. Thanks for the suggestions! We will edit the caption to make the figure self-contained.

[1] Zhang, Jianyu, et al. "Memory mosaics." arXiv preprint arXiv:2405.06394 (2024).

---

Q7: Long-term memory mask is randomized during training while fixed at inference. Visualization of long/short-term memories.

The motivation for creating long/short-term memories comes from the "flat-then-peak" attention distribution in Figure 1 (right). To achieve context length extrapolation, we need to allocate the "flat" part (position-invariant) to long-term memory, and the "peak" tail to short-term memory. As such, only the size of long-term memory increases during context length extrapolation.

Because we do not know the exact boundary between "flat" and "peak" parts, we use a randomized long-term memory mask length to further impair long-term memory from learning the tail of sequence. Note that on average this randomization of mask length equals a linearly-decayed importance.

We will compare and share the long/short-term attention distribution (visualization) using this randomized mask in the final version, explaining the motivation behind it and showing its benefits.

---

Q8: Add the number of shots in Table 1.

Thanks for pointing it out. Following the setups in Deepseek-V3 [8] and Llama3 [9], we set the number of shots for common benchmarks as shown in Table C. We will clarify the #shots setups in the final version.

Table C: Number of shots for common benchmarks.

---

Q9: Suggestions on varying the number of shots for the common benchmarks in Table 1?

First, we observe in [8-9] that people do not vary the number of shots for these common benchmarks. We conducted a study to explore different numbers of shots and found that performance on these common benchmarks (e.g., MMLU) usually does not increase with more shots. Note that few-shot learning is expected to see improved performance with additional shots.

Furthermore, Deepseek-R1 [7] even observed a decreased performance as the number of shots increased. In their original words: "When evaluating DeepSeek-R1, we observe that it is sensitive to prompts. Few-shot prompting consistently degrades its performance. Therefore, we recommend users directly describe the problem and specify the output format using a zero-shot setting for optimal results."

These results reveal that these tasks are not classic "few-shot learning". The shots in these tasks provide the information of "problem format", rather than contextual knowledge. Therefore, this paper, as well as [8-9], doesn't vary the number of shots for these common benchmarks.

Thanks for the suggestions. We will add this clarification in the final version.

[7] Guo, Daya, et al. "Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning." arXiv (2025).
[8] DeepSeek-AI. "DeepSeek-V3 Technical Report". (2024).
[9] Grattafiori, Aaron, et al. "The llama 3 herd of models." arXiv (2024).

---

Minor suggested edits.

Thanks for the suggestions on editing! We will fix them in the final version.

---

Summary

Thank you for your thoughtful reviews!

Q1-Q5 raise concerns mainly based on an incorrect assumption --- namely that "this work focuses on model efficiency or model serving". We provide details clarification to state our goal clearly, along with additional experimental results. These concerns should be well addressed after clarifying the goal of this work.

Q6-Q9 ask for details about the model, motivations, and experimental setups. We provide comprehensive explanations and experimental details.

Does the response address your concerns?

Dear reviewer YUYJ,

Thank you for your thoughtful reviews and suggestions regarding model efficiency. Although this work does not focus on model efficiency or model serving, your suggestions help clarify the goal of this paper. We provide a brief summary of our responses in the Summary paragraph, outlining our answers to questions Q1–Q9.

We hope the response addresses your concerns. We also understand that a mismatch regarding the assumed goal of the paper may impair the review-response process. Thus, please feel free to let us know if you have any further concerns/questions!

Thanks for your reply! For model efficiency issue, as stated in Q4: "Therefore, the training cost of Memory Mosaics v2 and baseline Transformer models remains similar. We do not observe any significant training or inference overhead.". In short, model efficiency is not the overhead of Memory Mosaics v2, compared with Transformer.

It is the highly-optimized implementation (such as, flash-attention and fused kernels) in Transformer that help transformer gain ~20% training and inference speed. One can, of course, optimize the Memory Mosaics v2 implementation via replacing FlexAttention by customized triton/cuda kernels to fill this gap. (Note that current Memory Mosaics v2 use FlexAttention[5] implementation to due to its flexibility). However, this optimization of implementation is not the main focus of this paper. This is a normal case in computer science. For example, Transformer was published in 2017, while a efficient implementation, flash-attention, was created in 2022.

For the term "scale in title "Memory Mosaics at Scale", there are many potentional problems to consider, including model efficiency, when scaling up a model. For example, some ideas work well on small data/model-size but fail on large data/model-size, a certain advantage of a method at small scale may also disappear at large scale, etc. Of course, training cost is also a key issue to consider during scaling up (thanks for your suggestion). The term "at scale" in the title needs to consider all these problems, not only efficiency. This paper optimizes the implementation until "model efficiency is not the overhead of Memory Mosaics v2". After that, we stop over-optimizing the implementation and focus on other problems. We will clarify this point in the final version.

Dear reviewer, we thanks for your helpful suggestion on model efficiency! We will state in the final version that this paper doesn't focus on the optimization of implementations (flash-attention, fused kernels). In current implementation, model efficiency is already not the overhead of Memory Mosaics v2, compared with Transformer. We will also consider the HELMET suite in the final version (e.g. removing the overlap).

We hope the answer address your concerns. We appreciate it if you can update scores based on these clarifications and new results.

Thank you for your thoughtful reviews and useful suggestions! It is a pleasure to discuss these questions.

---

W1 & Q1: Lacks analysis on computational overhead compared to transformers. Do these architectural changes incur any computational overhead during inference? How does the performance compare under the same amount of FLOPs?

Thank you for pointing out the need for computation analysis of Memory Mosaics v2 and transformers. We have now included #FLOPs and #parameter analysis for Memory Mosaics v2 and transformers across all three evaluation dimensions. The results demonstrate that Memory Mosaics v2 achieves better performance under comparable FLOPs/parameters or fewer FLOPs/parameters under comparable performance.

• persistent-knowledge storage and retrieval: Table 2 in our paper shows that after removing long-term memory, memory moasics v2 contains fewer parameters (8.3B vs 8.8B) and Flops (15.6B vs 16.7B) than Transformer. Despite this reduction, Memory Mosaics v2 still achieves a comparable performance as Transformer. Flops/token is estimated at context length 256 via [1] approach.

• New-knowledge storage and retrieval: This paper compares memory mosaics v2 and transformer on two scales in Table 4, allowing us to study the relationship between #params/#FLOPs and performance. Table A summarizes the #params, #FLOPS, and performance of these models. By plotting a curve between #params/#FLOPs and performance using the data in Table A, one can easily observe that Memory Mosaics v2 outperforms Transformer by over 10% under any #params or #FLOPs configuration. Flops/token is estimated at context length 32k via [1] approach. | model | # params | # Flops per token | # ruler question-answering (32k task-length)| |---|---|---|---| |transformer small | 1.5B |4.4B| 22.1| |transformer large | 8.8B|25.2B |41.1| |Memory Mosaics v2 small |1.6B|4.7B|53.4| |Memory Mosaics v2 large |9.9B|27.4B|53.4|
  Table A: Ruler question-answering performance of transformer and Memory Mosaics v2 at different parameters and FLOPs scales. Memory Mosaics v2 outperforms transformer by over 10% under the same #params or #FLOPs.

• In-context learning: Similar to the comparison in New-knowledge storage and retrieval, Table B summarizes the #params, #FLOPs/token, and #performance comparison on in-context learning tasks (TacRed 10-shots). Again, by plotting a curve between #params/#FLOPs and performance using the data in Table B, one can easily observe that Memory Mosaics v2 outperforms Transformer by over 10% under the same #params or #FLOPs configurations. (A similar conclusion also holds for other in-context learning tasks and other numbers of shots). | model | # params | # Flops per token | # in-context learning (tacred 10-shots)| |---|---|---|---| |transformer small | 1.5B |4.4B| 23.0| |transformer large | 8.8B|25.2B |50.0| |Memory Mosaics v2 small |1.6B|4.7B|45.5| |Memory Mosaics v2 large |9.9B|27.4B|63.0|
  Table B: In context learning performance of transformer and Memory Mosaics v2 at different parameters and FLOPs scales. Memory Mosaics v2 outperforms transformer by over 10% under the same #params or #FLOPs.

Thanks for the useful suggestions! We will add figures for these comparisons and analyses in the final version.

[1] Adam Casson. Transformer flops. 2023

---

Q2: How would this model work on generative tasks and natural language modeling?

The in-context learning and new-knowledge storage & retrieval tasks are evaluated in a generative manner. Specifically, we allow Memory Mosaics v2 to generate texts freely. Then check whether the generated text correctly answers the question.

Memory Mosaics v2 handles generative tasks and natural language modeling similarly to transformers. First Memory Mosaics v2 pre-fill prompts (e.g., a question), store KV-cache of the prompt, and then generate text using this stored KV-cache.

Meanwhile, other related techniques in transformers, such as KV hashing and quantization, can also be directly applied to Memory Mosaics v2.

---

Summary

Thanks for your thoughtful reviews!

Q1 highlights the need for computational analysis (e.g., FLOPs). We address this point by adding FLOPs and parameter analysis on all three evaluation dimensions.

Q2 raises the concern about the text generation process of Memory Mosaics v2. We provide clarification, showing that the generation process of Memory Mosaics v2 is similar to that of transformers. Many evaluation tasks in this paper are assessed in the generative manner.

We hope the response address your concerns!

Thank you for your thoughtful reviews! After reading the weaknesses W1-W3, we find that all of them can be effectively addressed by improving clarity and providing more experimental evidence.

---

W1: Compared to a standard Transformer, the proposed model requires more parameters to achieve comparable performance.

Actually, Memory Mosaics v2 achieves better performance under comparable #FLOPs/#parameters, or fewer #FLOPs/#parameters under comparable performance across all three evaluation dimensions.

• Persistent-knowledge storage and retrieval: Table 2 in our paper shows that after removing long-term memory, memory moasics v2 contains fewer parameters (8.3B vs 8.8B) and Flops (15.6B vs 16.7B) than Transformer. Despite this reduction, Memory Mosaics v2 still achieves a comparable performance as Transformers.

• New-knowledge storage and retrieval: This paper compares memory mosaics v2 and transformer on two scales in Table 4, allowing us to study the relationship between #params/#FLOPs and performance. Table A summarizes the #params, #FLOPs, and performance of these models. By plotting a curve between #params/#FLOPs and performance using the data in Table A, one can easily observe that Memory Mosaics v2 outperforms Transformer by over 10% under any #params or #FLOPs configuration. (To the ease of reading, we will use a figure to showcase this result in the final version.) | model | # params | # Flops per token | # ruler question-answering (32k task-length)| |---|---|---|---| |transformer small | 1.5B |4.4B| 22.1| |transformer large | 8.8B|25.2B |41.1| |Memory Mosaics v2 small |1.6B|4.7B|53.4| |Memory Mosaics v2 large |9.9B|27.4B|53.4|
  Table A: Ruler question-answering performance of transformer and Memory Mosaics v2 at different parameters and FLOPs scales. Memory Mosaics v2 outperforms transformer by over 10% under the same #params or #FLOPs.

• In-context learning: Similar to the comparison in New-knowledge storage and retrieval, Table B summarizes the #params, #FLOPs, and #performance comparison on in-context learning tasks (TacRed 10-shots). Again, by plotting a curve between #params/#FLOPs and performance using the data in Table B, one can easily observe that Memory Mosaics v2 outperforms Transformer by over 10% under the same #params or #FLOPs configurations. | model | # params | # Flops per token | # in-context learning (tacred 10-shots)| |---|---|---|---| |transformer small | 1.5B |4.4B| 23.0| |transformer large | 8.8B|25.2B |50.0| |Memory Mosaics v2 small |1.6B|4.7B|45.5| |Memory Mosaics v2 large |9.9B|27.4B|63.0|
  Table B: In-context learning (tacred 10-shots) performance of transformer and Memory Mosaics v2 at different parameters and FLOPs scales. Memory Mosaics v2 outperforms transformer by over 10% under the same #params or #FLOPs.

Even though W1 is not a genuine weakness of this paper, it helps improve the writing of the paper. Thank you for pointing it out! We will add these results (using figures), as well as clarifications, in the final version.

---

W2.1: How does the capacity scale with the number of Mosaic v2 parameters?

The previous answer in Tables A and B showcases how the capacity scales with the number of Mosaic v2 parameters and FLOPs, demonstrating that Memory Mosaics v2 outperforms Transformer by over 10% under the same #params or $FLOPs. We will add these results in the final version, addressing this concern.

---

W2.2: Clarify where the proposed memory structure provides its main advantages.

The proposed long/short-term memories structure in Memory Mosaics v2 is built based on the "flat-then-peak" attention distribution observed in Memory Mosaics [1] (Figure 1 in this paper). To achieve context length extrapolation, we need to allocate the "flat" part (position-invariant) to long-term memory, and the "peak" tail to short-term memory. This ensures that only the size of long-term memory increases during context length extrapolation.

For handling in-context learning tasks (few-shot learning), a fair treatment is required for different shots/examples/demonstrations, regardless of their position in a sequence. The position-invariant nature of long-term memory provides this fairness.

In contrast, transformer attention lacks this "flat" attention distribution (Figure 1 left in our paper). Therefore, transformers struggle with context length extrapolation and perform poorly on in-context learning tasks.

We will visualize the attention distribution of Memory Mosaics v2 with long/short-term memories in the final version, addressing the concern of long/short-term memories.

[1] Zhang, Jianyu, et al. "Memory mosaics." arXiv preprint arXiv:2405.06394 (2024).

---

W3: Include experimental evidence to support the claim that "Memory compression algorithms are weak in New-knowledge storage & retrieval and in-context learning tasks."

Weakness of memory compression algorithms in New-knowledge storage & retrieval:

The RULER paper [4] has already studied the weaknesses of model compression methods, such as Mamba and RWKV, via experimental comparison (Figure 4). Here is a summary:

In the original words of the RULER paper, "Finally, we show that non-Transformer architectures, such as RWKV and Mamba, still lag behind Transformer by large margins on RULER."

Note that the two question-answering RULER tasks used in our paper are notably more challenging than other "needle in a haystack" tasks in RULER, like multi-/single-key retrievals. The information entropy of these multi-/single-key retrieval tasks is extremely low, making these tasks easy to compress, allowing many models to easily achieve >99% performance (see Appendix E and F in RULER [4]).

Weakness of memory compression algorithms In-context learning:

[5] has already experimentally shown the weakness of model compression methods, such as Mamba and RWKV, in in-context learning (Note that this paper uses the same in-context learning data as in [5].). Here is a summary:

Memory compression algorithms, such as Mamba and RWKV, completely fail on the in-context learning tasks.

We will refer these results to experimentally support the claim (weakness of memory compression methods) in the final version.

[4] Hsieh, Cheng-Ping, et al. "RULER: What's the Real Context Size of Your Long-Context Language Models?." arXiv preprint arXiv:2404.06654 (2024).
[5] Li, Tianle, et al. "Long-context llms struggle with long in-context learning." arXiv preprint arXiv:2404.02060 (2024).

---

Q1: Perspective on memory compression algorithms

First, we agree that the memory and computation efficiency are significant aspects of memory compression algorithms. Our gated recurrent key feature extractor is also inspired by memory compression algorithms.

In your perspective, you mentioned this: "One could interpret many memory compression algorithms [...] as enhancing the long-term context understanding of the contextual associative memory, ultimately enabling more effective retrieval from the persistent memory." Note that memory compression algorithms don't contain a true "long-term context" analogous to either long-term memory in Memory Mosaics v2 or attention in transformer. (Otherwise, these algorithms would not be considered as memory "compression".)

Due to the absence of a mechanism capable of storing contextual information over large amounts of new information, memory compression algorithms face difficulties with tasks like new-knowledge storage & retrieval and in-context learning.

For persistent-knowledge storage and retrieval tasks, we agree that memory compression algorithms perform well. Actually, rich literature (e.g., Mamba, RWKV, xLSTM, etc) has already experimentally verified this point. In our paper, Table 2 (memory mosaics v2 large without long-term memory) also supports this point. Note that "memory mosaics v2 large without long-term memory" is a kind of fixed-sized sliding window, and thus acts like a memory compression algorithm.

Thanks for the suggestion! We will clarify that "memory compression algorithms (Mamba, RWKV, xLSTM, etc) perform well on persistent-knowledge storage and retrieval" in the final version, using rich experimental evidence from literature.

---

Summary

W1–W2 point out issues related to clarification. We have revised the text to improve clarity, added experimental details, and reorganized the experimental results accordingly.

W3 highlights the lack of experimental evidence supporting our claims about the weaknesses of memory compression algorithms. Thanks for pointing it out! We now refer to the extensive experimental evidence provided in [4] and [5] to support our argument.

Q1 raises a question regarding the perspective on memory compression algorithms. We have included a corresponding discussion in the revised version.

Does the response address your concerns?

Dear reviewer jbCa,

Thank you for your helpful suggestions. We've added additional clarifications and experimental results to address your concerns. We are eager to hear whether these additions adequately resolve them. Please don’t hesitate to let us know if you have any further questions! We appreciate it if you can update scores based on these clarifications and new results.

Thank you for your thoughtful reviews and valuable suggestions! It was a pleasure to answer your questions.

---

Q1: Could the same benefits be achieved without changing the attention mechanism, but by simply enforcing long/short-term masks on traditional transformer attention?

The long/short-term memory components in Memory Mosaics v2 are built based on the nice "flat-then-peak" attention distribution in Memory Mosaics (Figure 1 in our paper). Unfortunately, transformer attention lacks this nice attention distribution, due to its position-encoding and asymmetric query-key in transformers. Therefore, simply adding a long/short-term mask cannot help transformer attention to achieve the same in-context learning benefits.

This paper builds upon Memory Mosaics [1], which has already shown that the position-encoding (as well as corresponding asymmetric query-key) in transformer attention results in a position-dependent attention distribution (Figures 14 & 15 in [1]). This position-dependent behavior impairs transformer's ICL capabilities, as shown in Figures 8 & 9 of [1]. Simply adding long/short-term masks does not address this fundamental issue and thus fails to improve ICL performance.

We will visualize the attention distribution of MMv2 with long/short-term masks in the final version, addressing the concern of applying long/short-term masks to transformers.

[1] Zhang, Jianyu, et al. "Memory mosaics." arXiv preprint arXiv:2405.06394 (2024).

---

Q2: Motivation and Performance benefit of adaptive smoothing, gated recurrent key computation, long-/short-term memory?

Training Memory Mosaics v2 Large and Transformer Large on 1 trillion tokens requires over 4 GPU years with H100 GPU. Comprehensive ablations at this scale are not feasible within weeks. Fortunately, we can better address the concern using conceptual and practical evidence.

These three components are conceptually necessary and practically helpful:

• Adaptive Smoothing (bandwidth): In memory-based methods like kernel smoothing, the optimal bandwidth depends on the number of examples, and must be adaptive. Rich literature [2-3] supports this need for an adaptive bandwidth (see lines 89-91 for details). Furthermore, the adaptive bandwidth in MMv2, $\beta_0 + \beta_1 * n^{\alpha}$, degenerates to a constant bandwidth as in MMv1 when $\alpha = 0$, providing a chance to learn a constant bandwidth.
  Practically, MMv2 learns positive $\alpha$ (0.2~0.4) after training, achieves a significantly lower training loss and better performance on the common benchmarks (persistent-knowledge storage and retrieval).

• Gated recurrent key feature extractor: the gated recurrent key feature extractor is time-variant, enabling the model to extract semantic meanings without being constrained by sequence structure. An example is shown in line 99, where “tom-and-jerry” and “tom---and---jerry” share the same semantic meaning despite differing in sequence structure. Rich literature in recurrent-style networks (e.g., RWKV, mamba) has experimentally verified the benefit of the time-variant feature extractor (line 100).
  Practically, this gated recurrent key feature extractor helps MMv2 (small scale) on common benchmarks (+0.5%) and ruler benchmarks (+8%) when trained on 4k evaluated on 32k.

• Long/short-term associative memories: The motivation and benefits of long-/short-term memory have been discussed in previous answers.

We will include these clarifications and results in the final version.

[2] Trevor Hastie, Robert Tibshirani, Jerome Friedman, et al. The elements of statistical learning, 2009.
[3] E. García-Portugués. Notes for Nonparametric Statistics. 2024.

---

Q3: Long-term memory mask is randomized during training while fixed at inference. Visualization of long/short-term memories.

The motivation for creating long/short-term memories comes from the "flat-then-peak" attention distribution as shown in Figure 1 (right plot). To achieve context length extrapolation, we need to allocate the "flat" part (position-invariant) to long-term memory, and the "peak" tail to short-term memory. As such, only the size of long-term memory increases during context length extrapolation.

Because we do not know the exact boundary between "flat" and "peak" parts, we use a randomized long-term memory mask length to further impair long-term memory from learning the tail of a sequence. Note that on average this randomization of mask length equals to a linearly-decayed importance.

We will compare and share the long/short-term attention distribution (visualization) using this randomized mask in the final version, explaining the motivation behind it.

---

Q4: Formal equation for long vs short memories.

Thanks for the valuable suggestion! Following Eq. 4, we express long/short memories as:

short-term memory: $f(k; \{(k_{max(t-h,1)}, v_{max(t-h,1)}) ... (k_{t-1}, v_{t-1})\})$

long-term memory: $f(k; \{(k_1, v_1) ... (k_{t-m}, v_{t-m})\})$, where $t-m >= 1$

There are no other differences besides the mask of long/short-term memories. We will add these equations in the final version.

---

Q5: Weakness of Mamba and RWKV on RULER benchmark

Weakness of memory compression algorithms on Ruler:

The RULER paper [4] has already studied the weaknesses of model compression methods, such as Mamba and RWKV, via experimental comparison (Figure 4). Here is a summary:

In the original words of the RULER paper, "Finally, we show that non-Transformer architectures, such as RWKV and Mamba, still lag behind Transformer by large margins on RULER."

Note that the two question-answering RULER tasks used in our paper are notably more challenging than other "needle in a haystack" tasks in RULER, like multi-/single-key retrievals. The information entropy of these multi-/single-key retrieval tasks is extremely low, making them easy to compress, allowing many models to easily achieve >99% performance (see Appendix E and F in RULER [4]).

Weakness of memory compression algorithms on in-context learning:

[5] has already experimentally shown the weakness of model compression methods, such as Mamba and RWKV, in in-context learning (Note that our paper uses the same in-context learning datasets as in [5].). Here is a summary:

Table 2: Weakness of Mamba and RWKV on in-context learning benchmark (TacRed classification).

We will refer these results to experimentally support the claim (weakness of memory compression methods) in the final version.

[4] Hsieh, Cheng-Ping, et al. "RULER: What's the Real Context Size of Your Long-Context Language Models?." arXiv preprint arXiv:2404.06654 (2024).
[5] Li, Tianle, et al. "Long-context llms struggle with long in-context learning." arXiv preprint arXiv:2404.02060 (2024).

---

Q6: The goal of this work

Page 3 Footnote 5 claims: This work is neither a linearization of attention nor attention efficiency.

The goal of this paper is to develop new architectures that can better learn new tasks. We neither focus on the performance of persistent-knowledge storage and retrieval (many common benchmarks), nor model efficiency on model serving, nor throughput, nor memory efficiency. Even though they can be future research directions. The three evaluation dimensions in our paper are, therefore, designed to assess the new-tasks-learning capability.

Thanks for the valuable suggestion! We will emphasize our goal more clearly in the final version.

---

Summary

Q1–Q3 and Q6 concern clarification and explanation issues that could potentially impair readability and understanding. Thanks for pointing them out! We have provided more detailed explanations to address these concerns.

Q4 highlights the need for formal equations to make the paper scientifically sound; we have added the relevant equations accordingly.

Q5 points out the need for experimental evidence to support our claim regarding the "weakness of memory compression methods." To address this, we refer to the rich experimental evidence provided by the Ruler benchmark [4] and the in-context learning benchmark [5]. Note that these two benchmarks are the ones used in this paper.

We hope our responses sufficiently address your concerns!

Dear reviewer vHo6,

Thank you for your thoughtful reviews. We provide a brief summary of our responses in the Summary paragraph, outlining our answers to questions Q1–Q6.

We hope our answers sufficiently address your concerns. However, please feel free to let us know if you have any further questions! We appreciate it if you can update scores based on these clarifications and new results.