Neural Attention Search
Authors: Difan Deng
Marius Lindauer
Venue: NeurIPS 2025 poster ()
Keywords: Sparse Transformer
KV Cache Compression
Neural Architecture Search
Primary area: deep_learning
TL;DR: We present Neural Attention Search (NAtS), an end-to-end learnable sparse transformer

Abstract: We present Neural Attention Search (NAtS), an end-to-end learnable sparse transformer that automatically evaluates the importance of each token within a sequence and determines if the corresponding token can be dropped after several steps. To this end, we design a search space that contains three token types: (i) Global Tokens will be preserved and queried by all the following tokens. (ii) Local Tokens survive until the next global token appears. (iii) Sliding Window Tokens have an impact on the inference of a fixed size of the next following tokens. Similar to the One-Shot Neural Architecture Search approach, this token-type information can be learned jointly with the architecture weights via a learnable attention mask. Experiments on both training a new transformer from scratch and fine-tuning existing large language models show that NAtS can efficiently reduce the KV cache and the inference costs for the transformer-based models while maintaining the models' performance.

Neural Attention Search
Difan Deng Leibniz University Hannover
Marius Lindauer
Leibniz University Hannover
d. deng@ai. mni-hannover. de
L3S Research Center
m l.iidauerraa... uni hhanoverr de
Abstract
We present Neural Attention Search (NAtS), an end-to-end learnable sparse trans- former that automatically evaluates the importance of each token within sequence and determines if the corresponding token can be dropped after several steps. To this end, we design search space that contains three token types: (i) Global Tokens will be preserved and queried by all the following tokens; (ii) Local Tokens survive until the next global token appears; and (iii) Sliding Window Tokens have an impact on the inference of a fixed size of the next following tokens. Similar to the One-Shot Neural Architecture Search approach, this token-type information can be learned jointly with the architecture weights via a learnable attention mask. Experiments on both training new transformer from scratch and fine-tumie ane-tuning exist- ing large language models show that NAIS can efficiently reduce the KV cache size and the inference costs for the models while maintaining the models' performance.
Introduction
The ability to understand and infer from long-context information is crucial for many tasks such
as long document summarization [94] and question answering [20] 48), code generation [35] [66] or multi-round dialogues (9I]. Thanks to the ability to query the information from any position of the historical sequence, transformer-basee large language models 14
42 extend their
context length up to millions of tokens.
However, querying information from historical sequences requires complexity of O(L2) w.r.t. the input sequence length L. KV caching could reduce this time complexity to (L) by storing all the historical KV values. Nevertheless, with the increasing model size of recent LLMs, even the OLL) time-wise and memory-wise complexity could become a bottleneck during inference time.
Indeed, not all the tokens in sequence are equally important 46 Many of them are redundant and do contribute the final output. Humans can recognize this information without pre-defined fixed rules and summarize or discard the context information into much smaller content. Transformers could also learn this ability implicitly: Many tokens in the attention map might only have very low weights [96] and only have little influence on the final predictions. However, as the transformer learns this information implicitly, we might not know how the important tokens would distributed the context. Selecting these tokens and recognizing the attention distributions might require extra human experts' knowledge by either looking at the attention maps [27 (88 95| 96 or applying specific fixed rules 13 16 17 31 84 Since this knowledge is already contained in the transformer models, we could also ask the model to evaluate the importance of each token and learn to predict the optimal type for the given input tokens automatically.
Unlike prior works that rely on human expertise or predefined rules identify important tokens [15 27 28
96], we propose novel approach to evaluate the importance of each token by assigning different roles to each of the tokens. For example, some tokens will be preserved until the end, while other tokens might only survive for a short amount of time. These roles measure
54
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

the importance of each token and determine if it would survive within the next few tokens. Rather than pre-defining a set of fixed rules for each token, we ask the model to learn this information automatically. Finding the optimal role for each token is similar to the process of neural architecture search [23 1 where an optimal architecture is searched by the optimizer for given task. Thereby, in this work, we jointly optimize the choice of each token and the model weights by constructing learnable attention mask. Our approach is implicitly motivated by the one-shot neural architecture search approach [22] (55) 68), where the model parameters and architecture parameters are jointly optimized during the search process. However, our approach, Neural Attention Search (NAIS), searches for the optimal token roles jointly with the attention weights.
Our contributions are as follows:
1. We propose Neural Attention Search (NAtS), an end-to-end learnable sparse attention framework that automatically searches for the optimal roles for each token in the input sequence.
2. We introduce different token roles in our search space that can be later combined to construct a learnable attention mask and then jointly optimized with the model weights in NAIS,
3. We show that NAIS could efficiently reduce the KV cache required during inference time while maintaining most of the models' performance.
By automatically learning to focus on the most relevant information, NAtS paves the way for more efficient and scalable inference with LLMs in long-contex: applications.
Background and Related Work
2.1 Attention maps for transformers
Transformers 77] computes the correlation between different tokens by mapping the input sequences into three variables, Q. K, and V. An attention map A is first computed by pairing Q and K and then scaled and normalized with a softmax function. Finally, this attention map A is multiplied by V. Additionally, mask M is attached to the Attention Map to guide the attention to focus only certain tokens. This can be expressed as
QKT
A
O softmax(A + Madd)V
Vhhead
The additive attention mask M
(-inf,00 0} controls if the transformer needs to construct the correlation between Q, and Kj. A widely used mask Madd used in transformers is causal mask (79), where the upper parts of the mask are filled with inf; each token can only query the token information before it.
Computing the attention map of sequence with length L requires a complexity of O(L)) for both time and memory. This complexity can be lowered to O Leaahh)) during inference time, with Leaehe being the length of the KV cache. However, as expense, the cost of storing the KV cache gradually becomes unnegligible with the increasing size of large language models and the context length they could handle. Therefore, lots of work has been proposed to reduce the storage and computation overheads during the inference time [52].
To reduce the computational costs of an LLM on long context information, LLMLingua (66 trains another smaller model to compress the input texts. Many other studies work on reducing the KV cache size by deciding which tokens to evict or preserve with the past attention maps. Recent works such as H2O (96), Scissorhand (58), FastGen [1], KeyFormer [6]. SnapKV [54], PyramidKV ([5), A (27). ChunkKV (57), and CriticalKV [28 all apply different rules to identify the importance of each token and remove all the remaining tokens. Additionally, instead of simply removing the unimportant tokens, we can also merge them into the existing tokens [8]] 95).
Other works do not aim at reducing the KV cache size, but only select a subset of KV caches to compute the attention outputs, including Quest (76), TokenButler (9) AttentionPredictor [87], and XAttention [86]. However, these works still need to maintain the entire KV cache the GPUs This still results in huge GPU memory usage overhead. To alleviate this issue, some other works like Inf 8 HIP [5T), and ShodwKV (75] offload the KV cache to CPUs and only load the KV caches that are relevant for the current token predictions to GPUs.

---

While these existing methods offer various ways to reduce KV cache size, they often rely on inflexible predefined rules and potentially inaccurate heuristics based on past attention maps. NAtS, in contrast, introduces a novel approach that learns to dynamically assign token roles during inference, enabling a more adaptive and efficient use of the KV cache. By treating token role assignment as an optimization problem, NAtS leverages principles from neural architecture search to jointly optimize token roles and model weights, leading to a more flexible and powerful attention mechanism.
2.2 Sparse Attention
Sparse attention works on computing only a fraction of the attention map to reduce the computation costs and KV cache sizes with either pre-defined fixed patterns. including Sparse Transformers [17],
Longformer [13,, StreamingLLMM [84], and DuoAttention [85] Alternatively, this can be adjusted by special embedding tokens such as Longcoder [35 and SepLLM [16]. Additionally, one can also determine the types of sparse attention with a set of proxy values, including MInference [43], FlexPrefill [50]. SeerAttention (3). SpargeAttenntion [93], and XAttention [86]. However, these approaches aim to approximate the attention output of each layer from a pre-trained dense transformer that only involves the QK matrices. In contrast, NAIS learns the token roles directly from the final loss function and does not necessarily need to approximate the existing attention output. Hence, NAtS can be applied to both training from scratch or approximating the outputs from existing model. Additionally, the existing approaches can only be applied to approximating the attention maps during the pre-filling stages and still need preserve the entire KV cache during the decoding stage, making them less profitable for reasoning models where the model needs to generate lots of tokens during the decoding stage [2 In contrast, NAIS only checks the roles for each token and does not necessarily need to link to the actual attention map values. Hence, it can be applied to both pre-filling and decoding stages while keeping the KV cache sizes to low level.
In contrast to the fine-tuning free approaches above, we can also fine-tune the target model to achieve the desired sparsity. This includes Adaptively Sparse Attention [10] and Dynamic Memory Compression [64). However, both approaches rely on expensive cumulative productions during training time, making their approach less applicable to long-context scenarios. LLMLingua2 [66] trains another model to compress the target texts. Additionally, Kim et al. [46 asks the model to learn drop the tokens with lower attention scores. Landmark Attention
663 inserts landmark
tokens after certain time steps and applies these landmark tokens to summarize the information of the
tokens before that landmark token and recover. SPARSEK (60] only selects constant number of KV pairs for each query by introducing a differentiable SparseK operator. Mixture-of- only passes a limited number of tokens to each transformer layer. Native Sparse Attention [92] and MoBA [6]] all select a subset of the KV values to compute the attention masks and therefore reduce the computational overhead. However, these approaches still require preserving the entire KV cache values and only reducing computational costs during decoding stages.
[72]
MoA 29 proposes search for the optimal sparse attention for different heads, which share similar ideas to NAIS. However, MoA's search space only contains streaming LLM. This head-level search space of MoA is much smaller than our token-level search space, which might restrict the model's expressibility under even smaller budgets.
2.3 Neural Architecture Search
Designing a neural architecture for a specific task might require a lot of trial and error. Neural Architecture Search (NAS) automates this process by searching in pre-defined search space (24
Previous work on NAS mainly focused on searching within discrete search space by sampling a new
architecture from the search space, evaluating its performance, and updating the optimizers [98 99
One-Shot NAS (68) approaches instead share the same weights of the operators with respect to all the architectures in the search space. This allows to jointly optimize the architectures and weights. DARTS (55) and GDAS [22) further relax the discrete search space into continuous values to optimize the architecture parameters with gradient descent. The One- OnetShot..e NAS approach allows the optimizers to efficiently search for the optimal architecture within a huge search space. Similarly, NAtS has multiple options for each token as the search space and is able to search for the optimal token types jointly with the model weights. However, unlike One-Shot NAS approaches that consider the optimization problem as a bilevel optimization problem and optimize the model weights and architecture weights alternatively, we optimize the token state information and model weights within
3

---

one forward and backward pass. This is similar to mixture-of-experts (MOE) model 26 73 However, instead of distributing data across all experts uniformly, we only select one expert for each token and assign all data to that expert.
3 The NAIS Approach
In this section, we first introduce all the candidate token types in our search space. We then show that we can construct a learnable attention mask with the choice of each token type. Finally, we can efficiently reduce the KV cache size by dropping the unnecessary tokens during inference.
3.1 Search Space Design
Not all the tokens in a sequence are equally important. Just like a paragraph is composed of multiple sentences, sequence can be divided into multiple sub-sequences containing different information; some tokens might only be required within these sub-sequences. Hence, we design a search space for each token's 's role within the sequence and ask the model to automatically learn the optimal role for each token in the sequence.
We first define Global Tokens tokens containing important information that need to be preserved for the following predictions. Liu et al. [58] and Zhang et al. (96) showed that only a small fraction of the tokens contribute to most of the attention scores for self-attention computation. These tokens need to be preserved for models to recall the global information that helps to predict the next token. In vanilla transformer models, all the tokens are Global Tokens.
Global Tokens will not be evicted during inference time. Therefore, we should maintain as few Global Tokens as possible to ensure inference efficiency. Each Global Token should not only preserve the information within that position. Ideally, it should also be able summarize the information from previous sequences [16 35 Hence, split the entire sequence with the Global Tokens into multiple sub-sequences, with each sub-sequence ending with one Global Token. Each Global Token only needs to summarize the information from its sub- uub-sequences and the previous Global Tokens.
Local Tokens only survive until the next Global Token appears. Therefore, models will have the full attention within each sub-sequence to summarize the local sub-sequence information into the Global Token which is located at the end of the sub-sequence. Meanwhile, the model will be sparse within the input sequence. Local Tokens are considered as tokens that only provide lower-level information that helps the model to understand the local context information, but might not help further outside this context. This provides the model with flexible interface to control its sparsity based on the input context information.
Only the Global Tokens and Local Tokens might control the sparsity at a low granularity level. E.g., assuming that one input sequence is highly localized, each token only has high correlation with itself or a few neighboring tokens. In this case, they are all similar and are assigned with the same token type. However, none of the Global Token and Local Token could sparsify this attention map efficiently: if all the tokens are classified as Local Tokens, the input sequence will only be considered as one single subsequence, and all the Local Tokens will be equivalent to the Global Tokens.
Hence, we introduce Sliding Window Token. Sliding Window Tokens will only be preserved for the next W time steps and were previously considered one of the most popular sparse attention 96).
approaches
In contrast to other causal attention maps, Figure T(d) illustrates an exemplary attention mask constructed by the choices of different token types. In this case, we define the sliding window size as 4. Token 1, 4, 10 act as Global Tokens; Tokens 2, 6, 7, 8, 11 are Local Tokens; Token 3, 5, 9, 11 are Sliding Window Tokens. The Global Tokens splits the entire sequence into three subsequences that end at 4, 10, and the last index, respectively. Hence, Token 2 will only be required by Token 3 and 4. This rule applies the same for Token 6, 7, and 8, where they only interact with the tokens within the same subsequence. For the sliding window tokens, only the next three tokens could query their information, regardless of whether these tokens belong to the same sub-sequence. Only 6 out of 12 tokens are involved during inference time to make the next token prediction.
Combining different token types could already cover most of the sparse attention variables. For instance, assigning Global Tokens to the first few tokens and Sliding Window Tokens to all the other

---

Local Tokona
(a)
(b)
(c)
(d)
Figure I: A comparison between different casual attention maps. (a) The full attention map, where
each token is connected to the tokens hefore it. [(b) the local attention with sliding windows 3, every token will only get access to the information the tokens ahead. [[c] Longformer, besides the local attention, the first, 6th and 9th tokens are the pre-defined global tokens. T(d) NAtS dynamically
learns the optimal role for each token and constructs learnable mask based on the tokens' roles.
tokens results in streaming LLMs 84 and MoA 20 If we further all the tokens in some heads as Global Tokens, we will get DuoAttention 83 Assigning different tokens as Global Tokens results in the Vertical patterns [43] that appear many Token eviction strategies [54] 58]. Since different
layers or heads might have different numbers of global tokens, patterns like pyramidky (15] and adaky (27) can be easily derived from our search space. For a transformer network with N layers and H attention heads that needs to process one sequence of length L, given the three token types defined above, there are 3
possible configurations our search space. prohibitive to check all the combinations in such a huge search space, with the increasing sequence length that model receives. Here, we propose an end-to-end approach search for the optimal token types together with the network weights jointly.
3.2 Searching for the Optimal Token Types
Searching for the optimal token types within sequence is similar to searching for the optimal architectures for a given task in neural architecture search [24 81]. Following GDAS [22), we apply the Gumbel-Softmax trick 41 62] to sample from the discrete search space. The Gumbel-Softmax
trick allows the gradient to be backpropagated through the discrete choice of the token types.
Specifically, we first use linear layer (which we call Attention Score Layer) that maps each the input tensor for an attention layer E IRd to the likelihood for each option: a E R(H+Nqee) = Linear(x)) (x). where the number of KV heads and Nogts is the number of options in the search space. The type for each token is then sampled from a Gumble-Softmax function based on the likelihood values. Given that the number of options Nopus normally much smaller than the model head dimensions, this linear layer only introduces minimal overhead compared to the QKV linear layers.
We now construct learnable attention mask M with series of sampled token states. However, the additive mask in Eq. | will take - inf values, resulting in invalid gradient information. Hence, we the token information construct multiplicative attention mask Af {0, 71
e mul Σ, A is 3 M
(2)
The attention mask columns for Global Tokens and Sliding Window Tokens can be directly constructed since they will survive for fixed number of steps. However, the mask for Local Tokens M,,, is controlled by both the distribution from Local Tokens and Global Tokens as Local Tokens will survive until the next Global Token appears. In other words, to make M,, (j i) a valid value, no Global Token should appear between i and j.
For the sake of brevity, will use M instead of mul in the following part of this paper.
5

---

Formally, the attention masks can be created as follows:
M 1
(3)
if iijsitW W
M
(4)
(0
ifj>itW
M
II G.) n=
where W is the sliding window size and G, a Global Token at Position n. We then construct the attention masks based on the type of each token. After that, we mask out the upper triangular part of the mask to ensure its causality.
In practice, we first collect the index of the next global token 1,-- == min((klk >= in G2 1})
and rewrite Eq. 5 as:
[1
if 1 Gun if Seen
M
(6)
During the backward process, we collect the column-wise gradient values from our attention masks and apply these gradients to update the parameters from our Attention Score Layer. Different from other approaches that approximate the attention maps or attention outputs from existing model 9 22 877. NAtS can be learned directly from the target labels, and thus, can be applied for both pre-training a new sparse transformer from scratch or fine-tuning an existing transformer. To control the attention map's sparsity, we introduce small regularization value X that is directly applied to the gradient for Global Token and Local Tokens to encourage more Sliding Window Tokens. This regularization value À can be considered as a value similar to soft threshold, pushing the tokens whose contributions in the attention maps outside the scope of the sliding window sizes with lower levels towards Sliding Window Tokens. Hence, we could efficiently reduce the corresponding overall computational costs.
We show how the backward gradients are computed and how the values control the attention map sparsity in the appendix C
These rules are then integrated FlashAttention 18 19 to avoid explicitly computing the attention only need to collect
masks during the forward pass. In addition to the transformer computation, the next Global Token indices IGns (with a complexity of ((N) and then mask out the attention map values with the masks defined above. At the same time, we also skip all the computation blocks that do not contain any valid values (i.e., all the items in the mask of that block are 0) during the online attention process. Further details can be found in the appendix D
3.3 Efficient Inference with Different Token Types
NAtS can be applied to both pre-filling and decoding stages. During the pre-filling stage, we first compute the attention values for all the global and local tokens. We then compute attention map values for the sliding window tokens separately to fully utilize the parallelism of GPU architectures.
During the decoding stage, we dynamically map the input feature maps to the corresponding token types and discard the tokens no longer required by the following tokens once a new token arrives. The Sliding Window Tokens only survive for fixed number of time steps. We preserve queue in the KV cache space that only stores the most recent tokens and mask out the non- Sliding Window Tokens: when new tokens come, we place them after the tail of the queue to replace the tokens older than W.
Similar to the vanilla transformer forward process, when new tokens arrive, we concatenate them with the existing KV caches, generating new masks and computing the attention output. After that, we start our post-update process: we first check the state of each token to decide if we can remove them or keep them the cached values. Since different heads might disagree on the types of the same token, we record both the sizes for Global Tokens (Sizeg) and Local Tokens (Sizer.) for all the heads. New Sliding Window Tokens do not change these sizes since they will always be placed separately. However, when a new Global Token for any head arrives, we remove all the Local Tokens from the

---

corresponding heads and place the new Global Token right after the existing Global Tokens and then update our Size; and Size, accordingly. The same strategy is applied when new LocalT okens arrive: we place them at the end of the Local Tokens and increase the number for Local Tokens. A detailed updating process can be found under Appendix F
Figure Sesureeiiluree iillustrates an exemplary case to update the KV caches in NAIS with sliding window size of 3. The first two places store the most recent tokens and use mask to mask out the non-Sliding Window Tokens (yellow tokens). Since Token 8 for Head I is Sliding Window Token and Global Token for Head 2. We first move both tokens to the beginning of our cache to replace the old one. After that, we drop Token 8 in Head 1 since it has already moved to the sliding window caches. Then, since Token in Head 2 is Global Token, we drop all local tokens after the last Global Token (Token 1). Hence, we place Token after Token and remove Tokens and 6. The Sizeg is then updated from [5, 3] to [5, 4| and the Sizer, is updated from [0, 2] to [0, O]: both new tokens are merged into the existing tokens, we do not need the extra space to store them.
Bliding Window
Tokknss
Itwaldd Slkcing
Window Tosent
Ciistoii Troeess
Local Tokers
Figure An example of how caches are up- dated within NAIS when new tokens arrive with a model containing two heads. The two rows represent different heads.
4 Experiments
We implement NAIS based on the Flash Attention 2 [8] implementation on triton In our experiments, we first train NAIS parameters jointly within a transformer model from scratch. We then apply NAIS to fine-tune a large language model to show that NAIS could efficiently reduce the KV cache size required while maintaining the model performance. The codes for NAIS implementation and experimental designs can be found under:
com /ttp////ttaaeeeerrrraattettteeea
4.1 Training a Transformer From Scratch
We first apply NAtS to train GPT2 small style (69 transformer model with 128M parameters from scratch on the PG-19 Language Modeling Benchmark (70) with four Nvidia H100-PCIe GPUs and evaluate it on the test sets of PG19 with context length of 1024. Further details can be found under the appendix
As baseline, we train another dense transformer model under the same hyperparameter setting. During inference time, we compare NAIS with the following baselines besides the full Transformer: (i) Streaming LLM [84] only preserves the first few starting and the most recent few tokens for future prediction. (ii) H2O
[66] first computes the attention map and only preserves the tokens with the top-k attention scores. H2O and Streaming LLM are training-fre frree approaches that control the sparsity with pre-defined hyperparameters during inference time. In contrast, NAIS controls this sparsity by the sparsity hyperparameter value X. Hence, we train multiple models with different A. However, in our experiments, we observe that the attention map sparsity values converge much faster than the model loss. We could quickly estimate the attention map sparsity within the first 10,000 iterations to check if this sparsity value satisfies the required sparsity and early-stop the runs that do not satisfy the requirements [40].
We provide different hyperparameters for H2O and Streaming LLM for different sparsity (with the sliding window size of 32, 64, 128, 256, plus the same number of HH for H2O and 64,
170
153
128, 256. 512, plus 8 Sink tokens for Stream- ing LLM). This results in corresponding KV budget fraction ranging from 6.25% to 50%. Meanwhile, we train NAtS with the following X: 0, le 9, 9, le 8, 1e 7. We run each experiment with three different seeds.
15.0
14.0
5.10
KV Cocha Budget (5)
Figure 3: Perplexity vs KV Cache size under dif- ferent sparsity settings À on the PG19 dataset.
Figure 3 shows the perplexity and the corre- sponding standard deviation (as error bars) of
7

---

the PG19 test sets from different model settings across different seeds. The x-axis indicates the fraction of KV caches applied to generate the last token in the input sequence. As expected, larger À results in smaller KV cache fraction, while the performance is relatively stable across different seeds. Interestingly, even if À is set as O, the KV cache budget can still be reduced to around 35% with slightly lower perplexity compared to the full attention. This might indicate that the model also tends to sparsify the input information to focus more on the relevant parts from the input data [89].
H2O maintains nearly lossless perplexity with a budget of 50%. However, as the available budget decreases, NAtS's perplexity remains nearly the same until a KV cache size of around 10%. In contrast, H2O and Streaming LLM start to degenerate their performance starting from 25% budget allocation. The most sparse model in NAtS family only contains roughly 3% of the KV cache size (meaning roughly 30 cached tokens per layer on average) and achieves lower perplexity compared to H2O with 12.5% of the KV caches.
4.2 Fine-Tune a Large Language Model
We now apply NAIS to fine-tune an existing large language model 34 42 scenario. We construct new training dataset that follows the construction rules of LongBench as training set to fine-tune LLM models with NAIS. Some of LongBench's 's tasks are collected from the test sets of the previous benchmarks. Hence, we first collect the training sets from these benchmarks and construct these datasets following the data structure in LongBench. Additionally, we also add the passkey-retrieval dataset introduced in DuoAttention [83]. Overall, this dataset contains 7 000 instances with a maximal context length of 16 000. Further details can be found in Appendix F.2
in the long context
We only fine-tune the Attention Score Layer while keeping all the other parameters in the network fixed to approximate the original output from the corresponding base LLM. This approach is similar to AuoAtttnnion [85]. However, since want the model to capture the overall context information, we update Attention Score Layer with all the output from the full attention layer:
Lasstall
1 B 1=1==1
H
( full [
O NAss(?))
where B is the Batch Size. We only use the context information from the real-world dataset to ensure that the token information does not rely on specific prompts. However, the synthetic dataset might contain duplicated contexts, and hence, we only optimize the synthetic dataset with its corresponding labels [85] Additionally, we control the sparsity with the regularization value À instead of the additional loss item defined in DuoAttention, and therefore, only optimize NAIS with Laistal the
loss function. We fine-tune two long- cong-context models (Llama-3.1-882 In-t-ruttttt (34] and Mistral-7B- vo.3-Instruct (42) on two Nvidia H100 PCIe GPUs for one epoch using AdamW [47 59 with learning rate of 2e-3 with a warm-up from 2e-4 in the first 20% steps and reducing back to 2e-4 in the last 20% steps. We apply different À to allow for different sparsity. For fair comparison, we set the sliding window size W as 256, the same as DuoAttention.
In addition to the baselines in Section 4.1 we also evaluate other state-of-the-att art KV eviction strategies, including SnapKV [54). Adak (27], ChunkKV [57] PyramidKV (15). Critical V (28). These approaches rely a sliding window from the query matrix during the pre-filling stage to identify the important tokens. Unlike NAIS and the other approaches listed above that update the KV caches on the fly, these approaches only evict the KV cache once after the prefilling stage. Additionally, we evaluate DuoAttention (85) and MoA 29 Both approaches extend the existing streaming LLMS approach to provide different budgets for different attention heads. These two methods share the same idea as NAtS, where the importance of different KV caches should be learned instead of the pre-defined rules. However, instead of learning the importance of the head level, we aim at learning this information directly on the token level.
Since we only optimize the Attention Score Layer, the number of learnable parameters is Raaggrs dmodel Thheads Noptions: where Noptions is the number of options in our search space (in our case, it is 3). Hence, the size of the parameters that need to be stored is negligible (i.e., it only takes roughly 13MB on disk for each set of Attention Score Layer) compared to the weights of the LLM. Hence, users could pre-collect all these weights and apply the one that fits their compression requirement.
We evaluate all the baseline results on the RULER 38 and LongBench 12]. RULER a synthetic benchmark that evaluates a model's 's long-context capabilities across 4 task categories. LongBench
8

---

evaluates the model's 's capacity to understand information in different long context scenarios. Fol- lowing the setting from KV Press [65) we ask all the one-time KV eviction approaches evict the KV cache once the model receives all the context information before the question-related informa- tion arrives. For the other approaches, we eviet the KV caches dynamically when new tokens are generated.
Full
Duo
SLLM
H2O
Snap
Ada
Chunk
Critical
Pyradmid
MoA
NAIS
4k
95.06
94.25
59.37
83.03
54.90
70.03
69.19
87.52
56.41
44,63
95.32 (51%)
8k
93.22
92.42
48.58
81.51
54.10
71.43
67.20
89.01
54.55
19.80
93.30 (45%)
16k
90.53
88.24
50.97
65.55
54.38
76.32
67.07
88.03
58.17
15.16
90.01 (42%)
32k
85.95
84.65
45.68
11.48
47.13
72.48
55.88
82.88
48.05
12.79
86.16 (41%)
64k
84.01
83.54
51.40
0.07
69.61
78.10
71.05
81.93
70.49
10.86
84.09 (41%)
128k
73.83
73.11
46.95
0.00
39.26
64.49
41.61
70.53
40.02
11.05
74.77 (44%)
Full
Duo
SLLM
H2O
Snap
Ada
Chunk
Critical
Pyradmid
MoA
NAIS
4k
95.06
73.18
37.30
71.95
35.67
44.34
51.59
57.73
35.60
35.60
93.93 (24%)
8k
93.22
62.83
28.29
65.34
33.92
41.46
45.87
62.83
35.17
25.69
90.29 (19%)
16k
90.53
52.01
28.19
47.82
32.74
41.61
49.40
67.08
34.35
23.94
88.87 (16%)
32k
85.95
53.50
30.54
9.71
32.17
40.24
39.19
64.41
32.49
20.17
84.40 (15%)
64k
84.01
65.57
31.42
0.04
57.36
65.98
57.01
70.85
58.02
19.29
79.36 (15%)
128k
73.83
53.05
31.41
0.00
27.36
34.33
37.35
56.51
27.34
13.40
64.26 (16%)
Table 1: Ruler results with 50% KV budgets (top) and 25% KV budget (bottom) on LLama 3.1. All
Full models use the 100% KV budgets. We mark the actual KV budgets used by NAtS in the bracket. Models with the best performance besides the base Full Attention Models are bold
The results on the Ruler Benchmark with LLama 3.1 are shown in Table NAIS uses X == 3e - 7 for the 25% KV size level and x 5e 8 for the 50% KV size level. NAIS achieves nearly lossless scores across different input context lengths. When the KV budget drops to 25%, NAIS achieves 96% of the full attention scores with much smaller KV budget sizes, while the other baselines could no longer keep their performance with the decreased KV budgets. Results with can be found under Appendix G
For the LongBench tasks, we evaluate all the baselines with 50% and 25% of the KV cache sizes. We show the result with KV cache size of 25% for LLama3. in Table The model that is used in this task is a model with with (in many cases) smaller KV cache values. We provide further results, including results with Mistral-7B and NAIS trained with other parameters, in Appendix G
3e- - 7. NAIS achieves the best performance on most of the datasets
Full
Duo
SLLM
H2O
Snap
Ada
Chunk
Critical
Pyradmid
MoA
NAIS
NarrativeQA
31.35
22.62
27.83
21.07
27.08
29.02
26.27
30.48
28.02
19.19
30.53(13%)
Qusper
24.73
18.99
14.53
15.00
14.24
14.84
12.71
15.79
13.43
18.57
23.76(27%)
MuldfHeldQQA-en
29.46
25.36
16.47
18.24
17.72
19.31
17.44
17.72
14.75
28.94(24%)
MuldFieldQAAzh
60.01
54.95
33.94
29.14
33.77
35.34
33.49
36,30
32.76
21.23
61.18(22%)
HotpotQA
17.06
16.06
13.95
16.47
15.58
16.12
16.50
16.73
14.23
10.18
16.06/20%)
2WikiQA
16.64
16.18
13.53
11.28
13.45
13.17
13.41
13.89
12.36
11.42
17.05(23%)
Musique
11.59
8.41
8.96
11.13
10.81
10.68
11.05
11.25
9,40
5.93
11.44((19)
DuReader (th)
35.56
32.18
18.64
24.77
23.57
24.57
24.24
24.68
23.56
22.33
34.28(1756
GoyReport
34,30
27.63
28.20
26.90
28.51
28. 28
28.86
29.58
27.53
25.24
33.82(1996)
QMSum
23.30
22.28
20.67
20.39
21.41
22.05
21.27
22.27
21.98
19.59
23.11(17%)
MultNNwss
27.13
24.58
21.76
22.80
23.14
23.50
22.72
24.03
23.00
23,30
26.72(33%)
VCSUM (zh)
16.36
14.97
14.26
14.85
14.80
15.18
14.92
15.46
15.21
14.58
15.61(14%)
TREC
72.50
58.50
66.00
$5,00
53.50
$9.00
54,50
62.00
48.50
60,50
72.00(26%)
TriviaQA
91.15
82.81
90.69
89,98
91.00
91.55
90.72
90.97
90,92
71.51
91.61(22%)
MAMSum
43.72
40.14
41.97
41.98
43.59
43.11
42.32
43.79
43.15
42,30
44.21(16%)
LSHT
46.50
35.00
35.50
27.50
45.00
44.00
43,50
46.00
45.00
22.00
47,50 16%)
Passage Count
4.50
7.22
5.52
4.70
7.12(209))
RetiigeRetieeall-on
97.98
91.65
94.81
72.25
90.36
90.70
90.38
94.70
88.32
28.21
95.81(19%)
RussugeRetrieval-zh
77.99
55.22
28.97
52.13
73.06
78.74
55.95
80.15
74.66
24.78
78. 78.0%))%
LCC
54.10
54.27
52.38
54.52
55.40
54.93
54.12
54.26
54.14
53.54(38%)
RepoBench-P
51.39
57.26
49.65
55,54
51.80
52.44
52.87
52.57
52.41
53.35
53,
Table 2: LongBench Results with 25%
Budget Allocation. The Full model uses the 100% KV budgets.
The numbers in the brackets for NAtS the used value KV cache sizes. We bold the models with the best performance besides the base Full Attention Models.

---

Pre-filling Latency (8)
Decodingg Latency (ms)
Transformer
150
Wrnsforar
100
NAZS
NAIS
50
Pre fillifiling GPU Memory Usage (GB)
Decoding ng GPU Memoory Usage (GB)
60
40
Transformer
40
Transfomer
20
20
1050
100
200
300
400
500
600
700
100
200
300
4.00
500
600
700
Length
Length
Figure 4: Memory and latency usage during pre-filling (left) and decoding (right)
4.3
Latency Evaluation
We evaluate the latency and memory usage with the full attention Llama 3.1 with NAIS (a À 3e 7) during both pre-filling and decoding phases with different context lengths. The experiments implemented on a single Nvidia H100 PCIe GPU, and the model is stored as Bfloat16. We adapted RoPE (74) and RMSNorm kernel from FlashInfer [90] to accelerate the forward process. We apply chunked-prefillin
7 49 with a chunk size of 10 000 by dividing the input contexts into multiple smaller chunks reduce the peak memory usage from the intermediate activation values in the FFN layers. The full attention transformer quickly runs out of the GPU memory with a context length of around 200 000, while NAIS could efficiently reduce the pre-filling and decoding latency and memory usage, allowing NAtS to do inference to a context length up to 700 000, 3.5 times larger than the full attention transformer. For the context length of 200 000, NAIS consumes 2.24x less memory with 3.0x latency speed up during pre-filling and 2.6x less memory with a 1.38 decoding speed up.
5 Conclusion and Future Work
Efficiently managing the KV cache is crucial for deploying large language models in long-context applications. In this work, we propose NAtS, an approach that automatically optimizes the optimal roles of each token determine how long they should survive. By constructing learnable mask, NA1S learns to sparsify the attention map end-to-end. Our experiments show that NAIS uses much less KV caches compared to the State-of-the-an KV caching reduction approach. While NAtS shows promising results, future work could include exploration of further token roles and structured search spaces with hierarchies. Overall, believe that NAtS paves the way towards more efficient and scalable inference with LLMs.
6
Acknowledgement
The authors gratefully acknowledge the computing time provided to them on the nigh-performance computers Noctua2 at the NHR Center PC2 under the project hpe-prf-intexml. These are funded by the Federal Ministry of Education and Research and the state governments participating on the basis of the resolutions of the GWK for the national high performance computing at universities www.abr-verrin.dd//uuseeeeppp
Difan Deng was supported by the Federal Ministry of Education and Research (BMBF) under the project Al service center KISSKI (C0an1o.001822093C))..
References
[1] Proceedings of the International Conference Learning Representations (ICLR' 17), 2017.
Published online: iicrr.cc
[2] Proceedings of the International Conference on Learning Representations (ICLR' 199), 2019. Published online: iccr..co
[3] Proceedings of the 37th International Conference on Advances Neural Information Processing Systems (NeurIPS 233), 2023. Curran Associates.
10

---

[4] International Conference on Learning Representations (ICLR'24) 2024. Published online:
iclr.cc
[5] Forty-first International Conference on Machine Learning. 2024.
[6] M. Adnan, A. Arunkumar, G. Jain, P. Nair, Soloveychik, and P. Kamath. Keyformer: Kv cache reduction through key tokens selection for efficient generative inference. arXiv: 2403.09054 [cs.LG], 2024.
[7] A. Agrawal, Panwar, J. Mohan, Kwatra, Gulavani, and R. Ramjee. Sarathi: Efficient Ilm inference by piggybacking decodes with chunked prefills arXiv: 2308. 16369 [cs. CLJ, 2023.
[8] J. Ainslie, J. Lee-Thorp, M. de Jong, Y. Zemlyanskiy, F. Lebron, and S. Sanghai. GQA: Training generalized multi-query transformer models from multi-head checkpoints. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing (EMNLP). Association for Computational Linguistics, 2023.
[9] Y. Akhauri, AbouElhamayed, Y. Gao, C. Chang, N. Jain, and M. Abdelfattah. Tokenbutler: Token importance is predictable. arXiv: 2503.07518 CLs. 2025.
[10] S. Anagnostidis, D. Pavllo, L. Biggio, L. Noci, A. Lucchi, and T. Hofmann. Dynamic context pruning for efficient and interpretable autoregressive transformers. In Proceedings of the 37th International Conference on Advances in Neural Information Processing Systems (NeurlPS'2 233)
neu URL http://ppaaaeeeeeeeeeeeeetrre/aa
[11] Anthropic. The claude 3 model family: Opus, sonnet, haiku, 2024. URL
https: //www/wwwww//e cdn. anthropic com/ cttp:///cssoceaaaaaaafa003300000
Model_ Card Caar__arr__aaaeeeee Claude 3. pdf
[12] Y. Bai, X. Lv, J. Zhang, H. Lyu, J. Tang, Z. Huang. Z. Du, X. Liu, A. Zeng L. Hou, Y. Dong, J. Tang, and J. Li. LongBench: A bilingual, multitask benchmark for long context understanding. In L. Ku, A. Martins, and V. Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics, pages 3119-3137. Association for Computational Linguistics, 2024.
[13] Beltagy, M. E. Peters, and Cohan. Longformer: The long- long-document transformer. 055500000.05150 [cs. CLJ, 2024.
[14] Brown, Mann, Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry. A. Askell, S. Agarwal, A. Herbert- Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh. D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Lan- guage models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.-F. Balcan, and H. Lin, editors, Proceedings the 34th International Conference Advances in Neural Information Processing Systems (NeurIPS 20), pages 1877-1901. Curran Associates, 2020.
[15] Z. Cai, Y. Zhang, B. Gao, Y. Liu, T. Liu, K. Lu, W. Xiong. Y. Dong, B. Chang, Hu, and W. Xiao. Pyramidkv: Dynamic KV cache compression based on pyramidal information funneling. arXiv: 2406. 02069 [cs. CL), 2024.
[16] G. Chen, H. Shi, J. Li, Y. Gao, X. Ren, Y. Chen, X. Jiang, Z. Li, W. Liu, and C. Huang. Sepllm: Accelerate large language models by compressing one segment into one separator. arXiv: 2412.12094, [cs. CL), 2024.
[17] R. Child, S. Gray, A. Radford, and I. Sutskever. Generating long sequences with sparse transformers, arXiv:1904. 100.0900400099 [Gs. 2019.
[18] T. Dao. F: Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations (ICLR'24) iel 4. Published online: iclr.cc
11

---

[19] T. Dao, D. Fu, S. Ermon, A. Rudra, and C. Ré. FlashAttention: Fast and memory -efnnoreefffiienttt exact attention with io-awareness. S. Koyejo. S. Mohaaed, A. Agarwal, D. Belgrave, K. Cho. and A. Oh, editors, Proceedings of the 36th International Conference on Advances in Neural Information Processing Systems (NeurIPS 22), pages 16344-16359, Curran Associates, 2022.
[20] Dasigi, K. Lo, I. Beltagy, A. Cohan, N. Smith, and M. Gardner. A dataset of information- seeking questions and answers anchored in research papers. In Toutanova et al. 777, pages 4599 4610.
[21] DeepSeek-AL. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv:2501. [cs. CLJ, 2025.
[22] X. Dong and Y. Yang. Searching for a robust neural architecture four gpu hours. In Proceedings of the International Conference Computer Vision and Pattern Recognition (CVPR' 19). Computer Vision Foundation and IEEE Computer Society, IEEE, 2019.
[23] Elsken, J. Metzen, and F. Hutter. Neural Architecture Search. F. Hutter, L. Kotthoff, and J. Vanschoren, editors, Automated Machine Learning: Methods, Systems, Challenges, chapter 3, pages 63-77. Springer, 2019, Available for free at pottt://autone org/book
[24] T. Elsken, J. Metzen, and Hutter. Neural Architecture Search: A survey. Journal of Machine Learning Research, 20(55):1-21, 2019.
[25] A. Fabbri, I. Li, T. She. S. Li, and D. Radev. Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model. In J. Burstein, C. Doran, and T. Solorio, editors, Proceedings of the 2019 Conference the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1074-1084. Association for Computational Linguistics, 2019.
[26] W. Fedus, B. Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research (JMLR), 2023.
[27] Y. Feng, J. Lv, Y. Cao, X. Xie, and S. Zhou. Ada-kv: Optimizing KV cache eviction by adaptive budget allocation for efficient LLM inference. arXiv: 2407. [cs. CL), abs/2407. 1550, 2024.
[28] Y. Feng, Lv, Y. Cao, X. Xie, and S. Zhou. Identify critical kv cache in Ilm inference from an output perturbation perspective. arXiv: 2502. 03..5sc0...s.. CLJ, 2025.
[29] T. Fu, H. Huang, X. Ning, G. Zhang, B. Chen, T. Wu, H. Wang, Z. Huang, S. Li, S. Yan, G. Dai, H. Yang, and Y. Wang. Moa: Mixture of sparse attention for automatic large language model compression. arXiv: 2406. 14909 [cs. 2024.
[30] Y. Gao, Z. Zeng, D. Du, S. Cao, P. Zhou, J. Qi, J. Lai, H. So, T. Cao, F. Yang, and M. Yang. Seerattention: Learning intrinsic sparse attention in your llms. arXiv: 2410.13276 [cs.CL), CLJ, 2024.
[31] S. Ge, Y. Zhang, I. Liu, M. Zhang, J. Han, and Gao. Model tells you what to discard: Adaptive KV cache compression for llms. International Conference on Learning Representations (ICLR'24) icl []. Published online: iicrr.c
[32] Team Gemini and et al. Gemini: A family of highly capable multimodal models. vrXiv:2312.11800 (es.CL). CLJ, 2024.
[33] B. Gliwa, I. Mochol, M. Biesek, and A. Wawer. Samsum corpus: A suman-annotated dialogue dataset for abstractive summarization. arXiv: 1911.12237 [cs. CLJ, 2019.
[34] A. Grattafiori, A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, and et al. The Ilama 3 herd of models, 2024. URL httptttts///xii....t/ org/ wwtg/ba///aee abs/2407 22783
[35] D. Guo, C. Xu, N. Duan, J. Yin, and J. McAuley. Longcoder: A long-range pre- traited language model for code completion. A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning (ICMI.'23) 23), volume 202 of Proceedings of Machine Learning Research, pages 12098-12107 PMLR, 2023, URL wtttts::t/ppooeeeeeeeee
mlr ppess/v202/ v202/ v202/guo23/ hhml
12

---

[36] W. He, K. Liu, J. Liu, Y. Lyu, S. Zhao, X. Xiao, Y. Liu, Y. Wang. H. Wu, Q. She, X. Liu, Wu, and H. Wang. Dureader: a chinese machine reading comprehension dataset from real-world applications. In Proceedings of the Workshop on Machine Reading for Question Answering@AA @AI. 2018, Melbourne, Australia, July 19, 2018, pages 37-46, 2018. URL httpst /aaclanthology... org/ WWt/ V18-2605/
[37] X. Ho. A. Duong Nguyen, S. Sugawara, and A. Aizawa. Constructing A multi-hop QA dataset for comprehensive evaluation of reasoning steps. In Donia Scott, Núria Bel, and Chengqing Zong, editors, Proceedings of the 28th International Conference on Computational Linguisties, COLING 2020, pages 6609-6625. International Committee on Computational Linguistics, 2020.
[38] C. Hsich, S. Sun, S. Kriman, S. Acharya, D. Rekesh, F. Jia, Y. Zhang, and B. Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv: 2404.06654 [cs. CL.J, 2024.
[39] L. Huang, S. Cao, N. Nova Parulian, H. Ji, and L. Wang. Efficient attentions for long document summarization. In Toutanova et al. (77), pages 1419-1436.
[40] K. Jamieson and A. Talwalkar. Non-stochastic best arm identification and Hyperparameter Optimization. In Gretton and C. Robert, editors, Proceedings the Seventeenth International Conference on Artificial Intelligence and Statistics (AISTATS' volume 51. Proceedings of Machine Learning Research, 2016.
[41] E. Jang, S. Gu, and B. Poole. Categorical reparameterization with gumbel-softmax. In Proceedings of the International Conference Learning Representations (ICLR 17) icl (I] Published online: ccclr.cc
[42] A. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, D. Chaplot, D. de las Casas, and et al. Mistral 7b, 2023.
[43] H. Jiang, Li, C. Zhang, Q. Wu, Luo, S. Ahn, Han, A. Abdi, D. Li, C. Lin, Y. Yang. and L. Qiu. MInference 1.0: Accelerating pre-filling for long-contens LLMs via dynamic sparse attention. In Proceedings of the 38th International Conference on Advances in Neural Information Processing Systems (NeurIPS'244 24). Curran Associates, 2024.
[44] M. Joshi, H. Choi, D. Weld, and L. Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In R. Barzilay and M.- Y. Kan, editors, Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 1601-1611. Association for Computational Linguistics, 2017.
[45] Andrej Karpathy. NanoGPT.
cow www.o//aaaaaaaaaaaaoooooo http:///////ggggoooo/ooaaaaaaaaa nanoCPP 2022.
[46] S. Kim, S. Shen, D. Thorsley, A. Gholami, W. Kwon, J. Hassoun, and K. Keutzer. Learned token pruning for transformers. In A. Zhang and H. Rangwala, editors, Proceedings of the 28th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD 22), pages 784-794. ACM Press, 2022.
[47] D. Kingma and J. Ba. Adam: A method for stochastic optimization. In Proceedings of the International Conference on Learning Representations (ICLR'15), 15), 2015. Published online: iclr.cc
[48] T. Kociský, J. Schwarz, P. Blunsom, C. Dyer, K. Hermann, G. Melis, and F, Grefenstette. The narrativeqa reading comprehension challenge. Transactions of the Association for Computa- tional Linguistics, 6:317-328, 2018. URL https:////d/ org/ 10.
http:/aall_____000
[49] W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C. Yu, J. Gonzalez, H. Zhang, and I. Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, SOSP '23, page 6-16626. New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400702297.
[50] X. Lai, J. Lu, Y. Luo, Y. Ma, and X. Zhou. Flexprefill: A context-aware sparse attention mechanism for efficient long-sequence inference. International Conference on Learning Representations (ICLR 25), 2025. Published online: iclr.co
13

---

[51] H. Lee, G. Park. Y. Lee, J. Suh. J. Kim, W. Jeong, B. Kim, H. Lee, M. Jeon, and S. Hwang. training-free sub-quadratic cost transformer model serving framework with hierarchically pruned attention. arXiv: 2406.09827 CI.) 2024.
[52] H. Li, Y. Li, A. Tian, T. Tang, Z. Xu, X. Chen, N. Hu, W. Dong, Q. Li. and L. Chen. A survey on large language model acceleration based on kv cache management. arXiv: 2412.1944, [cs.Al], 2024.
[53] Xin Li and Dan Roth. Learning question classifiers. In 19th International Conference on Compulational Linguistics, COLING, 2002.
[54] Y. Li, Y. Huang, B. Yang, B. Venkitesh, A. Locatelli, H. Ye, T. Cai, P. Lewis, and D. Chen. Snapkv: LLM knows what you are looking for before generation. arXiv: 2404. 14469 (Ls.CL), 2024.
[55] H. Liu, K. Simonyan, and Yang. DARTS: Differentiable architecture search. Proceedings of the International Conference on Learning Representations (ICLR'19) 19) icl [2]. Published online: iclr.cc
[56] T. Liu, C. Xu, and J. McAuley. Repobench: Benchmarking repository-level code auto- completion systems, 2024. URL WRLLhttss///arti wwwp/rgg///aaaaeeeeeeee 03091 Published online: iclr.cc
[57] X. Liu, Z. Tang, P. Dong, Z. Li, B. Li, X. Hu, and X. Chu. Chunkkv: Semantic-preserving kv cache compression for efficient long-context Ilm inference. arXiv: 2502.00299(cs.CL CLJ, 2025.
[58] Z. Liu, A. Desai, F. Liao, W. Wang, V. Xie, 7. Xu, A. Kyrillidis, and A. Shrivastava. Scis- sorhands: Exploiting the persistence of importance hypothesis for LLM KV cache compression at test time. In Proceedings of the 37th International Conference on Advances in Neural Infor- mation Processing Systems (NeurIPS'23) neu
httttttt////eeeeeeeee wtwpiiiwwwwweeffooo net/forum?
id id=JZfg6wcigg
[59] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In Proceedings of the International Conference on Learning Representations (9C icl 2 Published online:
iclr.ec
[60] C. Lou, Z. Jia, Z. Zheng, and K. Tu. Sparser is faster and less is more: Efficient sparse attention for long-ran range transformers. arXiv:2406. 16747, [cs. CLJ, abs/2406.16747, 2024. URL https:://oi..
org/10. 48550/ arXiv. 2406 16747
[61] E. Lu, Z. Jiang, Liu, Y. Du, T. Jiang, C. Hong, S. Liu, W. He, E. Yuan, Y. Wang. Z. Huang, H. Yuan, S. Xu, X. Xu, G. Lai, Y. Chen, H. Zheng, J. Yan, J. Su, Y. Wu, N. Zhang, Z. Yang, X. Zhou, M. Zhang, and J. Qiu. Moba: Mixture of block attention for long-context llms. arXiv: 23102.1389 [Cs. 2025.
[62] C. Maddison, Mnih, and Y.Teh. The concrete distribution: continuous relaxation of discrete random variables. In Proceedings of the International Conference on Learning Representa- tions (ICLR'17) 17) icl URK.jtttp://opprrvvvis httpap/s//pppaerrreeeaaee//r//as online: iclr.cc
Published
[63] A. Mohtashami and M. Jaggi. Random-access infinite context length for transformers. In Proceedings the 37th International Conference on Advances in Neural Information Processing Systems (NeurIPS'23) 23) neu (3. URL HtRL/https:/// htttps://openrevier: http://pwww.peeeeeeeeeea/aaaaaaa nett/forrumP?k-TThhh44Wvy
[64] Nawrot, A. Lancucki, M. Chochowski, D. Tarjan, and E. Ponti. Dynamic memory compres- sion: Retrofitting 1lms for accelerated inference. In Forty-first International Conference Machine Learning iem (5)
[65] Nvidia. kvpress. http::///gihuub..cm////DD/
2022.
[66] Z. Pan, Q Wu, H. Jiang, M. Xia, X. Luo, Jue Z., Q. Lin, V. Rühle, Y. Yang. C. Lin, H. Zhao, I.. Qiu, and D. Zhang. LImlingua-2: Data distillation for efficient and faithful agnossi prompt compression. arXiv: 2403.12966 [cs. CLJ, 2024.
14

---

[67] PG19. https://wwws wttp://bwew.gttbbbrggggg.. org/
[68] H. Pham, M. Guan, B. Zoph, Q. Le, and J. Dean. Efficient Neural Architecture Search via parameter sharing. Dy and A. Krause, editors, Proceedings of the 35th International Conference Machine Learning (ICML'I8) volume 80. Proceedings of Machine Learning Research, 2018.
[69] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, Mishkin, J. Clark, et Learning transferable visual models from natural language supervision. arXiv:2103.000220
[cs. CVJ, 2021.
[70] J. Rae, A. Potapenko. S. Jayakumar, C. Hillier, and T. Lillicrap. Compressive transform- ers for long-range sequence modelling. Proceedings of the International Conference on Learning Representations (ICLR (2), 2020. URL httpp:://penneevee..neee/ee/eeee
SylKikSYDH Published online: iclr.ce cc
[71] Y. Rao, W. Zhao, B. Liu, J. Lu, J. Zhou, and C. Hsieh. Dynamicvit: Efficient vision transformers with dynamic token sparsification. In M. Ranzato, A. Beygelzimer, K. Nguyen, P. Liang, J. Vaughan, and Y. Dauphin, editors, Proceedings of the 35th International Conference on Advances Neural Information Processing Systems (NeurIPS: Curran Associates, 2021.
[72] D. Raposo, S. Ritter, B. Richards, T. Lillicrap, Humphreys, and A. Santoro. Mixture-of-depths: Dynamically allocating compute in transformer-based language models. arXiv: 2404.02255 [cs.LG], 2024.
[73] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Le, G. Hinton, and J. Dean. Outrageously large neural networks: The sparsely-gate gated mixture-of-experts layer. In Proceedings of the International Conference on Learning Representations (ICLR' 17) icl []. Published online: iclr.cc
[74] Su, M. Ahmed, Y. Lu, S. Pan, W. Bo, and Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, page 127063, 2024. URL httpp::::/// 10. 1016/ neucom. 2023. 127063
[75] H. Sun, L. Chang, W. Bao, S. Zheng, N. Zheng, X. Liu, H. Dong, Y. Chi, and B. Chen. Shadowky: Kv cache in shadows for high-throughput long-context Ilm inference. arXiv: 2410.21465 [cs. 2024.
[76] J. Tang, Y. Zhao, K. Zhu, G. Xiao, B. Kasikci, and S. Han. QUEST: query-aware sparsity for efficient long-context LLM inference. Forty-first International Conference on Machine Learning icm [
[77] K. Toutanova, A. Rumshisky, L. Zettlemoyer, D. Hakkani-Tür, I. Beltagy, S. Bethard, R. Cot- terell, T. Chakraborty, and Zhou, editors. Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Tech- nologies, NAACL-HLT 2021, Online, June 6-11, 2021, 2021. Association for Computational Linguistics.
[78] H. Trivedi, N. Balasubramanian, T. Khot, and A. Sabharwal. Musique Multihop questions via single-hop question composition. Trans. Assoc. Comput. Linguistics, 10:539-554, 2022.
[79] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, I. Jones, A. Gomez, L.. Kaiser, and I. Polo- sukhin. Attention is all you need. In I. Guyon, U. von Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Proceedings of the 31st International Conference Advances in Neural Information Processing Systems (NeurIPS'17) Curran Associates, Inc., 2017.
[80] Z. Wan, X. Wu, Zhang, Y. Xin, Tao, Z. Zhu, X. Wang, Luo. J. Xiong. L. Wang, and M. Zhang. D2o: Dynamic discriminative operations for efficient long-context inference of large language models. arXiv: 2406. 13035 [cs. CL). 2024.
[81] C. White, M. Safari, R. Sukthanker, B. Ru, Elsken, A. 7ela, D. Dey, and F. Hutter. Neural architecture search: Insights from 1000 papers. 2001002301 [cs.LG), 2023.
15

---

[82] H. Wu, M. Zhan, H. Tan, Z. Hou, D. Liang, and L. Song. VCSUM: A versatile chinese meeting summarization dataset. In Anna Rogers, Jordan Boyd- Graber,,, and Naoaki Okazaki. editors, Findings of the Association for Computational Linguistics: ACI. 2023, pages 6065-6079. Association for Computational Linguisties, 2023.
[83] C. Xiao, P. Zhang, X. Han, G. Xiao, Y. Lin, Z. Zhang. Z. Liu, and M. Sun. Infllm: Training-fres long-context extrapolation for Ilms with an efficient context memory. arXiv: 2402.04617 [cs.CL], 2024.
[84] G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis. Efficient streaming language models with attention sinks. arXiv:2309. 17453 [cs. CLJ, 2023.
[85] G. Xiao, J. Tang, J. Zuo, J. Guo, S. Yang, H. Tang, Y. Fu, and S. Han. Duoattention: Efficient long-context IIm inference with retrieval and streaming heads. arXiv:2410. 10819 [cs.CL], 2024.
[86] R. Xu, G. Xiao, H. Huang, J. Guo, and S. Han. Xattention: Block sparse attention with antidiagonal scoring. arXiv: 2503. 16428 [cs. CL), 2025.
[87] Q. Yang, J. Wang. X. Li, Wang. C. Chen, L. Chen, X. Yu, W. Liu, J. Hao, M. Yuan, and B. Li. Attentionpredictor: Temporal pattern matters for efficient Ilm inference. arXiv: 2502,04077 [cs.CL), 2025.
[88] Z. Yang, P. Qi, S. Zhang. Y. Bengio, W. Cohen, R. Salakhutdinov, and C. D. Manning. Hot- potqa: A dataset for diverse, explainable multi-hop question answering. In Ellen Riloff, David Chiang, Julia Hockenmaier, and uchi Tsujii, editors, Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 2369-2380. Association for Computational Linguistics, 2018.
[89] T. Ye, L. Dong, Y. Xia, Y. Sun, Y. Zhu, G. Huang, and F. Wei. Differential transformer, 2024.
[90] Z. Ye, L. Chen, R. Lai, W. Lin, Y. Zhang, S. Wang, T. Chen, B. Kasikei, V. Grover, A. Krishna- murthy, and L. Ceze. Flashinfer: Efficient and customizable attention engine for Ilm inference serving. arviv: 2501.01005 [Cs.DC), 2025.
[91] Z. Yi, J. Ouyang. Y. Liu, T. Liao, Z. Xu, and Y. Shen. A survey on recent advances in Ilm-based multi-turn dialogue systems. arXiv: 2402.18013 [cs.CL], CLJ, abs/2402.18013, 2024.
[92] Yuan, H. Gao, D. Dai, J. Luo, L. Zhao, Z. Zhang, Z. Xie, Y. Wei, L. Wang, Z. Xiao, Y. Wang. C. Ruan, M. Zhang, W. Liang, and W. Zeng. Native sparse attention: hardware-alignee and natively trainable sparse attention. arXiv: 2502. / 1089 [cs. Cts) 2025.
[93] Zhang, C. Xiang, H. Huang, J. Wei, H. Xi, Zhu, and J. Chen. Spargeattn: Accurate sparse attention accelerating any model inference. arXiv: 2502.18133 [cx.LG), 2025.
[94] T. Zhang, F. Ladhak, E. Durmus, P. Liang. K. McKeown, and T. Hashimoto. Benchmarking large language models for news summarization. arXiv: 2301. 13848 [cs.CL], 2023.
[95] Y. Zhang, Y. Du, G. Luo, Y. Zhong, Z. Zhang, S. Liu, and R. Ji. Cam: Cache merging for memory-efficient Ilms inference. In Forty-first International Conference Machine Learning net/forunm Ld-LCTmpp11165
iem 5 URL /ttpps://ppenreiiiii
[96] Z. Zhang, Y. Sheng, T. Zhou, T. Chen, I.. Zheng, R. Cai, 7. Song, Y. Tian, C. Ré, C. W. Barrett, 7. Wang, and B. Chen. H2O: heavy-hitter oracle for efficient generative inference of large language models. In Proceedings of the 37th International Conference on Advances in Neural Information Processing Systems (NeurIPS'22 23) neu 3. URL Settps//openrret
net/
[97] M. Zhong, D. Yin. T. Yu, A. Zaidi, M. Mutuma, R. Jha, A. Awadallah, A. Celikyilmaz, Y. Liu, x. Qiu, and D. Radev. Qmsum: A new benchmark for query- mumti-domain meeting summarization. In Toutanova et al. (77), pages 5905-5921.
[98] B. Zoph and Q. Le. Neural Architecture Search with reinforcement learning. In Proceedings of the International Conference on Learning Representations (ICLR' icl D Published online: iicr.coo cc.
16

---

[99] B. Zoph, V. Vasudevan, J. Shlens, and Q. Le. Learning transferable architectures for scalable image recognition. In Proceedings of the International Conference on Computer Vision and Puttern Recognition (CVPR' 18). Computer Vision Foundation and IEER Computer Society, IEFE, 2018
17

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors their evaluation. While "[Yes] " is generally preferable to "[No]" it is perfectly acceptable to answer " "[No] provided proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in a binary way, acknowledge that the true answer often and write a justification to elaborate. All supporting evidence can appear either in the main paper the supplemental material, provided in the appendix. If you answer [Yes] question, in the justification please point to the section(s) where related material for the question can be found.
nuanced, please just use your best judgment
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In this paper, we present NAIS in Section 3
The experimental results in show that strong performance and efficiency in different benchmarks.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Yes, we discuss the limitations in Section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
18

---

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specing well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification:
This paper does not include theoretical results
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: The detailed experiments set up and hyperparameter configurations are de- scribed in Section 4 and F The search space is described in Section and the algorithm is
described in Section D We also provide the codes under dateerhttps::.. //t///ttttt thub.. com/ automl/
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
19

---

If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g.. with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g.. registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Yes, all the datasets used in this paper open-sourced. We further provide a script on how to generate the training set for fine-tuning an LLM with NAIS under:
https://gittaei /ttp///gttttbttt// www.coo.caa.a.. www.aa////eeuaaaeentataa
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ /nips.cc/public/gu.dee http://gguuiee//eeeuuubbiiiiiiii wwwpannisssiooooooooico for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
20

---

6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g.. data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: Yes, we provide the detailed experiments setup and hyperparemter configura- tions in Section 4 F
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question:
Does the paper report error bars suitably and correctly defined or other appropriate
information about the statistical significance of the experiments?
Answer: [No]
Justification: All the experiments in this paper run only once with fixed random seed.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes, we report the hardware and resources we used to train (4 Nvidia H100 PCIe GPUs for 18 hours) /fine-tune the model (2 Nvidia H100 PCIe GPUs in 8 hours) in Section 4 and
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
21

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Yes, we have reviewed the NeurIPS Code of Ethics. NAIS is a methodology paper and does not have ethical concerns.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Yes, we have discussed the broader impact in Section B
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has socictal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper poses such risks.
Guidelines:
22

---

The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Yes, the data set and models used in our paper are open sourced and we cite the original paper/github that provides the data and model in this paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of license (e. CC-BY 4.0) should included for each asset. For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperssiiipssiiiiaea... paperssithcode thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of derived asset (if has changed) should be provided. this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
This paper does not release new assets
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details the dataset/code/moddll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
23

---

Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [Yes]
Justification: The experiments parts involve fine tuning an existing LLM model. We used model.
NAtS fine tune a Meta-Llama-3.1 InS-tsrrutt and
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non non-standard components.
Please refer to our LLM policy for what should or should be described.
24

---

A
Limitations
foundation While NAtS shows promising results that efficiently reduce the KV cache sizes and our search space already contains most of the KV cache eviction strategies. However, the construction of the attention mask still column-win wise oriented and mostly focuses on the vertical style sparse attention [43]. Hence, this work does not involve the block-sparse attention and slash attention that widely observed in other frameworks [43 (5)] Which in turn might result in a more sophisticated searching and inference process, related to the value of the Q matrix. One potential future direction would be to introduce a more structured search space that allows for the construction of more sparse attention types.
longer easily drop the desired tokens since they also be
B Broader Impact
LLMs are widely applied to different fields nowadays. However, the cost for LLM to store the KV cache and predict the next token is still huge, given the (L) computation and memory costs of full Attention Models. This prevents further adaptation of LLMs (and other transformer- models with properties similar to those described before) because of high energy consumption and limited context windows. NAtS achieves a substantial reduction in KV cache size with minimal impact model performance, outperforming existing state-of-the-art approaches. This increased efficiency can enable the deployment of larger, more powerful language models on resource-constraingg devices and facilitate the development of new applications that rely long-context understanding, such as advanced conversational AI, comprehensive document summarization, and complex code generation. By making long-context processing more accessible, NAIS has the potential to accelerate progress in natural language processing and related fields. Nevertheless, it does not solve other inherent problems of LLMs such as hallucinations.
C Details on Backward Propagation
C.1 Gradients for Attention Masks
To compute the gradients for M, set (A, M) e @ M; then the gradient for M is:
80 80 0g
(8)
0M 0g dm
0g
(9)
0M
Ög
A o M
(10)
0A
In Eq. and 10 show the gradient for M is the same the value that 80/0A is supposed to be
no mask is applied. Since we have J.,j = Ar. @ Mij- Let's set S, = E, 9i,j and Pt.j := 96.3 Then have dP dOVT. Therefore,
S:
dg.; = (diag( S
-
ddP,
(11)
S..
Combining Eq. with Eq. 10 and Eq. provides the same gradients as the vanilla softmax function with additive masks:
d.,. = (diag( AA
A
p
(12)
St. St.
dA dM M
= (diag - P. PT ddP
(13)
Since dg. required to compute the gradient for A;; and always needs to be computed. We can directly use this information to compute the gradients for the attention Mask M.
25

---

Following FlashAttention [9], we define D, do then
dMi,j = eAus (dPij D:)
(14)
S::
dAi.j P.j(dP:j D:)
(15)
and dAA is computed by Eq- 13 After that, needs to be recomputed anyway in the flash attention's 's backward process, this only results in little
backpropagate dA dq and du Since M
computational overhead.
However, in practice, we cannot guarantee an upper bound for en with Ma,5 since S S;
Ej E,s @ Mig- As result, lim,(4.,))ss (aا = co when Mig = O. Hence, we first clip S.
within (0, 1) and then compute dMM with the clipped value:
P = min( Ai.
1)
(16)
S.
dM Pdy D.)
(17)
C.2 Details computing the gradients for computing the token gradients info da
The gradients towards each token are collected through the column sum of each value weighted by the corresponding attention masks:
da. Σ dMds x M.
(18)
Where (G, L, )W) is the discretized token type. Intuitively, this shows the model's preference over short-range or long-range correlations: is quite close to j. then all the a will receive the same gradient information. However, if the model wants create a long-range correlation with >> j. only Global Tokens will receive the gradient information. The network will therefore prefer to classify the corresponding tokens global tokens.
Eq. shows that the Global Token controls the local mask size. Therefore, the gradients for Global Token should also be regularized by the gradient information from the previous tokens:
OM
i-1
II (1-Ga) == 1
(19)
aG
In cases where G₁ is O, Equation 19 the negative value of Equation 5 However, for the case where
G₂ 1. this equivalent the local mask values where is no longer set as a Global Token. This requires us find the index the next global token Gaea and the last global token I GGum where / is the length of the local sequence that ends at k.
This gradient information will then be collected and subtracted from the computed Global Tokens gradient values:
Σ
ML m.n x dMMm,,
if G. = 0,
mSi
Cre
da
G-
(20)
Σ
dMmmn
if ifG; = 1,
Guel 2m>=i
f>n2 Gra
d da da.
(21)
Intuitively, this gradient term da." G- checks if the new Global Token needs to be inserted into the current sub-sequence (when is 0) or we should remove the current Global Token to enlarge the current sub-sequence (when G is 1). We further illustrate this process in the appendix.
Figure 5 illlustrates an example of this. Token is global token, and we search for its next Global Token (which is Token 10 in this example). Assuming that want to change Token 4 to another
26

---

role, the regions within the boundary (orange ones) are those tokens that are influenced by this swtching: given that Token longer becomes Global Token, Token 2 and will be part of larger subsequence that ranges from I to 10, and the attention maps within the orange region should not be masked out. Intuitively, this gradient measures the regret that we made in order to switch one Global Token into Local Token. A similar idea can be found in the red region. Assuming that we want to switch Token 7 to a Global Token, then Tokens 5,6 will be split into a new subsequence and we will no longer connect them with Tokens 8, 9, 10 since they belong to two different sub-sequences after the switch. Hence, this value in the red regions measures the regret we mistakenly classify a Global Token non-Global Tokens.
In practice, the da G values for Global Token can be easily collected by looking at the last and next G index with scan function to collect the gradients to the corresponding positions. However, computing da, G- with G; O requires the gradient information from all the past token mask information. This might be computationally prohibitive in practice, since the valid values contained in Equation 20 are sparse, we ignore the cases for G; = 0 and only compute the dense gradients for the case where G =.
Negalive Grader Valuess for Glebal Tokan 4
Negatiee Gradient Values for Glotal Token 7
10
10
Figure 5: The gradient term da, G- for Token and 7
C.3 Optimizing for sparser Attention Maps
The sparse regularization term À is directly applied to the corresponding gradients for the Global Tokens and Local Tokens. Intuitively, if the column-wise wise sum of the gradients for each iteration is smaller than A. the attention maps of that column would require further update gradients updates towards a sparser transformer, they will be push the projection layer to classify those tokens as sliding window tokens. This ensures that the global tokens and local tokens actually require certain amount of column-wise attention map values to keep them activated. Hence, the sparse regularization term A could also be considered values smaller than that will be encouraged to be filtered out. This value should therefore penalize the number of unmasked tokens for each row. While the number of unmasked tokens for Global Token and Local Tokens in column are L i and I Guaa - with 1, G.ues defined in Section 3.2 have:
soft threshold for the attention map values, where attention map
Hence, we
b-i L Gra
dG Gaparce X
(22)
dG X X L
(23)
Combining Eq. 18 20
22 and Eq. 23 we have:
c = Σ dMi,j x M x Messual + dG dGGoparee da. G-
(24)
dat = Σ dMi.j x M x Mossual + dG
(25)
da sw Σ dMa,3 x MSW sw x M
Mensual
(26)
with MG. MSW ,M,,, defined in Eq. 3. and 4 and M
M.gsuul is a casual attention mask.
D Integrating NAIS into FlashAttention
Flashattention [8] continuously loads Q. K, V values to the SRAM and only computes the attention maps within the loaded blocks. NAIS loads the corresponding token states information a and con- struct the masks M dynamically with Equation collects the row-wise gradients for computing dQ and column-wise gradients to compute dK and dV. Hence, the gradients da can be computed together with dK and d'V. We show how NAIS can be
During the backward process, Flashattention
27

---

Algorithm NAIS forward pass on FlashAttentionz we mark the NAIS related operations with blue
Require: Matrices Q K. E BLxd token states a E RLxNNp indices of the next token states IGum E ER in HBM, block sizes Be, Br, sliding window size W
1: Divide Q into T,
blocks QI....; ......QQT, of size B, x d each, and divide K, V, a, I
B
into T. =
blocks K4.....Ky.. VI...., Vr., of size Be x d, 01,,..., OT, of size B, x
B2
Nopts- and Guu 1,.... I
Gieu of size B, each.
2: Divide the output O Rxd into T, blocks Og...,, Or, of size B, each, and divide the logsumexp L into T, blocks ......: LT. of size B, each.
3: for S ST do
4: Load Q, from HBM to on-chip SRAM.
5: On chip, initialize 0 B,)xx xd E RB (0) (0) RB, m (0) (-00) ER. Load 0j and ISm, , to SRAM and check if the current block contain any valid values according to Equations 6 if Do Compute then
6: for 1 j S do
7:
8:
9:
Construct attention mask M.j; with ag. IGuu j. and W with Equations 4 Load K,, V, from HBM to on-chip SRAM.
10:
On chip, compute S()) = Q.K3 E R B.B..
11:
chip, compute m. m. max(m 1-11) row.max(S (S) E B. P1D) = ex rowsum(()))) R.B, m3() MMM E B. (pointwise), d()1 m! 1-m 5
12:
(5) ((-1)
On chip, compute o = diag(em) (mm -10 + end if
13:
14:
15:
end for
On chip, compute O; diag( ) (T(
On chip, compute L, m; (T.) + log e (T.))
17:
Write O to HBM as the i-th block of Ó.
19:
Write L; to HBM as the i-th block of L.
20: end for
21: Return the output O and the logsumexp L.
We highlight the difference between the vanilla Flashattention2 and Flashattention2 with NAIS in blue. Compared with vanilla Flashattention models, NAIS adaptively skips the blocks that do not have valid mask values to avoid unnecessary computation. These invalid blocks will only be applied to update the gradients for da.
integrated into Flashattention2 18] in Algorithms I and 2
E Detailed Inference Process
In Section 3.3 we showed that NAIS could efficiently update the KV cache values and drop the unnecessary tokens. Here we provide a detailed pseudocode for this updating process.
As shown in Algorithm 3 once receive type. Depending on the new token type, we either: 1. append the new KV values after the next global tokens and remove the remaining local tokens 2. append the new KV values after the existing tokens 3. put the new KV values to the tail of the sliding window token queues (located at the beginning of the KV cache values)
KV cache pair, first check its corresponding
F Experiments Details
Here we discuss further details in our experiments. The codes for NAIS can be found under https:/ http//ggtttuu..../eee thub. comm///ttomm//eeeerlaaattttttttt
28

---

Algorithm 2 NAIS backward pass on FlashAttention2,, we mark the NAIS related operations with blue
Require: Matrices Q K. V, O, dO R2xdd token states RLxN iinicess of the next token states [Gus E Re in HBM, vector I. E TR in HBM, block sizes Bg, By, sliding window size W.
1: Divide Q into T,
N Br
blocks QI...., QT. of size B, x each, and divide K, V, a, G
in to Tc
blocks K I,...., Kr, and Vi...., Vr.,, of size Be d,
or, of size
Be
B, x Nopts and [ 1......... I of size B, each.
2: Divide O into T, blocks Oj....: Or, of size B, x each, divide dO into T, blocks do....,,do, dOr, of size B, x d each, and divide L into T, blocks Li...., Lr, of size B, each.
3: Initialize dQ= N in HBM and divide it into T, blocks dQ dQr, of size B, x d
each. Divide dK, dV E RR in to T. blocks dK dKr, and dV1, 1..... dVr. Te of size Be x each. Divide da E R Nog in to dai..... day, of size B, x Nopts each.
4: Compute D = rowsum(dO dO @ O) E R (pointwise multiply), write D to HBM and divide into T, blocks Di...., Dr, of size B, each.
5: for S j - Tg do
6: Load Kj, V j. 0j and IGeu, from HBM to on-chip SRAM.
7: Initialize dk, = (0) (0)B.xd d. dV, (0) daj (0) B. x Nogs on SRAM.
8:
for 1 < < T. do
9:
Load Q O,, dO,, dQ,, Li, D; from HBM to on-chip SRAM. On chip, compute 5 Q.K B.B.
11:
On chip, Construct attention mask M., with aj, I y, and W with Equations 3
On chip, compute 4 exp(S.) L,) E RR B. and P(() -P(()) M. Check if the current block contains any valid value if Do compute then
12:
13:
14:
On chip, compute dV, - dv, + ( do, E BRB.d
15:
16:
On chip, compute dP : d BR
17:
On chip, compute ds p(") @ (ap (
D:) E B.
Load dQ, from HBM to SRAM, then on chip, update dQ, +- dQ, +d. K, E ERBxxd and write back HBM.
18:
19:
On chip, compute dK, + dK, + ds end if
Q E BBRd
20:
On chip, clip pr(") to (0.1) and compute ds (f)
p( @ (dp! D2) E B.
21:
22:
On chip, update da, with Equations 24 25 26
23:
end for
24: Write dK,, j. and da, HBM.
25: end for
26: Return dQ, dk, dV.
E.1 Detailed Hyperparameter Setting for GPT2-small Training
Following the setting from NanoGPT 45 the GPT-2 style small has 12 layers and 12 heads with a
hidden dimension of 768. Instead of the learnable position encoding, we apply rotary embeddings [74]
to each transformer layer. The PG-19 dataset contains books extracted from Project Gutenberg 667
with about 2B tokens in the training sets. We train all models with context length of 1024 and batch size of 480 (using gradient accumulation). We train them for 600 000 iterations evaluate them on the test sets of PG19 with context length of 1024. Training one model can be finished in 16 hours with Nvidia H100 PCEi GPUs
F.2 Collecting the Fine-tune Training Set
To fine-tune NAIS LLMs, we collect the training datasets from different tasks for real-world tasks:
29

---

Algorithm 3 NAIS KV Cache Updating Process
Require: KV values K, V T sliding window size W, RLXNou Existing KV
cache K CCchh VCache E TR [ xdheade, number of global tokens natobal E TR number of
local tokens nocall R current tail of the sliding window queue
1: for 0 i < do
2: for 0 S h 11 do
3:
remove the tail of the sliding window queue: K&fache 0, V 0 a; is "lobalToken then
4:
move K Vi the (nk (nglobel)th >th position of the KV cache values: KCache
5:
n nglobal h
Ki VCache "A nglobal Vi,h
global global + 1, Rih focal update number of global and local tokens ni Tip
6:
7:
remove all the KV cache values after nit else if ag is LocalToker then move Ki,,, Vi, to the (n giobal + local )th position of the KV cache values: ngglobal +nckerat . = Viih- ni Jooal
global, KCachee O, Cachee = 0, for j > n
global
8:
9:
Caacee
ng +nefort = Ki,,,
10:
update number of local tokens nhh else
n +
11:
12:
move Ki,,, Van to the (t)th position of the KV cache values: K Ki VCache 4n h Vi
13:
end if
14: generate masks for the corresponding valid values
15: move the queue tail the next value: t (t 1) mod W
16: end for
17: end for
Multi-Document QA: HotPorQA
88
2WikiMultihopQA
37
MuSiQue
78 and
DuReader (zh) [36]
Single-Document QA: NarrativeQA
48
and Qasper 20
Summarization: GovReport [39] QMSum 197] MultiNews 25 and VCSUM (zh)
82
Few-shot Leraning: TREC [53] TriviaQA [44] and SAMSum 33
Code Completion: LCC 33 and RepoBench-P (56]
We also construct the synthetic passkey-retrieval dataset introduced in DuoAttention 85
This
dataset is generated by embedding multiple random passkeys in different locations with a long context. The model will then be asked to recall this passkey information.
For all the few-shot learning datasets, following Bai al. [1]. we randomly concatenate multiple question-answer pairs into one single extended context as one piece of data. The number of con- catenated samples for the TREC dataset ranges from [10, 100). This value is [2, 6] for TriviaQA and [10, 50] for SAMSum. Additionally, for the datasets that do not have enough context length (e.g., the DuReader dataset), we also merge multiple documents into one piece of data (in our case, this value is 4).
We do not collect all the data whose length goes beyond threshold ensure that the context can be fitted into our model. Additionally, we collect at most 500 instances in each dataset since some datasets might not contain data. the end, the real-world dataset contains 6436 data instances. We add another 564 instances from the synthetic dataset. In the end, our dataset contains 7000 instances. Finetuning model on this dataset for one epoch takes roughly 8 hours with Nvidia H100 PCIe GPUs.
30

---

Full
Duo
SLLM
H2O
Snap
Ada
Chunk
Critical
Pyradmid
MoA
NAIS
4k
93.68
92.63
58.89
5.60
46.50
57.92
59.33
74.06
44.00
55.73
93,68 (49%)
8k
91.29
89.94
59.21
2.82
38. 12
49.30
53.53
83.12
37.66
41.98
90.76 (43%)
16k
89.85
87.51
55.18
3.00
38.59
52.80
63.46
82.83
38.05
34.31
89.09 (41%)
32k
81.24
78.81
47.63
2.70
49.32
66.20
71.91
76.23
46.16
27.92
80.39 (42%)
Full
Duo
SI.LM
H2O
Snap
Ada
Chunk
Critical
Pyradmid
MoA
NA1S
4k
93.68
52.44
42.96
3.26
33.72
38.79
42.74
46.11
33.66
33,26
93.18 (29%)
8k
91.29
40.96
36.54
2.37
31.36
36.67
42.01
46.73
28.64
25.47
89.97 (24%)
16k
89.85
35.42
33.38
1.92
30.30
37.48
42.84
57.17
28.17
25.55
87.69 (21%)
32k
81.24
40.00
32.76
1.70
35.17
44.46
60.56
61.55
31.84
24.88
78.60 (21 %))
Table 3: Ruler results with 50%
KV budgets (top) and 25% KV budget (bottom) on Mistral-7B.
Instruct-v0.3, All Full models use the 100% KV budgets. We mark the actual KV budgets used by NAIS in the bracket.
G Further Experimental Results
G.1 Results on Ruler Benchmark with Mistral
Table the results of different optimizers on the Dystral-7B-Instruct--o model. In this case, the À for 25% and 50% budgets are 1e 6 and 3e respectively. We found that to achieve the same sparsity level, Mistral-7B requires a larger sparse regularization value À compared to the Llama 3.1 model. This might indicate that the attention map values in Mistral are more evenly distributed, and thus we need a larger À to force more tokens to only focus on the local information. The maximal context length of Mistral 7B is 32k. We only evaluate the models with maximal length of 32k. The result is consistent with the results shown in Table NAtS consistently outperforms the other baselines with smaller budgets.
G.2 Further results on LongBench
We show additional results on fine-tuning LLM on the LongBench dataset here. Table 4 Isowws the results with 25% KV budgets on I--strrttttt with À 6. The result confirms our conclusion that NAIS outperforms the other baselines in most datasets under similar budget level (and many times with even smaller budgets).
Full
Duo
SLLM
H2O
Snup
Ada
Chunk
Critical
Pyradmid
MoA
NAIS
NarrativeQA
29.21
16.49
27.56
23.65
22.70
24.34
21.22
27.72
21.97
21.67
26.7(9%)
Qusper
41.34
14.58
29.16
18.95
23.08
22.99
21.21
28.85
23.54
19.72
41.46/21%)
MuldFFeldQQAeen
52.51
24.83
31,88
31.52
38.56
38.11
34.11
45.18
35.34
28.80
52.36(18%)
MultFiclddQAzzh
58.03
29.15
30.47
29.21
31.42
32.30
32.40
38.
29.67
23.20
54.13(1658
HotporQA
49.62
42.10
46.93
47.48
45.94
39.64
15%)
2WikiQA
40.01
27.83
35.11
30.43
33.22
33.14
32.03
38.24
33.96
27.02
37.05(18%)
Musique
28.44
13.03
22.65
16.99
22.45
21.28
22.62
25.01
23.67
16.22
27. 97.88(149)
DuReader (2h)
34.87
26.51
18.69
20.66
25.34
27.55
28.45
27.04
24.82
19.11
35.25(12%)
GovReport
34.94
22.09
27.44 .4
27.11
28.31
27.94
28.94
29.33
26.94
25.58
33.39(
QMSum
25.67
17.25
22.11
20.84
21.43
21.95
21.95
22.77
22.39
19.53
23.70(13%)
MultiNews
27.85
23.41
22.01
23.03
23.77
23.62
23.34
24.37
23.89
22.55
27.35(28%)
VCSUM (2h)
16.34
13.58
1.5.41
15.13
15.28
15.38
14.69
16.10
14.96
14.87
16.06(119)
TREC
75.50
52.00
70.50
61.50
58.50
59.00
57.00
67.50
58.00
63.00
73.50(20%)
TriviaQA
88.89
84.85
88,92
86.60
88.89
89.39
89.19
88,98
89.31
86,94
89,80(179)
SAMSum
47.29
42.15
45.44
45.16
46.40
46.72
46.51
46.67
46.58
43.90
46.91(13%)
LSHT
39.75
17.50
29.00
22.00
38.50
39.00
39,50
39.25
37.50
17.25
5(11()
Passage Count
5.50
3.50
6.00
6.00
6.00
8.00
5.50
6.00
6.50
5.00
4.5(14%)
RussageRetrrevalleen
98.00
62.50
86.50
63.00
91.50
94.00
87,50
89.50
88.00
17.25
96.00/1396
PassageRettievalllt
96.50
34.50
71.00
16.00
92.00(14%)
LCC
53.02
50.87
$1.22
52.77
55.67
55.67
53.76
56.17
55.46
51.08
53.8
RepoBench-P
56.83
48.93
$4,80
55.74
56.82
57.06
56.75
56.75
56.40
53,59
56.7(23%)
Table LongBench Results with 25% Budget Size on Mistral-7B..
Tables aand show the evaluation results on LongBench with 50% budgets. Despite having fewer KV cache budgets in all the datasets, NAIS still achieves comparable performance on the LLLam3-8B model and better results on the Mistral model and generally performs comparable to the results with the full attention transformers.
31

---

Full
Duo
SLLMM
H20
Snap
Ada
Chunk
Critical
Pyradmid
MoA
NAIS
NarrativeQA
31.35
29.40
28.11 18.18
20.15
30.37
30.04
29.49
32.03
29.18
13.92
31.39(29%)
Qasper
24.73
20.37
18.15
18.09
18.44
19.42
17.24
22.13
18.54
22.25
25.17(45%)
MuldFieldQA-enn
29.46
26.70
17.89
22.50
23.21
24.62
22.55
26,30
22.95
14.24
28.94(419)
MulifieldQAAzhh
60.01
60.31
42.19
38.88
47.66
50.23
44.54
56.22
46.50
23.91
61.27(37%)
HotporQA
17.06
18.84
16.73
15.61
16.55
16.29
17.22
15.83
15.86
8.85
17.28(37%)
2WíkiQA
16.64
16.30
14.28
11.71
15.82
15.38
14.67
17.01
14.81
10.58
16.84(419)
Musique
11.59
10.75
10.12
11.19
10.63
11.89(36%)
DuReader (zh)
35.56
33.29
18.16
27.32
29.22
30.82
28.49
32.81
27.93
22.80
33,36(32%)
GovReport
34,30
32,99
30.84
29.46
31.34
31.51
31.92
33.09
30.54
25,45
34.23(36%)
QMSum
23.30
23.89
21.43
20.72
22.08
22.74
22.22
22.89
22.92
19.76
22.6(33%)
MultiNews
27.13
26, 29
24.77
25.10
25.27
25.48
25,08
25.81
25.10
24.86
27.17(51%)
VCSUM (2h)
16.36
15.70
15.29
15.03
15.70
15.77
15.59
15.89
15.32
14.67
16.24(28%)
TREC
72.50
72.50
71.00
65.50
60.00
66.50
63.50
70.50
59.50
64.00
43.50(445)
TriviaQA
91.15
90.41
91.44
90.25
91.47
90.97
91.63
91.14
90.98
81.17
42.14(40%)
SAMSum
43.72
42.68
43.41
42.57
44.21
43.44
43.00
43.41
43.91
39.64
42.61(30%)
LSHT
46.50
46.50
40.00
34.50
45.00
46.00
45.00
46.50
45.00
24.00
45.5(31%)
Passage Count
6.63
6.67
6.73
2.21
6.33
6.07
6.03
4.74
8.04
1,80
9.78(37%%
PassageRetrieval-e
97.98
98.55
96.67
93.54
95.98
97.88
95.83
98.12
96.10
17.96
97.32(36%)
Pussageeettieevllzt zhh
77.99
75.58
43.61
67.83
79.66
78.27
72.38
78.22
79.48
19.72
78. 17((4))
LCC
54.10
56.22
$3.10
55.61
54.08
53.13
54.23
52.33
$1.87
52.28
53.27(55%%
RepoBench-P
51.39
57.54
55.17
52.49
52.32
52.38
48.46
52.68(47%)
Table 5: LongBench Results with 50% Budget Size LLama 3.1 8B
Full
Duo
SLLM
H2O
Snap
Ada
Chunk
Critical
Pyradmid
MoA
NAIS
NurrativeQA
29.21
26.95
28.76
24:00
24.65
25.40
25.66
28.20
24.41
23.58
28.1(31%)
Qusper
41.34
35.65
35.80
27.88
31.68
31.94
31.38
38.10
31,50
28.60
41.75(46%)
MuldFieldQAAcn
52.51
52.50
37.35
40.15
46.12
48.52
44.73
52.20
45.70
35.52
52.79(42%)
MuldFieldQQA-zh
58.03
55.37
36.43
39.43
40.55
40.92
43.70
47.93
38.24
27.48
56.75(38%)
HotpotQA
49.62
52.94
47.66
47.24
49.34
49.57
49.51
47.20
$1.26
39.64
50.73(39%)
2WikiQA
40.01
38.09
37.27
36.75
40.24
37.49
40.77(42%))
Musique
28.44
29.45
27.24
20.61
26.43
25.77
27.28
28.54
25.85
16.84
28, 28.64(38%)
DuReader (zh)
34.87
36.44
18.63
26.82
30.29
32.18
32.72
32.
29.71
20.79
34, 7333%)
GoyReport
34.94
32.07
30.79
30.58
31.43
31.55
32,16
32.39
25.80
34.23 38%)
QMSum
25.67
24.20
22.91
22.48
23.16
24.36
23.47
25.12
24.11
21.08
25,94(37%)
MultiNews
27.85
26.97
25.18
25.62
25.97
26.24
25.83
26.26
25.56
24.39
27. 47(4%)
VCSUM (zh)
16.34
15.58
16.24
15.80
16.42
15.73
16.57
14.13
16.55(31%)
TREC:
75.50
74.50
74.00
66.00
66.50
68.50
69.00
74.50
64.50
69.00
75.00(499))
TriviaQA
88.89
87.37
89.06
88.36
88.71
88.81
88.69
88.08
89.55
87.05
89. 89.3(41%)
SAMSum
47.29
44.85
47.30
45.82
47.19
47.30
46.77
47.31
46.85
45.29
46.87(35%)
LSHT
39.75
37.50
33.00
31,50
39.25
39.25
39.25
40,25
38,00
21.00
38. 75(255)
Passage Count
5.50
6.00
5,50
4.00
4.00
5.50
5.00
5.00
4.00
4.50
7.50K39%)
Retrivaeen
98.00
99.00
89.50
81.00
97.50
98.50
97.00
98.00
98.50
33,00
98. 0(359)
PassageRetrieaal-hh
96.50
96.50
$1.50
76.00
89.75
95.50
92.00
97.00
96.50
17.50
96.0(35%)
LCC
53.02
53.22
52.94
54,40
54.50
54.14
53.79
54.63
53.78
$2.04
53. 53.66(56%)
RepoBench-P
56.83
55.89
55.47
57.55
56.76
57.05
55.68
56,38
56.36
54.67
57.38(52%)
Table 6: LongBench Results with 50%
Budget Size on Mistral-7B.
H
Ablation Study
H.1
Sparse Regularization Term A
We first study the impact of sparse regularization terms X. This value controls the efficient KV values cached the model and thus the model performance. The result shown in Table We underline the results that are better than the optimal baselines with 25% budgets, and bold the results that are better than the optimal baselines with 50%. Despite that NAIS in Table (NV ALS 3e - 7) used more than 25% overall KV budgets for some tasks, here show N AtS could still outperform many of the corresponding optimal basclines with an even smaller KV budget.
H.2 Sliding Window Length W
Another important hyperparameter for NAIS is the sliding window size W. We apply different sliding window sizes W (64, 128, 256, 512) to fine-tune the Llama3-8B model (with À = 3e 77).
The result is shown in Table 8 Overall, all the approaches perform similarly. However, a smaller sliding window size generally results in an overall larger KV cache size. A reduced sliding window size would force the model to apply more Global Tokens and Local Tokens construct the mid min-range
32

---

Full
NAIS Ic-6
NAIS 5c-7
NAIS 3c-7
NAIS lc-7
NAIS 5c-8
NarrativeQA
31.4
27.05(5%)
29.19(958)
30.53(33%)
31,39(298)
31.94(19%)
Qasper
24.7
11.39
28.6(21%)
23.76(27%)
25.17)45%)
(6.97
MultiFicldQA-cn
29.5
13%
27.41(18%)
28.94(24%)
29.7(549%) (8.94(419)
MultrFieldQA-zh
60.0
51.73(14%)
56.64(1%%)
61.18(22%)
61.27(37%)
(1987
HotporQA
17.1
15.47(10%)
17,44 15%%
16.06(20%)
17.28(37%)
16.65(49%)
2WikiQA
16.6
16,44 35.37(13%)
17.05(23%)
(4
16.75(53%)
Musique
11.6
10.78(9%)
9.87()33%
(1.42
(1.89(36%)
11.31
DuReader (zh)
35.6
30.06(9%)
(3%)
34,28( 17%)
33236
(454))
GevRepurt
34.3
30.82(10%)
(1.47)14%)
34.23(36%) 13%%
34.8 50%)
QMSum
23.3
23.37(8%))
23.09 (2%)
23.11
22.6(33%)
22.74(45%)
MultiNews
27.1
26.45(22%)
26.57.28%)
26.72(335%
27.17(51%)
(7.11
VCSUM (zh)
16.4
15.5607%)
(5.61(10%)
15.61(14%)
16.24(28%)
15.99
TREC
72.5
68.001456
70.5(20%)
72.0(26%)
73.5(44%)
73.0(56%))
TriviaQA
91.2
91.38( 2%%
90.32(17%)
92.61(22%
92.14(40%)
(1.64
SAMSum
43.7
43.98(8%)
44.0(12%)
44.21(16%)
42.61(30%)
43.84(419) 9
LSIIT
46.5
46 46.008%)
44.0012%)
47.5016%)
45.5(31%)
45.5(43%)
Passage Count
6.6
2.56(10%)
6.88(149))
7.12(20%)
9.78(37%)
8.82(50%)
Relriigeelrreevlleen
98.0
17(9))
82.65(14%)
95.81(19%)
97.32(369)
96.93(48%)
PassigeRetrieval-zh
78.0
33. 14(12%))
64.77(16%%
78.5(20%)
78.17(34%)
77.7(46%)
LCC
54.1
55.46(24%)
54.53(31%)
53,54(38))
54.4(699) 53.27(57%)
RepoBench-P
51.4
50.78(14%)
5185(21%)
53.48(28%)
52.68(47%)
52.36(59%)
Table 7: Ablation Study of Sparse Regularization values À for LLama3.1-8h
Full
NAIS
NAIS
NAIS 256
NAIS 512
NarrativeQA
31.4
21.61(33%)
30.76(15%)
30.53(13%)
28.96 28.96(12%)
Qusper
24.7
24.1(28%) 17.92(3%))
23.76(27%)
28.24(29%)
MulriFieldQA-cn
29.5
19.37(34%)
29.02(26%)
28.94(24%)
25.87(25%)
MultFicldQAAzh
60.0
35.49(35%)
60.37(22%)
61.18(22%)
58.08(26%)
HotpotQA
17.1
13.99(33%)
16.41(23%)
16.06(20%))
17.44(18))
2WikiQA
16.6
11.52(34%)
16.69(26%)
17.05(23%)
16.48(24%)
Musique
11.6
9.42(33%)
11.62(22%))
11.42(18%)
10.89(16%)
DuReader (zh)
35.6
22.26(34%)
35.75(19%))
34.28(17%)
33.91(16%)
GovReport
34.3
26.63(34%)
33.89)
33.82(19%))
33.05(19%)
QMSum
23.3
20.46(33%)
23.52(1958)
23.11(17%)
22.66(16%)
MultiNews
27.1
23.78(36%)
26.73(33%))
26.72(33%%)
26.46(40%)
VCSUM (zh)
16.4
13.8(34%)
16.1(15%)
15.61(14%)
15.9(15%)
TREC
72.5
46.5(34%)
73.0(31%)
72.0(26%)
70.5(25%)
TriviaQA
91.2
84.44(44))
01.89)66%)
91.61(22%)
91.66(21%)
SAMSum
43.7
38.97(34%)
43.62(18%)
44.21(169%)
44.47(17%)
LSHT
46.5
26.5(34%)
47.0(18%)
47.5(16%)
45.5(15%)
Passage Count
6.6
0.0(34%)
5.75(23%)
7.12(20%)
5.66(18%)
PassageRetrievalleen
98.0
7.13(33%)
96.71(23%)
95.81(199)
85.78(17%)
PussageRetrreevll-zt
78.0
5.15(34%)
81.64(21%)
78.5(20%)
70.66(21%)
LCC
54.1
39, 13(66))
54.92(39%)
53.54(38%)
14.15(44%)
RepoBench-P
51.4
36.87 36.87(34%)
51.89(32%)
53.48(28%)
51.34(26%)
Table 8: Ablation Study of Sliding Window Length W for Llama3
correlation since this distance cannot be covered by the sliding window tokens. However, as the number of sliding window sizes further increases, this compression rate might saturate since the remaining tokens might always require a long-range correlation whose distance is much larger than the sliding window size (e.g., these tokens might require the correlation between two tokens whose distances are larger than Ik or even more).
I KV Size Distributions
Here, we provide more KV distribution results with different datasets and hyperparameters on the longbench dataset.
Figures 6 and show the KV cache size distributions for Llama 3.1 8B on the Narrative QA and
Multi News Datasets. The models tend to preserve more KV caches in the shallower layers and ignore the KV caches in the intermediate layers. However, the model still preserves several KV caches near the output layers. Additionally, even in the same layer, the distributions of the KV cache sizes not
33

---

Unndd De Nartheo
Banod le Fnntrrateeo
lepsl
Teysl
200
mill
Figure 6: KV size of LLama Narrative QA Dataset
Figure 7: KV size LLama on Multi-News Dataset
evenly distributed in each layer. Some heads are always preferred while the others might be dropped as the context or À changes.
Micall led semeivon
Mizad Se- os Narssieeh
Moll
Wmill
Figure 8: KV size of Mistral on NarrativeQA Dataset
Misaal Hulliios
utaal 5o-3 Hudiies
Intil
Figure 9: KV size of Mistral on Multi-News Dataset
34

---

Figures 8 and show the KV cache size distributions from the Mistral model. Although the two models share some similar trends, e.g. both models tend to preserve more tokens in the shallower and the last few layers. However, compared to Llama, the KV caches for the Mistral model are more unevenly distributed and tend to gather towards some specific heads. Despite that, both Mistral models and LLaaa models are GQA models [8] and share similar architectures, the preserved token distributions can still be different. This shows that the KV importance distributions might not only depend the architectures, but are also closely related to the model weights, This highlights the importance of adapting different KV eviction strategies to different models.
alaut
Ind
Junlatt
Iuttla
Iu.
Ladg.. 5000
3000
Lad
Idl.daaull
[. Inth.h
alautaladu Lnl.bl.dd..
Figure 10: Number of Token Types for Multi-News Dataset
J Token Roles Distributions
Usmod toterr Nanatvook
hhtd dul
Lui
J.I.I.II....
L.
I.
I.L.Hal... J.I.I.L...l..
Figure 11: Number of Token Types for NarrativeQA Dataset
We also illustrate the token roles distribution Figure 10 and and oooo the Multi-News and Narrative datasets. Generally, most tokens are classified as either Global Tokens or Sliding Window Tokens, and only a few are considered Local Tokens. This might indicate that most attention operations still focus either long-range short-term correlations.
35


Reviews:
[{"review_id": "5XYDZEnzTq", "reviewer": "Reviewer_1gKy", "summary": "Neural Attention Search (NATS) introduces an end-to-end learnable sparse-attention framework that lets a transformer decide, during training, how long each token should remain in the key–-value cache.  Each token is classified as a global, Local, or Sliding-Window token via a Gumbel-Softmax search over a huge architectural space, and the chosen roles are encoded as a learnable attention mask. When applied to GPT-2-small (trained from scratch on PG-19) and to fine-tuned Llama-3-8B / Mistral-7B, NAtS cuts KV-cache size to 3 % – 50 % of full attention with negligible perplexity loss, enabling context windows up to 700 k tokens and 1.8×–2.6× speed-ups on an H100 GPUs", "questions": "- Could you report mean ± std over seeds (or another statistical test) for key metrics to substantiate that the observed ≤1 pp performance gaps are robust\n- does the per-token Gumbel sampling introduce communication overhead, and could the learned masks generalize across TP shards without retraining", "limitations": "no negative societal impact", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "strength:\n- the token-type search is optimized jointly with network weights and implemented inside FlashAttention 2, avoiding explicit masks. \n- Pushes single-GPU context length from 200 k to 700 k tokens with 2.24× memory savings during pre-fill.  \nweakness\n- The smallest lambda settings still incur up-to-34 % perplexity loss on certain tasks when KV budget falls below ~5%. \n- Search space is column-oriented; block-sparse and slash patterns are excluded", "was_revised": true, "final_justification": "Thanks for the authors for the responses, they largely resolve my concerns. I raised my score accordingly.", "rebuttal": "Dear reviewer 1gKy, we thank you for your constructive comments, \n\n> The smallest lambda settings still incur up-to-34 % perplexity loss on certain tasks when KV budget falls below ~5%.\n\nThanks for your comment, but we are not quite sure which result you refer to.\n\nIf we understand your comments correctly, you are referring to Figure 3 with the pre-training tasks, where the sparest model only requires a KV cache size of 2.3% to achieve a perplexity of 14.38. This perplexity value is only roughly 2.8% higher than the full transformer (13.99). Another candidate that requires 5.69% KV cache size achieves a perplexity of 14.08, only 0.6% higher than the full attention module. \n\nThe 34% perplexity loss might come from the highest value in Figure 3 (Streaming LLM with 12.5% KV cache size). In comparison, Streaming LLM receives a 19% perplexity loss with a KV cache size of 12.5% while H2O gets a 12% perplexity loss with KV cache size of  6.25%, both loss-wise and KV cache budget-wise worse than NAtS. \n\n> Search space is column-oriented; block-sparse and slash patterns are excluded\n\nWe agree that NAtS only incorporates the column-wise attention patterns. However, this decision also helps to efficiently reduce the KV cache size. The block-sparse and slash patterns must maintain the entire KV caches, as the valid values in the block-sparse and slash patterns need to be determined by the joint values of the Q, K, and V. Hence, we cannot easily discard KV caches, as they might be required by the following Q. Additionally, including the block-sparse and slash patterns might require additional checks for the types of the Q values, increasing the complexity during both training and testing phases from $O(L)$ to $O(L^2)$ as we need to check the types of both Q and KV values.\n\n> Could you report mean ± std over seeds (or another statistical test) for key metrics to substantiate that the observed ≤1 pp performance gaps are robust\n\nThanks for the advice. We have retrained some of the NAtS models ($\\lambda$ values with $0, 1e-8, 5e-9$). Due to the expense of training an LLM from scratch, we cannot afford the cost of training all the models from scratch with multiple seeds. The results are shown as follows:\n\n|             |   KV Budget | Perplexity     |\n|:------------|------------:|:---------------|\n| NAtS        |       0.338 | 13.980 (0.007) |\n| Transformer |       1     | 13.991 (0.007) |\n| H2O         |       0.5   | 14.048 (0.009) |\n| NAtS        |       0.093 | 14.074 (0.016) |\n| NAtS        |       0.057 | 14.104 (0.016) |\n| H2O         |       0.25  | 14.241 (0.012) |\n| H2O         |       0.125 | 14.704 (0.011) |\n| Streaming   |       0.5   | 14.730 (0.150) |\n| Streaming   |       0.25  | 15.431 (0.218) |\n| H2O         |       0.062 | 15.743 (0.027) |\n| Streaming   |       0.125 | 16.375 (0.269) |\n| Streaming   |       0.062 | 17.729 (0.339) |\n| Streaming   |       0.031 | 19.795 (0.444) |\n\nWe show the mean and std of the perplexity values with three repetitions in the table above. Despite that NAtS (with $\\lambda$ as 0) only requires around 33.8% of the required KV budgets, it achieves a lower perplexity compared to a full transformer with 100% KV budgets. Additionally, the table above is consistent with the results shown in Figure3: NAtS could achieve a much lower mean perplexity within the same budgets compared to the other baselines.  \n\n\n> does the per-token Gumbel sampling introduce communication overhead, and could the learned masks generalize across TP shards without retraining\n\nThe generated Gumbel samples only have a matrix of size $L \\times H \\times N_{ops}$ where L, H, and $N_{opts}$ are the sequence length, number of KV heads and number of operations, respectively. This is much smaller than the QKV tensors with size $L \\times H \\times D_{heads}$ where $D_{heads}$ is the head dimension. Additionally, we could also split head-wise the Gumbel samples (and their corresponding weights) across different GPUs, where each GPU only receives the Gumbel samples from the corresponding heads. The masks are generated for each head dynamically and independently within the computational kernels during both the training and validation process, and therefore could be directly applied across different TP shards without retraining."}, {"review_id": "rWywfN9M2C", "reviewer": "Reviewer_GRci", "summary": "The paper proposes Neural Attention Search (NAtS), a learnable sparse attention framework that assigns each token a discrete role—Global, Local, or Sliding Window—to control its lifespan and influence in attention computation. By jointly optimizing these token roles and model weights via a differentiable mechanism inspired by neural architecture search (NAS), the method aims to reduce KV cache size and inference cost, while maintaining the model's performance on long-context tasks.", "questions": "- How do you envision your method being applied to reasoning tasks that generate long sequences? \n\n- Can the framework be extended to support adaptive role reassignment during generation, or is it fixed once the role is predicted? What about for those multi-step tasks?", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n- The idea of casting sparse attention as a NAS problem is novel and provides a structured lens through which to learn token importance.\n\n- The method achieves both memory and latency reductions, which are crucial in long-context LLM inference and deployment.\n\nWeaknesses\n\n- While the NAS formulation is interesting, the underlying insight and motivation for using the NAS method is somewhat unclear. For example, other learning-based methods like SeerAttention directly learn the sparse attention mask—what are the practical trade-offs between learning token roles vs. learning masks? More discussion would help clarify the positioning.\n\n- The search space (Global / Local / Sliding Window tokens) feels hand-designed. It is hard to judge the validity of this space. For example, why are Local tokens dropped after a Global token? It would be helpful to explore or justify alternative or more flexible formulations.\n\n- The latency and memory results only compare with full attention, not with other sparse attention baselines (e.g., SeerAttention, MInference, DuoAttention). This makes it hard to judge the relative efficiency gains.\n\n- Section 3.3 (Efficient Inference) would benefit from a clearer and more formal algorithm description or pseudocode. The current writing is hard to follow and lacks clarity on KV cache update logic across heads and layers.", "was_revised": true, "final_justification": "I find the authors’ explanations regarding my concerns persuasive, and I have revised my score to 'weak acceptance.'", "rebuttal": "Dear reviewer GRci, \n\nWe thank you for agreeing with  NAtS’s efficiency and novelty. Regarding your concerns:\n\n\n> While the NAS formulation is interesting, the underlying insight and motivation for using the NAS method is somewhat unclear. For example, other learning-based methods like SeerAttention directly learn the sparse attention mask—what are the practical trade-offs between learning token roles vs. learning masks? More discussion would help clarify the positioning.\n\n\nBoth NAtS and SeerAttention propose learned sparse attention maps instead of fixed attention patterns. However, SeerAttention is trained to approximate the attention output of each layer from an existing dense transformer that only involves the QK matrix, whereas NAtS learns the token roles directly from the final loss function in an end-to-end manner that involves the QKV matrices. Hence, SeerAttention, as a post-hoc process, requires a dense transformer as a teacher model and can only be applied during the post-training stage. In contrast, NAtS does not necessarily require the teacher model and can be learned from scratch, allowing it to be used in both pre-training and post-training stages. \n\nAdditionally, SeerAttention still needs to maintain the full KV caches, as the decision of whether a block should be computed is determined by the gating values from both Q and K. There is no information telling the optimizer when a KV cache can be dropped. In contrast, NAtS evicts the tokens on the fly once the corresponding KV values are no longer required in the following attention computation process, and thus could efficiently reduce the GPU memory consumption. \n\nLast but not least, SeerAttention is only applied during the pre-filling stage, while NAtS can be used in both the pre-filling and decoding stages, as NAtS checks the types of each newly generated token and evicts non-global tokens with the corresponding rules for each token (as shown in Section 3.3) and thus could further reduce the computational costs for models with long outputs such as reasoning models. \n\nThank you for pointing this out. We will also describe the differences between our work and the existing sparse attention approach in our paper. \n\n\n\n> The search space (Global / Local / Sliding Window tokens) feels hand-designed. It is hard to judge the validity of this space. For example, why are Local tokens dropped after a Global token? It would be helpful to explore or justify alternative or more flexible formulations.\n\nGlobal tokens are the default settings in transformer models, while sliding window tokens are widely used in other sparse transformer models [1,2]. Therefore, we include them in the NAtS search space. These two token types can be considered as two extremes: global tokens will be kept forever, while sliding window tokens will only be kept for a fixed number of steps. However, in practice, we might also face cases where tokens need to be preserved longer than the existing sliding window size, but are no longer required after that. For instance, a piece of code might be important in its package but might not be necessary when its functions are imported into other packages. Since we cannot expect this distance to remain constant, as it depends on contextual information, we introduce local tokens whose lifetimes are learned during training. Hence, local tokens can be considered as a sliding window token with learnable sliding window sizes. Since this token type is an intermediate between global and sliding window tokens, we apply a global token as a signal that ends the lifetime of the local tokens. As we primarily focus on column-wise sparse attention, an alternative would be the introduction of sliding window tokens from different sliding window shapes, which could be covered by the local tokens.  We believe that these three roles are already sufficiently flexible to cover a wide range of different attention maps. \n\n\n\n> The latency and memory results only compare with full attention, not with other sparse attention baselines (e.g., SeerAttention, MInference, DuoAttention). This makes it hard to judge the relative efficiency gains.\n\nSince the other sparse attention baselines (SeerAttention, MInference) only work for the pre-filling stages and do not evict the KV cache, we mainly focus on the pre-filling latency here. To show the efficiency of NAtS during the pre-filling stage, we conduct experiments comparing the latency of NAtS against other baselines. With a sparse kernel designed for first computing attention value for the global/local and then the values for the sliding window tokens,  NAtS provides an overall lower latency compared to Duo Attention and Seer Attention across all context lengths, and better or comparable to MInference until 150k Context length. However, MInference can only be applied to accelerate the pre-filling stage and cannot provide further benefits towards memory reduction or decoding speed up, and thus cannot scale to a higher context length due to the GPU memory limits.\n\n| Prefilling Size |   FA2 |   MInference |   DuoAttn |   SeerAttn |   NAtS | \n|-------:|------:|-------------:|----------:|-----------:|-------:| \n|  10k |  0.69 |         1.53 |      1.11 |       0.48 |   0.59 | \n|  50k |  5.24 |         4.8  |      3.32 |       3.14 |   3.08 | \n| 100k | 16.43 |         7.99 |      8.03 |       8.14 |   6.98 | \n| 150k | 33.21 |        11.51 |     14.59 |      14.91 |  12.41 |\n | 200k | 56.98 |        15.74 |     22.1  |      24.15 |  19.22 |\n\n\n\n\n> Section 3.3 (Efficient Inference) would benefit from a clearer and more formal algorithm description or pseudocode. The current writing is hard to follow and lacks clarity on KV cache update logic across heads and layers.\n\nThanks for pointing this out. We will add pseudo-code to describe the process more precisely. The main message here is that we \n\n1. check the types of the KV cache from each head; \n\n2. Move the most recent sliding window token to the sliding window queues (located at the beginning of the KV caches); \n\n3. Check if global tokens are contained in the new tokens. If so, remove all local tokens from the corresponding heads and move the global tokens after the last global token of that head, and remove the remaining local tokens of that head.\n\n Thus, we could efficiently reduce the required KV caches and, furthermore, the corresponding computational times. \n\n\n\n> How do you envision your method being applied to reasoning tasks that generate long sequences?\n\n\nAs NAtS dynamically checks the token roles and evicts the local and sliding window tokens, regardless of whether the token belongs to the prefilling stages or is generated by the transformers, NAtS can be directly applied during the decoding stages. Additionally, since the NAtS token roles are independent of the Q values, following the KV cache update algorithm described above, NAtS can be directly applied to reasoning tasks without any further adjustment. \n\n\n> Can the framework be extended to support adaptive role reassignment during generation, or is it fixed once the role is predicted? What about for those multi-step tasks?\n\n\nOnce predicted, the token roles will be fixed. However, since we only update the projection layers (i.e., layers that map the input feature map values to the token types) with the next-token-prediction loss and do not introduce any auxiliary information about whether the corresponding token should be preserved, the roles of each token are assigned based on their impacts on all the potential following tokens and therefore should also provide further information in multi-step tasks. Additionally, since the token roles are not dependent on the Q values, we can ensure that the tokens preserved are also ready to be queried in the following steps, as the tokens that might only be used in the current step would be classified as local and sliding window tokens. Hence, despite the fixed token roles, NAtS can still be directly applied to multi-step tasks.\n\n[1] Child et al. Generating long sequences with sparse transformers\n\n[2] Xiao et al. Efficient Streaming Language Models with Attention Sinks"}, {"review_id": "sddx7a2FEf", "reviewer": "Reviewer_6LT7", "summary": "This paper proposes a neural attention search (NAtS) for three token types: global tokens, local tokens, and sliding window tokens. By learning an attention mask, NAtS is able to jointly learn token type information with architecture weights, similar to the single-shot neural architecture search method. The authors demonstrated in experiments training Transformer from scratch and fine-tuning existing large language models (LLMs) that NAtS can effectively reduce the KV cache size and reduce inference cost while maintaining model performance.", "questions": "NAtS can automatically learn token types. How does the ratio of global, local, and sliding window tokens change under different tasks and context lengths? Is this ratio change related to task characteristics (e.g., some tasks require stronger global understanding, while others focus more on local information)?\n\nNAtS treats token type search as a NAS problem. Compared with typical NAS methods, what are the characteristics of NAtS in terms of computational complexity or training stability during the search process? How does this \"optimizing token state information and model weights in a single forward and backward propagation\" approach reduce the computational overhead of NAS?\n\nIs there a more theoretical or in-depth analysis to explain how exactly λ guides the model to learn a sparse attention map? In practical applications, how do we determine the best λ value to strike a balance between performance and efficiency? Can we dynamically adjust λ based on the task or context?\n\nFigures 8 and 9 show the distribution of KV cache sizes across layers and heads, and mention that the KV cache distribution in the Mistral model is more uneven than that in Llama. Is this unevenness a feature of the Mistral model design itself (such as GQA/MQA), or does it reflect the specific sparsity patterns learned by NAtS on different model architectures? What guidance do these observations have for future research on sparse attention?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "+ The most notable advantage of NAtS is its ability to learn token importance end-to-end. By treating token role assignment as an optimization problem, NAtS leverages the principles of neural architecture search to jointly optimize token roles and model weights, avoiding the limitations of hard-coded patterns or heuristic rules commonly found in traditional sparse attention methods.\n+ NAtS can effectively reduce memory usage and accelerate reasoning in both the pre-filling and decoding stages, improving the efficiency and scalability of long text reasoning.\n+ NAtS is integrated into FlashAttention 2, which helps avoid explicit calculation of attention masks during forward propagation and can skip calculation blocks that do not contain valid values, further improving efficiency.\n- Although the search space of NAtS is flexible, the paper acknowledges that the construction of the attention mask is still column-oriented and focuses mainly on vertical-style sparse attention. Other important sparse attention types such as block-sparse attention and slash attention are not involved.\n- Compared with other training-free methods, NAtS requires training, which makes it difficult to apply it to large-scale models.\n- The paper mentions the time and number of GPUs required to train the Transformer model and fine-tune the LLM in NAtS, but there is a lack of in-depth analysis on the training complexity of the \"Attention Score Layer\" and whether this \"joint optimization\" method has significant advantages in training time or resource consumption compared to traditional NAS methods.\n- Although the paper conducted an ablation study of λ and W, a theoretical explanation or a more fine-grained empirical analysis of how λ precisely controls sparsity and its impact on model performance and KV cache size can be more in-depth. For example, why can the KV cache still be reduced to 35% when λ is set to 0? The impact of the sliding window size W on performance and sparsity is also worth further exploring its internal mechanism.\n- The appendix shows the distribution of KV cache sizes at different layers and heads. The reasons behind these distribution patterns or their deep connection with model performance and task types can be explained in more depth.", "was_revised": true, "final_justification": "Thank you for addressing some of my concerns. However, the comparison between methods that require training and those that do not still lacks fairness. I believe in and appreciate the potential and contributions of this work, but the authors may need to conduct more equitable experiments. Therefore, I have raised my scores for Clarity, Significance, and Originality.", "rebuttal": "Dear reviewer 6LT7, \n\nWe thank you for highlighting the flexibility and the efficiency of NAtS, regarding your concerns:\n\n> Other important sparse attention types such as block-sparse attention and slash attention are not involved.\n\nWe agree that NAtS only incorporates the column-wise attention patterns. However, this decision also helps to efficiently reduce the KV cache size. The block-sparse and slash patterns must maintain the entire KV caches, as the valid values in the block-sparse and slash patterns need to be determined by the joint values of the Q, K, and V. Hence, we cannot easily discard KV caches, as they might be required by the following Q. Additionally, including the block-sparse and slash patterns might require additional checks for the types of the Q values, increasing the complexity during both training and testing phases from $O(L)$ to $O(L^2)$ as we need to check the types of both Q and KV values, which might be pragmatic with the increased sequence length.\n\n> NAtS requires training, which makes it difficult to apply it to large-scale models.\n\nAs NAtS only trains the projection layers, which only provides a reduced memory overhead compared to fine-tuning a full model. Hence, we only need 2 H100 GPUs to fine tune an 8B LLM. \n\nAmong the existing training-free methods, H2O requires the explicit attention map values, which might requires lots of GPU memories and is not compatible with flash attention; Streaming LLM cannot get information in the middle of the context; SnapKV based approaches select the KV values based on the last Q values and thereby, relies heavily on the positions of Q values, which might lose essential information for scenarios like multi-round diagonal or reasoning models. Additionally, the pre-defined rules (such as the uniformly distributed token budgets for SnapKV, layerwise gradually decreased KV budgets for PyramidKV, or layerwise uniformly budget distributions for AdaKV) might not generalize to all the architectures and contexts. \n\nHence, since all these training-free approaches only make predictions based on the past information, they might fail to preserve the information that might be required by a longer context such as reasoning models or multi-round dialogue scenarios. \n\n> but there is a lack of in-depth analysis on the training complexity of the \"Attention Score Layer\" and whether this \"joint optimization\" method has significant advantages in training time or resource consumption compared to traditional NAS methods.\n\nThe Attention score layer is a linear layer that maps the input feature $X \\in R^{L \\times D}$ to $Logits~\\in R^{L \\ times (n_{ops}*h_{kv})}$ where $n_{opts}$ and $h_{kv}$ are the number of token types and number of KV heads respectively. This cost is minimal compared to the costs of computing the QKV values (roughly $n_{ops}/d_{head}$) where $d_{head}$ is the dimension of each head. Hence, overall, this operation also introduces a complexity of $O(L)$.\n\n The traditional NAS approach approximates a bilevel optimization process by updating the model weights and architecture parameters alternatively. Hence, within each updating iteration, only a small fraction of weights are updated. Since the number of parameters in the projection layer is much smaller than the number of parameters in the network, this might not be data-efficient. Hence, we propose to optimize the projection layers and the model parameters jointly to make full use of the training data, which is 2 times more data efficient than the bilevel optimization approaches in NAS frameworks. This approach is also applied in the mixture-of-expert models to optimize the gating parameters and model parameters jointly, showing its efficiency in training larger- scale models.  \n\n> a theoretical explanation or a more fine-grained empirical analysis of how λ precisely controls sparsity and its impact on model performance and KV cache size can be more in-depth.\n\nThe $\\lambda$ value is directly applied to the gradients of each logit for the token types. Hence, it can be considered as a regularization term that encourages more sliding window tokens. \n\nMore specifically, the gradients for each logit value are the column-wise sum of the corresponding attention mask values. As discussed in Appendix C, the gradients for the attention masks are the same as the gradients for each attention map value if no masks are applied (except for the normalization denominator). Hence, if the column-wise sum of the gradients for each iteration is smaller than $\\lambda$, indicating that the attention maps of that column requires no further update or gradients updates towards a sparser transformer, they will be push the projection layer to classify those tokens as sliding window tokens.  In contrast, if the column-wise tokens actually require larger attention map values, their corresponding gradients will be larger than $\\lambda$ and push the model towards global/local tokens. This ensures that the global tokens and local tokens actually require a certain amount of attention map values to keep them activated. Hence, this value could also be considered as a soft threshold for the attention map values, where attention map values smaller than that will be encouraged to be filtered out. \n\nThe reason that KV cache reduced to 35% when $\\lambda$ is set to 0 might be the result of the combination of Global Tokens and Local tokens, where local tokens encourage the correlations for the points that are closer to each other. Thus, we can drop those local tokens earlier if they might increase the model loss if those corresponding attention map values increase (as shown in Section C.2 and Figure 5). Additionally, some of the attention heads might only focus on the local dependencies. These heads might naturally surpass the attention map values with long-range correlations and therefore would prefer the sliding window tokens. \n\n We briefly discuss the impact of W and sparsity in Section G.2. Additionally, since the LLMA models use rope as their positional encoder, while the rope values decay as the distance between two data points increases, the sliding window size vs. sparsity might also be correlated to the rope parameters. \n\n\n> The appendix shows the distribution of KV cache sizes at different layers and heads. The reasons behind these distribution patterns or their deep connection with model performance and task types can be explained in more depth.\n\nNAtS learns the tokens roles only with the gradient information and without any prior information.  Overall, the distributions fit the observation from other works, such as the higher and intermediate layers are less important than the first few initial layers or the output layers. Additionally, only a partial fraction of heads are responsible for constructing long-term distributions, which are well supported by many architecture pruning works. However, as motivated in the paper, we are not trying to find a fixed pattern that works for any specific architectures or tasks, but are trying to ask the models to learn this information by themselves. \n\n\n> How does the ratio of global, local, and sliding window tokens change under different tasks and context lengths?\n\nOverall, the fraction of tokens changes across different tasks. For instance, on a RepoBench-P dataset, the fraction of global, local, and sliding window tokens is 23\\%, 2\\%, 75\\%, respectively. In contrast, for longer QA tasks, such as NarrativeQA, the fractions become 11\\%, 3\\%, 85\\%, respectively. Since the information density of codes is much larger than the information contained in natural languages. This further shows the importance of data-dependent budget allocation for different tasks \n\n> How does this \"optimizing token state information and model weights in a single forward and backward propagation\" approach reduce the computational overhead of NAS?\n\nPlease check our replies for weakness “whether this \"joint optimization\" method has significant advantages in training time or resource consumption compared to traditional NAS methods.”\n\n> Is there a more theoretical or in-depth analysis to explain how exactly λ guides the model to learn a sparse attention map? \n\nFor how $\\lambda$ influences the sparsity, please refer to our replies to “a theoretical explanation or a more fine-grained empirical analysis of how λ precisely controls sparsity and its impact on model performance and KV cache size can be more in-depth”\n\nRegarding the best $\\lambda$ values, since the sparsity values can already be observed in the early training stage (Line 274), we could determine if we would like to early-stop runs that cannot achieve the desired sparsity. Alternatively, we could also construct different $\\lambda$ values as different experts and consider the entire system as a mixture-of-expert system to propose the optimal $\\lambda$ values for the target tasks. Since $\\lambda$ controls the sparsity by directly modifying the weights for the token type logits, we could also control the $\\lambda$ values with the gradient information from the attention output $dO$, such as its norm values. However, for the sake of simplicity, we set a fixed value in our experiment.\n\n> Is this unevenness a feature of the Mistral model design itself (such as GQA/MQA), or does it reflect the specific sparsity patterns learned by NAtS on different model architectures? What guidance do these observations have for future research on sparse attention?\n\nBoth Mistral and LLama3 models are GQA models. Hence, this might be due to the different distribution of the two models parameters. For instance, most heads of mistral models might only focus on the local information with a sliding window: the attention maps of points that are far away to each other might not be large enough and can be removed; only a few heads are responsible for the global information. In LLama, more heads might contribute to the final predictions and thus we need to preserve the tokens from those heads."}, {"review_id": "N8aQp6bXUZ", "reviewer": "Reviewer_UM18", "summary": "throughout), Local Tokens (survive until the next global token), and Sliding Window Tokens (impact fixed-size windows). Using a Gumbel-Softmax trick and learnable attention masks, NAtS jointly optimizes token type assignments with model weights. The authors demonstrate NAtS's effectiveness through experiments on both training transformers from scratch and fine-tuning existing LLMs, showing significant KV cache reduction while maintaining performance.", "questions": "See weaknesses above.", "limitations": "See weaknesses above.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "### Strengths\n- Well-structured presentation. The paper is quite readable with clear explanations of the approach. \n- Strong empirical results. The experiments are thorough. They test on multiple benchmarks (RULER, LongBench), different models (LLama 3.1, Mistral-7B), various KV budget levels, and compare against many baselines (SnapKV, AdaKV, PyramidKV). NAtS consistently outperforms baselines. \n- Practical implementation details provided. The authors integrate their approach into FlashAttention (Algorithms 1-2 in the appendix), showing how to efficiently implement the sparse attention patterns. They also provide concrete latency/memory measurements in Figure 4.\n\n### Weaknesses\n- Limited novelty in core concepts. The idea of having different token types (global, local, sliding window) isn't particularly new - the paper itself mentions these patterns exist in StreamingLLM, Longformer, etc. The main contribution seems to be learning which tokens get which role, but this feels somewhat incremental. \n- Efficiency gains appear limited to very long contexts. Looking at Figure 4, the memory/latency benefits only become significant when context length exceeds 100k tokens. For more typical contexts (under 50k), the improvements seem marginal. This limits the practical impact since most real-world applications don't use such extreme context lengths.\n- Training overhead was not clearly discussed. While inference efficiency is well-studied, the paper doesn't clearly quantify the overhead of learning these token assignments during training. The additional Attention Score Layer parameters and Gumbel-Softmax sampling likely add computational cost, but this isn't analyzed.", "was_revised": true, "final_justification": "The rebuttal fixed all my concerns, so I raised my score from 4 to 5.", "rebuttal": "Dear reviewer UM18, \n\nwe thank you for your positive feedback regarding the efficiency of NAtS, regarding your concerns:\n\n> Limited novelty in core concepts. The idea of having different token types (global, local, sliding window) isn't particularly new - the paper itself mentions these patterns exist in StreamingLLM, Longformer, etc. The main contribution seems to be learning which tokens get which role, but this feels somewhat incremental. \n\nAlthough the token types are defined in previous works, finding the optimal token types for each token is still prohibitive, as the potential search space contains $3^{L×H ×N}$ options. Where $L, H, N$ are the sequence length, number of heads, and number of layers in the network. For a Llama model with 8 KV heads, 32 layers, even checking the type of one single token results in a complexity of $3^{256} \\approx 1.39\\times 10^{122}$ options. This value can even grow with the increased number of observed tokens. The predefined rules can only cover a tiny fraction of this search space. \n\nAdditionally, despite some other works such as SeerAttention[1] and Tokenbutler[2] also trying to learn a sparse attention map, they require a pre-trained model as a teacher model and cannot be learned end-to-end. In contrast, NAtS learns its token types directly from the gradient information, using either next-token prediction loss without any predefined rules. \n\nIn this paper, we show that: \n\n1. A learnable attention mask can be learned jointly with the transformer parameters towards a sparser transformer. \n\n2. The gradient of attention masks can be gathered to indicate the importance of each token and backpropagated towards their types. \n\n3. In the end, we achieve a token-level sparsity control that best trades off sparsity vs. model performance. \n\nFor a NAS problem [3], the search strategies are important as the search space design; hence, knowing how to search for the optimal token types should not only be considered incremental compared to the search space construction. \n\n\n> Efficiency gains appear limited to very long contexts. Looking at Figure 4, the memory/latency benefits only become significant when context length exceeds 100k tokens. For more typical contexts (under 50k), the improvements seem marginal. This limits the practical impact since most real-world applications don't use such extreme context lengths.\n\nThe efficiency of LLM inference is determined by many factors; for instance, the memory used to store the model parameters is much larger than the KV cache if the context length is not large enough. The cost of feedward layers cannot be negligible when the input context length is small. While for the decoding stage, memory bounds could be a bottleneck with the shorter content length.  Here, we mainly focus on the cases where the computational bounds are the bottleneck to highlight the efficiency of NAtS. \n\nHere, we also show the efficiency gains for the pre-filling (computationally dense) stages with a prefilling length smaller than 100k tokens. With a sparse kernel designed for first computing attention value for the global/local and then the values for the sliding window tokens, NAtS provides an overall $2.35 \\times$ lower latency compared to vanilla transformer within token budgets smaller than 100k\n\n|      |   10k |   20k |   30k |   40k |   50k |   60k |   70k |   80k |   90k |   100k |\n|:-----|--------:|--------:|--------:|--------:|--------:|--------:|--------:|--------:|--------:|---------:|\n| FA2  |    0.69 |    1.53 |    2.51 |    3.75 |    5.24 |    6.9  |    8.91 |   10.97 |   13.44 |    16.43 |\n| NAtS |    0.59 |    1.13 |    1.71 |    2.37 |    3.08 |    3.76 |    4.55 |    5.29 |    6.16 |     6.98 |\n\n\n> Training overhead was not clearly discussed. While inference efficiency is well-studied, the paper doesn't clearly quantify the overhead of learning these token assignments during training. The additional Attention Score Layer parameters and Gumbel-Softmax sampling likely add computational cost, but this isn't analyzed.\n\nThe Attention score layer is a linear layer that maps the input feature $X \\in R^{L \\times D}$ to $Logits~\\in R^{L \\ times (n_{ops}*h_{kv})}$ where $n_{opts}$ and $h_{kv}$ are the number of token types and number of KV heads respectively. This cost is minimal compared to the costs of computing the QKV values (roughly $n_{ops}/d_{head}$) where $d_{head}$ is the dimension of each head. Hence, overall, this operation also introduces a complexity of $O(L)$. The Gumbel-Softmax sampling only requires the token-level softmax and argmax operations across different types and thus does not introduce a complexity of $O(L)$ and is theoretically smaller than the overhead from the linear layers. The additional training overhead that NAtS brings to vanilla transformer for the pre-training tasks is only roughly 5% (0.93s per iteration for NAtS vs. 0.88s per iteration for vanilla transformer)\n\n\n[1] Gao et al. SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs\n\n[2] Akhauri et al. TokenButler: Token Importance is Predictable\n\n[3] Elsken et al. Neural architecture search: A survey."}]

Author rebuttal:

Meta-review: Neural Attention Search is an end-to-end learnable sparse transformer that automatically evaluates token importance to reduce KV cache and inference costs while maintaining performance. Key strengths include its novel casting of sparse attention as a Neural Architecture Search problem, end-to-end learnability, and integration into FlashAttention 2 for efficiency. It effectively reduces memory and accelerates reasoning in both pre-filling and decoding stages. Empirical results demonstrate strong performance across various benchmarks and models.

Initial weaknesses identified by reviewers involved the hand-designed nature of the search space, a perceived lack of novelty in core concepts, limited comparisons against other sparse attention baselines, and insufficient discussion of training overhead.

During the rebuttal, authors provided convincing responses. They clarified NAtS's minimal perplexity loss compared to baselines and its advantages over other sparse methods in terms of end-to-end learning and full-stage application. They provided additional latency comparisons against other sparse attention techniques and quantified the minimal training overhead. The authors also justified their search space design and the benefits of their joint optimization approach. All reviewers (1gKy, GRci, 6LT7, UM18) acknowledged that their concerns were addressed and consequently increased their scores. Based on the robust technical contributions and effective rebuttal, the paper is recommended for acceptance.
Decision: Accept (poster)