Dynamic token pruning for LLMs: leveraging task-specific attention and adaptive thresholds
摘要
Large language models (LLMs) have achieved state-of-the-art performance across a wide range of natural language processing (NLP) tasks. With their high inference computational costs, deployment is extremely challenging, especially to resource-constrained environments. Dynamic pruning methods, because they are efficient, are likely to assign uniform policies to all tasks and even forget task-specific knowledge and make optimal behavior complicated. To counter this limitation, we propose task-specific dynamic token pruning (TS-DTP), a novel optimization framework that reaches maximum efficiency for LLMs at inference without compromising task-specific performance and, in certain cases, improving upon it. TS-DTP utilizes task-specific knowledge to regulate the token selection process by applying task-specific attention weights and adaptive threshold learning. This approach enables better token importance decision-making through the dynamically adjustable pruning policy according to the downstream task need. It enables very high-grained control, keeping the meaningful contextual information, therefore promoting better performance compared to regular pruning methods. Experimental findings on a variety of NLP tasks (question answering, machine translation, sentiment analysis) validate that TS-DTP achieves extremely large reductions in computational expense and memory demands and similar or marginal gains in accuracy. Our findings are at the forefront of efficient and deployable LLM development and highlight the importance of task adaptation for optimal performance in low-resource settings.