Research map / Architectures and efficiency / Attention and transformers
Neural Attention Search
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Neural Attention Search (NAtS), a learnable sparse attention method that assigns each token a global, local, or sliding-window role via Gumbel-Softmax-searched masks, reducing KV cache while keeping performance.
Problem
Long-context LLM inference is bottlenecked by KV cache size and cost; existing KV eviction and sparse attention rely on fixed heuristics or approximate dense attention outputs.
Approach
Each token is assigned one of three roles (global, local, sliding-window) by a learned attention-score layer trained end-to-end through a learnable multiplicative attention mask, with Gumbel-Softmax sampling and a sparsity regularizer lambda, so token roles are learned from the task loss rather than predefined rules.
Key results
- On RULER with Llama-3.1-8B at 50% KV budget, NAtS matches or exceeds full attention at most lengths (e.g., 95.32 vs 95.06 at 4k), and at 25% budget keeps ~96% of full-attention scores.
- On PG-19 from-scratch GPT-2 small, NAtS with ~3% KV cache (about 30 tokens per layer) achieves lower perplexity than H2O at 12.5% budget.
- At 200k context, NAtS uses 2.24x less memory with 3.0x prefill speedup and 2.6x less memory with 1.38x decoding speedup; it can process up to 700k tokens on one H100.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use learned token-role attention with a 50% KV budget instead of full attention | Llama 3.1-8B · RULER (4k–128k) | score | No clear effect |
| Use learned token-role attention with a 25% KV budget instead of full attention | Llama 3.1-8B · RULER (4k–128k) | score | Mixed |
| Use NAtS rather than strongest listed non-full baseline | Llama 3.1-8B · LongBench (25% KV budget) | score | Mixed |
| Use NAtS rather than full attention at 200,000 context length | Llama 3.1-8B · Latency evaluation | prefill latency: 3.0× latency speed up | Helped |
| Use NAtS rather than full attention at 200,000 context length | Llama 3.1-8B · Latency evaluation | decoding latency: 1.38~× decoding speed up | Helped |
| Use NAtS rather than full attention | Llama 3.1-8B · Maximum context length | context length: 700,000; 3.5 times larger | Helped |
Surprising findings
- Even with the sparsity regularizer lambda set to 0, NAtS reduces KV cache to about 35% while slightly lowering perplexity versus full attention on PG-19.
- Fine-tuning only a small attention score layer (about 13MB) with a distillation loss on a 7000-instance dataset suffices to achieve near-lossless RULER performance.
Limitations
- Attention mask construction is column-wise and mostly covers vertical sparse patterns; block-sparse and slash patterns are not handled.
- All experiments run once with a fixed random seed, so no error bars are reported.
- Gradient for global-token removal ignores cases with G_i=0 for efficiency, making the gradient approximate.
Experimental setup
Datasets: PG-19 HotpotQA 2WikiMultihopQA MuSiQue DuReader NarrativeQA Qasper GovReport QMSum MultiNews VCSUM TREC
Benchmarks: PG-19 RULER LongBench
Metrics: perplexity accuracy/score KV cache budget fraction latency memory usage
Baselines: Full attention StreamingLLM H2O DuoAttention SnapKV AdaKV ChunkKV CriticalKV PyramidKV MoA
Builds on: FlashAttention-2 GDAS DuoAttention Streaming LLM
Closest papers by meaning
- Transformers Learn Faster with Semantic Focus (NeurIPS 2025 · similarity 0.89)
- Token Statistics Transformer: Linear-Time Attention via Variational Rate Reduction (ICLR 2025 · similarity 0.88)
- RAT: Bridging RNN Efficiency and Attention Accuracy via Chunk-based Sequence Modeling (NeurIPS 2025 · similarity 0.88)
- Rethinking Attention: Exploring Shallow Feed-Forward Neural Networks as an Alternative to Attention Layers in Transformers (Student Abstract) (AAAI 2024 · similarity 0.87)
- A Training-Free Sub-quadratic Cost Transformer Model Serving Framework with Hierarchically Pruned Attention (ICLR 2025 · similarity 0.87)
- Hardware-aligned Hierarchical Sparse Attention for Efficient Long-term Memory Access (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.