通过小大模型注意力映射，加速推理并减少显存占用
IAM: Efficient Inference through Attention Mapping between Different-scale LLMs
- 利用不同规模大模型间注意力矩阵的高相似性进行映射优化
- 预填充阶段加速15%，KV缓存使用减少22.1%且性能损失小
- 兼容现有缓存优化方法，适用于多系列模型
当前大语言模型在长上下文场景下面临显著资源消耗问题。尽管已有大量工作致力于提升推理效率，但多数方法仅利用模型内部稀疏性，未充分借助外部信息优化。本文发现不同规模大模型间的注意力矩阵具有高度相似性，由此提出IAM框架，通过小模型与大模型间注意力映射，实现双重优化：加速注意力计算并减少KV缓存占用。实验表明，IAM在不明显降低性能的前提下，使预填充阶段提速15%，KV缓存使用减少22.1%。跨多个模型系列的测试验证了其通用性。更重要的是，IAM与现有多种缓存优化方法正交，可作为当前高效推理工具箱中的灵活补充。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。