PromptDistill: Query-based Selective Token Retention in Intermediate Layers for Efficient Large Language Model Inference
| Title | PromptDistill: Query-based Selective Token Retention in Intermediate Layers for Efficient Large Language Model Inference |
| Publication Type | Conference Paper |
| Year of Publication | 2026 |
| Authors | Jin W., Song M., Pala T.D, Ken C.Y., Herremans D., Poria S. |
| Conference Name | 19th International Natural Language Generation Conference (INLG) |
| Date Published | 10/2026 |
| Abstract | As large language models (LLMs) tackle increasingly complex tasks and longer documents, their inference cost becomes a major bottleneck. To address this, we propose PromptDistill, a novel, training-free method that improves inference efficiency while preserving generation quality. PromptDistill identifies and retains the most informative tokens by leveraging early-layer attention, preserving their hidden states while reducing computation in later layers. Unlike previous methods such as H2O, SnapKV, and GemFilter, which perform compression only after full input processing, or ignore contextual dependencies, PromptDistill dynamically allocates computation to relevant tokens while maintaining global context. Experiments on LLaMA 3.1 8B, Phi 3.5 Mini, Qwen2 7B/14B, Phi 3 128k across LongBench, InfBench, and Needle in a Haystack demonstrate that PromptDistill achieves substantial efficiency gains with minimal quality loss. With a single-stage selection strategy, it outperforms GemFilter, H2O, and SnapKV, achieving an overall 1% to 5% performance improvement with faster inference. Additionally, we explore multi-stage selection, which further improves efficiency while maintaining strong generation performance. These results demonstrate the adaptability of LLM hidden representations and provide valuable insights into optimizing inference for long-context scenarios. Link to codes will be added after anonymous review. |
| URL | https://arxiv.org/abs/2503.23274 |