An Evolved Universal Transformer Memory

## Abstract

Abstract Prior methods propose to offset the escalating costs of modern foundation models by dropping specific parts of their contexts with hand-designed rules, while attempting to preserve their original performance.
We overcome this trade-off with Neural Attention Memory Models (NAMMs), introducing a learned network for memory management that improves both the performance and efficiency of transformers.
We evolve NAMMs atop pre-trained transformers to provide different latent contexts focusing on the most relevant information for individual layers and attention heads.
NAMMs are universally applicable to any model using self-attention as they condition exclusively on the values in the produced attention matrices.
Learning NAMMs on a small set of problems, we achieve substantial performance improvements across multiple long-context benchmarks while cutting the model’s input contexts up to a fraction of the original sizes. We show the generality of our conditioning enables zero-shot transfer of NAMMs trained only on language to entirely new transformer architectures even across input modalities, with their benefits carrying over to vision and reinforcement learning.
Our source code is available at https://github.com/SakanaAI/evo-memory .

## 1 Introduction

Figure: Figure 1: NAMMs use evolution to optimize the performance of LMs by pruning their KV cache memory. Evolved NAMMs can be zero-shot transferred to other transformers, even across input modalities and task domains.
Refer to caption: extracted/5936268/figures/visualizations/Figure1_small_c.png

Transformer architectures have become the golden standard in deep learning, with ubiquitous applications in the design of modern foundation models, exhibiting exceptional performance and scalability (Achiam et al., ; Das et al., ; Team et al., ; Dosovitskiy et al., ; Chen et al., ; Brohan et al., ; Gur et al., ).
The outputs of a transformer are exclusively conditioned on a recent context of input tokens, which for language models (LMs) generally correspond to a window of preceding words.
Thus, addressing the challenge of extending this context window is critical to enable tackling long-range tasks and is currently a focal area of research (Huang et al., ).
However, long contexts also immediately impact training and inference costs, with modern foundation models being increasingly resource-hungry and expensive. Many recent methods proposed to partially offset these costs by studying how to heuristically quantify the importance of each token stored in the model’s latent memory, i.e., stored in its Key-Value (KV) cache. Then, by simply evicting the least important tokens with hand-designed strategies, they have shown early success at reducing memory size while limiting performance losses (Luohe et al., ).

Our research aims to go beyond these hand-designed strategies as we hypothesize that shaping the latent memory KV cache of transformers entails new opportunities to improve their capabilities in downstream tasks.
One widely evidenced example in support of our hypothesis is the effectiveness of hand-crafted input context modifications through prompt engineering (Liu et al., ), even allowing foundation models to learn in-context entirely new skills at test time (Brown et al., ).
Furthermore, unlike prompt engineering, directly managing the memory of transformers enables the provisioning of distinct contexts to each latent level independently, such that individual layers and attention heads can focus on the most relevant information for their specific needs.

Motivated by these considerations, we propose Neural Attention Memory Models (NAMMs), introducing a new class of networks trained with evolution to learn an efficient memory system that maximizes the downstream performance of pre-trained transformers.
Evolution inherently overcomes the non-differentiability of memory management operations with binary outcomes (selecting tokens to preserve/discard) which renders gradient-based optimization incompatible.
Our efforts are inspired by the key role that natural evolution played in shaping human memory, which analogously appears to selectively incorporate and actively prune information based on its lifelong usefulness (Sherry & Schacter, ; Nairne & Pandeirada, ; Frankland & Bontempi, ).

**Table 1: Summarized NAMMs performance in language modeling (top) and zero-shot transfer settings (bottom)**
| Model/Eval | LongBench | InfiniteBench | ChouBun |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Performance | Cache size | Performance | Cache size | Performance | Cache size |  |
| Base model | 28.86 (1.00) | 32768 (1.00) | 1.05 (1.00) | 32747 (1.00) | 21.21 (1.00) | 12099 (1.00) |
| H2O | 28.37 (0.99) | 8192 (0.25) | 1.05 (1.00) | 8193 (0.25) | 19.86 (0.94) | 8292 (0.69) |
| L2 | 27.42 (1.00) | 8192 (0.25) | 1.63 (1.55) | 8193 (0.25) | 18.93 (0.89) | 8292 (0.69) |
| NAMMs | 29.33 (1.11) | 8155 (0.25) | 11.00 (10.45) | 13192 (0.40) | 24.44 (1.15) | 9895 (0.82) |
| Model/Eval | Llama 3 70B | Computer Vision | Reinforcement Learning |  |  |  |
| Performance | Cache size | Performance | Cache size | Performance | Cache size |  |
| Base model | 35.22 (1.00) | 10107 (1.00) | 43.84 (1.00) | 7039 (1.00) | 29.04 (1.00) | 3000 (1.00) |
| H2O | 34.17 (0.97) | 6662 (0.66) | 41.97 (0.96) | 4479 (0.64) | 28.70 (0.99) | 2048 (0.68) |
| L2 | 33.50 (0.95) | 6662 (0.66) | 41.45 (0.95) | 4479 (0.64) | 27.91 (0.96) | 2048 (0.68) |
| NAMMs | 34.70 (0.99) | 8365 (0.83) | 44.38 (1.01) | 5100 (0.72) | 31.73 (1.09) | 2434 (0.81) |

Our NAMMs are conditioned on features entirely constructed from the attention matrix, making them universally applicable to any transformer-based architecture.
Learning NAMMs atop a pre-trained Llama 3 8B model (Dubey et al., ), we not only obtain efficiency benefits, with substantial reductions in the number of retained tokens in the KV cache, but also exceed the performance of the full-context model with notable margins. We validate these findings across 36 different tasks from LongBench (Bai et al., ), InfiniteBench (Zhang et al., ), and ChouBun(^1^11ChouBun is the pronunciation of “{CJK}UTF8ipxm長文”, literally translating to “long text” in Japanese.), a new Japanese benchmark designed to assess long-context capabilities beyond the common English and Chinese. These results mark a clear contrast with the aforementioned hand-designed strategies that appear to inevitably trade off efficiency for performance, in line with their stated purpose.

Furthermore, we show that the generality of our parameterization enables zero-shot transfer of NAMMs trained on three natural language tasks to entirely new transformer models.
In particular, we obtain further performance and efficiency improvements not only when using the evolved NAMMs with other LMs of increased size, but also transformers with entirely different architectures concerned with new input modalities, for problems such as vision and reinforcement learning.
In a nutshell, our main technical contributions can be summarized as the following:

- •
We introduce NAMMs, a novel memory evolution framework that adds a new dimension to optimizing transformer models without altering their powerful architectures.
- •
We design and successfully train NAMMs on top of pre-trained transformer models, obtaining both performance and efficiency gains on several long context language tasks.
- •
We show NAMMs, trained only on language tasks, can be transferred zero-shot to any other transformers, retaining benefits across different input modalities and task domains.

We share our full training code to facilitate future advances in foundation models through memory.

## 2 Background and preliminaries

Attention and transformers. Transformers are neural network architectures designed specifically for efficiently processing input sequences.
These models take as input a stream of tokens (e.g., embeddings of words, image patches, robotic states, etc.) and, produce a set of latents with the same length within their layers.
Multi-headed dot product attention (Vaswani et al., ), or simply self-attention, characterizes modern transformers, facilitating effective information sharing across token representations.
The attention layer conducts a set of parallel computations, each known as an attention head, mapping tokens to query, key, and value vectors $\in\mathbb{R}^{d}$. These vectors are organized along the sequence dimension in the matrices $Q$, $K$, and $V$, and the layer’s output is computed as:

$$ $\textit{attention}_{M}(Q,K,V)=AV,\quad\text{where, }\quad A=\textit{softmax} \left(M\times\frac{QK^{T}}{\sqrt{d}}\right).$ (1) $$

Here, $M$ represents an optional mask multiplying the attention matrix $A$, usually enforcing an auto-regressive conditioning such that each token cannot attend to its future. An interpretation of the attention layer comes from the elements of the attention matrix $A^{j}_{i}$, i.e., the dot products between each key $i$ and query $j$ normalized along the column dimension.
Intuitively, each of these values can be understood as the relative importance of token $i$ in processing the input representation of token $j$.

Frequency-based feature extraction. An established canonical technique to pre-process one-dimensional non-stationary signals is the Short-Time Fourier Transform (STFT) (Allen & Rabiner, ). This technique has seen plenty of applications for feature extraction concerning audio, biomedical, seismic, and many more kinds of modalities. The STFT performs a time-convolution of a signal, shifting each convolutional window to the frequency domain through a discrete Fourier transform, producing a spectrogram representation of the original input. We use $\omega^{t}\in\mathbb{R}^{N+1}$ to denote the fixed-sized vector produced at each timestep $t$, where the $N$ frequencies span from zero up to the Nyquist frequency (half the original sampling rate).
Mathematically, the $m$-th frequency from an STFT for time $t$ is extracted from an input vector $v\in\mathbb{R}^{T}$ as:

$$ $\omega^{t}[n]=\sum_{t^{\prime}=0}^{T}v[t^{\prime}]w[t-t^{\prime}]e^{\frac{-n \pi t}{N}}.$ (2) $$

Here, the convolutional filter of the SFTF is defined by the product of a finite-length window function $w$ with each exponential term in the Fourier transform. A popular choice for $w$ is the Hann window (Oppenheim, ), employing a smooth decay at its edges which helps minimize the overestimation of the magnitudes of the higher frequencies in $\omega$ due to spectral leakage (Harris, ).

## 3 Neural Attention Memory Models

Figure: Figure 2: Schematic depiction of our Neural Attention Memory Model design. We extract features from a spectrogram over the attention values of the KV cache tokens (left), which we reduce via an element-wise exponential moving average (EMA) operation (center). These features are fed to our memory model’s networks with fully connected (FC) and cross-token BAM connections (right).
Refer to caption: extracted/5936268/figures/visualizations/BAM_summary_fc.png

An immediate limitation of transformers is the quadratic costs associated with computing the attention matrix $A$.
To partially address this issue, during auto-regressive generation, the latents for the keys and values of the tokens generated at the previous steps are usually stored in what is referred to as the KV cache.
This object can be regarded as being analogous to the memory of the transformer, which now, at each step, only needs to compute the query, key, and value of the latest token and perform attention over a horizontal vector by exploiting causal ordering. In this section, we describe the feature extraction, architecture, and optimization of NAMMs, which have been designed to act on the KV cache to improve both the performance and practicality of this powerful class of models.

### 3.1 Attention spectrograms for model-agnostic feature extraction

The feature extraction framework of NAMMs is designed to be agnostic to the parameterization of the base transformer they are applied for.
In particular, we build a representation for each token in the current KV cache memory from its corresponding column vector in the attention matrix $A_{i}$.
To meaningfully compress this unbounded vector signal, we process it via an STFT with a fixed-sized Hann window (Figure 2, left).
This operation produces a spectrogram representation of the attention columns $\omega_{i}^{t}$, representing the frequencies with how the queries attend to each of the stored key tokens (indexed by $i$) on a compressed time-axis (indexed by $t$).
Thus, this representation exposes precisely the knowledge of how each token’s relative importance varies across all past queries in a compact form factor, discarding all other information specific to the learned transformer weights.

As NAMMs rely only on the attention values for their input, they are universally applicable to any layer producing an attention matrix.
This property is crucial, enabling us to avoid learning individual memory models for the different layers of a transformer, thus, greatly limiting the number of total optimized parameters. Furthermore, it also allows efficient training on top of smaller foundation models for targeted problems, and later transferring the resulting models zero-shot at test-time to larger architectures and arbitrary applications.

### 3.2 Memory model design and cross-token communication

Figure: Figure 3: Our backward mask makes each token attend exclusively to its future relatives in the KV cache.
Refer to caption: extracted/5936268/figures/visualizations/bam_mask_causal.png

NAMMs parameterize a small neural network $m_{\phi}$ to output a scalar selection score $s_{i}=m_{\phi}(\omega_{i}^{1:T})$ for each $i^{th}$ token in the KV cache.
First, to obtain a consistent input dimension, we reduce the attention spectrogram into a smaller feature vector $\omega_{i}$ by compressing the time-axis via an element-wise exponentially moving average (EMA: $\omega_{i}=\sum_{t}\gamma^{t}\omega^{t}_{i}$; Figure 2, center).
We then append positional encodings and feed the vector $\omega_{i}$ to the memory model’s network $m_{\phi}$ to produce the score $s_{i}$.
Finally, we evict from the KV cache memory all latent tokens with $s_{i}<0$, effectively treating the problem as a binary classification task.
We repeat this process with a fixed interval, every set number of new input tokens, $n_{\mathrm{up}}$.

Backward attention memory models (BAM). For the design of $m_{\phi}$, we posit that sharing information from all tokens in memory could be key for assessing their importance.
A particularly motivating scenario in LMs arises when considering the case of repeated words or sentences, where learning a diversity measure that compares different tokens would allow preventing redundancies in the KV cache.
Corroborating this intuition, even from a biological perspective, memory formation and retention appear to adhere to models of neuronal competition (Han et al., ).

Based on these considerations, we design the backward attention memory architecture (BAM) for parameter-efficient sharing of information while making use of the powerful inductive biases enabled by the masked self-attention operation.
In particular, we implement $m_{\phi}$ via an initial self-attention layer with a counter-causal mask $\hat{M}$, which we refer to as backward (Figure 3). This design serves to introduce a purposeful asymmetric relationship, allowing to distinguish between older and newer tokens. We then output $s_{i}$ from a final linear operation:

$$ $\displaystyle o_{i}=\textit{attention}_{\hat{M}}(K_{\Omega},V_{\Omega},Q_{ \Omega}),\quad s_{i}$ $\displaystyle=\textit{linear}(o_{i}),$ (3) $$

where $K_{\Omega},V_{\Omega},Q_{\Omega}$ are the key, value, and query matrices from all feature vectors $\omega_{i}$ in memory.
Using BAM to tackle the previous motivating scenario, only the representation for older tokens would be potentially affected by the presence of newer duplicates.
Thus, just by learning a simple diversity metric within self-attention, backward masking would provide the memory model with the potential to preserve only the most informed occurrence of each token without risking discarding any information in its entirety (since the score for the latest instance of each repeated token would be independent of its past).

### 3.3 Training NAMMs with incremental evolution

Figure: Figure 4: Mean and standard deviation over the CMA-ES population batch performance (left), together with the performance of the learned mean parameter on each task (right).
Refer to caption: extracted/5936268/figures/results/training_curves_bam.png

We evolve the network weights of our NAMMs to directly optimize the performance on a subset of long-context language modeling tasks from LongBench (Bai et al., ). As we share a single $m_{\phi}$ across all layers, even with our largest NAMM we only evolve about 4000 total parameters.
We use the seminal CMA-ES optimization algorithm (Hansen, ) and apply NAMM atop a Llama 3 8B base model (Dubey et al., ) with a context extended from 8192 to 32768 tokens via NTK-aware positional interpolation (bloc97, ).
Due to the inference costs of LMs with long inputs, we sample a subset of different prompts from each task in each generation and propose training in an incremental fashion: starting from a single task, and adding additional tasks at later training stages.
Empirically, we found both these choices to provide effective regularization, improving generalization (see Appendix C).
The performance of modern LMs on LongBench varies considerably across tasks, and even across different task prompts.
Hence, instead of using the raw scores, we opt to maximize normalized performance relative to the vanilla base model’s stored evaluation performance on each same subset of prompts, retaining all tokens in its KV cache memory.

We choose three tasks from different LongBench categories across both English and Chinese where the Llama 3 base model seems to particularly struggle: PassageRetrieval-en, DuReader, and NarrativeQA.
We evolve our NAMM for 300 generations in its first incremental phase, 250 in its second, and 120 in its third. We diminish the number of generations to counteract the increasing costs with each additional phase and make more efficient use of computational resources.
At the end of each phase, we resume from the best previous checkpoint. We provide training curves of our main backward-attention model in Figure 4, showing the average and standard deviation of the normalized batch performance across the population (left), together with the normalized per-task and average performance on all samples of the optimized mean from CMA-ES (right).
We refer to Appendix A for additional architectural and optimization details, together with the set of hyper-parameters. We also provide additional statistics and training curves for other memory model designs in Appendix C.

## 4 Experimental Results

In this section, we evaluate and analyze evolved NAMMs as compared to full-context transformers and two recent hand-designed methods for KV cache management: H2O (Zhang et al., ) and L2 (Devoto et al., ). We compare each method in terms of absolute and normalized performance and also provide the resulting average cache size recorded at the end of each prompt. We first consider three long-context language modeling benchmarks spanning 36 diverse tasks in three languages, using the same Llama 3 8B base transformer from training. Then, we evaluate the capabilities of zero-shot transferring NAMMs to other unseen transformers and task domains. In particular, we not only consider transfer to larger LMs, but also transformers with tokens constructed from modalities other than language. Across all these settings, we also compare BAM with a simpler 2-layer MLP architecture and provide summarized results after every stage of incremental evolutions. We refer to Appendix C additional evaluations, ablation studies, and all learning curves. Lastly, we perform targeted qualitative analysis to understand the behavior of our new memory framework.

### 4.1 Long-context language understanding

**Table 2: NAMMs evaluation on LongBench. The normalized performance (in brackets) is calculated using the base model with full cache. The tasks used for NAMM’s training are highlighted in gray.**
| Model/Task id | Single-Doc QA | Multi-Doc QA | Summarization |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 1-1 | 1-2 | 1-3 | 1-4 | 2-1 | 2-2 | 2-3 | 2-4 | 3-1 | 3-2 | 3-3 | 3-4 |  |
| Base model | 10.38 (1.00) | 12.79 (1.00) | 22.60 (1.00) | 21.31 (1.00) | 10.41 (1.00) | 12.67 (1.00) | 7.54 (1.00) | 25.86 (1.00) | 29.34 (1.00) | 23.93 (1.00) | 0.92 (1.00) | 2.66 (1.00) |
| H2O | 8.75 (0.84) | 13.07 (1.02) | 22.11 (0.98) | 21.62 (1.01) | 10.28 (0.99) | 12.40 (0.98) | 7.20 (0.95) | 26.58 (1.03) | 28.56 (0.97) | 23.98 (1.00) | 0.88 (0.96) | 2.25 (0.85) |
| L2 | 8.83 (0.85) | 13.13 (1.03) | 22.22 (0.98) | 21.79 (1.02) | 9.97 (0.96) | 12.15 (0.96) | 5.88 (0.78) | 24.96 (0.97) | 28.05 (0.96) | 23.28 (0.97) | 1.15 (1.25) | 1.52 (0.57) |
| NAMM (Ours) | 9.14 (0.88) | 12.63 (0.99) | 21.94 (0.97) | 21.34 (1.00) | 9.71 (0.93) | 11.63 (0.92) | 6.98 (0.93) | 20.58 (0.80) | 28.78 (0.98) | 24.39 (1.02) | 1.04 (1.13) | 3.63 (1.36) |
| Model/Task id | Few-shot Learning | Synthetic | Code | Overall |  |  |  |  |  |  |  |  |
| 4-1 | 4-2 | 4-3 | 4-4 | 5-1 | 5-2 | 5-3 | 6-1 | 6-2 | All tasks | Test tasks | Cache size |  |
| Base model | 73.00 (1.00) | 89.45 (1.00) | 46.54 (1.00) | 40.00 (1.00) | 1.48 (1.00) | 12.18 (1.00) | 28.80 (1.00) | 69.09 (1.00) | 65.17 (1.00) | 28.86 (1.00) | N/A | 32768 (1.00) |
| H2O | 73.00 (1.00) | 90.03 (1.01) | 46.48 (1.00) | 34.00 (0.85) | 2.18 (1.47) | 9.93 (0.82) | 27.76 (0.96) | 69.37 (1.00) | 65.44 (1.00) | 28.37 (0.99) | N/A | 8192 (0.25) |
| L2 | 66.41 (0.91) | 84.92 (0.95) | 45.78 (0.98) | 34.38 (0.86) | 3.13 (2.11) | 11.00 (0.90) | 28.68 (1.00) | 73.45 (1.06) | 55.20 (0.85) | 27.42 (1.00) | N/A | 8192 (0.25) |
| NAMM (Ours) | 73.00 (1.00) | 89.81 (1.00) | 46.35 (1.00) | 40.00 (1.00) | 3.04 (2.05) | 27.55 (2.26) | 28.60 (0.99) | 69.53 (1.01) | 66.35 (1.02) | 29.33 (1.11) | 1.07 | 8155 (0.25) |

Longbench. In Table 2, we provide results across all LongBench tasks (Bai et al., ).
Our NAMM yields concrete improvements to the Llama 3 8B transformer both when considering the full set or exclusively the held-out set of test tasks that were not used for evolution, with improvements of 11% and 7% respectively. At the same time, our NAMM also yields efficiency side benefits, notably reducing the context-extended KV cache size. Instead, while both H2O and L2 produce even smaller cache sizes, they both come with some performance costs - in line with their stated objective of retaining rather than improving the original full-context performance. These results indicate how two different hand-designed extremes, either retaining all tokens or aggressively dropping them, come with their own set of downsides for either efficiency or performance. On the other hand, NAMMs are the only methods making consistent improvements from the base model across both axes, highlighting how end-to-end evolutionary optimization can open new orthogonal directions beyond what is feasible with manually-designed heuristics.

**Table 3: NAMMs evaluation on InfiniteBench. The normalized overall performance (in brackets) is calculated using the average performance of the base model with full cache.**
| Model/Task name | Retrieval | Dialogue | Novel | Math | Code | Overall |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Ret.PassKey | Ret.Number | Ret.KV | En.Dia | En.Sum | En.MC | En.QA | ZH.QA | Math.Find | Code.Run | Code.Debug | All tasks | Cache size |  |
| Base model | 0.00 | 0.00 | 0.00 | 1.00 | 7.73 | 0.00 | 1.05 | 1.79 | 0.00 | 0.00 | 0.00 | 1.05 (1.00) | 32747 (1.00) |
| H2O | 0.00 | 0.00 | 0.00 | 1.50 | 5.38 | 0.00 | 1.01 | 1.71 | 1.71 | 0.25 | 0.00 | 1.05 (1.00) | 8193 (0.25) |
| L2 | 0.00 | 0.00 | 0.00 | 1.00 | 5.41 | 0.44 | 0.83 | 2.59 | 7.43 | 0.25 | 0.00 | 1.63 (1.55) | 8193 (0.25) |
| NAMM (Ours) | 11.86 | 11.86 | 1.80 | 1.00 | 14.91 | 36.24 | 8.78 | 17.67 | 10.57 | 1.75 | 4.57 | 11.00 (10.45) | 13192 (0.40) |

InfiniteBench. In Table 3, we provide results across the InfiniteBench tasks (Zhang et al., ). In this benchmark, the average prompt length is close to 200K tokens making it extremely challenging, especially for LMs that were not expensively finetuned for very long context understanding. In fact, as reported by Zhang et al. (), even GPT4 (Achiam et al., ) cannot exceed a performance of 1% on some of its problems. In line with these results, the full-context Llama 3 together with H2O and L2 obtain near-zero performance on most tasks. Instead, our NAMM provides outstanding improvements, bringing overall benchmark performance from 1.05% to 11%.
We also observe that while our NAMM’s memory size is larger than for LongBench, it is considerably lower in relation to the base model’s (now only 40%). This result suggests that NAMMs emergently learned a scalable memory strategy, forgetting redundant and detrimental information at an increasing rate with longer contexts without requiring the hand-designed hard cache limits enforced by L2 and H2O.

**Table 4: NAMMs evaluation on ChouBun.**
| Model/Task | Extractive QA | Summarization | Overall |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| JA.WikiQA | JA.EdinetQA | JA.CorpSecQA | JA.CorpSecSum | All tasks | Cache size |  |
| Base model | 22.91 (1.00) | 28.34 (1.00) | 11.83 (1.00) | 21.75 (1.00) | 21.21 (1.00) | 12099 (1.00) |
| H2O | 20.76 (0.91) | 26.39 (0.93) | 10.42 (0.88) | 21.87 (1.01) | 19.86 (0.94) | 8292 (0.69) |
| L2 | 19.60 (0.86) | 24.06 (0.85) | 8.23 (0.70) | 23.83 (1.10) | 18.93 (0.89) | 8292 (0.69) |
| NAMM (Ours) | 21.34 (0.93) | 28.61 (1.01) | 14.64 (1.24) | 33.15 (1.52) | 24.44 (1.15) | 9895 (0.82) |

ChouBun. Our new benchmark focuses on tasks designed exclusively in Japanese, a novel language unseen during NAMMs training. We hope this benchmark might itself be a valuable contribution to the research community, allowing the assessment of long-context capabilities in multilingual LLMs beyond the already-ubiquitous English and Chinese. We provide further benchmark statistics, details about task composition, together with evaluation metrics for a wider range of popular LLMs in Appendix B. In Table 4, we report our results evaluating NAMMs. Once again, we observe a clear contrast with prior hand-designed methods. While integrating either H2O or L2 leads to notable performance drops, NAMMs provides substantial improvements, with overall performance up by 15% from the full-context Llama 3 8B base model.

### 4.2 Zero-shot transfer across architectures and modalities

**Table 5: NAMMs evaluation on LongBench with a Llama 3 70B model. The normalized performance (in brackets) is calculated using the base model with full cache.**
| Model/Task id | Single-Doc QA | Multi-Doc QA | Summarization |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 1-1 | 1-2 | 1-3 | 1-4 | 2-1 | 2-2 | 2-3 | 2-4 | 3-1 | 3-2 | 3-3 | 3-4 |  |
| Base model | 9.38 (1.00) | 13.84 (1.00) | 24.99 (1.00) | 17.78 (1.00) | 11.73 (1.00) | 14.26 (1.00) | 8.11 (1.00) | 26.43 (1.00) | 13.13 (1.00) | 24.55 (1.00) | 23.20 (1.00) | 10.08 (1.00) |
| H2O | 8.80 (0.94) | 13.48 (0.97) | 25.02 (1.00) | 18.44 (1.04) | 12.36 (1.05) | 14.32 (1.00) | 8.15 (1.01) | 26.22 (0.99) | 13.37 (1.02) | 24.50 (1.00) | 23.20 (1.00) | 9.22 (0.91) |
| L2 | 8.57 (0.91) | 13.40 (0.97) | 24.70 (0.99) | 17.94 (1.01) | 12.77 (1.09) | 13.85 (0.97) | 7.13 (0.88) | 25.74 (0.97) | 12.78 (0.97) | 23.21 (0.95) | 23.35 (1.01) | 8.45 (0.84) |
| NAMM (Ours) | 9.13 (0.97) | 13.53 (0.98) | 24.25 (0.97) | 17.82 (1.00) | 11.45 (0.98) | 13.76 (0.96) | 8.34 (1.03) | 21.79 (0.82) | 12.66 (0.96) | 24.21 (0.99) | 23.56 (1.02) | 8.62 (0.86) |
| Model/Task id | Few-shot Learning | Synthetic | Code | Overall |  |  |  |  |  |  |  |  |
| 4-1 | 4-2 | 4-3 | 4-4 | 5-1 | 5-2 | 5-3 | 6-1 | 6-2 | All tasks | Test tasks | Cache size |  |
| Base model | 78.00 (1.00) | 92.43 (1.00) | 48.67 (1.00) | 45.50 (1.00) | 22.50 (1.00) | 75.37 (1.00) | 33.89 (1.00) | 74.60 (1.00) | 71.19 (1.00) | 35.22 (1.00) | N/A | 10107 (1.00) |
| H2O | 77.50 (0.99) | 92.43 (1.00) | 48.33 (0.99) | 39.75 (0.87) | 18.12 (0.81) | 64.69 (0.86) | 33.89 (1.00) | 74.61 (1.00) | 71.09 (1.00) | 34.17 (0.97) | N/A | 6662 (0.66) |
| L2 | 76.50 (0.98) | 93.22 (1.01) | 46.15 (0.95) | 36.25 (0.80) | 16.98 (0.75) | 64.34 (0.85) | 36.28 (1.07) | 74.38 (1.00) | 67.43 (0.95) | 33.50 (0.95) | N/A | 6662 (0.66) |
| NAMM (Ours) | 78.50 (1.01) | 92.36 (1.00) | 48.49 (1.00) | 45.50 (1.00) | 19.07 (0.85) | 74.19 (0.98) | 34.28 (1.01) | 74.71 (1.00) | 72.42 (1.02) | 34.70 (0.99) | 0.99 | 8365 (0.83) |

Cross-scale adaptation. In Table 5, we provide results zero-shot transferring our NAMM from the Llama 3 8B to the Llama 3 70B model on LongBench. Across all tasks, we find performance to be very close to the full-context baseline with an overall gap of less than 1% even for the test subset. While NAMMs are not able to improve the overall full-context performance in this first transfer setting outside specific task categories (e.g., coding and few-shot learning), they still outperform both H2O and L2 baselines and retain a similar efficiency as with their original training transformer.

**Table 6: Evaluation on the LongVideoBench and MLVU benchmarks with Llava Next Video 7B.**
| Model/Task name | LongVideoBench | MLVU | All tasks | Cache size |
| --- | --- | --- | --- | --- |
| Base model | 43.45 (1.00) | 44.23 (1.00) | 43.84 (1.00) | 7039 (1.00) |
| H2O | 40.91 (0.94) | 43.03 (0.97) | 41.97 (0.96) | 4479 (0.64) |
| L2 | 40.84 (0.94) | 42.07 (0.95) | 41.45 (0.95) | 4479 (0.64) |
| NAMM (Ours) | 44.58 (1.03) | 44.18 (1.00) | 44.38 (1.01) | 5100 (0.72) |

Vision Language Understanding. In Table 6, we provide results zero-shot transferring to the computer vision domain, evaluating NAMMs with a Llava Next Video 7B model (Zhang et al., ) on LongVideoBench (Wu et al., ) and Multi-Task Long Video Understanding (MLVU) (Zhou et al., ). As when evaluated atop Llama 8B, our NAMM is the only method recording gains over the full-context base transformer in both benchmarks. Furthermore, we find that NAMMs learns to forget almost exclusively parts of redundant video frames rather than the language tokens describing the final prompt, even though they were never faced with such modality during training. This result validates that our NAMM recovered a domain-agnostic memory management strategy, further highlighting their flexibility.

**Table 7: Evaluation on D4RL with a Decision Transformer. The normalized performance (in brackets) is calculated using the base model with full cache.**
| Model/Task name | Hopper-v3 | Walker2d-v3 | HalfCheetah-v3 | Overall |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Medium | Med-Replay | Expert | Medium | Med-Replay | Expert | Medium | Med-Replay | Expert | All tasks | Cache size |  |
| Base model | 33.36 (1.00) | 18.37 (1.00) | 44.62 (1.00) | 68.21 (1.00) | 7.18 (1.00) | 38.98 (1.00) | 34.91 (1.00) | 5.06 (1.00) | 10.64 (1.00) | 29.04 (1.00) | 3000 (1.00) |
| H2O | 33.19 (1.00) | 17.86 (0.97) | 49.10 (1.10) | 67.63 (0.99) | 7.59 (1.06) | 40.03 (1.03) | 26.73 (0.77) | 4.46 (0.88) | 11.74 (1.10) | 28.70 (0.99) | 2048 (0.68) |
| L2 | 32.85 (0.98) | 17.96 (0.98) | 43.75 (0.98) | 65.47 (0.96) | 7.18 (1.00) | 40.64 (1.04) | 30.10 (0.86) | 4.76 (0.94) | 8.52 (0.80) | 27.91 (0.96) | 2048 (0.68) |
| NAMM (Ours) | 36.10 (1.08) | 18.86 (1.03) | 49.39 (1.11) | 70.87 (1.04) | 7.53 (1.05) | 50.02 (1.28) | 34.56 (0.99) | 5.90 (1.17) | 12.34 (1.16) | 31.73 (1.09) | 2434 (0.81) |

Reinforcement learning. In Table 7, we provide our zero-shot transfer results for the offline reinforcement learning setting, where we apply NAMMs atop a decision transformer (Chen et al., ) and consider the canonical continuous-control tasks from the D4RL benchmark (Fu et al., ). We find our NAMM improves the base transformer quite considerably in this domain across eight out of nine offline tasks with over 9% overall gains, opposing the performance loss of the other efficient baselines. We posit that since the nature of the decision transformer optimization is closely tied to behavior cloning, the ability to discard part of the context is likely to allow NAMMs to forget and avoid imitating previous mistakes autoregressively. In support of this hypothesis, we observed slightly higher average rewards in the transitions for the retained tokens (by 1.4%, 0.8%, and 12.3% for the Hopper, Walker2d, and HalfCheetah environments, respectively).

**Table 8: Summarized comparison of different NAMMs in language modeling (top) and zero-shot transfer (bottom)**
| Model/Eval | LongBench | InfiniteBench | ChouBun |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Performance | Cache size | Performance | Cache size | Performance | Cache size |  |
| Base model | 28.86 (1.00) | 32768 (1.00) | 1.05 (1.00) | 32747 (1.00) | 21.21 (1.00) | 12099 (1.00) |
| NAMM (MLP, s1) | 28.83 (1.05) | 7639 (0.23) | 3.08 (2.93) | 11329 (0.35) | 22.09 (1.04) | 9525 (0.79) |
| NAMM (MLP, s2) | 29.22 (1.07) | 8475 (0.26) | 4.00 (3.80) | 13031 (0.40) | 22.06 (1.04) | 9815 (0.81) |
| NAMM (BAM, s1) | 28.91 (1.05) | 7951 (0.24) | 10.14 (9.63) | 11173 (0.34) | 22.73 (1.07) | 9569 (0.79) |
| NAMM (BAM, s2) | 29.25 (1.07) | 8267 (0.25) | 9.78 (9.29) | 12789 (0.39) | 24.05 (1.13) | 9867 (0.82) |
| NAMM (BAM, s3) | 29.33 (1.11) | 8155 (0.25) | 11.00 (10.45) | 13192 (0.40) | 24.44 (1.15) | 9895 (0.82) |
| Model/Eval | Llama 3 70B | Computer Vision | Reinforcement Learning |  |  |  |
| Performance | Cache size | Performance | Cache size | Performance | Cache size |  |
| Base model | 35.22 (1.00) | 10107 (1.00) | 43.84 (1.00) | 7039 (1.00) | 29.04 (1.00) | 3000 (1.00) |
| NAMM (MLP, s1) | 34.11 (0.97) | 7930 (0.78) | 40.44 (0.92) | 584 (0.08) | 29.30 (1.01) | 1993 (0.66) |
| NAMM (MLP, s2) | 34.29 (0.97) | 8445 (0.84) | 40.39 (0.92) | 713 (0.10) | 29.58 (1.02) | 2834 (0.94) |
| NAMM (BAM, s1) | 34.11 (0.97) | 7947 (0.79) | 41.52 (0.95) | 723 (0.10) | 30.44 (1.05) | 2009 (0.67) |
| NAMM (BAM, s2) | 25.20 (0.72) | 8276 (0.82) | 44.63 (1.02) | 4948 (0.70) | 31.53 (1.09) | 2534 (0.84) |
| NAMM (BAM, s3) | 34.70 (0.99) | 8365 (0.83) | 44.38 (1.01) | 5100 (0.72) | 31.73 (1.09) | 2434 (0.81) |

NAMMs comparison. In Table 8, we provide summarized results comparing NAMMs with either BAM or the simpler MLP architecture at the end of each stage of incremental evolution. First, we note that even the MLP NAMM after stage 1 impressively improves performance across all language benchmarks. Additionally, performance sees near-monotonic improvements with each additional stage of incremental evolution in both language and zero-shot transfer settings. Comparing our implementations, the performance benefits from the memory models with BAM appear consistently superior to the MLP. Moreover, on ChouBun. we observe that the performance with BAM sees a notable upswing after the second stage of incremental training, which might be associated with the introduction of another ideogram-based language in the training set.(^2^22The DuReader task, used in the second stage of incremental training, uses the Chinese language.) The same improvement not occurring with the MLP-based NAMMs might be further evidence of architectural performance saturation, highlighting the effectiveness of our main implementation.

### 4.3 Understanding Neural Attention Memory Models

Figure: Figure 5: Memory size and token oldness as recorded for each layer in the base model (top) and for each task in LongBench (bottom). We normalize these statistics per task using either their average across all task prompts (top) or the mean sample length (bottom).
Refer to caption: extracted/5936268/figures/visualizations/analysis_per_layer_task.png

Influence of layer depth.
We begin analyzing NAMMs by focusing on the final amount of retained tokens and their oldness(^3^33We define oldness of a retained token as the number of new queries since its introduction in the KV cache.).
At the top of Figure 5, we provide these normalized metrics as a function of layer depth.
Interestingly, our learned NAMM does not appear to affect the KV cache uniformly, retaining visibly more and older tokens for some of the early-middle layers of the base transformer.
One possible interpretation of our results, complementing recent analysis (Wendler et al., ), is that these layers might be particularly important for processing and aggregating information over longer contexts, thus requiring larger memories than the rest.

Influence of task structure.
At the bottom of Figure 5, we instead provide these metrics while varying the source task, this time normalized by the average prompt lengths shown in green. Our results illustrate an inverse correlation between normalized memory size and prompt length (with a Pearson coefficient of -0.84), further confirming our earlier observations of sub-linear memory growth and favorable scaling to longer contexts. Additionally, we observe that in the code completion tasks (with task id 6-1 and 6-2) NAMMs learn to preserve visibly more tokens relative to their average prompt lengths. This result appears intuitively consistent with the higher information density in code, leaving room for less redundancy as opposed to natural language.

Figure: Figure 6: Qualitative inspection of the text from decoding the retained tokens in the KV cache after applying NAMMs. We compare the behavior of NAMMs across the layers with the highest (left) and lowest average retained tokens (right), for either a natural language (top) or coding task (bottom).
Refer to caption: extracted/5936268/figures/visualizations/analysis_tok_pruning.png

Selected qualitative examples. We qualitatively find these analyzed trends by inspecting the text corresponding to the forgotten tokens for a few selected prompts. In particular, we consider the layers with the highest and lowest average retained tokens (15 and 24), for tokens from either a natural language or coding task (PassageRetrieval-en, id 5-1, and RepoBench-P, id 6-2). As shown in Figure 6, for early-middle layers, NAMMs tend to focus on retaining global information such as the task preamble and key words throughout the text. Instead, for later layers, NAMMs seem to forget many of these tokens, whose information has likely been already incorporated in the previous layers, allowing the transformer to focus more on tokens with more detailed local information. Furthermore, in coding tasks, we find that the pruned tokens are mostly contiguous, corresponding to whitespace, comments, and whole segments of boilerplate code. This is in contrast to natural language tasks, where NAMMs appear trying to exploit some of the grammatical redundancies of the English syntax often dropping specific tokens mid-sentences.

Figure: Figure 7: Density plot of gradient magnitudes for each score with respect to all memory tokens (left), together with a qualitative analysis extracting slices from three tokens (center) and computing the dot products of the gradients with the scored-token’s feature vector (right).
Refer to caption: extracted/5936268/figures/visualizations/analysis_bam_fig_noSP.png

Backward attention cross-token interactions.
To analyze cross-token interactions learned through our BAM architecture, we record the gradients of each token score $s_{i}$ with respect to all input features $v_{j}$ for all tokens in memory after storing 1024 tokens, i.e., for $j=1,2,\dots,1024$. We denote these quantities as $\nabla g_{i}^{j}=\partial s_{i}/\partial v_{j}$.
We provide a qualitative visualization of our results on the PassageRetrieval-en task for a randomly selected layer and prompt in Figure 7. On the left subplot, we provide a visualization of the squared magnitudes $(\nabla g_{i}^{j})^{T}\nabla g_{i}^{j}$ for each combination of tokens (either scored or attended upon in BAM, i.e., indexed by $i$ or $j$). Here, the effects of the backward mask are clearly visible, allowing tokens to exclusively attend to later ones, where $i>j$. Predictably, these magnitudes mostly peak on the subplot’s diagonal, indicating the self-influence that each token’s features have on its corresponding output score. However, there are also notable exceptions, as shown in the center subplot, where we overlap three slices from our left surface plot corresponding to the gradients of the first, together with the highest and lowest-scored tokens in memory (respectively indexed by $i=$0, 292, and 800). We provide additional directional information of each gradient vector from these slices in the right subplot, where we take its dot product with the scored token’s own feature vector $(\nabla g_{i}^{j})^{T}v_{i}$. After the first notable spike, at $i=j$, most other dot-product spikes with the largest magnitudes consistently have negative values. Hence we can logically deduce that the scores of these tokens would benefit from pushing the representations of future tokens away from their own. This result appears to validate the hypothesis that BAM learns a mechanism for cross-token competition, incentivizing diversity and promoting tokens covering unique frequencies in the attention spectrogram.

Additional analysis. We provide additional analytic results in Appendix D. For instance, we compare the generated responses before and after the introduction of NAMMs in a very long context task, and show the sensitivities of the token scores for each input feature. These results illustrate how NAMMs learn to overcome different failure modes of long context LMs, and that they mainly rely upon high-frequency components of the attention matrices, further evidencing the need to go beyond simple strategies and the potential of end-to-end learning for token-level memory systems.

## 5 Related works

Devoto et al. () and Yao et al. () try to identify the least important tokens to evict using heuristics such as L2 magnitude and entropy. Alternative strategies include considering simple statistics from the attention matrix (Liu et al., ; Oren et al., ; Zhang et al., ). Ge et al. () and Li et al. () build on these ideas by applying multiple strategies based on matching specific attention patterns. Unlike this prior work, our approach uniquely employs a black-box model to learn token-level memory management and shows potential for improving both the performance and efficiency of transformers. We refer to App. E for references and connections to the wider literature, including efficient architectures, memory, and evolution.

## 6 Discussion and future work

This work introduced Neural Attention Memory Models, providing a new framework to enhance the performance of transformers while significantly reducing memory footprint.
By evolving NAMMs on top of pre-trained LMs, we demonstrated their effectiveness across diverse long-context tasks in three languages, significantly surpassing previous hand-designed KV cache eviction frequently hindering performance, and the original model relying on costly full-context conditioning.
Our carefully designed approach also enabled NAMMs, trained solely on language tasks, to achieve zero-shot transferability across architectures, input modalities, and task domains.
This work has only begun to explore the design space of our memory models, which we anticipate might offer many new opportunities to advance future generations of transformers.
In this regard, we believe NAMMs should not be viewed as a replacement for gradient-based optimization, but rather an orthogonal framework that could be combined and alternated with parameter fine-tuning. Such an extension has the potential to unlock efficient long-context training, drawing parallels to the iterative process of learning and evolution that shaped human memory.

## 7 Author contributions

Edoardo Cetin initiated the project, led the design and implementation of NAMMs, and provided major contributions to writing. Qi Sun designed and implemented the zero-shot transfer experiments with Llama 3 70B and Llava Next Video 7B, and provided contributions and feedback to writing. Tianyu Zhao devised and implemented ChouBun, and provided contributions and feedback to writing.
Yujin Tang coordinated the project, gave key advice for the design of NAMMs, and provided major contributions to writing.

## Acknowledgements

The authors would like to thank David Ha and Llion Jones for providing valuable discussions during the early stages and feedback while drafting the text. This paper is based on results obtained from a project, JPNP20017, subsidized by the New Energy and Industrial Technology Development Organization (NEDO).

## Appendix A Implementation details

### A.1 Model specifics and NAMMs execution

We evolve our Neural Attention Memory Models on top of a context-extended Llama 3 8B (Dubey et al., ) base model. In particular, we employ the NTK-aware positional interpolation strategy (bloc97, ) to extend the context by four times from 8192 to 32768. Unlike prior strategies that require further gradient fine-tuning to avoid performance collapse (Chen et al., ), NTK-aware positional interpolation has been shown to produce sensible results even when applied zero-shot. In case the length of a task prompt still exceeds 32768 we perform mid-sentence cropping (Xiao et al., ; Jin et al., ), as standard in long-context LM evaluation (Bai et al., ; Zhang et al., ).

When applying NAMMs, we only affect the execution of the base model with a fixed frequency, once every $n_{up}=512$ steps. When feeding longer prompts to our model, we simply split the tokens into $n_{up}$-sized chunks. We note that due to modern frameworks being bound primarily by memory constraints, input-splitting in itself has minimal effects on running time, with similar approaches being already performed under the hood by established kernel procedures (Dao et al., ).

### A.2 Feature extraction and architecture details

**Table 9: NAMMs hyper-parameters used for training and evaluation. The omitted CMA-ES hyper-parameters can be obtained by following the recommended default calculation by Hansen ().**
| NAMMs hyperparameters |  |
| --- | --- |
| Spectrogram window size $n_{w}$ | 32 |
| Spectrogram window stride $s_{w}$ | 16 |
| Spectrogram window type | Hann |
| Spectrogram EMA reduction coefficient $\gamma$ | $0.99^{16}$ |
| Positional features | 8 |
| NAMMs execution delay | 512 |
| NAMMs non-linearity | ReLU |
| Optimization hyperparameters, notation from Hansen () |  |
| Evolution algorithm | CMA-ES |
| Elite ratio | 0.5 |
| Mean coefficient $c_{m}$ | 1 |
| Initial step size $\sigma$ | 0.65 |
| Samples batch size per-task | 64 |
| Population size | 32 |
| Task for incremental stage 1 | PassageRetrieval-en |
| Task for incremental stage 2 | DuReader |
| Task for incremental stage 3 | NarrativeQA |
| BAM-specific |  |
| Hidden dimensions | 16 |
| Use bias | True |
| Masking strategy | counter-causal |
| Number of attention layers | 1 |
| Number of final linear layers | 1 |
| Use residual connections | True |
| Use multiplicative interactions | True |
| MLP-specific |  |
| Hidden dimension | 25 |
| Number of hidden layers | 2 |
| Use residual connections | True |

Our new feature extraction framework is a key component for enabling the transfer properties of NAMMs. In practice, we extract the attention spectrogram from the real-valued attention matrix using a Hann window of size $n_{w}=32$, resulting in just seventeen complex-values frequencies that we convert to real numbers by simply taking their magnitude, yielding each $\omega_{i}^{t}\in\mathbb{R}^{17}$. We use a stride of half the window size $s_{w}=16$, producing $n_{T}=n_{up}/s_{w}=32$ frequency representations over the time axis of the attention matrix from the latest chunk of $n_{up}$ queries, $\omega_{i}^{1},\dots,\omega_{i}^{n_{T}}$. Thus, we reduce these frequency representations over the time axis via an element-wise exponentially moving average operation. We note that our EMA does not only consider the $n_{T}$ representations computed for the frequency of each token in the $n_{up}$-sized chunk of the latest queries, but also the discounted EMA at the previous execution step or our memory for each retained token, denoted $\omega_{i}^{\prime}$. Thus, each of our reduced spectrogram representations reflects the full history of previous attention values:

$$ $\omega_{i}=\left(\sum^{n_{T}}_{t=1}\gamma^{t-1}\omega^{t}_{i}\right)+\gamma^{n _{T}}\omega_{i}^{\prime},$ (4) $$

where we use $\gamma$ to denote the EMA’s discount factor. To expedite learning the weights of our architecture, we ensure all spectrogram features have unit variance at initialization across our training data, using the statistics of the base Llama 3 model computed on the first task employed in incremental learning (PassageRetrieval). Finally, we also concatenate a small eight-dimensional sinusoidal positional embedding using the oldness of each token, i.e., the amounts of new queries observed since its introduction in the KV cache.

Figure: Figure 8: Schematic depiction of the components of our Neural Attention Memory Models, denoted $m_{\phi}$, parameterized with our BAM architecture. The spectrogram representation of each token, denoted $\omega_{i}$, is processed by an attention layer followed by a simple linear operation to output its relative score. Backward masking introduces asymmetry, ensuring that each token can only attend to its future relatives.
Refer to caption: extracted/5936268/figures/visualizations/bam_app_det.png

Our backward-attention memory network processes these representations by directly first applying the self-attention layer employing the counter-autoregressive backward masking introduced in Section 3, designed to facilitate asymmetric interactions between tokens in memory. The output of self-attention is then fed to a single final linear layer to obtain the final score. We employed a few important additional design choices following some preliminary testing. First, motivated by efficiency considerations, we use a single head within our attention mechanism and no layer normalization. Second, our attention layer produces outputs that are twice the dimensionality of the spectrogram features. These outputs are integrated back into the main network before the final linear layer via both residual and multiplicative interactions. We provide a schematic depiction of our minimal architecture in Figure 8. Through our minimalist design choices, our full network comprises only just over four thousand learnable parameters, a negligible amount, orders of magnitudes lower than even a single layer in modern transformers.

### A.3 Zero-shot transfer

For our zero-shot transfer experiments, we consider a Llama 3 transformer with 70B parameters (Dubey et al., ), a Llava Next Video transformer with 7B parameters (Zhang et al., ), and a decision transformer (Chen et al., ) with about 1M parameters. For our 70B experiments, we follow the exact same setup as when evaluating our 7B Llama model used in training. For our video-language model, we extract $12\times 12$ image tokens from 48 uniformly sampled frames, 6912 in total. We also slightly shift the selection score threshold by 5, to counteract the lower number of total tokens and get a comparable average cache size to the L2 and H2O baselines. We adapt the code and follow the standardized experimental setup from Li et al. (). For the reinforcement learning experiments, we encode each state, action, and return-to-go into separate tokens and do not apply any restrictions or modifications to our standard NAMM LM setup. We average the performance collected over 20 random seeds to account for the stochasticity of the initial state in the Gym Mujoco environments (Brockman et al., ). Our RL experiments adapt the checkpoints and setup provided by Beeching & Simonini ().

## Appendix B ChouBun details

**Table 10: Statistics of ChouBun. Lengths are counted by tokens produced by Llama 3 tokenizer.**
| Statistics | Extractive QA | Summarization | Overall |  |  |
| --- | --- | --- | --- | --- | --- |
| JA.WikiQA | JA.EdinetQA | JA.CorpSecQA | JA.CorpSecSum | All tasks |  |
| Number of documents | 20 | 20 | 30 | 30 | 70 |
| Number of QA pairs | 200 | 390 | 150 | 30 | 770 |
| Number of reference answers | 1 | 1 | 1 | 5 | 1 or 5 |
| Document length max. | 13027 | 10152 | 85981 | 85981 | 85981 |
| Document length mean | 10131 | 8994 | 26220 | 26220 | 13317 |
| Document length min. | 8196 | 6825 | 5640 | 5640 | 5640 |
| Answer length max. | 40 | 208 | 30 | 140 | 208 |
| Answer length mean | 7 | 11 | 8 | 80 | 21 |
| Answer length min. | 1 | 1 | 1 | 55 | 1 |

**Table 11: Performance of a wider range of LLMs on the ChouBun benchmark.**
| Model/Task name | Extractive QA | Summarization | Overall |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| JA.WikiQA | JA.EdinetQA | JA.CorpSecQA | JA.CorpSecSum | All tasks | Max. length |  |
| mistralai/Mistral-7B-v0.1 | 8.68 | 8.34 | 16.25 | 10.50 | 10.94 | 32768 |
| rinna/llama-3-youko-8b | 16.68 | 12.23 | 17.03 | 22.27 | 17.05 | 8192 |
| meta-llama/Meta-Llama-3-8B | 14.58 | 14.77 | 16.86 | 22.84 | 17.27 | 8192 |
| meta-llama/Llama-2-7b-hf | 16.77 | 9.92 | 20.86 | 21.97 | 17.38 | 2048 |
| 01-ai/yi-6b-200k | 30.36 | 23.64 | 38.09 | 21.11 | 28.30 | 200000 |
| elyza/Llama-3-ELYZA-JP-8B | 20.77 | 21.45 | 35.59 | 40.21 | 29.50 | 8192 |

The ChouBun benchmark is created to assess the generalization ability of NAMMs to a new language (Japanese), but we hope it will also serve as a standard benchmark for Japanese LLMs. The benchmark is composed of two task categories — extractive QA and abstractive summarization — and four tasks as follows.

- •
JA.WikiQA is an extractive QA task about 20 randomly sampled articles from the 20240429 dump of Japanese Wikipedia(^4^44[https://dumps.wikimedia.org/other/cirrussearch/](https://dumps.wikimedia.org/other/cirrussearch/)). Each article corresponds to 10 QA pairs, and there are 200 QA pairs in total.
- •
JA.EdinetQA is an extractive QA task based on 20 security reports from EDINET(^5^55[https://disclosure2.edinet-fsa.go.jp/](https://disclosure2.edinet-fsa.go.jp/)). The EDINET security reports are in CSV format, which makes them less human-readable. Nevertheless, we choose not to convert the format because the conversion process per se is non-trivial, and using a CSV-style text input helps us evaluate a model’s capability of understanding structured data. The total number of QA pairs in JA.EdinetQA is 390.
- •
JA.CorpSecQA is another extractive QA task based on 30 security reports downloaded from three corporation websites (MUFG(^6^66[https://www.mufg.jp/ir/report/security_report/](https://www.mufg.jp/ir/report/security_report/)), NTT(^7^77[https://group.ntt/jp/ir/library/results/](https://group.ntt/jp/ir/library/results/)), and Toyota(^8^88[https://global.toyota/jp/ir/library/securities-report/](https://global.toyota/jp/ir/library/securities-report/))). We extract texts from original file in PDF format. There are 150 QA pairs in total.
- •
JA.CorpSecSum is an abstractive summarization task based on the same data of JA.CorpSecQA. Each document corresponds to one data point, and we collect 5 reference summaries for each data point.

Collecting human annotations for long-text tasks is challenging, therefore we use synthetic QA pairs and summaries. In particular, we prompt various LLMs(^9^99gpt-4o-2024-05-13, gpt-4o-mini-2024-07-18, gpt-4-turbo-2024-04-09, and claude-3-5-sonnet-20240620) to generate multiple question-answer pairs or summaries for each document. Different instructions are designed for the two tasks and they are shown in Figure 9. To improve the reliability of the synthetic data, we ensure that every answer in extractive QA tasks is a text span presented in its corresponding source document. In Table 10, we provide the statistics of the benchmark.

We use F1 score and ROUGE score for evaluation in the extractive QA tasks and summarization task, respectively. Reference text and hypothesis text are pre-tokenized by the MeCab tokenizer(^10^1010[https://github.com/polm/fugashi](https://github.com/polm/fugashi)). A wider range of LLMs’ performance on the ChouBun benchmark is presented in Table 11.

Figure: Figure 9: LLM prompts for generating synthetic QA pairs and summaries in ChouBun.
Refer to caption: x1.png

## Appendix C Additional results

### C.1 Performance across incremental stages and architectures

We provide additional results and analysis to the summarized one, complementing Section 4, with the detailed performance across different NAMMs, evaluating the best checkpoints after each stage of incremental training stage, and ablating the BAM architecture with an MLP.

**Table 12: NAMMs evaluation on LongBench (Bai et al., ). The normalized performance (in brackets) is calculated using the base model with full cache. The aggregate test task performance of NAMMs models is taken by averaging the normalized scores on the tasks not used for incremental evolution.**
| Model/Task id | Single-Doc QA | Multi-Doc QA | Summarization |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 1-1 | 1-2 | 1-3 | 1-4 | 2-1 | 2-2 | 2-3 | 2-4 | 3-1 | 3-2 | 3-3 | 3-4 |  |
| Base model | 10.38 (1.00) | 12.79 (1.00) | 22.60 (1.00) | 21.31 (1.00) | 10.41 (1.00) | 12.67 (1.00) | 7.54 (1.00) | 25.86 (1.00) | 29.34 (1.00) | 23.93 (1.00) | 0.92 (1.00) | 2.66 (1.00) |
| NAMM (MLP, s1) | 7.60 (0.73) | 12.74 (1.00) | 22.74 (1.01) | 21.08 (0.99) | 9.58 (0.92) | 12.24 (0.97) | 6.48 (0.86) | 19.41 (0.75) | 27.76 (0.95) | 23.61 (0.99) | 0.95 (1.03) | 3.44 (1.29) |
| NAMM (MLP, s2) | 6.76 (0.65) | 12.77 (1.00) | 23.74 (1.05) | 20.56 (0.96) | 9.69 (0.93) | 12.21 (0.96) | 6.93 (0.92) | 22.40 (0.87) | 27.30 (0.93) | 24.20 (1.01) | 1.72 (1.87) | 2.78 (1.05) |
| NAMM (BAM, s1) | 5.77 (0.56) | 12.76 (1.00) | 22.94 (1.02) | 21.55 (1.01) | 9.47 (0.91) | 12.21 (0.96) | 6.51 (0.86) | 18.73 (0.72) | 28.06 (0.96) | 23.97 (1.00) | 1.01 (1.10) | 4.00 (1.50) |
| NAMM (BAM, s2) | 7.08 (0.68) | 12.70 (0.99) | 22.21 (0.98) | 21.50 (1.01) | 9.94 (0.95) | 12.21 (0.96) | 7.13 (0.95) | 20.34 (0.79) | 28.87 (0.98) | 23.84 (1.00) | 0.92 (1.00) | 3.94 (1.48) |
| NAMM (BAM, s3) | 9.14 (0.88) | 12.63 (0.99) | 21.94 (0.97) | 21.34 (1.00) | 9.71 (0.93) | 11.63 (0.92) | 6.98 (0.93) | 20.58 (0.80) | 28.78 (0.98) | 24.39 (1.02) | 1.04 (1.13) | 3.63 (1.36) |
| Model/Task id | Few-shot Learning | Synthetic | Code | Overall |  |  |  |  |  |  |  |  |
| 4-1 | 4-2 | 4-3 | 4-4 | 5-1 | 5-2 | 5-3 | 6-1 | 6-2 | All tasks | Test tasks | Cache size |  |
| Base model | 73.00 (1.00) | 89.45 (1.00) | 46.54 (1.00) | 40.00 (1.00) | 1.48 (1.00) | 12.18 (1.00) | 28.80 (1.00) | 69.09 (1.00) | 65.17 (1.00) | 28.86 (1.00) | N/A | 32768 (1.00) |
| NAMM (MLP, s1) | 73.00 (1.00) | 89.48 (1.00) | 46.80 (1.01) | 37.50 (0.94) | 2.46 (1.66) | 23.98 (1.97) | 28.46 (0.99) | 69.75 (1.01) | 66.40 (1.02) | 28.83 (1.05) | 1.01 | 7639 (0.23) |
| NAMM (MLP, s2) | 74.00 (1.01) | 88.64 (0.99) | 46.04 (0.99) | 41.50 (1.04) | 1.53 (1.03) | 25.94 (2.13) | 29.78 (1.03) | 69.80 (1.01) | 65.23 (1.00) | 29.22 (1.07) | 1.02 | 8475 (0.26) |
| NAMM (BAM, s1) | 73.00 (1.00) | 89.81 (1.00) | 46.70 (1.00) | 38.75 (0.97) | 2.19 (1.48) | 25.14 (2.06) | 28.51 (0.99) | 69.50 (1.01) | 66.51 (1.02) | 28.91 (1.05) | 1.00 | 7951 (0.24) |
| NAMM (BAM, s2) | 73.00 (1.00) | 90.03 (1.01) | 46.85 (1.01) | 42.00 (1.05) | 2.35 (1.59) | 24.69 (2.03) | 28.46 (0.99) | 69.65 (1.01) | 66.57 (1.02) | 29.25 (1.07) | 1.04 | 8267 (0.25) |
| NAMM (BAM, s3) | 73.00 (1.00) | 89.81 (1.00) | 46.35 (1.00) | 40.00 (1.00) | 3.04 (2.05) | 27.55 (2.26) | 28.60 (0.99) | 69.53 (1.01) | 66.35 (1.02) | 29.33 (1.11) | 1.07 | 8155 (0.25) |

**Table 13: NAMMs evaluation on InfiniteBench (Zhang et al., ). The normalized overall performance (in brackets) is calculated using the average performance of the base model with full cache.**
| Model/Task name | Retrieval | Dialogue | Novel | Math | Code | Overall |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Ret.PassKey | Ret.Number | Ret.KV | En.Dia | En.Sum | En.MC | En.QA | ZH.QA | Math.Find | Code.Run | Code.Debug | All tasks | Cache size |  |
| Base model | 0.00 | 0.00 | 0.00 | 1.00 | 7.73 | 0.00 | 1.05 | 1.79 | 0.00 | 0.00 | 0.00 | 1.05 (1.00) | 32747 (1.00) |
| NAMM (MLP, s1) | 0.00 | 10.00 | 0.00 | 3.00 | 7.27 | 3.93 | 1.57 | 4.26 | 0.57 | 0.00 | 3.30 | 3.08 (2.93) | 11329 (0.35) |
| NAMM (MLP, s2) | 10.17 | 11.86 | 0.00 | 2.50 | 7.48 | 3.06 | 1.58 | 4.10 | 1.71 | 0.00 | 1.52 | 4.00 (3.80) | 13031 (0.40) |
| NAMM (BAM, s1) | 9.49 | 9.83 | 1.80 | 0.50 | 14.36 | 37.12 | 8.95 | 16.20 | 5.71 | 1.50 | 6.09 | 10.14 (9.63) | 11173 (0.34) |
| NAMM (BAM, s2) | 11.86 | 11.86 | 1.80 | 1.00 | 14.62 | 35.37 | 8.96 | 15.45 | 0.57 | 1.75 | 4.31 | 9.78 (9.29) | 12789 (0.39) |
| NAMM (BAM, s3) | 11.86 | 11.86 | 1.80 | 1.00 | 14.91 | 36.24 | 8.78 | 17.67 | 10.57 | 1.75 | 4.57 | 11.00 (10.45) | 13192 (0.40) |

**Table 14: NAMMs evaluation on the new ChouBun benchmark. The normalized performance (in brackets) is calculated using the base model with full cache.**
| Model/Task name | Extractive QA | Summarization | Overall |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| JA.WikiQA | JA.EdinetQA | JA.CorpSecQA | JA.CorpSecSum | All tasks | Cache size |  |
| Base model | 22.91 (1.00) | 28.34 (1.00) | 11.83 (1.00) | 21.75 (1.00) | 21.21 (1.00) | 12099 (1.00) |
| NAMM (MLP, s1) | 21.60 (0.94) | 26.81 (0.95) | 10.34 (0.87) | 29.60 (1.36) | 22.09 (1.04) | 9525 (0.79) |
| NAMM (MLP, s2) | 20.76 (0.91) | 26.30 (0.93) | 11.86 (1.00) | 29.32 (1.35) | 22.06 (1.04) | 9815 (0.81) |
| NAMM (BAM, s1) | 19.19 (0.84) | 28.85 (1.02) | 14.36 (1.21) | 28.51 (1.31) | 22.73 (1.07) | 9569 (0.79) |
| NAMM (BAM, s2) | 20.75 (0.91) | 28.46 (1.00) | 14.55 (1.23) | 32.45 (1.49) | 24.05 (1.13) | 9867 (0.82) |
| NAMM (BAM, s3) | 21.34 (0.93) | 28.61 (1.01) | 14.64 (1.24) | 33.15 (1.52) | 24.44 (1.15) | 9895 (0.82) |

Extended language modeling results. We report our results for LongBench, InfiniteBench, and ChouBun in Tables 12, 13, 14. First, we note that even training on a single task with our simple MLP architecture impressively improves performance across all benchmarks. Additionally, performance across benchmarks sees near-monotonic further improvements with each stage of our incremental evolution recipe. Comparing our implementations, we note that the performance benefits from the memory models with backward attention are consistently superior to the fully connected variant in both initial stages of incremental training, empirically validating our hypothesis about the importance of global KV cache information for determining the importance of each token. Lastly, on ChouBun. we observe that the performance with BAM sees a notable upswing after the second stage of incremental training, which might be associated with the introduction of another ideogram-based language in the training set.(^11^1111The DuReader task, used in the second stage of incremental training, uses the Chinese language.) The same improvement not occurring with the MLP-based NAMMs might be further evidence of architectural performance saturation, highlighting once again the effectiveness of our main implementation design.

**Table 15: NAMMs evaluation on LongBench (Bai et al., ) with a Llama 3 70B model. The normalized performance (in brackets) is calculated using the base model with full cache. The aggregate test task performance of NAMMs models is taken by averaging the normalized scores on the tasks not used for incremental evolution. The tasks on which NAMMs are trained are highlighted with a gray background.**
| Model/Task id | Single-Doc QA | Multi-Doc QA | Summarization |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 1-1 | 1-2 | 1-3 | 1-4 | 2-1 | 2-2 | 2-3 | 2-4 | 3-1 | 3-2 | 3-3 | 3-4 |  |
| Base model | 9.38 (1.00) | 13.84 (1.00) | 24.99 (1.00) | 17.78 (1.00) | 11.73 (1.00) | 14.26 (1.00) | 8.11 (1.00) | 26.43 (1.00) | 13.13 (1.00) | 24.55 (1.00) | 23.20 (1.00) | 10.08 (1.00) |
| NAMM (MLP, s1) | 6.94 (0.74) | 13.82 (1.00) | 24.27 (0.97) | 17.60 (0.99) | 10.83 (0.92) | 14.17 (0.99) | 7.89 (0.97) | 20.81 (0.79) | 13.09 (1.00) | 23.30 (0.95) | 23.28 (1.00) | 8.66 (0.86) |
| NAMM (MLP, s2) | 7.88 (0.84) | 13.71 (0.99) | 23.27 (0.93) | 18.18 (1.02) | 11.41 (0.97) | 14.11 (0.99) | 8.07 (0.99) | 21.75 (0.82) | 14.28 (1.09) | 24.48 (1.00) | 22.00 (0.95) | 8.99 (0.89) |
| NAMM (BAM, s1) | 7.31 (0.78) | 13.75 (0.99) | 24.51 (0.98) | 17.78 (1.00) | 10.82 (0.92) | 14.08 (0.99) | 7.59 (0.94) | 19.27 (0.73) | 13.89 (1.06) | 23.71 (0.97) | 23.41 (1.01) | 8.87 (0.88) |
| NAMM (BAM, s2) | 3.57 (0.38) | 13.86 (1.00) | 23.02 (0.92) | 18.71 (1.05) | 4.94 (0.42) | 13.32 (0.93) | 1.90 (0.23) | 17.74 (0.67) | 10.39 (0.79) | 20.45 (0.83) | 23.18 (1.00) | 8.13 (0.81) |
| NAMM (BAM, s3) | 9.13 (0.97) | 13.53 (0.98) | 24.25 (0.97) | 17.82 (1.00) | 11.45 (0.98) | 13.76 (0.96) | 8.34 (1.03) | 21.79 (0.82) | 12.66 (0.96) | 24.21 (0.99) | 23.56 (1.02) | 8.62 (0.86) |
| Model/Task id | Few-shot Learning | Synthetic | Code | Overall |  |  |  |  |  |  |  |  |
| 4-1 | 4-2 | 4-3 | 4-4 | 5-1 | 5-2 | 5-3 | 6-1 | 6-2 | All tasks | Test tasks | Cache size |  |
| Base model | 78.00 (1.00) | 92.43 (1.00) | 48.67 (1.00) | 45.50 (1.00) | 22.50 (1.00) | 75.37 (1.00) | 33.89 (1.00) | 74.60 (1.00) | 71.19 (1.00) | 35.22 (1.00) | N/A | 10107 (1.00) |
| NAMM (MLP, s1) | 78.00 (1.00) | 92.28 (1.00) | 48.37 (0.99) | 43.50 (0.96) | 20.76 (0.92) | 68.66 (0.91) | 33.89 (1.00) | 74.58 (1.00) | 71.68 (1.01) | 34.11 (0.97) | 0.99 | 7930 (0.78) |
| NAMM (MLP, s2) | 77.00 (0.99) | 91.93 (0.99) | 48.60 (1.00) | 44.75 (0.98) | 17.17 (0.76) | 70.21 (0.93) | 36.18 (1.07) | 74.72 (1.00) | 71.30 (1.00) | 34.29 (0.97) | 0.99 | 8445 (0.84) |
| NAMM (BAM, s1) | 77.50 (0.99) | 92.46 (1.00) | 48.24 (0.99) | 45.00 (0.99) | 17.32 (0.77) | 69.87 (0.93) | 33.89 (1.00) | 74.58 (1.00) | 72.40 (1.02) | 34.11 (0.97) | 0.99 | 7947 (0.79) |
| NAMM (BAM, s2) | 74.50 (0.96) | 51.45 (0.56) | 39.73 (0.82) | 15.00 (0.33) | 5.86 (0.26) | 13.35 (0.18) | 34.29 (1.01) | 73.81 (0.99) | 61.91 (0.87) | 25.20 (0.72) | 0.79 | 8276 (0.82) |
| NAMM (BAM, s3) | 78.50 (1.01) | 92.36 (1.00) | 48.49 (1.00) | 45.50 (1.00) | 19.07 (0.85) | 74.19 (0.98) | 34.28 (1.01) | 74.71 (1.00) | 72.42 (1.02) | 34.70 (0.99) | 0.99 | 8365 (0.83) |

**Table 16: Evaluation on the LongVideoBench and MLVU benchmarks with Llava Next Video 7B. The normalized performance (in brackets) is calculated using the base model with full cache.**
| Model/Task name | LongVideoBench | MLVU | All tasks | Cache size |
| --- | --- | --- | --- | --- |
| Base model | 43.45 (1.00) | 44.23 (1.00) | 43.84 (1.00) | 7039 (1.00) |
| NAMM (MLP, s1) | 39.64 (0.91) | 41.24 (0.93) | 40.44 (0.92) | 584 (0.08) |
| NAMM (MLP, s2) | 41.06 (0.94) | 39.72 (0.90) | 40.39 (0.92) | 713 (0.10) |
| NAMM (BAM, s1) | 41.06 (0.94) | 41.98 (0.95) | 41.52 (0.95) | 723 (0.10) |
| NAMM (BAM, s2) | 45.03 (1.04) | 44.23 (1.00) | 44.63 (1.02) | 4948 (0.70) |
| NAMM (BAM, s3) | 44.58 (1.03) | 44.18 (1.00) | 44.38 (1.01) | 5100 (0.72) |

**Table 17: NAMMs evaluation on D4RL (Fu et al., ) using a Decision Transformer model (Chen et al., ; Beeching & Simonini, ). The normalized overall performance (in brackets) is calculated using the average performance of the base model with full cache.**
| Model/Task name | Hopper-v3 | Walker2d-v3 | HalfCheetah-v3 | Overall |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Medium | Med-Replay | Expert | Medium | Med-Replay | Expert | Medium | Med-Replay | Expert | All tasks | Cache size |  |
| Base model | 33.36 (1.00) | 18.37 (1.00) | 44.62 (1.00) | 68.21 (1.00) | 7.18 (1.00) | 38.98 (1.00) | 34.91 (1.00) | 5.06 (1.00) | 10.64 (1.00) | 29.04 (1.00) | 3000 (1.00) |
| NAMM (MLP, s1) | 33.01 (0.99) | 18.39 (1.00) | 38.09 (0.85) | 70.82 (1.04) | 7.25 (1.01) | 44.61 (1.14) | 35.64 (1.02) | 5.05 (1.00) | 10.87 (1.02) | 29.30 (1.01) | 1993 (0.66) |
| NAMM (MLP, s2) | 33.48 (1.00) | 19.24 (1.05) | 30.07 (0.67) | 73.22 (1.07) | 7.95 (1.11) | 48.21 (1.24) | 33.59 (0.96) | 5.81 (1.15) | 14.67 (1.38) | 29.58 (1.02) | 2834 (0.94) |
| NAMM (BAM, s1) | 35.02 (1.05) | 18.24 (0.99) | 45.95 (1.03) | 69.33 (1.02) | 7.91 (1.10) | 44.45 (1.14) | 34.25 (0.98) | 5.12 (1.01) | 13.68 (1.29) | 30.44 (1.05) | 2009 (0.67) |
| NAMM (BAM, s2) | 35.18 (1.05) | 18.79 (1.02) | 48.08 (1.08) | 71.97 (1.06) | 7.70 (1.07) | 49.74 (1.28) | 35.67 (1.02) | 5.78 (1.14) | 10.82 (1.02) | 31.53 (1.09) | 2534 (0.84) |
| NAMM (BAM, s3) | 36.10 (1.08) | 18.86 (1.03) | 49.39 (1.11) | 70.87 (1.04) | 7.53 (1.05) | 50.02 (1.28) | 34.56 (0.99) | 5.90 (1.17) | 12.34 (1.16) | 31.73 (1.09) | 2434 (0.81) |

Extended zero-shot transfer results. We report our extended zero-shot transfer results for the 70B model and the offline RL setting in Tables 15, 16, and 17. We see the benefits from NAMMs again increase as we incorporate backward attention, and with each stage of incremental training to a similar extent as with the language modeling tasks. These results further highlight the potential benefits of scaling up the architecture of our memory model and increasing the number of incremental stages. To this end, given the generality of our parameterization, an interesting unexplored approach could be to incorporate different base models and input modalities during evolutionary training, something that would substantially increase problem diversity to obtain an even more robust transfer behavior.

### C.2 Training curves with fully-connected NAMMs

Figure: Figure 10: Mean and standard deviation over the CMA-ES population batch performance (left), together with the performance of the learned mean parameter on each task (right) for the training of the MLP NAMM.
Refer to caption: extracted/5936268/figures/results/training_curves_mlp.png

In Figure 10, we provide training curves of our Neural Attention Memory Model using a simple MLP architecture rather than backward attention, evaluated in Section 4. In the left sub-plot, we show the average and standard deviation of the normalized batch performance across the population, while in the right sub-plot, we show the normalized per-task and average performance on all samples of the optimized mean from CMA-ES. When compared with the BAM training curve from Figure 4, we note a few interesting differences, although its evaluation performance on the full LongBench benchmark is lower across both incremental phases (see Table 2), both its population batch performance and the CMA-ES full-task performance on the training sets are either comparable or slightly higher than BAM’s. This dichotomy appears to indicate that cross-token interactions might provide a better inductive bias, mitigating the overfitting potential of NAMMs.

### C.3 Evolution of memory size during training

Figure: Figure 11: Final memory size of NAMM parameterized by the learned mean of CMA-ES for both the BAM (left) and the MLP implementations (right).
Refer to caption: extracted/5936268/figures/results/training_curves_cs.png

In Figure 11, we provide training curves for the evolution of the memory size collected at the end of each task prompt of our NAMMs. On the left and right subplots, we provide results for the BAM and MLP implementations, respectively. For both architectures, we find that the memory size generally increases with training. This result suggests that NAMMs might learn to recognize additional valuable tokens as training progresses, enabling the corresponding performance improvements on the training tasks. Hence, they might indicate that there is some degree of a trade-off between the efficiency and performance of NAMMs. However, we note that both models are trained only for performance maximization, without any incentive to be more conservative. To this end, exploring regularization strategies to make NAMMs aware of deployment costs is an interesting direction for future work to obtain tailored sweet spots to cater to instance-specific resource constraints.

### C.4 Incremental training ablation

Figure: Figure 12: Mean and standard deviation over the CMA-ES population batch performance (left), together with the performance of the learned mean parameter on each task (center) and its final memory size for the NAMM trained without incremental evolution.
Refer to caption: extracted/5936268/figures/results/training_curves_abl.png

**Table 18: NAMMs incremental learning (IL) ablation evaluation on LongBench (Bai et al., ). The No IL baseline is trained from scratch on both the PassageRetrieval-en and DuReader tasks, the same employed during the second stage of incremental learning.**
| Model/Task id | Single-Doc QA | Multi-Doc QA | Summarization |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 1-1 | 1-2 | 1-3 | 1-4 | 2-1 | 2-2 | 2-3 | 2-4 | 3-1 | 3-2 | 3-3 | 3-4 |  |
| Base model | 10.38 (1.00) | 12.79 (1.00) | 22.60 (1.00) | 21.31 (1.00) | 10.41 (1.00) | 12.67 (1.00) | 7.54 (1.00) | 25.86 (1.00) | 29.34 (1.00) | 23.93 (1.00) | 0.92 (1.00) | 2.66 (1.00) |
| NAMM (BAM, s1) | 5.77 (0.56) | 12.76 (1.00) | 22.94 (1.02) | 21.55 (1.01) | 9.47 (0.91) | 12.21 (0.96) | 6.51 (0.86) | 18.73 (0.72) | 28.06 (0.96) | 23.97 (1.00) | 1.01 (1.10) | 4.00 (1.50) |
| NAMM (BAM, s2) | 7.08 (0.68) | 12.70 (0.99) | 22.21 (0.98) | 21.50 (1.01) | 9.94 (0.95) | 12.21 (0.96) | 7.13 (0.95) | 20.34 (0.79) | 28.87 (0.98) | 23.84 (1.00) | 0.92 (1.00) | 3.94 (1.48) |
| NAMM (BAM, s3) | 9.14 (0.88) | 12.63 (0.99) | 21.94 (0.97) | 21.34 (1.00) | 9.71 (0.93) | 11.63 (0.92) | 6.98 (0.93) | 20.58 (0.80) | 28.78 (0.98) | 24.39 (1.02) | 1.04 (1.13) | 3.63 (1.36) |
| NAMM (BAM, no IL) | 6.46 (0.62) | 12.72 (0.99) | 22.87 (1.01) | 21.22 (1.00) | 9.91 (0.95) | 11.77 (0.93) | 5.61 (0.74) | 18.94 (0.73) | 27.63 (0.94) | 22.60 (0.94) | 0.91 (0.99) | 1.75 (0.66) |
| Model/Task id | Few-shot Learning | Synthetic | Code | Overall |  |  |  |  |  |  |  |  |
| 4-1 | 4-2 | 4-3 | 4-4 | 5-1 | 5-2 | 5-3 | 6-1 | 6-2 | All tasks | Test tasks | Cache size |  |
| Base model | 73.00 (1.00) | 89.45 (1.00) | 46.54 (1.00) | 40.00 (1.00) | 1.48 (1.00) | 12.18 (1.00) | 28.80 (1.00) | 69.09 (1.00) | 65.17 (1.00) | 28.86 (1.00) | N/A | 10107 |
| NAMM (BAM, s1) | 73.00 (1.00) | 89.81 (1.00) | 46.70 (1.00) | 38.75 (0.97) | 2.19 (1.48) | 25.14 (2.06) | 28.51 (0.99) | 69.50 (1.01) | 66.51 (1.02) | 28.91 (1.05) | 1.00 | 8205 |
| NAMM (BAM, s2) | 73.00 (1.00) | 90.03 (1.01) | 46.85 (1.01) | 42.00 (1.05) | 2.35 (1.59) | 24.69 (2.03) | 28.46 (0.99) | 69.65 (1.01) | 66.57 (1.02) | 29.25 (1.07) | 1.04 | 8521 |
| NAMM (BAM, s3) | 73.00 (1.00) | 89.81 (1.00) | 46.35 (1.00) | 40.00 (1.00) | 3.04 (2.05) | 27.55 (2.26) | 28.60 (0.99) | 69.53 (1.01) | 66.35 (1.02) | 29.33 (1.11) | 1.07 | 8409 |
| NAMM (BAM, no IL) | 73.00 (1.00) | 89.28 (1.00) | 46.43 (1.00) | 38.75 (0.97) | 2.49 (1.68) | 29.28 (2.40) | 28.46 (0.99) | 69.80 (1.01) | 64.77 (0.99) | 28.79 (1.03) | 0.98 | 8457 |

We provide a full set of ablations results for our incremental training strategy, training a Neural Attention Memory Model with the BAM architecture from scratch on both the PassageRetrieval-en and DuReader tasks, as employed during the second stage of incremental learning. We evolve this Neural Attention Memory Model for 360 consecutive generations and provide training curves in Figure 12. In the left sub-plot, we show the average and standard deviation of the normalized batch performance across the population, in the center sub-plot, we show the normalized per-task and average performance on all samples of the optimized mean from CMA-ES, and on the right subplot we show the corresponding memory size. Furthermore, in Table 18, we provide the full LongBench evaluation results for this baseline, also showing our original incremental model’s performance for ease of comparison. Interestingly, the non-incremental NAMM obtained a notably higher score on the training tasks with a normalized performance of 1.57, in contrast to the normalized performance of 1.41 achieved by the best checkpoint from the second incremental training stage. Yet, outside the PassageRetrieval-en and DuReader tasks, its performance is notably inferior and very close to the original performance of the base model. These results appear to indicate that the usefulness of incremental training goes beyond the faster evolution provided by reducing the number of evaluation prompts to assess performance and that this strategy plays an important role in regularizing evolution and making Neural Attention Memory Models effectively generalize to new tasks.

## Appendix D Additional analysis

Figure: Figure 13: Distribution of gradient magnitudes for the token scores with respect to all the seventeen features in our attention spectrogram representations. In the rightmost-lower subplot, we also compare the total magnitudes of the frequency information with the recency information in the positional embeddings.
Refer to caption: extracted/5936268/figures/visualizations/analysis_per_feature.png

### D.1 Sensitivity to attention frequencies and positional encodings

We analyze the magnitudes of the gradients of the token scores $s_{i}$ with respect to each dimension in the token feature vectors. This procedure quantifies how varying each dimension in our attention spectrogram representation locally affects the output score of NAMMs, thus, providing a heuristic measure of its relevance (since scores determine which tokens get discarded). In Figure 13, we plot the distribution of magnitudes for all the seventeen features up to the Nyquist frequency (0 0 to $16$) in the attention spectrogram. All frequency distributions seem to cover a wide range of values, with each mean being close to the global mean, seemingly indicating NAMMs learn to make use of all available spectrogram information for at least some of the tokens. Additionally, we note that many of the higher frequencies have distributions with higher means and larger tails than the ‘ground frequency’ at dimension 0. Furthermore, as shown in the rightmost-lower subplot, NAMMs appear visibly less sensitive to recency information provided by the concatenated positional embeddings, with a lower total influence than frequency information on token scores. Overall, these observations seem to further validate the importance of going beyond simple hand-designed methods solely based on token recency and the sum of the attention values, which has so far been considered a strong established recipe for KV cache management (Oren et al., ; Zhang et al., ; Ge et al., ; Devoto et al., ).

### D.2 InfiniteBench results comparison

On the InfiniteBench tasks, our NAMM achieve particularly outstanding improvements over the base model and other baselines, with an over ten-fold score increase (from 1.05% to 11%). However, we note that even with NAMMs, the performance of Llama 3 8B still lags considerably behind the performance of powerful LMs designed specifically for long-context problems, as reported in Zhang et al. (). Nonetheless, on the En.Sum task, concerned with the summarization of fictitious novels, we find our main NAMM brings the performance of the context-extended Llama 3 from 7.73 to 14.91 even slightly beyond GPT4’s (14.73). While this performance is still low in absolute terms(^12^1212InfiniteBench tasks are scored in a range between 0 and 100.), such a result appears quite notable and suggests that improvements from NAMMs are orthogonal in nature to the ones brought by architectural improvements and scaling, which, by themselves, might be insufficient to address the challenges brought by long and noisy contexts.

Figure: Figure 14: Qualitative examples comparing the ground produced responses by Llama3 with and without our NAMM memory, together with GPT4, on two prompts from the En.Sum task part of InfiniteBench.
Refer to caption: extracted/5936268/figures/visualizations/infinibench_examples.png

We qualitatively inspect the effects of NAMMs on En.Sum by comparing example answers generated by Llama 3 with and without our memory models, together with examples generated by GPT4. As illustrated in Figure 14, we find both the Llama and GPT models to incur several failure modes, producing answers that entirely miss the objective of the original task. For instance, the context-extended Llama 3 often gets stuck in generation loops continuously repeating part of sentences without coherent structure. Instead, the GPT answers appear to forego summarizing the text and rather attempt to continue the provided passage, by generating end-of-text tokens or even roleplaying some of the characters. However, while introducing NAMMs appears to avoid many instances of these failure modes, we find the summarization of the memory-augmented Llama 3 still displays many imperfections such as misspelling character names (left) or lacking much depth by being extremely concise (right).

## Appendix E Extended related works

Similar to our NAMMs implementation, memory management through token eviction has been explored mostly to reduce memory constraints and enable querying LMs with longer contexts (Luohe et al., ).
Commonly, strategies entail simply cropping input prompts to a shorter length, often more effective when done from the middle rather than the ends (Xiao et al., ; Jin et al., ).
More advanced, several heuristic strategies have been proposed to identify and evict the least important tokens in the KV cache, selectively pruning it to a fixed size for each layer.
These strategies assess token relevance using metrics like L2 magnitude (Devoto et al., ) or entropy (Yao et al., ), or analyze statistics from the attention matrix, such as value magnitude or cumulative sums (Liu et al., ; Oren et al., ; Zhang et al., ).
Building on these ideas, Ge et al. () and Li et al. () apply multiple strategies simultaneously, choosing the best fit for each layer by matching them with specific attention patterns.
However, unlike previous work, our approach uniquely employs a black-box model to learn KV cache management, aiming to enhance efficiency and boost performance.

Many other methods to reduce memory consumption, affecting the KV cache, are mostly orthogonal and likely complementary to our approach.
For instance, MQA (Shazeer, ) and GQA (Ainslie et al., ) propose merging different attention heads during the training of LLMs, either fully or partially, to improve deployment-time throughput.
Brandon et al. (), pushed these strategies further, attempting to merge heads even across different layers.
GQA is commonly employed in many modern LMs, including the LLama 3 family of models which we use to train and evaluate NAMMs on language tasks (Dubey et al., ).
Furthermore, several methods have looked at KV cache compression through either quantization of the keys and values (Hooper et al., ; Dong et al., ; ) or even the whole hidden states (DeepSeek-AI et al., ).
Similarly to the aforementioned prior work concerning KV cache pruning, these methods considered mainly hand-designed strategies, such as employing different quantization rates based on heuristically recognizing important tokens.
We note that using evolution to optimize for which channels to merge or compress could also yield new interesting unexplored approaches, combining these orthogonal directions with some of the principles introduced by NAMMs.

There has also been much research interest in exploring new architectures to explicitly model components of a memory system or to address key challenges of reasoning over longer contexts.
For instance, past work has looked at incorporating neural models of memory within neural networks by implementing different reading and writing operations - either directly replacing their layers (Weston et al., ; Sukhbaatar et al., ), or introducing new auxiliary components (Rae et al., ; Lample et al., ).
In relation to transformers, more recent works have been proposed rethinking the ingredients of the self-attention operation, mostly in the context of LMs.
These works looked at either efficient linear approximation to self-attention to overcome quadratic costs (Beltagy et al., ; Katharopoulos et al., ; Wang et al., ; Peng et al., ), or introducing new kinds of persistent tokens and storage to extend information propagation (Dai et al., ; Munkhdalai et al., ; Hwang et al., ).
However, as also noted by Dao et al. (), none of these methods and approximations have managed to replace standard approaches so far.
We take a different approach that can be integrated in a zero-shot manner even without any fine-tuning.

Lastly, methodologically related to NAMMs, there have been other prior methods making use of evolution for or with transformer models.
For example, Tang & Ha () also trained a small attention-based model through evolution, exploiting the inherent parameter efficiency behind these operations.
Furthermore, So et al. () proposed using evolution to meta-optimize the basic building of transformers via neural architecture search, while Akiba et al. () focused on evolving different merging strategies across layers belonging to LMs with different capabilities.
As for these works, we note that evolution plays a critical role for NAMMs, allowing us to directly optimize for target performance and overcome the inherent non-differentiability underlying our new framework.

## Appendix F Limitations and future extensions

### F.1 Exploring the design space of Neural Attention Memory Models

In this work, we introduced Neural Attention Memory Models and showed their efficacy and potential to improve the performance and efficiency of transformers, even when evaluated zero-shot for unseen architectures and domains. However, given the novelty of our framework, we note that our design choices were mostly motivated by simplicity and practicality rather than quantitative empirical evidence. Thus, there is an extremely large design space in terms of the implementation, training, and deployment of these models that should be explored beyond this work, which is likely to yield further improvements.

For instance, while our current feature extraction, based on computing the spectrogram of the attention matrix, enables capturing global frequency information about the attention values of each token, it might fall short of modeling local information with enough granularity. This hypothesized limitation inherently comes from a few design choices we made with the purpose of limiting the input size and corresponding parameter count of our memory models. In particular, our spectrogram features only consider the real components of a short-time Fourier transform with a small Hann window of size thirty-two. Thus, we only provide NAMMs information about a relatively limited number of thirty-two frequencies, losing any notion of the phase of the attention matrix that would be captured by the full complex-valued Fourier coefficients. Consequently, the representations of tokens with high attention values for entirely non-overlapping queries occurring with the same frequency would be indistinguishable to our models. Moreover, our exponentially moving average reduction over the time dimension of the spectrograms provides an additional layer of heavy compression inevitably trading off expressivity for simplicity.

To partially address these concerns, an alternative design we explored entailed delaying the initial element-wise exponentially moving average reduction. Concretely, this involved computing $T$ different scores, feeding $m_{\phi}$ all feature vectors $\omega_{i}^{t}$ for $t=1,2,\dots,T$, across the attention spectrogram’s compressed time axis, only then reducing the resulting scores $s^{1:T}_{i}$ via EMA. While, in principle, this alternative ordering would allow for additional expressivity without adding to the parameter count, in practice, when evaluated with an initial version of the simple 2-layer MLP model, we found no significant performance difference and opted for the former lighter option. However, introducing cross-token interactions with the improved BAM design and further scaling is likely to introduce a need of re-evaluating this choice.

One further limitation comes from the current reliance on the exact values of the attention matrix. This reliance precludes NAMMs training from making use of fast kernel algorithms developed to accelerate inference by foregoing materializing attention values (Dao et al., ). While the main focus of this work has been to introduce NAMMs and display its potential to improve transformers across different domains, more scalable parameterizations and efficient backend integrations remain exciting open challenges for future research.

### F.2 Improving long-context sparse retrievals

**Table 19: NAMMs evaluation on the canonical Needle In A Haystack task (Kamradt, ). The normalized performance (in brackets) is calculated using the base model with full cache.**
| Model/Task name | Needle prompt length | Overall |  |  |  |
| --- | --- | --- | --- | --- | --- |
| 0-10000 | 10001-20000 | 20001+ | All prompt lengths | Cache size |  |
| Base model (full cache) | 8.87 (1.00) | 7.53 (1.00) | 3.50 (1.00) | 6.32 (1.00) | 32768 |
| NAMM (BAM) | 9.00 (1.02) | 4.80 (0.64) | 3.05 (0.87) | 5.36 (0.85) | 10208 |
| NAMM (BAM), $\gamma=0.9999^{s_{w}}$ | 9.00 (1.02) | 5.33 (0.71) | 3.45 (0.99) | 5.68 (0.90) | 10347 |

One notable example exemplifying some of the aforementioned limitations, comes from the canonical Needle In A Haystack task (Kamradt, ), which has been used to qualitatively evaluate LLMs for their ability to remember sparse information over long noisy horizons. We provide results on this task using the best-performing NAMM after three stages of incremental training with the BAM architecture, averaging evaluation scores provided by a GPT-4 model (Achiam et al., ) across different prompt ranges, consistently with Bai et al. (). As shown in Table 19, while NAMMs do not manage to exceed the overall performance of the base model, they still provide some notable efficiency gains. However, looking more closely at the score distribution across different prompt length ranges we observe an unexpected trend that is in contrast with the rest of our results on other benchmarks. In particular, while our NAMM obtains slightly higher than the base model for prompts with a size less than 10000, it seems to increasingly struggle with longer prompts.

After comparing the spectrogram features extracted for the different prompts, our explanation for these results highlights one current failure mode of the current implementation. In particular, the Needle In a Haystack task is constructed such that the model is tasked to remember some important information introduced at the beginning of the prompt, and later followed by completely unrelated ‘filler’ text. Hence, the attention scores and the corresponding spectrogram features for the tokens containing the relevant information are forcibly sparse, being high only at the very beginning of the prompt. Yet, since the evaluated NAMM reduces these features over the time axis of the spectrogram with an EMA coefficient of $\gamma=0.99^{s_{w}}$, all the frequency information regarding these tokens will be inevitably overwritten. To empirically validate our theory we provide results simply raising the EMA coefficient from $\gamma=0.99^{s_{w}}$ to $\gamma=0.9999^{s_{w}}$. Since our NAMMs was never actually trained with this higher coefficient, we note that this change effectively brings the input features out-of-distribution. Nonetheless, as shown in the final row of Table 19, the larger coefficient still manages to improve performance on the longer prompts by enabling the preservation of the frequency components from the target ‘needle’ over a longer horizon. These findings suggest that future NAMM designs should consider higher EMA reduction coefficients or, potentially, even directly learning this parameter with evolution in addition to the NAMM’s network weights.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. Paper seems to imply they transform the attention map after applying the causal mask. Is this true? If so, what is STFT doing exactly? Is it doing anything meaningful? 
2. How important is it to use the  fourier transformation to the signals instead of the original attention map?
3. I found the intuition for the gradient analysis section fairly confusing, and I am not sure what the takeaway is or why the quantities they measure matter. I wish this was better motivated, or it could be moved to the Appendix.
soundness: 3
strengths: The algorithmic design seems sound and seems to perform well based on their results. It’s interesting that the auxiliary model can then transfer to other models and modalities. I liked some of their ablation studies, such as showing that NAMM drops more tokens from later layers than earlier layers. This seems to imply that there is much more to learn from attention scores than what manually-designed algorithms do.
summary: They propose a learned strategy for maintaining the size of your KV cache. An auxiliary model utilizes the attention map to produce a score of how important each key-value pair is for predicting future tokens. This model is small (1-layer transformer). They train this auxiliary model using the attention maps of Llama 3 8B for long-context language tasks, but show that this auxiliary model is compatible with other models (Llama 3 70B) and can transfer to other modalities (video). The KV cache management reduces the memory of the cache, but performance does not decrease by much, sometimes improving upon using the full cache.
weaknesses: There are a couple problems with the paper presentation (missing citations) + rigor of their experiments. I highly recommend that these problems are addressed for the score to be raised. 


There are some problems with the flow of the paper due to many of the details being relegated to the Appendix. 
- The related works paragraph in the main body is a bit too short, and poses NAMM to be the first method to propose a learned black-box strategy. However, previous works already do this. For example, Pyramid-BERT [1] does exactly that, and DMC [2] proposes a continual pre-training objective. 
- Please include limitations of the method in the main body or refer to the Appendix section where you do somewhere. 
- Limitations should also address the computational complexity of this method. How slow is NAMM compared to other methods? 
- There should be an algorithm block in the main body detailing exactly how NAMM is used during inference. The written description is too vague. Especially include details such as the 512-chunking that is detailed in the Appendix. 
- There are no details of exactly how the auxiliary model is trained in the main body or the appendix, except for a sentence mentioning that the auxiliary model optimizes the performance of 3 benchmarks. Please detail the exact objective somewhere. 


Weaknesses in experiments:
- They should report the computational complexity of NAMM, how much longer is it to do inference using NAMM versus other baselines? In other words, how expensive is it to generate scores from the auxiliary model for each attention layer? 
- A result that they particularly emphasize in this paper is that models do better on long-context benchmarks with NAMM cache management over utilizing the full KV cache. That makes sense, the auxiliary model seems to pick out the most important tokens as the aux model is _trained on long-context benchmarks itself_. However, I worry that the significance of these results is not rigorously tested and may simply be because they do not use any fine-tuning strategy to adapt the  base model to long-context tasks, as it is standard to do in practice.
  - All the reported results apply NAMM on top of a Llama-3-8B model that they adapt to use long-context using “NTK-aware positional interpolation,” a zero-shot method that can be found in a Reddit post. I am not aware of the validity or the limitations of this approach. They should also apply NAMM on top of standard long-context models already out there, such as Mistral 7b with 32k context. Does NAMM still do significantly better? 
   - NAMM is only compared against L2 and H2O, but it should also compare against other finetuning strategies e.g. DMC, and strategies that also utilize the attention scores explicitly such as FastGen [3].
- The degree to which each baseline prunes the KV cache varies significantly, so it is hard to tell whether NAMM has better performance on benchmarks simply because it does not prune the cache to the degree the other baselines do. It would be good to actually control the hyperparameters of each method, such as the threshold at which a key-value pair is pruned, such that performance between methods is compared for some fixed cache size. Ideally, there should be a plot depicting this with x axis as cache size and y axis as performance. 


1. Pyramid-BERT: Reducing Complexity via Successive Core-set based Token Selection https://arxiv.org/pdf/2203.14380
2. Dynamic Memory Compression: Retrofitting LLMs for Accelerated Inference 
3. Suyu Ge,  et al.   Model tells you what to discard:  Adaptive KV cache compression for LLMs.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Please see the weaknesses.
soundness: 3
strengths: 1. The idea behind the contribution is simple and easy to understand.
2. NAMMs can be zero-shot transferred to other Transformers, which is quite interesting.
3. The benchmarks are diverse (but not comprehensive as discussed in weakness 2).
summary: Transformers are known to suffer from quadratic memory blow-ups. This paper attempts to make Transformers more efficient by using Neural Attention Memory Models (NAMMs). This involves pruning the KV cache memory using NAMMs.
weaknesses: 1. Pruning is lossy, meaning that it results in losing some of the information in the KV cache memory. For example, during the pruning process, you may very well discard information that the model could have found useful later on.
2. KV cache memory pruning (more broadly KV cache management) has been investigated in the last couple of years, and while the authors compare against some of them, like H2O and L2, they do not discuss or compare against some of the other relevant work. Some of the relevant work (in language modelling) include:
   - Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs, ICLR 2024
   - CacheGen: KV Cache Compression and Streaming for Fast Large Language Model Serving, SIGCOMM 2024

I appreciate that comparing against so many related works can be cumbersome. However, it is crucial that the authors discuss the related work and either compare their work against them or discuss why they are not comparable.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: The cache size is important for complexity. However, in practice, what we actually care about is the empirical runtime and memory complexities. What is the runtime and memory needed to process new tokens compared to that of the baselines (e.g., H2O)?

Evolution is typically expensive to train using. What are the computational resources used to train NAMMs in terms of the runtime and memory?

Considering that language models get significantly larger than 8B, it is important that the idea generalizes to larger settings. How would the resources needed to train and perform inference scale with the size of the base models?

When pruning, the size of the cache size affects the performance significantly. In practice, we may want to select an optimal cache size that trades off performance and computational resources. In this paper, only 1 cache size was chosen for each dataset. How was this cache size determined for each model? And why are the results reported for only a single cache size? 

Relatedly, in H2O, plots are shown with the x-axis being the cache size (ratio) and y-axis being the performance. As a result, we can see the performance trade-off with respect to the cache size. Could you elaborate on why the results were not shown the same way? And could you include these plots?

Could you include details in the Appendix regarding the various datasets (e.g., Long Bench, InfiniteBench, LongVideoBench, MLVU, D4RL etc...)?

The paper runs experiments on D4RL using Decision Transformers. However, the reported results are significantly worse than what is reported in many papers (e.g., original DT paper, Online DT, Aaren, etc..). Could you elaborate on why this is the case? Could you add details regarding the experimental setup that resulted in these numbers?

For the LongBench results, the overall numbers are reported as "All tasks". However, a subset of the tasks are used for training NAMM. Reporting "All tasks" seems unfair for evaluating generalization. Instead, could you report separate aggregate scores for (1) the tasks used in training and (2) the held-out tasks not used in training, to more clearly show generalization performance."



Interestingly, using NAMM improves over even the base model in several cases. What is your intuition as to why there's such a large improvement in some cases?
soundness: 2
strengths: The performance gains using NAMMs over H2O are significant. 

Improving the performance using NAMMs over the base model is impressive.
summary: The paper proposes Neural Attention Memory Models (NAMMs), a network learned via evolution (CMA-ES) that can be applied to Transformers, reducing the number of tokens needed (KV-cache) at inference time. 
Results show that NAMMs outperform recent methods for pruning the KV-cache and occasionally improves the performance of the base model.
weaknesses: H2O has some advantages over NAMMs. Notably, H2O does not require any new training. In contrast, NAMMs are trained using evolution. As such, NAMMs may require significantly more compute. The paper, however, lacks information regarding the computational resources needed by the method.

Furthermore, in practice, we care about the cache size and inference computational complexity. It is unclear how (1) the performance of NAMMs fare when varying the cache size and (2) the complexity (runtime and memory) of applying NAMMs at inference time.

Several details are missing to understand if the idea will generalize to other models. In this paper, only Llama 3 8B was considered. However, in contrast, the H2O paper considered several architectures (e.g., OPT,
LLaMA, and GPT-NeoX), showing their idea generalizes across architectures. However, in this case, it's unclear how NAMMs performance will generalize across different language models. 

The results on Decision Transformer (DT) are odd considering that several papers (e.g., the original DT and Online DT papers) showed results that are significantly better than what is reported in Table 7. 

Benchmarks in the paper are mentioned by name (e.g., Long Bench, InfiniteBench, LongVideoBench, MLVU, D4RL etc...) without concrete details (e.g., types of tasks in the benchmarks' datasets). It is unclear to readers who are unfamiliar with the various fields.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1 Removing tokens with importance scores <0 implies that the number of removed tokens may vary each time. 

 2.How can the consistency of dimensions be guaranteed? Without consistent dimensions, subsequent training would be impossible. 

3.Does this require an additional constraint? I feel this point was not clearly explained in the paper.

4.In my understanding, compression is performed before assessing the importance of tokens. Would it be possible to compare the effects without compression? It seems that compression also has its advantages.

5.Could you list specific learning rates, parameter settings, and other details? Publishing the experimental code would also be beneficial.
soundness: 2
strengths: 1. The paper is well-structured with a clear logical flow, making it easy for readers to follow. Figures are appropriately placed alongside the relevant text, contributing to a harmonious overall layout. The experimental results are effectively highlighted and processed, allowing readers to easily grasp the key points. This is a significant strength of the paper.

2. The NAMMs method proposed in the paper achieves a reduction in model parameter size, which facilitates training.

3. The architecture introduced by the authors is straightforward and easily reproducible. Extensive experiments have been conducted to substantiate its effectiveness, further validating the claims made in the paper.
summary: This paper focuses on optimizing the memory management of Transformers. The authors introduce the NAMMs (Neural Attention Memory models) method, which employs trainable parameters to learn the importance of token column vectors. By removing less important tokens, this approach reduces resource requirements and enhances efficiency. The effectiveness of the proposed method is validated through experiments conducted on multiple datasets.
weaknesses: 1 My understanding is that the method involves the fusion of column vectors for all tokens (3.1 169）. However, if the dimensionality is too high, could this lead to substantial additional computational time? Please provide some computational complexity analysis or empirical runtime comparisons.

2 Removing tokens with importance scores <0 implies that the number of removed tokens may vary each time. How can the consistency of dimensions be guaranteed? Without consistent dimensions, subsequent training would be impossible. Does this require an additional constraint? I feel this point was not clearly explained in the paper. (How is the following issue implemented？）
 a.How they handle varying numbers of removed tokens in practice
 b.If there are any mechanisms to ensure dimension consistency
 c.Whether any additional constraints are used 

3 In my understanding, compression is performed before assessing the importance of tokens. Would it be possible to compare the effects without compression? It seems that compression also has its advantages. Please conduct an ablation study specifically comparing performance with and without the compression step.

4.Could you list specific learning rates, parameter settings, and other details? Publishing the experimental code would also be beneficial. (Please provide the following materials.)
 a.A table of hyperparameters used for each experiment
 b.Any details on hyperparameter tuning processes
 c.Information on where to access their code or if they plan to release it This would give readers a clearer picture of how to reproduce the results.

Meta-review: Post rebuttal, all of reviewers vote for acceptance. The AC checked all the materials and concurs that the paper has done a valuable exploration of using evolution algorithms to compress the KV cache based Transformer memory for efficient decoding. The paper received initial mixed reviews, but the concerns have been thoroughly addressed during the rebuttal period, with authors paying great efforts providing new results and re-writing the draft. With these efforts, the paper can be accepted. Please incorporate necessary changes in the final version.

Decision: Accept (Poster)