Elamur Research⌕
☰
home / papers / 2026 / Octopus: Gated Selective Attention for Memory-Bounded Long-Context Inference in Large Language Models

Octopus: Gated Selective Attention for Memory-Bounded Long-Context Inference in Large Language Models

Paper · 2026 · Chien Van Nguyen, Ryan A. Rossi, Linh Ngo Van et al.

Summary

The KV cache of a Transformer grows linearly with context length, so long-context inference becomes memory-bound. Subquadratic models and linearized Transformers such as LoLCATs and Lizard fix memory use, but they compress history lossily and lose ground on reasoning. Octopus takes a different approach. It keeps standard softmax attention in a pretrained model and adds a small MLP gate to every layer. The gate predicts a utility score for each KV head and token, and this score is added as a log-space bias to the attention logits. Because of this design, the method works with optimized kernels like FlashAttention-2 and FlexAttention.

Training has two stages, and the backbone stays frozen in both:

  • Stage 1: the gates are distilled to match the teacher's attention outputs, with a sparsity regularizer.
  • Stage 2: the gates are fine-tuned on next-token prediction.

At inference, the cache holds sink tokens, a recent window, and the highest-scoring tokens, all within a fixed budget. On GSM8K with a 128-token budget, Octopus beats linearized baselines by more than 36 points and even outperforms its full-cache teacher. This suggests that learned selective retention can also act as a noise filter.

Key points

  • GSM8K (5-shot) results at budget B=128:
    • Octopus: 75.97%
    • Full-cache teacher: 71.56%
    • Lizard: 39.42%
    • LoLCATs: 33.05%
  • Parameter overhead: the gating module adds under 0.06% parameters, and only the gates are trained.
  • Runtime flexibility: the cache budget can be changed at runtime without retraining. Experiments covered GSM8K and MMLU.
  • Gating behavior: utility scores are head-specific, and they favor content-bearing entities (e.g., "Neil", "Armstrong") over stop words.
  • Limitations:
    • A fine-tuning phase is required, unlike training-free methods such as H2O and SnapKV.
    • Quality is capped by the teacher model.
    • Evicted tokens cannot be recovered.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓University of OregonEugene · United States23 works · 3,272 citations🏢Adobe SystemsSan Jose · United States35 works · 3,711 citations🎓Hanoi University of Science and TechnologyHanoi · Vietnam10 works · 89 citations🎓Trường Đại học Khoa học và Công nghệ Hà NộiHanoi · Vietnam8 works · 88 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksComputer Vision and Pattern Recognition · 2,869 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,320 works

Cite

Cite as Chien Van Nguyen, Ryan A. Rossi, Linh Ngo Van, Franck Dernoncourt and Thien Huu Nguyen. 2026. Octopus: Gated Selective Attention for Memory-Bounded Long-Context Inference in Large Language Models. Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers).

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.acl-long.1631

@inproceedings{nguyen2026octopus,
  title = {{Octopus: Gated Selective Attention for Memory-Bounded Long-Context Inference in Large Language Models}},
  author = {Chien {Van Nguyen}, Ryan A. Rossi, Linh Ngo Van, Franck Dernoncourt and Thien Huu Nguyen},
  booktitle = {Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
  year = {2026},
  doi = {10.18653/v1/2026.acl-long.1631},
  url = {https://doi.org/10.18653/v1/2026.acl-long.1631}
}

References

35 references, 17 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, and 1 others. 2023. Gpt-4 technical report. arXiv preprint arXiv:2303.08774.
  2. Art of Problem Solving. AIME Problems and Solutions. Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, and 1 others. 2024. Longbench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd an…
  3. Akhiad Bercovich, Itay Levy, Izik Golan, Mohammad Dabbah, Ran El-Yaniv, Omri Puny, Ido Galil, Zach Moshe, Tomer Ronen, Najeeb Nabwani, Ido Shahaf, Oren Tropp, Ehud Karpas, Ran Zilberstein, Jiaqi Zeng, Soumye Singhal, Alexander Bukharin, Yian Zhang, Tugrul Konuk, and 114 others. 2025. Llama-nemotron:…
  4. Ngoc Bui, Shubham Sharma, Simran Lamba, Saumitra Mishra, and Rex Ying. 2026. Cache what lasts: Token retention for memory-bounded KV cache in LLMs. In The Fourteenth International Conference on Learning Representations. find
  5. Zefan Cai, Wen Xiao, Hanshi Sun, Cheng Luo, Yikai Zhang, Ke Wan, Yucheng Li, Yeyang Zhou, Li-Wen Chang, Jiuxiang Gu, and 1 others. 2025. R-kv: Redundancy-aware kv cache compression for reasoning models. In The Thirty-ninth Annual Conference on Neural Information Processing Systems. find
  6. Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Yucheng Li, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Junjie Hu, and 1 others. 2024. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling. arXiv preprint arXiv:2406.02069.
  7. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, and 1 others. 2021. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168.
  8. Tri Dao. 2023. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691.
  9. Tri Dao and Albert Gu. 2024. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21060. find
  10. Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. 2023. Model tells you what to discard: Adaptive kv cache compression for llms. arXiv preprint arXiv:2310.01801.
  11. Daniel Goldstein, Eric Alcaide, Janna Lu, and Eugene Cheah. 2025. Radlads: Rapid attention distillation to linear attention decoders at scale. arXiv preprint arXiv:2505.03005. find
  12. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, and 1 others. 2024. The llama 3 herd of models. arXiv preprint arXiv:2407.21783.
  13. Albert Gu and Tri Dao. 2023. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752. find
  14. Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2020. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300.
  15. Disen Lan, Weigao Sun, Jiaxi Hu, Jusen Du, and Yu Cheng. 2025. Liger: Linearizing large language models to gated recurrent structures. In Forty-second International Conference on Machine Learning.
  16. Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. 2024. Snapkv: Llm knows what you are looking for before generation. Advances in Neural Information Processing Systems, 37:22947–22970.
  17. Jiaheng Liu, Dawei Zhu, Zhiqi Bai, Yancheng He, Huanxuan Liao, Haoran Que, Zekun Wang, Chenchen Zhang, Ge Zhang, Jiebin Zhang, and 1 others. 2025. A comprehensive survey on long context language modeling. arXiv preprint arXiv:2503.17407.
  18. Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. 2023. Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time. Advances in Neural Information Processing Systems, 36:5… find
  19. Chien Van Nguyen, Ruiyi Zhang, Hanieh Deilamsalehy, Puneet Mathur, Viet Dac Lai, Haoliang Wang, Jayakumar Subramanian, Ryan A Rossi, Trung Bui, Nikos Vlassis, and 1 others. 2025. Lizard: An efficient linearization framework for large language models. arXiv preprint arXiv:2507.09025.
  20. Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. 2024. You only cache once: Decoder-decoder architectures for language models. Advances in Neural Information Processing Systems, 37:7339– 7361. find
  21. Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. 2024. Quest: Query-aware sparsity for efficient long-context llm inference, 2024. URL https://arxiv.org/abs/2406.10774. ↗
  22. Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. 2023. Stanford alpaca: An instruction-following llama model. https://github.com/tatsu-lab/stanford_alpaca. ↗
  23. Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, and 1 others. 2023. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971.
  24. Chien Van Nguyen, Huy Huu Nguyen, Thang M Pham, Ruiyi Zhang, Hanieh Deilamsalehy, Puneet Mathur, Ryan A Rossi, Trung Bui, Viet Dac Lai, Franck Dernoncourt, and 1 others. 2024. Taipan: Efficient and expressive state space language models with selective attention. arXiv preprint arXiv:2410.18572. — [arXiv:2410.18572](https://arxiv.org/abs/2410.18572) find
  25. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances in neural information processing systems, 30.
  26. Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, and 1 others. An empirical study of mamba-based language models, 2024. URL https://arxiv.org/abs/2406.07887. ↗
  27. Junxiong Wang, Daniele Paliotta, Avner May, Alexander Rush, and Tri Dao. 2024. The mamba in the llama: Distilling and accelerating hybrid models. Advances in Neural Information Processing Systems, 37:62432–62457. find
  28. Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. 2023. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453. find
  29. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, and 1 others. 2025. Qwen3 technical report. arXiv preprint arXiv:2505.09388.
  30. Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. 2024. Gated linear attention transformers with hardware-efficient training. In Forty-first International Conference on Machine Learning. find
  31. Tianzhu Ye, Li Dong, Yuqing Xia, Yutao Sun, Yi Zhu, Gao Huang, and Furu Wei. 2024. Differential transformer. arXiv preprint arXiv:2410.05258. — [arXiv:2410.05258](https://arxiv.org/abs/2410.05258) find
  32. Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. 2019. HellaSwag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791–4800. Association for Computational Linguistics. find
  33. Michael Zhang, Simran Arora, Rahul Chalamala, Benjamin Frederick Spector, Alan Wu, Krithik Ramesh, Aaryan Singhal, and Christopher Re. 2025. LoL-CATs: On low-rank linearizing of large language models. In The Thirteenth International Conference on Learning Representations. find
  34. Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, and 1 others. 2023. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36:346… find
  35. Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, and 1 others. 2023. Pytorch fsdp: experiences on scaling fully sharded data parallel. arXiv preprint arXiv:2304.11277. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/nguyen-2026-octopus-gated-selective-attention-memory/