TokenButler: Token Importance is Predictable

• URL: http://arxiv.org/abs/2503.07518v1
• Date: Mon, 10 Mar 2025 16:41:14 GMT
• Title: TokenButler: Token Importance is Predictable
• Authors: Yash Akhauri, Ahmed F AbouElhamayed, Yifei Gao, Chi-Chih Chang, Nilesh Jain, Mohamed S. Abdelfattah,
• Abstract summary: Large Language Models (LLMs) rely on the Key-Value (KV) Cache to store token history, enabling efficient decoding of tokens.<n>Prior research has shown that only a small subset of tokens contribute meaningfully to each decoding step.<n>We introduce TokenButler, a high-granularity, query-aware predictor that learns to identify these critical tokens.
• Score: 8.514853311344458
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Large Language Models (LLMs) rely on the Key-Value (KV) Cache to store token history, enabling efficient decoding of tokens. As the KV-Cache grows, it becomes a major memory and computation bottleneck, however, there is an opportunity to alleviate this bottleneck, especially because prior research has shown that only a small subset of tokens contribute meaningfully to each decoding step. A key challenge in finding these critical tokens is that they are dynamic, and heavily input query-dependent. Existing methods either risk quality by evicting tokens permanently, or retain the full KV-Cache but rely on retrieving chunks (pages) of tokens at generation, failing at dense, context-rich tasks. Additionally, many existing KV-Cache sparsity methods rely on inaccurate proxies for token importance. To address these limitations, we introduce TokenButler, a high-granularity, query-aware predictor that learns to identify these critical tokens. By training a light-weight predictor with less than 1.2% parameter overhead, TokenButler prioritizes tokens based on their contextual, predicted importance. This improves perplexity & downstream accuracy by over 8% relative to SoTA methods for estimating token importance. We evaluate TokenButler on a novel synthetic small-context co-referential retrieval task, demonstrating near-oracle accuracy. Code, models and benchmarks: https://github.com/abdelfattah-lab/TokenButler

Related papers

• Learning to Attribute with Attention [75.61481181755744]
  We propose treating attention weights of different attention heads as features. This way, we can learn how to effectively leverage attention weights for attribution. Our resulting method, Attribution with Attention (AT2), reliably performs on par with approaches that involve many ablations.
  arXiv  Detail & Related papers  (2025-04-18T15:36:28Z)
• Token Pruning in Multimodal Large Language Models: Are We Solving the Right Problem? [19.35502303812707]
  Multimodal large language models (MLLMs) have shown remarkable performance for cross-modal understanding and generation, yet still suffer from severe inference costs.<n>Recently, abundant works have been proposed to solve this problem with token pruning, which identifies the redundant tokens in MLLMs and then prunes them to reduce the computation and KV storage costs.<n>In this paper, we answer these questions one by one, providing insights into the design of future token pruning methods.
  arXiv  Detail & Related papers  (2025-02-17T07:05:36Z)
• Stop Looking for Important Tokens in Multimodal Language Models: Duplication Matters More [18.928285521147057]
  We show that importance is not an ideal indicator to decide whether a token should be pruned.<n>We propose DART (Duplication-Aware Reduction of Tokens), which prunes tokens based on its duplication with other tokens.<n>Experiments demonstrate that DART can prune 88.9% vision tokens while maintaining comparable performance.
  arXiv  Detail & Related papers  (2025-02-17T06:56:28Z)
• HashAttention: Semantic Sparsity for Faster Inference [91.54218318798603]
  HashAttention is a principled approach casting pivotal token identification as a recommendation problem.<n>It efficiently identifies pivotal tokens for a given query in this Hamming space using bitwise operations.<n>It can reduce the number of tokens used by a factor of $1/32times$ for the Llama-3.1-8B model with LongBench.
  arXiv  Detail & Related papers  (2024-12-19T02:34:15Z)
• More Tokens, Lower Precision: Towards the Optimal Token-Precision Trade-off in KV Cache Compression [71.42818367729573]
  In large language models (LLMs), the memory usage of KV cache has become a critical bottleneck during inference.<n>The mainstream KV compression methods, including KV pruning and KV quantization, primarily focus on either token or precision dimension separately.<n>In this paper, we comprehensively investigate the token-precision trade-off in KV cache compression.
  arXiv  Detail & Related papers  (2024-12-17T09:20:31Z)
• HashEvict: A Pre-Attention KV Cache Eviction Strategy using Locality-Sensitive Hashing [32.62377392686119]
  We introduce HashEvict, an algorithm that uses locality-sensitive hashing (LSH) to compress the KV cache.<n>HashEvict can compress the KV cache by 30%-70% while maintaining high performance across reasoning, multiple-choice, long-context retrieval and summarization tasks.
  arXiv  Detail & Related papers  (2024-12-13T06:00:27Z)
• Multi-Stage Vision Token Dropping: Towards Efficient Multimodal Large Language Model [45.01871133425388]
  We propose Multi-stage Token Dropping (MustDrop) to measure the importance of each token from the whole lifecycle. MustDrop reduces about 88.5% FLOPs on LLaVA with a compression ratio of 92.2% while maintaining comparable accuracy.
  arXiv  Detail & Related papers  (2024-11-16T13:45:33Z)
• Training-Free Exponential Context Extension via Cascading KV Cache [49.608367376911694]
  We introduce a novel mechanism that leverages cascading sub-cache buffers to selectively retain the most relevant tokens.<n>Our method reduces prefill stage latency by a factor of 6.8 when compared to flash attention on 1M tokens.
  arXiv  Detail & Related papers  (2024-06-24T03:59:17Z)
• Get More with LESS: Synthesizing Recurrence with KV Cache Compression for Efficient LLM Inference [78.65321721142624]
  We focus on a memory bottleneck imposed by the key-value ( KV) cache. Existing KV cache methods approach this problem by pruning or evicting large swaths of relatively less important KV pairs. We propose LESS, a simple integration of a constant sized cache with eviction-based cache methods.
  arXiv  Detail & Related papers  (2024-02-14T18:54:56Z)
• How can objects help action recognition? [74.29564964727813]
  We investigate how we can use knowledge of objects to design better video models. First, we propose an object-guided token sampling strategy that enables us to retain a small fraction of the input tokens. Second, we propose an object-aware attention module that enriches our feature representation with object information.
  arXiv  Detail & Related papers  (2023-06-20T17:56:16Z)
• Adaptive Sparse ViT: Towards Learnable Adaptive Token Pruning by Fully Exploiting Self-Attention [36.90363317158731]
  We propose an adaptive sparse token pruning framework with a minimal cost. Our method improves the throughput of DeiT-S by 50% and brings only 0.2% drop in top-1 accuracy.
  arXiv  Detail & Related papers  (2022-09-28T03:07:32Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.