Title: Neural Attention Search
Venue: NeurIPS 2025
Accepted: yes
Keywords: Sparse Transformer

KV Cache Compression

Neural Architecture Search
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nthroughout), 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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.\n\nQUESTIONS:\nSee weaknesses above.\n\nLIMITATIONS:\nSee weaknesses above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n+ 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.\n\nQUESTIONS:\nNAtS 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

comments: 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

STRENGTHS AND WEAKNESSES:
strength:
- the token-type search is optimized jointly with network weights and implemented inside FlashAttention 2, avoiding explicit masks. 
- Pushes single-GPU context length from 200 k to 700 k tokens with 2.24× memory savings during pre-fill.  
weakness
- The smallest lambda settings still incur up-to-34 % perplexity loss on certain tasks when KV budget falls below ~5%. 
- Search space is column-oriented; block-sparse and slash patterns are excluded

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
- 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
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 2
DATE: 2025-07-12