An Evolved Universal Transformer Memory
Authors: Edoardo Cetin
Qi Sun
Tianyu Zhao
Yujin Tang
Venue: ICLR 2025 Poster (2025)
Keywords: Transformers
Evolution
Memory
KV cache
attention
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose a new memory system designed for transformers able to achieve performance improvements in language tasks, and zero-shot transferring to different architectures and input modalities.

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.

Published as a conference paper at ICLR 2025
AN EVOLVED UNIVERSAL TRANSFORMER MEMORY
Edoardo Cetin, Qi Sun, Tianyu Zhao, Yujin Tang Sakana AI, Japan {edo, qisun, tianyu, angiintanggssakana ai
ABSTRACT
Prior methods propose to offset the escalating costs of modern foundation mod- els by dropping specific parts of their contexts with hand-designed rules, while attempting preserve their original performance. We overcome this trade-off with Neural Attention Memory Models (NAMMs), introducing learned network for memory management that improves both the performance and efficiency of transformers. We evolve NAMMs atop pre-trained transformers to provide differ- ent latent contexts focusing on the most relevant information for individual layers and attention heads. NAMMs universally applicable to any model using self- attention as they condition exclusively on the values the produced attention ma- trices. Learning NAMMs small set of problems, formance improvements across multiple long-conten benchmarks while cutting the model's input contexts up fraction of the original sizes. We show the gen- erality 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 https:: gtttbbgggttt/tt///////// com/ http///aaaaa//-eeeeaaaaae
achieve substantial per-
INTRODUCTION
Transformer architectures have be- come the golden standard in deep learning, with ubiquitous applica- tions in the design of modern foun- dation models, exhibiting exceptional performance and scalability (Achiam
cool
Language
cool
very
Evolution
Evolution
Vision
2023 Das et al.] 2023 Team et al.
2023 Dosovitskiy et al. 2020
Chen et al. [202aa] Brohan
al.
2023 Gur et al.) (023). The out- puts of a transformer are exclusively conditioned on recent context of in- put tokens, which for language mod- els (LMs) generally correspond window of preceding words. Thus, addressing the challenge of extend- ing this context window is critical enable tackling long-range tasks and is currently a focal area of re- search (Huang et al., (023). However, long contexts also immediately impact training and inference costs, with modern foundation models being increasingly hsuurryehungrry 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. 2024).
RL
Figure 1: NAMMs use evolution to optimize the perfor- mance 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.
Our research aims to go beyond these hand- desi-nd strategies as we hypothesize that shaping the latent memory KV cache of transformers entails new opportunities to improve their capabilities in

---

Published as a conference paper at ICLR 2025
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. (023)) 3) even allowing foundation models to learn in-context entirely new skills test time (Brown et al. 2020).
Furthermore, unlike prompt engineering, directly managing the memory of transformers enables the provisioning of distinct contexts 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), intro- ducing a new class of networks trained with evolution to learn an efficient memory system that maxi- mizes the downstream performance of pre-trained transformers. Evolution inherently overcomes the non non-differentiabiliyy of memory management operations with binary outcomes (selecting tokens to preserve/discard) which renders gradient- baseed 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 1987 Nairne& & Pandeirada,
2010
Frankland & Bontempi 2005).
Our NAMMs are conditioned on features entirely constructed from the attention matrix, making them universally applicable to any architecture. Learning NAMMs atop pre-trained Llama 3 8B model (Dubey et al.) 2024), 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
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
transformer-based
Base model
28.86
32768
1.05
32747
21.21
12099
H2O
28. auto
8192 1229
1.05
8193 (0:2)
19.86 3u40
8292 0199
(1i0)
8192 0229)
1.63 (U9))
8193 (022)
18.93 1.90
8292 (us))
FastGen NAMMs
27.88 auso
9538 1::4
1.42
23016 mom
18.93 3u:00
8616 ou/))
29.33 (L.))
8155 (233)
03.4)
13192 0.49)
24.44 (.15)
9895 (1s)
ModeUEval
Llama 3 70B
Computer Vision
Reinforvement Learning
Performance Cache size Performance Cuche size Performance Cache size
model
35.22 (100))
10107 c100
(100)
70.39 (100)
29.0 (1.0
3000 (1:00
H2O
34,17
6662 ma
41.97 a0s0
4479
28.70 70
2048 ov9
33.50 over)
6662 (199
41.45 avrr)
4479 (166))
27.91 9140
2048 (19))
validate these findings across 36 different tasks from LongBench (Bai
NAMMs
34,70
8365
44.38 (1.r
5100
31.73
2434 o1st)
et al.) 2023). InfiniteBench (Zhang et al. 2024a), and ChouBurl 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 three natural language tasks entirely new transformer models. In particu- lar, 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 architec- tures concerned with new input modalities, for problems such as vision and reinforcement learning. In nutshell, our main technical contributions can be summarized as the following:
We introduce NAMMs, 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, ob- taining 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.
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 set of latents with the same length within their layers. Multi-headed dot product attention (Vaswani et al. 2017). or simply self-attention, characterizes modern transformers, facilitating effective information sharing across
CChouBun is the pronunciation of E
x", literally translating to "long text" in Japanese.
2

---

Published as a conference paper at ICLR 2025
Spectrogram extraction
EMA reduction
Score computation
A
connection
8AM connection
xy² y xx
xy
m& (w) (w)
0.1
1
0.4
0.2
0.4
0.0
0.2
0.0
0.1
0.5
0.6
mm (as1)
my
0.0
ii
10
0.0
0.2
0.1
0.2
1
0.0
0.0
0.2
0.3
0.0
0.1
0.3
0.2
0.9
m (w)(+1)
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-win 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).
token representations. The attention layer conducts set of parallel computations, each known as an attention head, mapping tokens to query, key, and value vectors E R.. These vectors are organized along the sequence dimension in the matrices Q, K, and V, and the layer's output is computed as:
A = softmax (M QKT Vd
attention (Q, K, AV, where,
Here, M represents optional mask multiplying the attention matrix A, usually enforcing 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 AJ, i.e., the dot products between each key i and query 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- feature extraction. An established canonical technique to pre-process one- dimensional non-stationary signals is the Short-Time Fourier Transform (STFT) (Allen & Rabiner, 1977]. 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 signal, shifting each convolutional window to the frequency domain through discrete Fourier transform, producing a spectrogrum representation of the original input. We use w' E IR N+1 to denote the fixed-sized vector produced each timestep f, where the N frequencies span from zero up to the Nyquist frequency (half the original sampling rate). Mathematically, the nith frequency from an STFT for time is extracted from an input vector E E as:
w/ [m]]
Σ ultTw(t l'le TR
Here, the convolutional filter of the SFTF is defined by the product of finite-length window functiunnttion w with each exponential term the Fourier transform. A popular choice for w is the Hann window (Oppenheim] 1999), employing a smooth decay at its edges which helps minimize the overestimation of the magnitudes of the higher frequencies in w due to spectral leakage (Harris 1978).
3 NEURAL ATTENTION MEMORY MODELS
immediate limitation of transformers the quadratic costs associated with computing the atten- tion 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 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.
horizontal vector by exploiting causal ordering. In this section, we describe
3

---

Published as a conference paper at ICLR 2025
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 representation for each token in the current KV cache memory directly from its corresponding unmodified column vector in the attention matrix A. 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 w representing the frequencies with how the queries attend to each of the stored key tokens (indexed by i) on 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 compact form factor, discarding all other information specific to the learned transformer weights.
As NAMMs rely only 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- est-time to larger architectures and arbitrary applications.
3.2 MEMORY MODEL DESIGN AND CROSS-TOKEN COMMUNICATION
Causal atternioon maak
NAMMs parameterize a small neural network mø to output a scalar selection score s; (w!!") for each 8th token in the KV cache. First, to obtain consistent input dimension, we reduce the atten- tion spectrogram into smaller feature vector wi by compressing the time-axis via an element-wist exponentially moving average (EMA: Wi- = E, r'wf; Figure 2 center). We then append po- sitional encodings and feed the vector W; to the memory model's network m& to produce the score s, Finally, we evict from the KV cache memory all latent tokens with s; < O, effectively treating the problem as a binary classification task. We repeat this process with fixed interval, every set number of new input tokens, Rup.
P
"Bockward" ntnn mask
Backward attention memory models (BAM). For the design of mg. we posit that sharing information from all tokens in memory could be key for assessing their importance. A particularly motivat- ing scenario in LMs arises when considering the case of repeated words or sentences, where learning diversity measure that com- pares different tokens would allow preventing redundancies in the KV cache. Corroborating this intuition, even from biological per- spective, memory formation and retention appear to adhere to mod- els of neuronal competition (Han et al., 2007).
Ceweeoe Mop
Most recurt beyy
Our backward mask makes each token attend exclusively to its future rela- tives in the KV cache.
Figure
Based on these considerations, we design the backward attention memory architecture (BAM) for parameter-efficient sharing of in- formation while making use of the powerful inductive biases en- abled by the masked self-attention operation. In particular, plement m& via initial self-attention layer with a counter-causal mask M, which refer as backward (Figure This design serves to introduce a purposeful asymmetric relationship, allowing to distinguish between older and newer tokens. We then output s; from final linear operation:
im-
0 = attention so (Kn, Va, Qn), 8; = linear((x))
where Kn, Vo, Qn the key, value, and query matrices from all feature vectors wi 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).
4

---

Published as a conference paper at ICLR 2025
CMA-EB perforaace (sampled prompta)
CMA-ES
Best phase checkpoint
Best phase checkpoint
Dulleader
HerrativaQA
Normateed parforerca
Incremental phase
Incremental phase
100
200 400
500
500
100
303
400
500
genanationa
Figure 4: Mean and standard deviation over the CMA-ES population batch performance (left), to- gether with the performance of the learned mean parameter on each task (right).
In practice, when applying NAMMs, we only af-
Algorithm NAMMs
fect the KV cache of the base model with a fixed frequency, once every Trup steps. When feeding longer prompts to our model, we simply split the tokens into -Sied chunks. We summarize the full execution pipeline of NAMMs in Algorithm We refer to Appendix A and our shared code for additional implementation details and discussion.
Input: KV cache Ieeratonn i. Stride size Updane interval Atsetion matri A (ltest Ray queries). Past STFT w
1:
it
Pap == then
use NAMMs every Payy steps or column A,
2:
5t in
do
3:
STFT wf for O, ReT
Bu nr Bup / EMA reduction
4:
Reduce wa (2e 7 ww + T""
5:
mig (w;) Output: KV cache where >
apply NAMMs, Eq. p update the KV cache
3.3 INCREMENTAL EVOLUTION
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 al. 2023). As we share a single m across all layers, even with our largest NAMM we only evolve about 4000 total parameters. We use the seminal CMA-ES optimization algorithm (Hansen) 2006) and apply NAMM atop a Llama 3 8B base model (Dubey et al. 2024) with a context extended from 8192 to 32768 tokens via NTK-aware positional interpolation (bloc97 2023). Due to the inference costs of LMs with long inputs, we sample 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 generaliza- tion (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 max- imize 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. Using evolution, we note that our training loop simply corresponds to running inference NAMMs atop the base, requiring no expensive backpropagation
dedicated hardware.
We choose three tasks from different LongBench categories across both English and Chinese where the Llama base model seems to particularly struggle: PassageRetrieval-en DuReader, and Nar- rativeQA; optimizing the normalized exact match, ROUGE-L, and F1 metrics, respectively. 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 cach phase, 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 perfor- mance 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 three recent hand-designed methods for KV cache management: H20 (Zhang et al.] 2024c)
5

---

Published as a conference paper at ICLR 2025
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.
Son QA
Docti-Doc QA
Summarization
Model/Task id
1-1
1-2
1-3
1-4
2-1
2-2
2-3
2-1
3-1
3-2
3-3
3-1
Base model
10.38 (1m0 cl 22.60 (1)) 21:31 10.41 u.n
12,67
7.54 (ua) 25,86
0:30 29.34
23.93
0.92
2.66
H2O
8.75 man 13.07 (10 22. 22.11 ove 21.62 c1o 10.28 n09
12.40 y0e:
7:20 am 26.58
n.40 28.56
pm 23.98
D.88 m99
2.25 ns
1.2
8.83
13,13
22.22
21.79
9.97
12.15
5.88
24.96
nm 28,05
23.28
1,15
1.52
FastGien
10.49 asa
13.09 (Jo) 21.85 ovn) 21.57 cum 10.28 (00)
11.50 200
6.77
auy 17.04
0:50 28.98
(290 23.53 0
0.85 (.00)
3.02 0:30
NAMM (Ours)
9.14
12.63
m09 21.94 21.34
9.71
11.63 200 6.98
20.58
28.78
24.39
0.04
3.63
Model/Task ld
Few-shot Learning
Synthetic
Code
Overall
4-3
5:3
tasks Cache
Base model
73.00 c3an 89:45 46.54 (1R) 40,00 curs
1,48 core
12.18 pore 28.80 n. 69.09
65.17
1:90
86
N/A
10107 (:00)
H20
73.00 (1am) 90.03 cun 46.48 (10)) 34.00 (van)
2.18 (1i)
9.93 ovats 27.76 yiw 69.37
0.30 65.44
28.37 (.99
N/A
6662 (160
L.2
66.41 049 84. 92 mon 45 78 34. $4.38 009
3.13
11:00
28.6
73.45
0:00
55,20
(145) 27.42 t0
N/A
6662
FastGen
73.00 (1a0) mon (10)) (100
10.23 yaar: 27.01 w*1 69.34
64.50
(:99 22,88 (3
N/A
9538 (2
NAMM (Ours)
73.00 (3a0) 89.81 (100) 46.35 (1/00) 40.00 (uon)
3.04 (200)
27.55 (2:0) 28.60
(009) 69.53
0:30 66.35
(182)
29.33 (..1
1.07
8409 (230)
and L.2 (Devoto et al. 2024), and FastGen (Ge et al., 2024). 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 span- ning 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 transform- ers with tokens constructed from modalities other than language. Across all these settings, we also compare BAM with simpler 2-layer MLP architecture and provide summarized results after every stage of incremental evolutions. We refer to Appendix [ additional evaluations (e.g., transferring NAMMs to Mistral LM), ablation studies (e.g.. comparing different architectures and input fea- tures), and all learning curves. Lastly, we conclude the Section with a targeted qualitative analysis, aimed at understanding the behavior of our new memory framework.
4.1 LONG-CONTEXT LANGUAGE UNDERSTANDING
22
Longbench. In Table we provide results across all LongBench tasks (Bai et al. 2023) and in Fig- ure [ we provide a summarized comparison varying the maximum cache size of H2O and L2 (we pro- vide similar analysis for FastGen in Figure 9 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, H2O, 12, and Fastgen all come with performance costs which notably grow the smaller their cache sizes in line with their stated objective of retaining rather than improving the original full-context performance. These results emphasize the inevitable tradeoff induced by prior hand-designed methods, able to obtain efficiency gains but at increasing performance costs due to their lossy heuristics. On the other hand, we find NAMMs successfully provide a paradigm shift, yielding consistent improvements from the base model across both per- formance and efficiency axes by learning to discard unhelpful information, highlighting how end- to-end evolutionary optimization can open new orthogonal directions beyond what is feasible with manually- designed heuristics.
Arergepertomaa
IDC
MAVE (Ound)
Bose mosel IFull cache)
2046
6090
5192
28304
3270R
Madnum cache site
Figure 5: Comparing NAMM with H2O and L.2 while varying the cache size.
we provide results across the InfiniteBench tasks (Zhang et al. 2024a). 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. (2024)), even GPT4 (Achiam et al. 2023) cannot exceed performance of 1% on some of its problems. In line with these results, the coll-context Llama 3 together with H2O and 1.2 obtain near-zero performance on most tasks. Instead, our NAMM provides outstanding
InfiniteBench. In Table

---

Published as a conference paper at ICLR 2025
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 Bosaey Neteer RctKV
Ev.Dia
ExSiev
ELMC
E.
ZB
Maik Cisisfann
ColteDolay
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 (100
32747 (
H2O
0.00
0.00
0.00
1.50
5.38
0.00
1.71
1.71
0.25
0.00
1,05
8193
L.2
0.00
0.00
0.00
1,00
5.41
0.44
0.83
2.59
7.43
0.25
0.00
1.63
8193 (0231
FastGen
0.00
0.00
5.82
1.31
1.38
5,43
0.25
1.42
23016 (0:0)
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 0944)
13192 y0:rs
Table 5: NAMMs evaluation on LongBench with a Llama 3 70B model. The normalized perfor- mance (in brackets) is calculated using the base model with full cache.
Single-Doe QA
Multi-Doc QA
Summarization
Model/Task id
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 c.
13.84.0m 24.99 cux 17.78
11.73 nos
14.26 pse
8.11 mas
26,43 u.0 13.13 u.00 24.55 c. 23. 33.20
10.08 0.80
H20
8.80 osn 13.48 mom 25,02
18.44 (om
12.36 (m)
14,32 pum
8.15 0:0
26.22 p: 13.37 paa)) 24.50
in 23.20
9.22 (0:
L2
8.57 m09
40 mom 24.70
nos: 17.94 cum
12.77 a.
13.85 sun
7,13 auo
25.74 mmm 12.78 nam 23.21 (5 23.35 cpm
8:45 12
NAMM (Ouns)
9:13 msh
13.53 mns 24
17.82 (10 11.45 avs) 15.76 8.34 (1an) 21.79 pm 12,66 24,21
23.56 (100)
8.62
Few-shot Learning
Synthetle
Code
Overall
Model/Task id
4-1
4-2
4.3
4-4
5-1
5.2
5-3
6-1
6-2
tasles
Test tasks
Cache size
Base model
78.00 78.00 cmn 92.45 92.43.110s 48.67 cum 45.50 cum
22.50 a.m 75.37 now 33.89 n:00 74.60 u.n0 71.19
0.00 36.22 0.09
N/A
10107 0.50
H2O
77.50 a.m. 92.43 quo) 48.33
39.75 sum
18.12 (un)) use 33.89 n.0 74.61 0.30 71.09
0.00
34.17 (m
N/A
6662 (240
L2
76.50 m:9) 93.22 cum 46. oun: 36 nom 16.90 aus 61.34 mo: 36.28 um 74,38 0.30 67.43 (230)
33.50 13:
N/A
6662 (2:0)
NAMM (Ouns)
78.50
36
45.50
a00) 74.19 y0e 34.28
0:00 (1a0) 34.70
8365
improvements, bringing overall benchmark performance from 1.05% to 11% We also observe that
while our NAMM's memory size larger than for LengBench, it is considerably lower in relation to the base model's 's (now only 40%). This result suggests that NAMMs emergently learned scalable memory strategy, forgetting redundant and detrimental information at an increasing rate with longer contexts without requiring the hand-de hard cache limits enforced by 1.2 and H2O.
ChouBun. Our new benchmark focuses on tasks designed exclu- sively in Japanese, a novel language unseen during NAMMs training. We hope this benchmark might itself be valuable contribution to the research community, allowing the assessment of long-context capabilities in multilingual LLMs beyond the already--ubiuitous 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.I 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.
Table 4: Extractive QA
NAMMs evaluation on ChouBun. Summarization
Overall
Model/Task
JA. WiKIA
EdinerQA JA, CorpSecCAA
JA. CorpSecSion
All tasks Cache size
Base model
22.91 po:
28.34 (1:00
11.83 (100
21.75 (100
21.21 (10 12099 (100
H2O 1.2
20.76
26.39 mas
10:42 puso
21.87 (1.0
19.86 mo4) 8292 mmm
19.60 ps:
24.06 man
8.23 mmm
23.83 c.mm
18.93 mam 8292 mmm
FastGen
23.83 (1.04)
8.23 0.39
19.60 caso
24.06
oueo
8616 o070
NAMM (Ours) 21.34
28.61
14.64 (.30
33.15
24.44 (.15) 9895 nvs
4.2 ZERO-SHOP TRANSFER ACROSS ARCHITECTURES AND MODALITIES
Cross-scale adaptation. In Table 5 we
Table 6: Evaluation on the Long Bide and
provide results zero-shot transferring our
MLVU benchmarks with Llava Next Video 7B.
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.
Model/Task name LongVideoBench
MLVU All tasks Cache size
Base model
43.45
44.23
43.84
7039 (180
H2O
40.91
43.03 (m 41.97
4479 (054)
1.2
40.84 pag
42,07 m. 41.45
may
4479 mm)
NAMM (Ours)
44.58 (1.00)
44.18
(1188..00 00 44,38 nan 5100 mm7
7

---

Published as a conference paper at ICLR 2025
Table 7: Evaluation on D4RL with Decision Transformer. The normalized performance (in brack- ets) is calculated using the base model with full cache.
Hopper-va
Walker2d.w3
HallCheetahow3
Overall
Model/Task name
Medium
Med- Replay
Expert
Mediam
Med-Replay
Expert
Medion
Med-RRplay
Exert
All tasks
Cache size
Base model
33.36 0.90
18.37 (ues 0.00) 21
cues
7.18
38.98
5,06 t100
10.64 cue
29.04 u.n 3000 0.90
H20
33.19 0.90
17.86 oom)
49,10 p.
67.63 oum
7.59 0.00
40.03
ua9 26.73 em
4.46 12:00
11.74 cum
28.70 msn 2048 12:0
L2
32.85
43.75
(3:5))
65.47 oue
7.18 ci
40.64
30,10
8.52 mom 27.91 1399) 20418 11
NAMM (Ours)
36.10 (La0)
18.86 (.00)
49.39 0.30
70.87
(L00))
7.53 0.00
50.02 (120
34,56
5,90 (1.7))
12.34 (L16)
31.73 (.a)) 2434 (130
Vision Language Understanding, In Tablel6 we provide results zero-shot transferring to the com- puter vision domain, evaluating NAMMs with a Llava Next Video 7B model (Zhang et al., 2024b) on Long VideoBeecBBhh (Wu et al. 2024) and Multi-Task Long Video Understanding (MLVU) (Zhou al 2024). 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 describ- ing 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.
al.
Reinforcement learning. In Table 7 we provide our zero-shot transfer results for the offline re- inforcement learning setting, where we apply NAMMs atop a decision transformer (Chen et al. 2021b) using the spen-sourced models from Beeching & Simonini (2022) pre-trained on the canon- icll the continuous-contro tasks from D4RL (Fu et al., 2020p. 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 mis- takes 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).
NAMMs comparison. In Table g provide summarized results com- paring NAMMs with either BAM or the simpler MLP architecture at the end of each stage of incremen- tal evolution. First, we note that even the MLP NAMM after stage impressively improves performance across all language benchmarks. Ad- ditionally, performance sees near- monotonic improvements with each additional stage of incremental evo- lution in both language and zero-shot transfer settings. Comparing our im- plementations, the performance ben- efits from the memory models with BAM appear consistently superior to the MLP. Moreover, on ChouBun. we observe that the per- formance with BAM sees 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 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.
Table 8: Summarized comparison of different NAMMs in language modeling (top) and zero-shot transfer (bottom)
Model/Eval
LongBench
Infiniteeench
Choulun
Performance Cuche size Performance Cuche size Performance Cache size
model
28, pm
32768 (r
p.05
32747 pio
21,21 prim)
12899
NAMM (MLP, sl)
28,83 (u8
7639 p.30
3,08 (290)
11329 p25
22,09 1100
9525 0.70
NAMM (MLP, s2)
29.22
8475
4,00
13831 pub
12.0611000
9815
NAMM (BAM, sl)
28.91 pm
7051 mm
10.14 mub
11173 pso
22.73 (ror)
NAMM (BAM, x2)
29.25 (ut)
8267 m025
9.78 (129)
12789 1039 26.05 11:7)
9867 (.92)
NAMM (BAM, 63)
29.33
8155
11,00
13192
24.44
Mode//Eval
Llama 708
Computer Vision
Reinforcement Learning
Performance Cuche sine Performance Cuche size Performance Cache size
Base model
35,22
10107
43.
29,04
300)
NAMM (MLP, sl)
34.11 nch
7930 cm
40.44 yom
564 mmm
29.30
(M93
NAMM (MLP, s2)
34.29 mmt)
8845 0.5
40.39 mot
713 (.10)
29.58 (1/0)
2834 (m
NAMM (BAM. 81)
34.11
7947
41.52
723
30.44
2009
(BAM s2)
25.20
M276 ou
44.63 nxo
4948 (230)
31.53 (100)
NAMM (BAM, 53)
34.70 mm
8365 mm
44.58 p000
5100
31.73 mm)
2434
4.3 UNDERSTANDING NEURAL ATTENTION MEMORY MODELS
Influence of layer depth. We begin analyzing NAMMs by focusing the final amount of retained At the top of Figure 6 we provide these normalized metrics as a function
tokens and their o
Thh DuReader task, used in the second stage of incremental training, uses the Chinese language. We define oldness of a retained token as the number of new queries since its introduction in the KV cache.
8

---

Published as a conference paper at ICLR 2025
2.00
Layer memory
Memory tokens
1.00
1.00
0.00
Loyer depih
Layer mamory tia
Memory tokans didnas
Average task prompt langth
15600
652
10600
0.44
3500
0.96
5000
2380
3-2
3-2
3-3
14
24
3-2
2-3
3-2 34
54
5-2
LongBench Soak
Figure 6: 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 statisties per task using either their average across all task prompts (top) or the mean sample length (bottom). Layer #15, Head #1
Layer #24, Head Aii
<bogin Jextp
paragraphe Wilppedoa along
<lbogin bext parographs Wikipedia slonn
Possagefeeteeis en (Nenuro garggg
Paragrash d. At See 21, sne susmiied France, and doon dhe declared mestrnee of Love feont Caeee of Boween "hoonseur Diee), who was piino minater at the begivii living of the regu of Louds XV (1723-1726). Dusing his
Paragguahh ph Abdha 2535, rolumndd France and was
decaaeed muchess or Louce herm Dake of Bourtonn (MMoroeu ie Duc") une prine ninitro begiming the rogo of XV (1723 1728). During his the roval couth and enginesred marragee Lood Masermalel Vermandoot, recaled
mesiety are dbmmmaated rmps court engeseered mamape cour Lecacay mally
misstry she
Mario Lecoozyccka
suted muviale noviale-sausin sausing Beunion bevisvee his casile Chacilyy.
recalled
cival faledd Instond,
minister causing Bourboni be bamateed to casile Chaniily b.1
<begin lexti-Flease complece givan bokow:
<Bogin bext complete given Jolows.
class KasRef KasRefresh@gers www.rrrrrstttiooo implements Operation
public wwwfereengggrrttiee. deplements Operation
ReceBanch (Coding)
sezort kader 11. uttil sentalizon SeowoHTSe-taltzer
Import to uttll serializer 2oo-te5Sental:zee import- badombin wii weweiaaleene-eeeeeeeeee aoneeo--nooootteiggging www.deuennnnggiiaanenn www.eaaeoooeeeeeeienoo.o ioport kndent in perta He Sro-Seriniseerr
secont www.eeaeeeeeeeoe weseeiiaeeeeeee seseeeseeeeeeeee www.eaaaeee--eeeeeeeeeoo www.-eoeeeeeeiiiiieiiiooe wws.iiiiiiiietiiiiiieeeerr
seport knden sertaisertiseee steen www.Sorta freny
evorything local www.eeeeeeyyyyeee marages
montgges
everythling for this
Bouthor Jocnua Kissoon
éouthon Joshua cissoon
[-]
Figure 7: 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).
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. 2024), 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 6 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 com- pletion 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 opposed natural language.

---

Published as a conference paper at ICLR 2025
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 natural language or coding task (PassageRetrieval-en) id 5-1, and RepoBench-P. id 6-2). As shown in Figure 7 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.
Additional analysis. We provide additional analytic results in Appendix [D] For instance, we ana- lyze how the presence of each token in memory affects the scores of the other tokens, 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 show that NAMMs learn mechanisms for cross-token competition relying on high-frequency components of the atten- tion matrices and illustrate how they learn to overcome different failure modes of long context LMs, further evidencing the need to go beyond simple strategies and the potential of end-to-end learning for loken-level memory systems.
RELATED WORKS
Devoto et al. (0024) and Yao et al. (0024) try to identify the least important tokens to evict using heuristics such as L2 magnitude and entropy to improve efficiency. Alternative strategies include considering simple statistics from the attention matrix (Liu et al.] [2024b] Oren et al., 2024 Zhang et al. 2024c). Ge et al. (2024) and Li et al. (2024b) build on these ideas by applying multiple strate- gies based on matching specific attention patterns. Motivated by similar considerations, Nawrot et al.) (2024) proposed directly fine-tuning the original base transformer to compress the KV cache while minimizing a regularization loss to preserve the original base model's behavior and limit perfor- mance degradation. Complementary
method, MQA (Shazeer 2019) and GQA (Ainslie al.) 2023) propose merging attention heads during training to improve inference throughput. Similarly, also KV cache quantization is another orthogonal area where different hand-designed strategies have been proposed (Hooper et al., 2024; Dong et al. 2024a b), with even recent work empirically show- ing their direct compatibility with token eviction methods (Liu et al., 2024a). We note that, unlike this prior work, our approach uniquely learns black-box model to maximize performance through token-level memory management and shows potential for providing improvements to both the effec- tiveness and efficiency of transformers. We refer to App. including efficient architectures, memory, and evolution.
F for references and to the wider literature,
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 transfcrability across architectures, imput modalities, and task domains. While NAMMs do appear to provide benefits beyond what achieved with hand-designed strategies, we believe there is much room for improvement (e.g., see Limitations F). 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 extension has the potential unlock efficient long-contex training, drawing parallels to the iterative process of learning and evolution that shaped human memory.
10

---

Published as a conference paper at ICLR 2025
7 AUTHOR CONTRIBUTIONS
Edoardo Cetin initiated the project, led the design and implementation of NAMMs, and provided major contributions 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 writ- ing. 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 dur- ing 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).
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
Joshua Ainslie, James Lee-Thorp Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points. arXiv preprint arXiv:2305. 23.i::2300.13225,, 2023.
Takuya Akiba, Makoto Shing, Yujin Tang. Qi Sun, and David Ha. Evolutionary optimization of model merging recipes. arXiv preprint 4rXiv:2403.13187. 2024.
Jont B Allen and Lawrence R Rabiner. A unified approach to short-time fourier analysis and syn- thesis. Proceedings of the IEEE, 65(11):1558-1555 1977.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang. Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308. 2023.
Yushi Bai, Xin Lv, Jiajie Zhang, Yuze He, Ji Qi Lei Hou, Jie Tang, Yuxiao Dong, and Juanzi Longalign: A recipe for long context alignment of large language models. arXiv preprint MXX1:24011.18088 2024.
Edward Beeching and Thomas Simonini. Introducing decision transformers hugging face, Mar 2022. URL /ttp://eee//nggiooooeeeoo/////et trr.sooonntttaaafffommr....
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05//55 2020.
bloc97. NTK-Aware Scaled RoPE allows LLaMA models to have extended (8k+) con- text size without any fine-tuning and minimal perplexity degradation. 2023. URL https://wwwt reddit. Laamm///LooLLLLMa//aacacccc scaled_rege recedddree___aaa____ alope_allowss 1wt.lalaaaaaaaaaaa htw.LLLaaM//mmmne///////////ttt/ to__tve/hh______
William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, and Jonathan Ragan Kelly. Reducing transformer key-value cache size with cross-layer attention. arXiv preprint arXiv:2405.1229: 2024.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint 16Xiv:1606.015:5 2016.
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Xi Chen, Krzysztof Choroman- ski, Tianli Ding, Danny Driess, Avinava Dubey, Chelsea Finn, et al. Rt-2: Vision- carion--anguag--attto models transfer web knowledge robotic control. arXiv preprint arXiv:2307.1588 2023.
11

---

Published as a conference paper at ICLR 2025
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901,, 2020
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:1 5084-15097, 2021a.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34: 5084-15097, 2021b.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306. 2023.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdi- nov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint 10010000000000000 2019.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances Neural Information Processing Systems, 35:16344-16355 2022.
Abhimanyu Das, Weihao Kong, Rajat Sen, and Yichen Zhou. A decoder-only foundation model for time-series forecasting. arXiv preprint arXiv:2310. 10688, 2023.
DeepSeek-AL, Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Hanwei Xu, Hao Yang, Haowei Zhang, Honghui Ding, Huajian Xin, Huazuo Gao. Hui Li, Hui Qu, J. L. Cai, Jian Liang, Jianzhong Guo, Jiaqi Ni, Jiashi Li, Jin Chen, Jingyang Yuan, Junjie Qiu, Junxiao Song, Kai Dong, Kaige Gao, Kang Guan, Lean Wang, Lecong Zhang, Lei Xu, Leyi Xia, Liang Zhao, Liyue Zhang, Meng Li, Miaojun Wang, Mingehuan Zhang, Minghua Zhang, Minghui Tang, Mingming Li, Ning Tian, Panpan Huang, Peiyi Wang, Peng Zhang, Qihao Zhu, Qinyu Chen, Qiushi Du R. J. Chen, R. I.. Jin, Ruiqi Ge, Ruizhe Pan, Runxin Xu, Ruyi Chen, S. Li, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shaoqing Wu, Shengfeng Ye, Shirong Ma, Shiyu Wang, Shuang Zhou, Shuiping Yu, Shunfeng Zhou, Size Zheng. Т. Wang, Tian Pei. Tian Yuan, Tianyu Sun, W. L. Xiao, Wangding Zeng, Wei An, Wen Liu, Wenfeng Liang, Wenjun Gao, Wentao Zhang, X. Q. Li, Xiangyue Jin, Xianzu Wang, Xiao Bi, Xiaodong Liu, Xiaohan Wang. Xiaojin Shen, Xiaokang Chen. Xiaosha Chen, Xiaotao Nie, Xiaowen Sun, Xiaoxiang Wang, Xin Liu, Xin Xie, Xingkai Yu, Xinnan Song, Xinyi Zhou, Xinyu Yang. Xuan Lu, Xuecheng Su, Y. Wu, Y. K. Li, Y. X. Wei, Y. X. Zhu, Yanhong Xu, Yanping Huang, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Li, Yaohui Wang, Yi Zheng, Yichao Zhang, Yiliang Xiong, Yilong Zhao, Ying He, Ying Tang, Yishi Piao, Yixin Dong, Yixuan Tan, Yiyuan Liu, Yongji Wang. Yongqiang Guo. Yuchen Zhu, Yuduan Wang, Yuheng Zou, Yukun Zha, Yunxian Ma, Yuting Yan, Yuxiang You, Yuxuan Liu, Z. Z. Ren, Zehui Ren, Zhangli Sha, Zhe Fu, Zhen Huang, Zhen Zhang, Zhenda Xie, Zhewen Hao, Zhihong Shao, Zhiniu Wen, Zhipeng Xu, Zhongyu Zhang, Zhuoshu Li, Zihan Wang, Zihui Gu, Zilin Li, and Ziwei Xie. Deepseek-v2: A strong, economical, and efficient mixture-of-expeee language model. arXiv preprini 240.:22400.044444 2024.
Alessio Devoto, Yu Zhao, Simone Scardapane, and Pasquale Minervini. A simple and effective 1_2 norm-based strategy for kv cache compression. arXiv preprint arXiv:2406. 100006..1430 11400, 2024.
Harry Dong, Xinyu Yang, Zhenyu Zhang, Zhangyang Wang, Yuejie Chi, and Beidi Chen. Get more with less: Synthesizing recurrence with ky cache compression for efficient Ilm inference. arXiv preprint 20XX::24402.09388.. 2024a.
Shichen Dong, Wen Cheng, Jiayu Qin, and Wei Wang. Qaq: Quality adaptive quantization for Ilm kv cache. arXiv preprint arXiv:2403.04 04644, 2024b.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An
12

---

Published as a conference paper at ICLR 2025
image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010.11929, 2020.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024.
Paul W Frankland and Bruno Bontempi. The organization of recent and remote memories. Nature reviews neuroscience, 6(2): 19--100,, 2005.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data- data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 2020.
Suyu Ge. Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells
you what to discard: Adaptive KV cache compression for LLMs, In The Twelfth International Conference on Learning Representations, 2024. URRKattpss//ppeeeeiieee net forum??
id=uNrFppPMyy
Izzeddin Gur, Hiroki Furuta, Austin Huang, Mustafa Safdari, Yutaka Matsuo, Douglas Eck, and Aleksandra Faust. A real-world webagent with planning, long context understanding, and pro- gram synthesis. arXiv preprint arXiv:2307. 12856, 2023.
Jin-Hee Han, Steven A Kushner, Adelaide P Yiu, Christy J Cole, Anna Matynia, Robert A Brown, Rachael L Neve, John Guzowski, Alcino Silva, and Sheena A Josselyn. Neuronal competition and selection during memory formation. science, 316(5823);457-466 2007.
Nikolaus Hansen. The cma evolution strategy: a comparing review. Towards a new evolutionary computation: Advances in the estimation of distribution algorithms, pp. 75-102, 2006.
Fredric J Harris. On the use of windows for harmonic analysis with the discrete fourier transform. Proceedings the the IEEE, 66(1):51-83, 1978.
Coleman Hooper, Schoon Kim, Hiva Mohammadzadch, Michael W Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache quantization. arXiv preprint arXiv:2441..180 2024.
Xin Huang, Ashish Khetan, Rene Bidart, and Zohar Karnin. Pyramid-bert. Reducing complexity via successive core-set based token selection. arXiv preprint arXiv::203. 14380, 2022.
Yunpeng Huang, Jingwei Xu, Zixu Jiang, Junyu Lai, Zenan Li, Yuan Yao, Taolue Chen, Lijuan Yang, Zhou Xin, and Xiaoxing Ma. Advancing transformer architecture in long-context large language models: A comprehensive survey. arXiv preprint arXiv:2311.12355 2023.
Dongseong Hwang. Weiran Wang, Zhuoyuan Huo, Khe Chai Sim, and Pedro Moreno Mengibar. Transformerfam: Feedback attention is working memory. arXiv preprint (tt::2404.09977. 2024.
AQ Jiang, A Sablayrolles, A Mensch, C Bamford, DS Chaplot, D de las Casas, F Bressand, G Lengyel, G Lample, L Saulnier, et al. Mistral 7b (2023). arXiv preprint trXiv:2310.06825, 2023.
Hongye Jin, Xiaotian Han, Jingfeng Yang. Zhimeng Jiang, Zirui Liu, Chia-Yuar Chang, Huiyuan Chen, and Xia Hu. Llm maybe longlm: Self-extend Ilm context window without tuning. arXiv preprint 2201v2240..01325, 2024.
Greg Kamradt. Needleinahaystack: Doing simple retrieval from Ilm models at various context lengths to measure accuracy, 2024, URL https:://tttttti com/ http//kamrradt... /LLMTeLL. WVeedJeenAHaaystak///rreeee
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers mns: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5165. PMLR, 2020.
Guillaume Lample, Alexandre Sablayrolles, Mare' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys. Advances in Neural Information Processing Systems, 32, 2019.
13

---

Published as a conference paper at ICLR 2025
Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforce- ment learning: State-of-the-are pp. 45-73. Springer, 2012.
Bo Li, Peiyuan Zhang, Kaichen Zhang, Fanyi Pu, Xinrun Du, Yuhao Dong, Haotian Liu, Yuanhan Zhang, Ge Zhang, Chunyuan Li, and Ziwei Liu. Lmms-eval: Accelerating the development of large multimodal models, March 2024a. URL https://ggtt //t.tttttttttt//e// com/
Evolvi I. Lab/ mmss mmm--eal
Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Lim knows what you are looking for before generation. arXiv preprint arXiv:2404. 14469, 2024b.
Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig, Pre- train, prompt, and predict: A systematic survey of prompting methods in natural language pro- cessing. ACM Computing Surveys, 55(9))1-35,, 2023.
Yuhan Liu, Hanchen Li, Yihua Cheng. Siddhant Ray, Yuyang Huang, Qizheng Zhang. Kuntai Du, Jiayi Yao, Shan Lu, Ganesh Ananthanarayanan, et al. Cachegen: Kv cache compression and streaming for fast large language model serving. Proceedings of the ACM SIGCOMM 2024 Conference, pp. 38-56, 2024a.
Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for Ilm kv cache compression at test time. Advances in Neural Information Processing Systems, 36, 2024b.
Shi Luohe, Zhang Hongyi, Yao Yao, Zuchao, and Zhao Hai. Keep the cost down: review on methods optimize Iim's kv-cache consumption. arXiv preprint arXiv:2407.18003 2024.
Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: Efficient infinite context transformers with infini-attention. arXiv preprint arXiv:2404.07143, 2024.
James S Naime and Josefa NS Pandeirada. Adaptive memory: Ancestral priorities and the mnemonic value of survival processing. Cognitive psychology, 61(1):1-22, 2010.
Piotr Nawrot, Adrian Lańcucki, Marcin Chochowski, David Tarjan, and Edoardo M Ponti. Dy- namic memory compression: Retrofitting 1lms for accelerated inference. arXiv preprint arXiv:2403.09636, 2024.
Alan Oppenheim. Discrete-time signal processing. Pearson Education India, 1999.
Matanel Oren, Michael Hassid, Yossi Adi, and Roy Schwartz. Transformers are multi-state rnns. arXiv preprint arXiv:2401.06104 2024.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv:2103.02143 2021.
Jack Rae, Jonathan Hunt, Ivo Danihelka, Timothy Harley, Andrew W Senior, Gregory Wayne, Alex Graves, and Timothy Lillicrap. Scaling memory-augmente neural networks with sparse reads and writes. Advances in Neural Information Processing Systems, 29, 2016.
Noam Shazeer. Fast transformer decoding: One write-head all you need. arXiv preprint
arXiv:1911.02150,, 2019.
David F Sherry and Daniel L Schacter. The evolution multiple memory systems. Psychological review, 94(4):439, 1987.
David So, Quoc Le, and Chen Liang. The evolved transformer. International conference on machine learning, pp. 5877-5886. PMLR, 2019.
Sainbayar Sukhbaatar, Jason Weston, Rob Fergus, et al. End-to-end memory networks. Advances in neural information processing systems, 28, 2015.
14

---

Published as a conference paper at ICLR 2025
Yujin Tang and David Ha. The sensory neuron transformer: Permutation invoonnnnninaan neural networks for reinforcement learning. Advances in Neural Information Processing Systems, 34: 22574-22587, 2021.
Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, lean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312.
2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL wwweeefeee23544415
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXivv2006.04766 2020.
Chris Wendler, Veniamin Veselovsky, Giovanni Monea, and Robert West. Do llamas work in en- glish? on the latent language multilingual transformers. arXiv preprint 1000:2202..1088. 2024.
Jason Weston, Sumit Chopra, and Antoine Bordes. 140.3914 2014.
Memory networks.
arXiv preprint
Haoning Wu, Dongxu Li, Bei Chen, and Junnan Li. Longvideobench: A benchmark for long-context interleaved video-language understanding, 2024. URL
15754
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.74553 2023.
Yao Yao, Zuchao Li, and Hai Zhao. Sirllm: Streaming infinite retentive Ilm. arXiv preprint arXiv:2405. 2...:2205.2228. 2024.
Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, et al. Infinitebench: Extending long context evaluation beyond 100k tokens. arXiv preprint 24Xv:2240..17/11 2024a.
Yuanhan Zhang, Bo Li, haotian Liu, Yong jae Lee, Liangke Gui, Di Fu, Jiashi Feng, Ziwei Liu, and Chunyuan Li. Llava-next: A strong zero-shot video understanding model, April 2024b. URL wwideetttvid/////
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, al. H2o: Heavy-hitter oracle for efficient gen- erative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024c.
Junjie Zhou, Yan Shu, Zhao, Boya Wu, Shitao Xiao, Xi Yang, Yongping Xiong, Bo Zhang, Tiejun Huang, and Zheng Liu. Mlvu: A comprehensive benchmark for multi-task long video understanding. arXiv preprint urXiv:2406 04264, 2024.
15

---

Published as a conference paper at ICLR 2025
Table 9: NAMMs parammersseeeerss used for training and evaluation. The omitted CMA-EE hyper-
parameters can be obtained by following the recommended default calculation by Hansen 2006).
NAMMs hyperparameters
Spectrogram window size fla
32
Spectrogram window stride Sw
Spectrogram window type
Hann
Spectrogram EMA reduction coefficient
0.99
Positional features
8
NAMMs execution delay
512
NAMMs non- non-lincarity
ReLUU
Optimization hyperparameterss notation from Hansen 2006
Evolution algorithm
CMA-ES
Elite ratio
0.5
Mean coefficient Cm
1
Initial step size
0.65
Samples batch size per-task Population size Task for incremental stage
64
32
PassageRetrieval-en
Task for incremental stage
DuReader NarrativeQA
Task for incremental stage
BAM-specific
Hidden dimensions
16 Truc counter-causal
Use bias
Masking strategy
Number of attention layers
Number of final linear layers
Use residual connections
Truc
Use multiplicative interactions
True
MLP-specific
Hidden dimension
25
Number of hidden layers Use residual connections
2
True
A IMPLEMENTATION DETAILS
A.1 MODEL SPECIFICS AND NAMMS EXECUTION
We evolve our Neural Attention Memory Models top of a context-extended Llama 8B (Dubey et al. 2024) base model. In particular, we employ the NTK-aware positional interpolation strategy (bloc97] 2023) 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., (023). NTK-aware positional interpolation has been shown produce sensible results even when applied zero-shot. In case the length of task prompt still exceeds 32768 we perform mid-sentence cropping (Xiao et al. 2023 Jin et al. 2024), as standard in long-context LM evaluation (Bai et al. 2023 Zhang et al, 2024a).
When applying NAMMs, we only affect the execution of the base model with fixed frequency, once every Nap = 512 steps. When feeding longer prompts to our model, we simply split the tokens into -iied chunks. We note that due to modern frameworks being bound primarily by memory constraints, input-splitting in itself has minimal effects running time, with similar approaches being already performed under the hood by established kernel procedures (Dao et al. 2022).
A.2 FEATURE EXTRACTION AND ARCHITECTURE DETAILS
Our new feature extraction framework is 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 nhe 32, resulting in just seventeen complex-values frequencies that
16

---

Published as a conference paper at ICLR 2025
convert to real numbers by simply taking their magnitude, yielding each w/ R We use stride of half the window size &w = 16, producing ny = nup/Sw Sw = 32 frequency representations over the time axis of the attention matrix from the latest chunk of uup queries, , ,..". Thus, 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 nt representations computed for the frequency of each token in the uup Bup-Sized sized chunk of the latest queries, but also the discounted EMA at the previous execution step or our memory for each retained token, denoted w/- Thus, each of our reduced spectrogram representations reflects the full history of previous attention values:
Wi
+ TTTT
(4)
where we use y denote the EMA's discount factor. To expedite learning the weights of our ar- chitecture, we ensure all spectrogram features have unit variance at initialization across our training data, using the statistics of the base Llama model computed on the first task employed in incremen- tal learning (PassageRetrieval), Finally, we also concatenate 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. We provide an extended summarized pseudocode description of the execution pipeline in Algortihm (to complement Algorithm in the main text).
Algorithm NAMMs
Input: Attention matrix A, Hann window size nww Stride size sw. Update interval Rup Output: Memory score for each token
Initialize token score array S
2: Split imput tokens into chunks of size Thup
3: for each input chunk do
4: Extract attention matrix A₁ from the latest uup queries
5: for column in A do
6:
Apply STFT of window size Tap and stride 8u to An(:, Compute spectrograms w/ E C17 for = 1,...,
7:
D 1rT -Bupp 8m
8:
Calculate w; with Equation 4
9:
Update memory w/ 6 W)
10:
Normalize spectrogram features Wi
11:
Concatenate positional embedding to wi
12:
Update S with BAM predicted score m
return S
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 facilitate asymmetric in- teractions between tokens in memory. The output of self-attention is then fed to single final linear layer to obtain the final score. We employed few important additional design choices following some preliminary testing. First, motivated by efficiency
my (1))
m (w)-1)
m (0)-2)
Linear
Linear
Linear
Add & Multiply
Add & Multiply
Add & Multiply
Backward Attention
Backward Attention
Backward Attention
W
W
Wi-2
Spectrogram features
Spectrogram features
Spectrogram features
Figure 8: Schematic depiction of the components of our Neural At- tention Memory Models, denoted mæs parameterized with our BAM architecture. The spectrogram representation of each token, denoted wi processed by an attention layer followed by a simple linear oper- ation to output its rélative score. Backward masking introduces asym- metry, ensuring that each token can only attend to its future relatives.

---

Published as a conference paper at ICLR 2025
considerations, we 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 schematic depiction of our minimal architecture in Figure [ 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-SHOP TRANSFER
For our shot transfer experiments. we consider a Llama transformer with 70B parame- ters (Dubey et al.) (024), a Llava Next Video transformer with 7B parameters (Zhang et al. 2024b), and decision transformer (Chen al. 2021b) with about 1 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 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. (024a)). For the reinforcement learning experiments, we encode each state, action, and return-to- into separate tokens and do not apply any restrictions modifications to our standard NAMM LM setup. We average the perfor- mance collected over 20 random seeds to account for the stochasticity of the initial state in the Gym Mujoco environments (Brockman et al.) 2016). Rather than re- re-training a decision transformer from scratch, our RL experiments adapt the open-sourced checkpoints and implementation provided by Beeching & Simonini (2022). We would like that note that on some task-dataset combinations of D4RL, these checkpoints appear to yield lower performance than what was reported in the original decision transformer paper (e.g.. Walker pre-trained on medium medium-expert data) Chen et al. (2021b). However, we do not believe these differences should affect our conclusions we used the same base model for all our memory management baselines.
A.4 EVOLUTIONARY OPTIMIZATION
As described in Section 3 we optimize NAMMs with the Covariance Matrix Adaptation Evolu- tion Strategy (CMA-ES) (Hansen, 2006). Being an evolutionary algorithm, CMA-ES does not re- quire any gradient information and can directly optimize black-box undifferentiable metrics. This property allows us to both optimize for the non-differentiable token selection task of our NAMMs and also maximize non-differentiable task performance metrics directly. In the case of the Long- Bench (Bai et al. 2023) tasks considered for training, these metrics correspond to exact match accuracy (PassageRetrievallog
ROUGE-L score (DuReader), and F1 score (NarrativeQA).
On high level, given neural network with P parameters, CMA-ES maintains mean vector c E IR" and a covariance matrix Σ PxP. Then, it repeats the following steps:
Sampling. CMA-ES generates population of neural networks, sampling their parameters from the multivariate normal N(u, Σ) distribution.
2. Evaluation. Each population candidate is evaluated to the objective function for the objec- tive function used.
3. Updating. By both selecting a subset of the population candidates and also weighting them based on their overall ranking the mean and covariance are updated towards higher- performing regions of the search space.
We provide the main pyper-parameters in Table and refer to either the work by Hansen (2006) or our shared code for the full implementation details.
A.5 FASTGEN IMPLEMENTATION AND TUNING
18

---

Published as a conference paper at ICLR 2025
We re-implemented the recent FastGen method proposed by Ge et al. (2024), which proposes to adopt hand-designed combination of different strategies targeted to retain tokens with high at- tention values, belonging to recent words, encoded from particular grammatical features (i.e., punctuation, 'special tokens"). particular, after observing the input prompt, FastGen performs a "profiling step" where the strategy able evict the most amount of tokens is selected such that:
A Al2 - T.
(5)
Here, A is the full-cache attention matrix, A is the "reconstructed" attention matrix re-calculated after performing softmax between each layer's queries and keys with masked-out entries for the keys evicted by the individual strategies. Furthermore, T is the main threshold hyper-parameter determining how aggressively FastGen is allowed to prune tokens even if resulting in degradation to the attention-rrcconsttuciion heuristic.
We note that, unlike our other baselines, FastGen is only directly compatible with language modeling tasks. This is because one of the main ways it differs from H2O is by preserving particular grammar- based tokens in some of its strategies (e.g., punctuation, special words, etc.). Thus, as this baseline was specifically designed for LMs rather than arbitrary transformers, did consider applying it in the 0-shot transfer settings, and only focused on Llama 8B.
We note that longer prompts (sometimes far beyond tens/hundreds of thousand tokens), for ef- ficiency consideration, we performed the profiling steps in our re-implementation after the first 4096 tokens any prompt exceeds this length. We also found to avoid losing too much performance over the base model on longer context tasks we had to retune its main 'threshold.' We selected T=0.999, this choice allowed FastGen to retain over 95% nor- malized performance while still discarding non-trivial portion of tokens on all LongBench, shown in Figure 9 Other than the main threshold for attention reconstruction. FastGen has two other main hyperparameters: the 're- cency ratio,' the 'attention ratio' determining the portion of most recent tokens or with the highest attention values to retain in its individual strategies. We set these parameerss to 0.3, following the paper's recommendation.
are dealing with much
3.10
NANH (Dum)
FastGen
Base model (Full cacho)
Aeerage pelfrmmtee
0.95
0.905
0.00
0.08
2910
7A28
8883
9538
Cachessar
Figure 9: Averaged normalized performance and cache size of FastGen compared NAMM over LongBench when varying the threshold pa- rameter T.
19

---

Published as a conference paper at ICLR 2025
Table 10:
Statistics of ChouBun. Lengths are counted by tokens produced by Llama 3 tokenizer.
Extractive QA
Summarization Overall
Statisties
JA. WikiQA
JA EdineeQAA JA.CorpSecCA
JA. CorpSeSeuu
All tasks
Number documents
20
20
30
30
70
Number of QA pairs
200
390
150
30
770
Number of reference answers
1
1
1 or
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
8
Answer length min.
Table 11: Performance of wider range of LLMs the ChouBun benchmark. Summarization
Extractive QA
Overall
Model/Task name
JA. WikQQA
EA CorpSeCCAA
JA.CorpSecSum
All tasks
Max. length
mistralai/Misn (
8.68
8.34
16.25
10,50
10.94
32768
rinna/ 1atra//lammm---yyuo-------
16.68
12.23
17.03
22.27
17.05
8192
meta 1lama/Meta alama 3 8B
14.58
14.77
16.86
22.84
17.27
8192
met meta-lllaa//Lamm. 7b hf
16.77
9.92
20.86
21.97
17.38
2048
01-al/y2-6b-200k alyz:a/lleaa-3-E111e..
30.36
23.64
38.09
21.11
28.30
200000
20.77
21.45
35.59
40.21
29,50
8192
B BENCHMARK DESCRIPTIONS
CHOUBUN DETAILS
The ChouBun benchmark is created to assess the generalization ability of NAMMs to new lan- guage (Japanese), but we hope it will also serve as standard benchmark for Japanese LLMs. The benchmark is composed of two task categories extractive QA and abstractive summarization and four tasks follows.
JA. WikiQA is extractive QA task about 20 randomly sampled articles from the 20240429 dump Japanese Wikipedia Each article corresponds to 10 QA pairs, and there are 200 QA pairs in total.
JA.Edere is an extractive QA task based on 20 security reports from EDINET 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 CSV-style text input helps us evaluate model's capability of understanding structured data. The total number of QA pairs in JA. JAEEdieeIQA is 390.
CorpSrA is another extractive QA task based on 30 security reports downloaded from three corporation websites (MUFQ) NTT and Toyota B). We extract texts from original file in PDF format. There 150 QA pairs total.
JA. CorpSecSSum is an abstractive summarization task based on the same data of JA. Mo.COA. Each document corresponds to data point, and collect reference summaries for each data point.
http:///roonnnt//////////a/////r
20

---

Published as a conference paper at ICLR 2025
Prompt for extractive QA
< D 20 E
#
- 3T \n
E 03 0T In In
3. 7 X
ESG, T TER In
### SEM 1 ###
[question_ 1JAn
### @ 1 ###
(answer 1](m -.An
### M 20 ###\n (question v
### @ 20 ###
[answer 20
Fo : \n ### x ##
[doc_text)
Prompt for abstractive summarization
1. # < 3 In
# - トレンド 3
3.
2 3 In
EX Fo7stt (5
### 3 1 ###
(summary_
.AAn
### #5 5 #
(summary 5/
FON с
### x T #
[doc_ text]
Figure 10: LLM prompts for generating synthetic QA pairs and summaries in ChouBun.
Collecting human annotations for long-tex text tasks is challenging, therefore we use synthetic QA pairs and summaries. In particular, we prompt various LLM to generate multiple question- answeer pairs or summaries for each document. Different instructions are designed for the two tasks and they are shown in Figure TO, To improve the reliability of synthetic data, we ensure that every answer in extractive QA tasks is text span presented in its corresponding source document. In Table TO we provide the statistics of the benchmark.
We use Fl 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 McCab tokenizer 100 A wider range of LLMs' performance the ChouBun benchmark is presented in Table II
'gpt 4opt-4o-202405-111
- upt-4-turbo-2024 04-09, and
el claude- 3-5 - 3-t------sscnntt-220-0 "https://githubb.com m/ fugashi
21

---

Published as a conference paper at ICLR 2025
B.2 BENCHMARKS SUMMARY
We provide a summary of the types of tasks and domains of the other benchmarks we considered for our experiments. We refer interested readers to the relative referenced papers for full details.
LongBench (Bai et al. 2023). This benchmark comprises 21 different tasks targeted to evaluate the
long-context capabilities of LMs. These tasks include both English and Chinese and come from ei- ther moditicd/subsampled versions of existing datasets or synthetic generation. The authors divided them in 6 categories, numbered with the prefixes 6: single-document QA, multi-document QA, summarization, few-shot learning, synthetic, and code. The tasks have a reported average length of 6711 English words and 13386 Chinese characters.
InfiniteBench (Zhang et al, 2024a). This benchmark comprises 12 different tasks designed to go beyond the existing benchmarks and push the limits in long-context L.s.. In fact, while popular prior long context benchmarks, including LongBench, focus on prompts of around 10K tokens InfiniteBench considers tasks with contexts beyond 100K tokens. These tasks again include both English and Chinese and come from either modified/subaampled versions of existing datasets or synthetic generation. The authors divided them into categories: retrieval, dialogue, novel, math, and code. We note some of these tasks considered extremely difficult, with even powerful proprietary LMs such as GPT4 not able to get above a performance 1%.
Long VinggBieeeench (Wu et al. 2024). This benchmark comprises 3763 curated long vidcos with subtitles. These videos are coupled with 6678 human-annotated questions focusing on 17 different categories. The benchmark is focused on what the authors refer as frame-specific 'reasoning' style questions. In particular, for these kinds of questions, video language models are tasked to respond 'referred queries' targeting particular parts of the whole video context.
Multi-task Long Video Understanding Benchmark (Zhou et al., 2024). This benchmark focuses on evaluating long-video understanding performance. It includes videos averaging 12 minutes in length up to 2 hours. The videos span different genres such as movies, documentaries, surveil- lance videos, ego-centric videos, games, and cartoons. In total, this benchmark comprises 2593 evaluation problems divided into 9 categories: topic reasoning, anomaly recognition, video summa- rization, needle question answering. ego reasoning, plot question answering, sub-scene captioning, action count, and action order. These problems are quite diverse including both multi-choice and generation-style questions for video language models.
D4RL (Fu et al., 2020). This benchmark focuses evaluating offline reinforcement learning agents (Lange et al. (0112). In particular, it provides pre-training datasets for different reinforce- ment learning tasks simulated through Mujoco based on OpenAI gym (Brockman al. 2016). The datasets are named based on the displayed agent skills (e.g., expert medium), and based on their prior learning experiences. Evaluation
inclusion of 'replay data' from the demonstrator agent's is then performed after pre-training by running the learned agents online in the respective environ- ments. We focus the most popular subset of this benchmark. involving continuous-contro tasks with three different agents: Hopper, HalfCheetah, and Walker-2d, evaluating the agent after pre- training on Expert, Medium, and Medium Replay data. Rather than re-training from scratch, we use the open- spuresd checkpoints from Chen et al. (2021b) and focus on the evaluation aspect of the benchmark.
22

---

Published as a conference paper at ICLR 2025
Table 12: NAMMs evaluation on LongBench Bai et al. 2023). 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.
Single-Doc QA
Multi-Dur QA
Summarization
Model/ Task ld
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 can
can
22.60 21.31 (1m
10.41 cum
12.67 (um)
25,86 (uo
aum
23,93 (100)
0.92 cue
2,66 (uo0)
NAMM (MLP, sI)
7,60
12.74 c3 22. 21.08 om 9.58 onty
12:24 omn 6,48 oum)
19.41
am 27.76 cm 23.61 205) 0.95 (:0)
3.44 0.29
NAMM (MI
6.76.x5
23.74 (4
12.21
msh
mm 24.20 a. 1.72
2.78 (m0
NAMM (BAM, st)
a.77
12.76
018 22.94 010 21.55 amt
9.17 adto
12:21 0m0y 6.51 000
18.73
007)) 28.06 om 23.97 0.0))
1.01 1.:0
4.00 0.59
NAMM (RAM, <2)
7,08
12:70
0:99 22.21 aso 21.50
9.94
12.21 000
7.13
20.34
0m9 28.87 20m 23.84 (1
0.92
3.94
NAMM (BAM,
0.14
12.63
(199 21.94 (45) 21.34 (100 9.71 009
11.63
6.98 019
20.58
28.78 pmm 24.39 (1.0)
1.04 p:1
3.63 (1.50)
Model/Task id
Few-shot Learning
Synthetic
Code
Overall
tasks
tacks Cache
Base model
73.00
89.45
c1.0 46.54
(mm 40.00
c.m
12.18 cum 28. 28.80 cum 69.09 cue 65,17 aum 28.86 (1:))
N/A
32768
NAMM (MLP,
73:00
89.45
(100
46.80
(100 37.50
039
2,46 (1m
23.98
(107) oon 69.75 (10) (ux)
28.83 (1o0)
7639 (12))
NAMM (MLP, $2) 74.00 amo 88.54
0:99 46.04 asn 41,50 (100
(30
25.94 c210) 29.78 cmm 69,80 cum 65.23 cum
29.22 (um
1.02
8475 o1211
NAMM (RAM, <1) 7s
46.70
38.75
2.19
cum (.o
7951
NAMM (BAM, x2)
a3.00
90.03
(180 46.85
(180) 42.00
2.35
(1:90
24.69
(2m0 28.46 oven 69.65 a.00 66.57 (um
29.25 (..9))
1.04
8267 aus
NAMM (BAM, x3)
a1.00
89.81
46.35
40.00
3,04
27.55
28.60
69.53
(1n 66.35
29.33
1.07
8155
Table 13: NAMMs evaluation on InfiniteBench (Zhang et al. 2024a). The normalized overall per-
formance (in brackets) is calculated using the average performance of the base model with full cache.
Retrieval
Dialogue
Novel
Math
Code
Overall
Model/Task name
Ret Roodey Nesder KV
Fx PNu
In Sues
EAMC
Bugs
24
Part Gode Ras Cist Debay
All tasks
Cache size
Base model
0.50
0.00
0.00
1.00
7.73
0.00
1.05
1.79
0.00
0.00
0.00
1,05
32747 pm
NAMM (MLP, s!)
0.00
10.00
0.00
3.00
7.27
3.93
1.57
4,26
0.57
0.00
3,30
3.08 (240)
11329 (233)
NAMM (MLP. x2)
10.17
11.86
7.48
3.06
1.58
4.10
1.71
0.00
1.52
c.00
13031 (0:4)
NAMM (BAM, s))
9.49
9.83
1.80
0.50
14.36
37.12
8.95
16.20
5.71
1,50
6.09
10.14 c1st
11173 (::4
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 01299
12789 (0:99
NAMM (BAM, 53)
11.86
11.86
1.80
1.00
14.91
36.24
8.78
17.67
10.57
1.75
4.57
11.00
13192
C ADDITIONAL RESULTS
C.I 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.
Extended language modeling results. We report our results for LongBench, InfiniteBench, and First, we note that even training on a single task with our simple MI.P architecture impressively improves performance across all benchmarks. Additionally, perfor- mance across benchmarks sees near-monotonic further improvements with each stage of our incre- mental 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 [I] The same improvement not occurring with the MI.P-based NAMMs might be further evidence of architectural performance saturation, highlighting once again the effectiveness of our main implementation design.
ChouBun in Tables I2
1 14
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 T6 and[7] We see the benefits from NAMMs again increase as we incorporate backward attention, and with each stage of incremental training to a sim- ilar 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.
"The DuReader task, used in the second stage of incremental training, uses the Chinese language.
23

---

Published as a conference paper at ICLR 2025
Table 14: NAMMs evaluation on the new ChouBun benchmark. The normalized performance (in
brackets) is calculated using the base model with full cache.
Extractive QA
Summarization
Overall
Mask name
WIRIQA EdinetQA JA. CorpSeeQA
JA. SorpSeSuum
All tasks
Cache size
Base model
22,91 (1.00
28.34 (1:00
11.83 01000
21.75 (um)
21.21 (.00 12099 (1000)
NAMM (MIPP; s1)
21.60 mn
26.81 0.99)
10.34 muss
29.60 (.90
(1 79
NAMM (MLP. s2)
20.76 019
26.30 quse
11.86 01009
29.32 (1201)
22.06 (.09 9815 (080)
NAMM (BAM. s!)
19.19
28.85
14.36
28.51 (.3)
22.73 c.mm 9569
NAMM (BAM. s2)
20.75 0.00))
28.46
14.55 0.22
32.45 1::4
24.05 (115)
9867
NAMM (BAM. s3)
21.34
01509
28.61 (1.0))
14.64 0.24
33.15 0.52)
24.44 (L15)
9895 00:03)
Table 15: NAMMs evaluation on LongBench (Bai et al. 2023) with Llama 3 70B model. The nor- malized 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
3-4
Base model
9,38.1.m
13.84
24.99
17.78 cm 11.73 am
14.26 amm
8.11 c.1
26.43 crm
13.13 24.55 (.:)) 23.20 p:9
10.08 (1o
NAMM (MLP,
6.34.
(108
24.27
17,60
090 10.83.09
14,17
7.89
13.09 (m 23.30 mos) 23.28
8.66
NAMM (MLP, x2)
0.8
13:71
da9 23:27 as0 18.18 (10) 11.41 907 14.11 2099
8.07 c09
21.75 (02) 14,28 nu 24.48 (100) 22.00 pm
8.99 aum
NAMM (RAM, sI)
13.75 0.99 24.51 aso 17.78
cm 10.82 mm 14,08 om
7.59
19.27
am 13.89 (m 23.71 amh 23.41 pm
8.87 avon
NAMM (BAM, x2)
3:57 0:30
13.86
(180 23.02 (4))) 18.71 (1s5) 4.94
13.52 (00
1.90
17.24
(uh) 10.39
pm 20.45 (m)) 23.
m.0
8.13 (ui)
NAMM (BAM, s3)
9.13
13.53
(.99 24.25 am 17.82 (100 11.45
acp
13.76
8.34 (um))
21.79
(02) 12.66 nom 24.21
a05 23.56 0.323
8.62 (000)
Few-shot Learning
Synthetic
Code
Overall
Model/Tash Taskk id
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
c1s0 can
cum cun 75.37 qum (um
74.60 (uo) 71.19 (um)
(1o0)
10107 (Ue))
NAMM (MLP, s!)
78.00
92.28 cm 48.57 osn 43.50 om 20.76 osm 68.66 om 33.89 cm 74.58 (1 71.68 com 34,11 s0n
0.99
7930 om
NAMM (MLP
a4,0
91,93 48.60 44,75
17.17
70.21
36.18
74.72
71.30
34,29 avo
8445
NAMM (BAM, x1) 77.50 am) 92.46
0100 48.24 (10 45,00 (m)) 17.32 (07) 69.87 (00 33.89 (.00 74.58 (100 72.40 (um 34,11 soo
0.99
7947 put
NAMM (BAM. <2)
74.50
a.0 51.45
0.30 39.73
15,00 e35 5.86
13.35
34.29
73.81
61.91
25.20
0.79
8276 euz
NAMM (BAM, x3) 78.50
auo 92.36
(100
48.49
(100)
45.50
(100
19.07
os9) 74.19
(oo 34.28
(10))
74.71
(00)) 72.42
(uo
34.70 (00
8365 (ui))
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-CONNECTER NAMMS
In Figure II 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 popula- tion, 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 F note few interesting differences, although its evaluation performance the full Long- Bench 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 better inductive bias, mitigating the overfitting potential of NAMMs.
C.3 EVOLUTION OF MEMORY SIZE DURING TRAINING
In Figure 12 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 gen- erally 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, note that both models are trained only for
24

---

Published as a conference paper at ICLR 2025
Table 16: Evaluation on the Long VideoBeenn 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, s
39.64 (0.11)
41.24 (0.93) 40.44 (092)
584 (0.05)
NAMM (MLP. s2)
41.06 (0.90)
39.72 (0.90) 40.39 (092)
713 (0.10)
NAMM (BAM, s!)
41.06 (0.91)
41.98 (0.95) 41.52 (095)
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. 2020) using a Decision Transformer model (Chen
jt al. 2021b [Beeching Simonini 2022). The normalized overall performance (in brackets) is calculated using the average performance of the base model with full cache.
Hopper-v3
Walker2d-r3
HalfCheetah-v3
Overall
Model/Task name
Repiay
Expert
Medium
Med ReplRy
Expert
Med-Roplay
Expert
Cache
Base model
33,36
18.37 (100)
44.62
68.21
7.18 u.r
38.95
34,91
5.06 (100)
10.54 n.m 29.04
3000
NAMM (MLP, s1) 33.01 am9
18.39.u.))
38.09
11:09 70.83 om
7.25 n.)
41.61
0.14 35.64
num
5.05 cma
10.87
um
1993 (1:9)
NAMM (MLP, s2) 33.48
19.24
30.07
73.23
7.95
48.21
33.59
5.81
14.67
29.58
2834 (130
NAMM (BAM. x1) 35.02 0.95
18.24 mm
45.95
p:00 69.33 0.00
7.91 p.
44.45
(1:0 34:25
you)
5.12
13.68 (.39 30.44 0.0
2009 (.67
NAMM (BAM, s2)
35.18
18.79
48.08
71.97
7,70
49.74
35.67
5.78
10.82
31.53
2534
NAMM (BAM, x3)
n.20
18.86 (1m
49.39
(.10) 70.87
0.00
7.53 p.95
50.02
(1:3)
34.56
yow.
5.90 (317)
12:34
0.16
31.73
cm
2434 (110)
performance maximization, without any incentive to be more conservative. To this end, exploring regularization strategies make NAMMs aware of deployment costs is interesting direction for future work to obtain tailored sweet spots to cater to instance-specific resource constraints.
C.4 INCREMENTAL. TRAINING ABLATION
We provide 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 13 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, show the normalized per-task and average performance on all samples of the optimized mean from CMA-ES, and on the right subplot show the corresponding memory size. Furthermore, in Table 18 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 notably higher score on the training tasks with a normalized performance of 1.57, in contrast to the normalized performance of 1. 1.41 achieved by the best checkpoint from the second incremental training stage. Yet, outside the RetraaeRetrieval-en and DuReader tasks, its performance is notably inferior and very close to the original performance of the base model. These results appear 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.
C.5 RUNNING TIMES AND MEMORY SAVINGS
We provide details about the efficiency and costs of NAMMs on top of the Llama 8B base model used for training. For our main experimental setup, we used rented cloud instances with Nvidia H100 GPUs, Intel Xcon Platinum 848 IC CPUs, and 1932GB of RAM. We performed model inference for each prompt on a single GPU, with batch size 1. During training, we used a single node with GPUs, distributing the evaluation of our population across 8 processes. However, we like to remark that since training NAMMs does not require any gradient computation, we were not restricted by any kind of hardware during training. In this regard, using inference-specializee resources beyond GPUs might provide considerable speedups and lower costs to ones employed in this work.
25

---

Published as a conference paper at ICLR 2025
Chic Pepulatio pample
CPA monn peloorannnn pewall)
Best phase chockpoint
Pesssppfhttrievlllnn
Dulloader
24-
1.4
3.-
Incremental phooe
160
350
600
500
206
380
400
Teal
Figure 11: 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.
CHA-DS memory stre (DAA)
OMM-ES rean mamoyy san (Mar)
12000
WWW
32009
Ww
30000
300
Best phase
Best phase
6000
Incremental phose
checkpoint
Pecresenntal phase
sccrementaa plase
2000
100
Figure 12: Final memory size of NAMM parameterized by the learned mean of CMA-ES for both the BAM (left) and the MLP implementations (right).
Table 20: Training times per generation of our fi-
Training. We collected the training time for each generation of NAMMs. As detailed in Section and Appendix A with the employed hyper-parameters, each generation consisted of running the base model for population size x task samples = 32 x 64 = 2048 prompts for each task. Thus, each incremental phase got linearly more expensive, with up to 2048 x 3 = 6144 NAMM evaluation in the final phase. These prompts were distributed across our processes balancing the number of tokens evaluated each. We also note that the average prompt length and the nature of each task (e.g., exact match, summa- rization, etc.) varied quite significantly LongBench, making their evaluation costs non-uniform.
nal NAMMs using our available distributed hard- ware.
Incremental phase
Time per generation (s)
Phase 1 (+PassugeRettieval-en))) Phuse 2 (+DuReader)
1597.62
4027.33
Phuse 3 (+NarrativeQAA
9848.31
Inference. We collected running times of NAMMs and our baselines in different settings. In par- ticular, these include both: Using samples from the full LongBench benchmark with an average length of 12099 2. Using only samples from LongBench selected to exceed the base transformer maximum length with an average length of 32641. Finally, we also record the running time of an ablated version of our NAMM run top of the base transformer that does not modify its KV cache, in order to disentangle the gains from the reduced memory and analyze the pure overheads from our model's 's execution.
As shown in Table 19 the running time overhead of our NAMM ablation that does not evict tokens is small when compared to the base model. Instead, the running time of NAMMs and the baselines
26

---

Published as a conference paper at ICLR 2025
CMM Riputaii
CMM
1300
Peooopefetrirrall
200
Dufleader
11003
1.7
2:00
1900
1.2
200
4.5
Figure 13: 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.
Table 18: NAMMs incremental learning (IL) ablation evaluation on LongBench (Bai et al. 2023 The No IL baseline is trained from scratch on both the PassageRetrieval-ee and DuReader tasks, the same employed during the second stage of incremental learning.
Model/Task id
Single- Doc QA
Multi-Doe Doc QA
Summarization
3-2
Base model
10.38 c1n
12.79 pm 22.60 mm 21.31 na
10.41 cas
12.67 psn 7,54 ma
25,86 us
29.34
can
23.93
0.92
2.66
NAMM (BAM, si)
5.77
12.76
22.94
21.55
9.47
12.21
6.51
18.73
28.06
23.97
1,01
4,00
NAMM (BAM, K2)
7.08 muo
12.70 wm
22.31 mm 21
9.94 o39
12:21 amn
7.13 mm
20.34
28.87
23.84
m00 0.92
3.94 n-4)
NAMM (BAM, s3)
9.14 oas
12.63 wsn
21.94 pom
21.34
9,71 680
11.63
6.98
20.58 a
28.78 e39
24.39
1.04
3,63 0.30
NAMM (BAM,
6.46
9.91
11.77
22.60
0.91
Model/Task id
Few-shot Learning
Synthetic
Code
Overall
4-1
4-2
4-3
4-4
5-1
5-2
5-3
5-1
6-2
All taoks
Test tanks
Cache size
Base model
73,00 cao 39:45
mm 46.54 nom 40.00
1,48 0100
12.18 (m 28,80 now 69.09 c9. 65.17 28.86 p1m
N/A
10107
NAMM (BAM, s])
73.00 cam 39.81 (100 46.70 nom 38.75
2.19 (146)
25.14 (200) 28.51 nom 69.50 c.n 06.51 (10) 28.91
1.00
8205
NAMM (BAM, 62)
73,00 90.03
46.85
24.69
1.04
8521
NAMM (BAM, s3)
73,00 cum 39.81 (100 46.35 nm 40.00 c:
3.04 (49)
21.35
(230) 28.00 nom 69.53 c.0 06.35 (18))
29.33 (.10
1.07
8409
NAMM (RAM, [I.)
71000 89.28
(100) 46.43 38.75
214
29.28
28.46
mom 69.80 64.77 s196. 28.79
0.98
8457
while evicting tokens is always inferior to the base model in all settings, and scales positively with longer prompts.
Furthermore, in Table 21 also reported estimated effects in peak GPU memory consumption, which were calculated from the peak KV cache sizes, together with the sizes of additional information (e.g., attention matrix) and models used by each method (again recorded on LongBench). We would like to note, however, that as the main objective of our work was provide performance benefits we did not particularly optimize our code for memory efficiency or speed. Thus, actual empirical savings with our shared implementation might differ from these calculated estimates. For instance, both our NAMMs and H20 baseline do not employ specialized kernels to replace FlashAttention (Dao et al. 2022).
Memory.
Table 21: Calculated peak storage cost and savings (MB) of NAMMs, together with the H2O and 1.2 baselines with default hyperparameters.
Model
KV cache storage Additional storage averheade Total savings
Base model (full cache)
65282
H2O
17408
5659
42215
1.2
17408
0
47874
NAMM
5668
32378
C.6 MISTRAL BASE MODEL AND FINETUNING NAMMS
We also analyzed an additional 0-shot transfer setting, this time applying NAMMs top of the Mistral 7B base model (Jiang et al. 2023). We considered 2 different setups: Zero-shot appli- cation, taking best NAMM model trained with the Llama 8B context-extended model. 2. Post
27

---

Published as a conference paper at ICLR 2025
Table 19: Running times from using NAMMs on top of a Llama 3 8B base model, together with the running time of the H2O and L2 baselines. NAMMs (full cache) indicates our NAMMs used without pruning the KV cache memory in order to show its total overhead without the gains from reducing the KV cache size.
Longbench (12099 average tokens per sample)
Method
Time per task sample
Base model (full cache)
4.0 (1.00)
NAMM (full cache)
4.54 (1.13)
H2O
3.78 (0.94)
L2
3.74 (0.93)
NAMM (Ours)
3.97 (0.99)
Selected max length samples (32641 average tokens per sample)
Method
Time per task sample
Base model (full cache)
15.52 (1.00)
NAMM (full cache)
17.6 (1.13)
H2O
11.79 (0.76)
L2
10.5 (0.68)
NAMM (Ours)
13.61 (0.88)
Table 22: NAMMs evaluation on LongBench using Mistral 7B v0.3 as base model. 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.
Single-Doc QA
Multi-Doc QA
Summarization
Model/Task id
1-1
1-2
1-3
1-4
2-1
2.2
2:3
2-4
3-1
3-2
3-3
34
model
(180)
28.45 (m0
pm
nm
11.31 0300
28.62 (uan 22,48 (10) 25.47 pin
12.41 avo))
H2O
7.96 0.36 7.87
am
28.98
16.47
som
11.83
mg
11.06
7.4 a.)
19.19
0.0) 25.54
can 20.55 mno 25.38 pm
12.19 avt
msb
quo
mm
059
0.59
030
0.49
mm
mu
oun
NAMM (0-sbot)
13.01 053) 10.30
(12)
28.98
usa
17.10
ama
11.34
mm
12:24 ma8
a.3
(4.1
25,22
csb 20.23 000
27.79
nmm
11,97 0000
NAMM (Finetunc)
9.82
10.36
28.40
17.38
11.72
12.11
7.28
19.04
25,77
21.74
27.72
12.52
Model/ Taskel//assiii id
Few-shot Learning
Synthetic
Code
Overall
4-1
4-2
4.3
4.4
5-1
5-2
5-3
6-1
All tasks
Test tacks
Cache size
Base model
76.50 cue 90.50 cum 36.55 uro 41.50 um
1.00 oo:;
32.17 mas 27.50 cas 65.91 cas 62.71 crm 30.33 (1/
NIA
10107 (100)
H2O
76.00
90.14
38.54
31,00
1.50
6.88
24.13 ase 66.25 (.m 64.46 cuao 26.27
NIA
6662
L2
72.50 ess 75.65 asso
33.91 wom
13.50 auo
0.50 wxx
6.00 a39
24.50 m. 61,46 50.11 man 26.27 ocm
NIA
6662 avea
NAMM (0-shot)
75.50
90.83 crm
40.84
36.25
1.73
26.04 nis 27.50
69 v1 64,60 caa
30.12
1.04
8518 oosn
NAMM (Finctunc)
cue 0.m mm
cm
1.0 cm 0.59) (1o (ua))
8654 pur
cross-model fine-tuning, running a small amount of additional evolutionary optimization comprising 20 generations using CMA-ES and the same training tasks used for Llama. We provide our full results and analysis in Table 22
Analogously to our other zero-shot transfer results provided in Section we find that NAMMs yield considerable benefits also when transferred to the Mistral model, overcoming the efficiency- performance tradeoff of hand-designed baselines. Furthermore, this analysis also shows that per- formance could be further improved by few finetuning generations after transferring to a different base models. While we did not investigate finetuning with our other transformers (e.g.. Llama 70B), believe these results highlight the potential of cheaply improving NAMMs already-remarkable zero-shot benefits, which we hope will be further explored in future work.
C.7 ATTENTION SPECTROGRAM FEATURES ABLATION STUDY
We examined ablating the attention spectrogram features produced by the STFT procedure and re- training NAMMS with two different alternatives:
28

---

Published as a conference paper at ICLR 2025
Table 23: NAMMs evaluation on LongBench ablating the STFT features. The normalized perfor- mance (in brackets) is calculated using the base model with full cache. The tasks used for NAMM's training are highlighted in gray.
Single-Doc QA
Multi-Doc QA
Summarization
Model/Task
1-2
1-3
1-4
2-2
2-3
2-4
3-2
3-3
3-4
Base model
10.38
22,60
10.41
25,86
29.34
23.03
0.92
(100
NAMM (BAM, x2)
7.08 mmm 12.70
cm 22.21 ass
cmm
9.94 dm
12.21 mm
7.13 mm
20.34
am 28.87
u00 23.84 110
0.92 mm00
3.94 0.40
NAMM (BAM, s3)
9.14 xum
12.63 am 21,94 (a)) 21,34
9.71 (14
11,03 000
0.98 000
20.58
26.78
w0b 24.39 (12
1/44 0.30
3,63 0.30
NAMM (RAD features)
10.09
12.93
21.56
9.65
8.26
18.00
3,4%
NAMM (Raw attention)
9.98 pie
12:77 cm 22.04 esso 21.48 9.72 am
12.08 mm
6.31 m.:0
22.61
mum 28.68 201
23.76 2000
0.87 p0m
c.
Model/Task id
Few-shot Leurning
Synthetic
Code
Overall
4-4
5-3
All
tasks
Cache size
Base model
73,00 p.m 50.4 45 on 46, 54
40.00 c1
c/4
12.18
28,90 num 69 09 (m 65.17 pm0 28,86 2100
NIA
10107 ci
NAMM (BAM.
73.00 90.03 auo 46.85 42,00
2.35
24.89
28.46 69
66.57
29.25
1.04
8521 (2:4
NAMM (BAM,s3)
73.00 nue 89.1 cum
3.0
27.55
0.30
28
(u0)
66.35
(10))
1,07
(.4))
NAMM (KAD features)
73.00 (um 80 48 on 46.64 c1 a
(. mz
21.15
27.76 mo 69 68 (m0) 56.31 (009)
28.71 (100
1.00
6000
NAMM (Raw altention) 73,00 59.48 46.66 39,50
(.4 2. 28.46
69.73
66.77
29.10
0.95
(223
1. The naive approach of using the raw attention values directly (cropped a fixed length) as input to NAMMs.
2. Substicuting the STFT features by constructing a "handcrafted" feature representation that simply includes three values: i. The sum of the attention values of each token. ii) The recency of each token. iii. The diversity of each token (computed by concatenating the keys and values represent each token and averaging the L2 distance all other tokens). We refer to this baseline as RAD.
We trained these baselines only for two incremental phases on the PassageRetrieval-en and Dureader tasks (thus, we also compared them with our original NAMM model after phase 2). Please refer to Table 23 for our results. Overall, find our baselines yield quite different behaviors, both underperforming our original NAMM design.
First, we find our naive baseline, taking as input the cropped attention value, is not able to improve over the full cache model when evaluated on the whole of LongBench. However, we note that its performance is strongly suggestive of the occurrence of overfitting, which we believe is to be expected as our memory model now only conditions on very high-frequency information that only considers the latest attention values.
the training task is significantly beyond the base model. Thus,
find this
Second, we find that our "handcrafted' Recency-Attentee Seet--tttttio--ivvetteee baseline is instead able to im- prove over the original model, but its improvements are only marginal. We find these results con- sistent with section D.2 of the extended analysis, which suggests that the behavior of NAMMs is considerably influenced by combination of different frequencies in the attention spectrogram which are lost by this approach.
29

---

Published as a conference paper at ICLR 2025
Grodlctt RoOn maggs scrcen cache tokenn (100)
Grudint sore mage tudes scroe RV cache tokana Hlg!
Grod due products acros cadle toherd Gyrmogg
Bandom taken Highest scared tolon.
377
10
-227
430 DV cadt
6
1300
Ofhe: 477 cotche Indcc
Figure 14: Density plot of gradient magniudees 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 is feature vector (right).
D
ADDITIONAL ANALYSIS
D.1 BACKWARD ATTENTION CROSS-TOKEN INTERACTIONS
We analyze the cross-token interactions learned through our BAM architecture by recording the gradients of each token score 8; with respect to all input features vU for all tokens in memory after storing 1024 tokens, i.e., for j = 1,2,.... 1024. We denote these quantities as:
3s:
(6)
Jug
We provide a qualitative visualization of our results on the PassageRetrieval-en task for randomly selected layer and prompt in Fig On the left subplot, we provide a visualization of the squared magnitudes (Vyy)) ) Vyg! for each combination of tokens (either scored or attended upon in BAM, i.c., indexed by i 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 '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 (Vg')" 04 After the first notable spike, at i = most other dot-product spikes with the largest magnitudes consis- tently 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.
D.2 SENSITIVITY ATTENTION FREQUENCIES AND POSITIONAL ENCODINGS
We analyze the magnitudes of the gradients of the token scores 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 its relevance (since scores determine which tokens get discarded). Figure I5 we plot the distribution of magnitudes for all the seventeen features up to the Nyquist frequency (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 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 subplot, NAMMs appear
frequency' at dimension 0. Furthermore, as shown in the rightmost- visibly less sensitive 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
30

---

Published as a conference paper at ICLR 2025
Memory sensiliily attonton froquancies
Prequecy 0
Prquercy
Frcuercy 2
frequeccy 3
mequercy
Pequensy'
0.52
Peqqurcy
hrqquecy 13
Presercy
hequency
hequency
0.30
Figure 15: 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 com- pare the total magnitudes of the frequency information with the recency information in the positional embeddings.
on token recency and the sum of the attention values, which has so far been considered strong established recipe for KV cache management (Oren et al. 2024 Zhang al. 2024c; Ge et al.
2024
Devoto et al. 2024).
D.3
INFINITBBENCH 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 I 1.05% to 11%). How- ever, 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, reported in Zhang et (2024a). Nonetheless, on the En.Sum task, concerned with summarization of fic- titious no we find our main NAMM brings the performance of the context-extended I Ilama from 7.73 to 14.91 even slightly beyond GPT4's (14.73). While this performance is still low in absolute terms such 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.
We qualitatively inspect the effects of NAMMs on En.Sum by comparing example answers gen- erated by Llama 3 with and without GPT4. As illustrated in Figure T6 we find both the Llama and GPT models to incur several fail- ure 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).
memory models, together with examples generated by
I'IniiittBeend tasks scored in range between and 100.
31

---

Published as a conference paper at ICLR 2025
Prompt io 27
Prompt id
fass Medeline is young arphan being relsed by Mire. Bruse her cruet, weathhy aunt A servant named Rhana na provides ls with some the fow kindinnn essca sho receives, be liling her etorie and singing songe her: Ono day, os purishment for imprisonaa lsea ists find suggosst
Reshan Makory is ainy-hhee-gg yeon-old Bllack man living in Sanmogga, New York, in 1841. He happlly married and has three dhikmren who hu adores, He worls as carpenter and is also a skilled viokinst who otten plays for hts frends.
Torget solution
fightinn
rod room, builying cousin Spencer Bruce,
Reshhr''s was into alevary out was
laager Herry Malory,
which Moddition Bruce dies. While
romains close friond Rashan's Samlly day, white milt visimat
leiing harself
sees uncle
Rosther temporerty Bulloving that
travaling for
thusisstioally
achool To ists delight, Bruce concura VnteCriccs the uewoad School, tai finds that her lile is from ityliic. The school's ls headmaster Mr. Bity a cruel, hypporticca, and sbusive mon. Eity studeets whilh using the
Mes, Bi that los
Roshhn doesn bell his with that he's
might whill WWashington
D c., Reshin becomes sicc and iosisa corsadisssnsse When Rashiin wakes he finds Nimeet shackled in dark cell Intrilashon discovere that he has been kichspped by a srave inader nemed Lorana Bitbii, who bastts Reshai every time ho instats that he's a free man.
preachse a dodrine of powerit and privetion school's funcs to provide a wealthy and opulent Hestyle for his own family.
[.]
[.]
was a ittie, and had been, was a inti,, and had been, was a ltfee. was (i2kk, and was lifie. was iitle, I was lille, I was a litle, and was iitle, was a infle. was a lifee was a lttie. was a litte was Itfe. was a litte, wwas litie,
ngton. the state of the fother was a stave in the state of the fother was
Lama3
s siave in the state the faiherr was a alave in the state of the fatherr was a slave
the father was a stave in the fother was stave the father was a slave in the
a wis a lttlo, loass iste, was Iifoe twa litic. was ittic, was lido.
kitt,, was ittin. twas
or was staeve was slave in the fathor was slave father was a slave father
fuiter was slave the ficher was atlae in the lather
(.)
[-]]
Liamo3 ND0/M
was Madalin, the daughter wealthy Englishman and Portuguese waman. father dled educaled school Yorkshiro, then live with undle sisters.
namator laist Madelins, story
young and sht tu
by forter
life the
Spenoor
North
[--]
o hotel in the Henry Mallor, depase and that hare been well acqueinted with and being duy
Sandy
sard county,
GPT4
"End J'ane Eyret by Charlotte Brone...
Rashan Matory, named in the annosed memoral al of Anno Malion, for upwards iftee years lest post that isid Ibashan Malory was born in the lown of Kingebury, in the county of Washington, and was resident of the countyy of Washington
[-]
Figure 16: Qualitative examples comparing the ground produced responses by Llama3 with and without our NAMM memory, together with GPT4, on two prompts from the Sum task part of InfiniteBench.
E EXTENDED RELATED WORKS
Similar to our NAMMs implementation, memory management through token eviction has been ex- plored mostly reduce memory constraints and enable querying LMs with longer contexts (Luohe eet al. 2024). Commonly, strategies entail simply cropping input prompts to a shorter length, of- ten more effective when done from the middle rather than the ends (Xiao et al.) 2023: Jin et al. 20244.. More advanced, several heuristic strategies have been proposed to identiffy and evict the least important tokens in the KV cache, selectively pruning it to a fixed size for each layer. These strate- gies assess token relevance using metrics like L2 magnitude (Devoto et al.] 2024) or entropy (Yao) jet al. 2024),, or analyze statistics from the attention matrix, such as value magnitude or cumulative sums (Liu et al. 2024b Oren et al.) 2024 Zhang et al. (024c). Building on these ideas, Ge et al. (2024) and Li et al. (2024b) apply multiple strategies simultaneously, choosing the best fit for each layer by matching them with specific attention patterns. Similar ideas where also explored in older work targeting encoder-decoder models, for instance, Huang et al. (2022) proposed more complex strategy for token selection based solving the core-set problem with parallelized greedy ap- proach. However, unlike previous work, our approach uniquely employs a black-box model to learn KV cache management in order to boost the base model's performance with improved efficiency coming as a free side benefit.
Many other methods to reduce memory consumption, affecting the KV cache, are mostly orthogonal and likely complementary to our approach. For instance, MQA (Shazeer 2019) and GQA (Ainslie) al., 2023) propose merging different attention heads during the training of LLMs, either fully or partially, to improve deployment-ting
throughput. [Bandon et al. (2024), pushed these strategies fur- ther, attempting to merge heads even across different layers. GQA is commonly employed in many modern LMs, including the LLaaa 3 family of models which we use to train and evaluate NAMMs on language tasks (Dubey et al. 2024). Furthermore, several methods have looked at KV cache compression through either quantization of the keys and values (Hooper et al. 2024 Dong et al.) 2024a|b) or even the whole hidden states (DeepSeek-Al et al. 2024). Similarly to the aforemen- tioned prior work concerning KV cache pruning. these methods considered mainly strategies, such as employing different quantization rates based on heuristically recognizing impor- tant tokens. We note that using evolution optimize for which channels merge compress could
32

---

Published as a conference paper at ICLR 2025
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 com- ponents of a memory system or to address key challenges of reasoning over longer contexts. For instance, past work has looked incorporating neural models of memory within neural networks by implementing different reading and writing operations either directly replacing their layers (We- 2015), or introducing new auxiliary components (Rae al. 2016 Lample et al. 2019). 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. 2020 Katharopoulos al. 2020 Wang et al. 2020; Peng et al. 2021 ) or introducing new kinds of persistent tokens and storage to extend information propagation (Dai et al. 2019; Munkhdalai et al,, 2024 (wang et al. 2024). However, as also noted by Dao et al. (2022), none 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.
ston al.) [004]] Sukhbaatar
Lastly, methodologically related to NAMMs, there have been other prior methods making use of evolution for or with transformer models. For example, Tang & Ha (202I) also trained small attention- based model through evolution, exploiting the inherent parameter efficiency behind these operations. Furthermore, So et al. 2019) proposed using evolution to meta-optimize the basic build- ing of transformers via neural architecture search, while Akiba et al. (2024) focused evolving different merging strategies across layers belonging to LMs with different capabilities. As for these works, we note that evolution plays critical role for NAMMs, allowing us directly optimize for target performance and overcome the inherent non-differentiabiliy underlying our new framework.
33

---

Published as a conference paper at ICLR 2025
Table 24:
NAMMs evaluation on the canonical Needle In A Haystack task (Kamradt) (024). The
normalized performance (in brackets) is calculated using the base model with full cache.
Needle prompt length
Overall
Model/Task name
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 (uam)
5.36 (0.5)
10208
NAMM (BAM), 7 0.9999
9.00 (1.02)
5.33 (0/1)
3.45 (199)
5.68 (190)
10347
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 poten- tial 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 empiri- cal 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 likely to yield further improvements.
For instance, while our current feature extraction, based on computing the spectrogram of the atten- tion matrix, enables capturing global frequency information about the attention values of each token, might fall short of modeling local information with enough granularity. This hypothesized limita- tion inherently comes from few design choices made with the purpose of limiting the input size and corresponding parameter count our memory models. In particular, our spectrogram features only consider the real components of short-time Fourier transform with 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 ini- tial element-win exponentially moving average reduction. Concretely, this involved computing T different scores, feeding m, all feature vectors w/ for 1, T, across the attention spec- trogram's 's compressed time axis, only then reducing the resulting scores s!- via EMA. While, in principle, this alternative ordering would allow for additional expressivity without adding to the pa- rameter 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 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 of fast kernel algorithms developed to accelerate inference by foregoing materializing attention values (Dao et al.) 2022). While the main focus of this work has been to introduce NAMMs and display its potential 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
One notable example exemplifying some of the aforementioned limitations, comes from the canoni- cal Needle In A Haystack task (Kamradt 2024), 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
34

---

Published as a conference paper at ICLR 2025
architecture, averaging evaluation scores provided by GPT-4 model (Achiam et al. 2023 across different prompt ranges, consistently with Bai et al. (2024). As shown in Table 24 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 observe unexpected trend that is in contrast with the rest of our results other benchmarks. In particular, while our NAMM obtains slightly higher than the base model for prompts with size less than 10000, 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 Haystack task is constructed such that the model is tasked to remember some important information introduced 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 y = 0.99*', all the frequency information regarding these tokens will be inevitably overwritten. To empirically validate our theory provide results simply raising the EMA coefficient from ? = 0.99* to y = 0. 0.9999*<. 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 the final row of Table 24 the larger coefficient still manages improve performance on the longer prompts by enabling the preservation of the frequency com- ponents from the target 'needle' over 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 's network weights.
35


Reviews:
[{"review_id": "6oZtgvIgz2", "reviewer": "Reviewer_PVCR", "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.", "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? \n2. How important is it to use the  fourier transformation to the signals instead of the original attention map?\n3. 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "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.", "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. \n\n\nThere are some problems with the flow of the paper due to many of the details being relegated to the Appendix. \n- 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. \n- Please include limitations of the method in the main body or refer to the Appendix section where you do somewhere. \n- Limitations should also address the computational complexity of this method. How slow is NAMM compared to other methods? \n- 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. \n- 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. \n\n\nWeaknesses in experiments:\n- 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? \n- 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.\n  - 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? \n   - 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].\n- 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. \n\n\n1. Pyramid-BERT: Reducing Complexity via Successive Core-set based Token Selection https://arxiv.org/pdf/2203.14380\n2. Dynamic Memory Compression: Retrofitting LLMs for Accelerated Inference \n3. Suyu Ge,  et al.   Model tells you what to discard:  Adaptive KV cache compression for LLMs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "a1cQs9ax3P", "reviewer": "Reviewer_tak4", "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.", "questions": "Please see the weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea behind the contribution is simple and easy to understand.\n2. NAMMs can be zero-shot transferred to other Transformers, which is quite interesting.\n3. The benchmarks are diverse (but not comprehensive as discussed in weakness 2).", "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.\n2. 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:\n   - Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs, ICLR 2024\n   - CacheGen: KV Cache Compression and Streaming for Fast Large Language Model Serving, SIGCOMM 2024\n\nI 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "r0B64RYRah", "reviewer": "Reviewer_J4vq", "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. \nResults show that NAMMs outperform recent methods for pruning the KV-cache and occasionally improves the performance of the base model.", "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)?\n\nEvolution is typically expensive to train using. What are the computational resources used to train NAMMs in terms of the runtime and memory?\n\nConsidering 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?\n\nWhen 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? \n\nRelatedly, 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?\n\nCould you include details in the Appendix regarding the various datasets (e.g., Long Bench, InfiniteBench, LongVideoBench, MLVU, D4RL etc...)?\n\nThe 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?\n\nFor 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.\"\n\n\n\nInterestingly, 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The performance gains using NAMMs over H2O are significant. \n\nImproving the performance using NAMMs over the base model is impressive.", "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.\n\nFurthermore, 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.\n\nSeveral 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,\nLLaMA, 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. \n\nThe 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. \n\nBenchmarks 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "I45H99FElB", "reviewer": "Reviewer_845X", "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.", "questions": "1 Removing tokens with importance scores <0 implies that the number of removed tokens may vary each time. \n\n 2.How can the consistency of dimensions be guaranteed? Without consistent dimensions, subsequent training would be impossible. \n\n3.Does this require an additional constraint? I feel this point was not clearly explained in the paper.\n\n4.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.\n\n5.Could you list specific learning rates, parameter settings, and other details? Publishing the experimental code would also be beneficial.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "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.\n\n2. The NAMMs method proposed in the paper achieves a reduction in model parameter size, which facilitates training.\n\n3. 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.", "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.\n\n2 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？）\n a.How they handle varying numbers of removed tokens in practice\n b.If there are any mechanisms to ensure dimension consistency\n c.Whether any additional constraints are used \n\n3 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.\n\n4.Could you list specific learning rates, parameter settings, and other details? Publishing the experimental code would also be beneficial. (Please provide the following materials.)\n a.A table of hyperparameters used for each experiment\n b.Any details on hyperparameter tuning processes\n 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)