Elamur Research⌕
☰
home / papers / 2026 / Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference

Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference

Paper · 2026 · Rei Taniguchi, Yuyang Dong, Makoto Onizuka et al.

Summary

Layer-wise token pruning reduces the KV cache during LLM inference. It keeps only a subset of important tokens beyond a certain Transformer layer. Existing methods such as FastKV and GemFilter fix this "selection layer" in advance, regardless of the task. A layer that suits easy QA can therefore fail badly on harder tasks like KV retrieval, where context and query look semantically alike.

The authors propose ASL (Adaptive Selection Layer), a training-free approach that picks the selection layer for each input during prefilling. ASL tracks how much the attention-based ranking of tokens shifts across consecutive layers. Once this rank variance falls below a user-set relative threshold, attention has settled on a stable set of tokens. ASL then selects those tokens a single time and passes only them to deeper layers. It can be paired with SnapKV to compress the decoding-stage cache, or placed inside GemFilter's two-pass scheme.

On InfiniteBench, RULER and NIAH, ASL beats FastKV, GemFilter and PyramidInfer on difficult tasks while keeping KV budgets and decoding speed comparable. It does so by trading some prefilling time for accuracy, which makes KV compression more robust across task difficulty.

Key points

  • The evaluation covers three benchmarks, NIAH (Needle in a Haystack), RULER, and InfiniteBench, with contexts reaching 256k tokens. The tested models are Llama-3.1-8B-UL and Qwen2.5-7B.
  • Baselines are the layer-wise pruning methods FastKV, GemFilter and PyramidInfer. ASL is also tested as a two-pass variant built into GemFilter (ASL_2pass).
  • KV budgets of 2048 and 8192 were tested. Peak memory stayed nearly the same as the baselines, so ASL's extra memory cost is negligible. Time per output token is also almost unchanged.
  • The relative variance threshold τ was swept from 0.2 to 0.6. Larger τ consistently lowers time to first token, while its effect on accuracy depends on the model. The authors recommend τ = 0.3.
  • Limitations:
    • Only one-shot methods were studied.
    • Only two LLMs were tested.
    • ASL was less competitive when combined with GemFilter on Llama-3.1-8B-UL.
  • Code is available at github.com/TANIGUCHIREI/ASL.

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

🎓The University of OsakaOsaka · Japan15 works · 517 citations🏢NECTokyo · Japan6 works · 546 citations🎓Nagoya UniversityNagoya · Japan8 works · 8 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 2Research subjects and languages the work is about.Language models & pretraining · 3,320 worksEfficient models & compression · 242 works

Cite

Cite as Rei Taniguchi, Yuyang Dong, Makoto Onizuka and Chuan Xiao. 2026. Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference. Findings of the Association for Computational Linguistics: ACL 2026.

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

@inproceedings{taniguchi2026adaptive,
  title = {{Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference}},
  author = {Rei Taniguchi, Yuyang Dong, Makoto Onizuka and Chuan Xiao},
  booktitle = {Findings of the Association for Computational Linguistics: ACL 2026},
  year = {2026},
  doi = {10.18653/v1/2026.findings-acl.48},
  url = {https://doi.org/10.18653/v1/2026.findings-acl.48}
}

References

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

  1. Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, and 1 others. 2024. PyramidKV: Dynamic KV cache compression based on pyramidal information funneling. arXiv preprint arXiv:2406.02069.
  2. Moses S Charikar. 2002. Similarity estimation techniques from rounding algorithms. In Proceedings of the thiry-fourth annual ACM symposium on Theory of computing, pages 380–388. find
  3. Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W Mahoney, Sophia Shao, Kurt Keutzer, and Amir Gholami. 2024. KVQuant: Towards 10 million context length LLM inference with KV cache quantization. Advances in Neural Information Processing Systems, 37:1270–1303.
  4. Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. 2024. Ruler: What’s the real context size of your long-context language models? arXiv preprint arXiv:2404.06654.
  5. Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H. Abdi, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. 2024a. Minference. https://github.com/microsoft/MInference/tree/main. ↗
  6. Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H. Abdi, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. 2024b. MInference 1.0: Accelerating pre-filling for long-context LLMs via dynamic sparse attention. Advances in Neural Information Proc…
  7. Weisheng Jin, Maojia Song, Tej Deep Pala, Yew Ken Chia, Amir Zadeh, Chuan Li, and Soujanya Poria. 2025. PromptDistill: Query-based selective token retention in intermediate layers for efficient large language model inference. arXiv preprint arXiv:2503.23274. find
  8. Dongwon Jo, Jiwon Song, Yulhwa Kim, and Jae-Joon Kim. 2025. FastKV: KV cache compression for fast long-context processing with token-selective propagation. arXiv preprint arXiv:2502.01068. — [arXiv:2502.01068](https://arxiv.org/abs/2502.01068) find
  9. Greg Kamradt. 2024. Needle in a haystack – pressure testing LLMs. https://github.com/gkamradt/ LLMTest_NeedleInAHaystack. ↗
  10. Wonbeom Lee, Jungi Lee, Junghwan Seo, and Jaewoong Sim. 2024. InfiniGen: Efficient generative inference of large language models with dynamic KV cache management. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pages 155–172. find
  11. 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.
  12. Akide Liu, Jing Liu, Zizheng Pan, Yefei He, Reza Haffari, and Bohan Zhuang. 2024a. MiniCache: KV cache compression in depth dimension for large language models. Advances in Neural Information Processing Systems, 37:139997–140031. find
  13. Songwei Liu, Chao Zeng, Lianqiang Li, Chenqian Yan, Lean Fu, Xing Mei, and Fangmin Chen. 2024b. FoldGPT: Simple and effective large language model compression scheme. arXiv preprint arXiv:2407.00928. — [arXiv:2407.00928](https://arxiv.org/abs/2407.00928) find
  14. Aurick Qiao, Zhewei Yao, Samyam Rajbhandari, and Yuxiong He. 2024. SwiftKV: Fast prefill-optimized inference with knowledge-preserving model transformation. arXiv preprint arXiv:2410.03960. find
  15. Ziran Qin, Yuchen Cao, Mingbao Lin, Wen Hu, Shixuan Fan, Ke Cheng, Weiyao Lin, and Jianguo Li. 2025. CAKE: Cascading and adaptive KV cache eviction with layer preferences. arXiv preprint arXiv:2503.12491. — [arXiv:2503.12491](https://arxiv.org/abs/2503.12491) find
  16. Zhenmei Shi, Yifei Ming, Xuan-Phi Nguyen, Yingyu Liang, and Shafiq Joty. 2024. Discovering the gems in early layers: Accelerating long-context LLMs with 1000x input token reduction. arXiv preprint arXiv:2409.17422.
  17. Prajwal Singhania, Siddharth Singh, Shwai He, Soheil Feizi, and Abhinav Bhatele. 2024. Loki: Low-rank keys for efficient sparse attention. Advances in Neural Information Processing Systems, 37:16692–16723. find
  18. Hanshi Sun, Li-Wen Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Xin Liu, Harry Dong, Yuejie Chi, and Beidi Chen. 2024. ShadowKV: KV cache in shadows for high-throughput long-context LLM inference. arXiv preprint arXiv:2410.21465. find
  19. Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. 2024. Quest: Query-aware sparsity for efficient long-context LLM inference. arXiv preprint arXiv:2406.10774.
  20. 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.
  21. Zihao Wang, Bin Cui, and Shaoduo Gan. 2024. SqueezeAttention: 2d management of KV-cache in LLM inference via layer-wise optimal budget. arXiv preprint arXiv:2404.04793. — [arXiv:2404.04793](https://arxiv.org/abs/2404.04793) find
  22. Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, and 3 others. 2020. Transfo…
  23. Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. 2024c. KIVI: A tuning-free asymmetric 2bit quantization for KV cache. arXiv preprint arXiv:2402.02750.
  24. 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
  25. Lingkun Long, Rubing Yang, Yushi Huang, Desheng Hui, Ao Zhou, and Jianlei Yang. 2025. SlimInfer: Accelerating long-context LLM inference via dynamic token pruning. arXiv preprint arXiv:2508.06447. find
  26. Dongjie Yang, XiaoDong Han, Yan Gao, Yao Hu, Shilin Zhang, and Hai Zhao. 2024. PyramidInfer: Pyramid KV cache compression for high-throughput LLM inference. arXiv preprint arXiv:2405.12532.
  27. Bohan Yu and Yekun Chai. 2025. EvolKV: Evolutionary KV cache compression for LLM inference. arXiv preprint arXiv:2509.08315. find
  28. Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Hao, Xu Han, Zhen Thai, Shuo Wang, Zhiyuan Liu, and 1 others. 2024. ∞ bench: Extending long context evaluation beyond 100k tokens. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume… find
  29. 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
  30. Xiabin Zhou, Wenbin Wang, Minyan Zeng, Jiaxian Guo, Xuebo Liu, Li Shen, Min Zhang, and Liang Ding. 2024. DynamicKV: Task-aware adaptive KV cache compression for long context LLMs. arXiv preprint arXiv:2412.14838. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/taniguchi-2026-adaptive-layer-selection-layer-wise/