IAM: Efficient Inference through Attention Mapping between Different-scale LLMs
This addresses efficiency issues for users deploying LLMs in resource-constrained environments, though it is incremental as it builds on existing optimization methods.
The paper tackles the problem of high resource consumption in large language models (LLMs) during inference, especially with long contexts, by introducing the IAM framework that uses attention mapping between different-scale LLMs to accelerate prefill by 15% and reduce KV cache usage by 22.1% without significant performance loss.
LLMs encounter significant challenges in resource consumption nowadays, especially with long contexts. Despite extensive efforts dedicate to enhancing inference efficiency, these methods primarily exploit internal sparsity within the models, without leveraging external information for optimization. We identify the high similarity of attention matrices across different-scale LLMs, which offers a novel perspective for optimization. We first conduct a comprehensive analysis of how to measure similarity, how to select mapping Layers and whether mapping is consistency. Based on these insights, we introduce the IAM framework, which achieves dual benefits of accelerated attention computation and reduced KV cache usage by performing attention mapping between small and large LLMs. Our experimental results demonstrate that IAM can accelerate prefill by 15% and reduce KV cache usage by 22.1% without appreciably sacrificing performance. Experiments on different series of models show the generalizability of IAM. Importantly, it is also orthogonal to many existing KV cache optimization methods, making it a versatile addition to the current toolkit for enhancing LLM efficiency.