arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2407.12866cs.CLcs.AIcs.LG
收藏

超越KV缓存:共享注意力机制提升大语言模型效率

Beyond KV Caching: Shared Attention for Efficient LLMs

Kyushu University(九州大学)

Bingli Liao, Danilo Vasconcellos Vargas

AI总结:

针对LLM推理效率瓶颈,提出跨层共享注意力权重的SA机制,利用注意力分布等向性减少计算量和KV缓存,在多种模型上以极小精度损失实现高效推理。

摘要:

大语言模型(LLM)的效率仍然是一个关键挑战,尤其是在计算资源受限的环境中。这些模型中传统的注意力机制虽然强大,但由于需要在不同层之间重新计算和存储注意力权重,因此需要大量的计算和内存资源。本文提出了一种新颖的共享注意力(Shared Attention, SA)机制,旨在通过直接跨多层共享已计算的注意力权重来提升LLM的效率。与以往专注于共享中间键值(Key-Value, KV)缓存的方法不同,我们的方法利用了高级LLM在预训练后观察到的注意力分布的等向性(isotropic)趋势,以减少推理期间所需的计算浮点运算次数(flops)和KV缓存的大小。我们通过实验证明,在各种LLM上实施SA仅导致标准基准测试上的极小精度损失。我们的研究结果表明,SA不仅节省了计算资源,还保持了稳健的模型性能,从而促进了在资源受限环境中部署更高效的LLM。

英文摘要:

The efficiency of large language models (LLMs) remains a critical challenge, particularly in contexts where computational resources are limited. Traditional attention mechanisms in these models, while powerful, require significant computational and memory resources due to the necessity of recalculating and storing attention weights across different layers. This paper introduces a novel Shared Attention (SA) mechanism, designed to enhance the efficiency of LLMs by directly sharing computed attention weights across multiple layers. Unlike previous methods that focus on sharing intermediate Key-Value (KV) caches, our approach utilizes the isotropic tendencies of attention distributions observed in advanced LLMs post-pretraining to reduce both the computational flops and the size of the KV cache required during inference. We empirically demonstrate that implementing SA across various LLMs results in minimal accuracy loss on standard benchmarks. Our findings suggest that SA not only conserves computational resources but also maintains robust model performance, thereby facilitating the deployment of more efficient LLMs in resource-constrained environments.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑