CASAK-V: Dynamic Sparse Attention and Adaptive KV-Cache Compression for Memory-Efficient Long-Context LLM Inference

Authors: Hamza Mohammed, Sai Chand Boyapati, Hang Yin
Primary area: foundation or frontier models, including LLMs
Keywords: large language models, sparse attention, kv-cache compression, long-context processing, meta-learning, adaptive algorithms, memory efficiency, inference optimization, on-device deployment, context-aware models, dynamic attention, transformer architectures, efficient natural language processing, machine learning systems, attention mechanisms, sparse computation, benchmarking, model compression, resource-constrained computing, edge ai, computational complexity, information retrieval, self-attention, transfer learning, deep learning, artificial intelligence, chunk-wise compression, pattern recognition

Abstract:
The emergence of long-context Large Language Models (LLMs) has triggered a rapid expansion of applications across various domains. However, these models remain inaccessible for on-device or on-premises deployments due to significant computational and memory challenges. The quadratic complexity of attention mechanisms and the substantial memory requirements of KV-caches, hinder adoption in resource-constrained environments. Current solutions, such as sparse attention mechanisms and KV-cache compression techniques, often rely on pre-observed patterns or context-independent, head-specific profiling strategies, which can compromise model accuracy, especially in long-context processing. This paper introduces Context-Aware adaptive Sparse Attention with Key-Value cache compression (CASAK-V), an inference-time approach that dynamically generates and applies head-specific sparse attention patterns. CASAK-V leverages a meta-learning framework to fine-tune a compact pre-trained vision-language encoder-decoder transformer for sparse pattern identification from per-layer attention scores. These patterns include fixed local windows, dynamic column stripes, block-sparse, and various other learned hybrid configurations. The technique additionally implements adaptive chunk-wise KV-cache compression using policies adapted from these layer-wise sparse configurations. To retain context-awareness, these configuration are dynamically adjusted during token generation, based on an attention map reconstruction heuristic. Our evaluations show that CASAK-V achieves minimal performance degradation on long-context benchmarks (LongBench), while reducing memory usage by 40% and delivering near-linear runtime complexity compared to full attention and caching. In summary, CASAK-V enables efficient long-context processing in memory-limited environments, extending the applicability of LLMs and facilitating their deployment in on-premises and on-device scenarios.

Decision: desk rejected
Reviewer scores: 0

Reviews and discussion:

Metareview: