Beyond KV Caching: Shared Attention for Efficient LLMs

📄 arXiv: 2407.12866v1 📥 PDF

作者: Bingli Liao, Danilo Vasconcellos Vargas

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-07-13

💡 一句话要点

提出共享注意力机制，提升资源受限场景下大语言模型效率

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 注意力机制 模型压缩 效率优化 资源受限环境

📋 核心要点

  1. 传统LLM注意力机制计算和存储注意力权重消耗大量资源，限制了其在资源受限环境中的应用。
  2. 论文提出共享注意力（SA）机制，通过跨层共享注意力权重，减少计算量和KV缓存大小。
  3. 实验表明，SA在保持模型性能的同时，显著降低了计算资源消耗，提升了LLM的效率。

📝 摘要（中文）

大型语言模型（LLM）的效率仍然是一个关键挑战，尤其是在计算资源有限的情况下。传统注意力机制虽然强大，但由于需要跨层重新计算和存储注意力权重，因此需要大量的计算和内存资源。本文提出了一种新的共享注意力（SA）机制，旨在通过直接跨多个层共享计算的注意力权重来提高LLM的效率。与之前专注于共享中间键值（KV）缓存的方法不同，我们的方法利用了预训练后高级LLM中观察到的注意力分布的各向同性趋势，以减少推理期间所需的计算量和KV缓存的大小。我们通过实验证明，在各种LLM中实施SA会导致标准基准上的精度损失最小。我们的研究结果表明，SA不仅节省了计算资源，而且保持了强大的模型性能，从而有助于在资源受限的环境中部署更高效的LLM。

🔬 方法详解

问题定义：大型语言模型在推理过程中，传统的注意力机制需要为每一层都计算和存储注意力权重，这导致了巨大的计算开销和内存占用，尤其是在长序列和深层模型中。现有的KV缓存优化方法虽然可以减少重复计算，但仍然需要存储大量的中间结果，限制了模型在资源受限设备上的部署。

核心思路：论文的核心思路是利用预训练后的大型语言模型中注意力分布呈现出的各向同性特性，即不同层之间的注意力分布具有一定的相似性。因此，可以通过共享注意力权重来减少冗余计算和存储，从而提高模型的效率。

技术框架：该方法的核心是共享注意力（SA）机制。在模型的某些层，不再独立计算注意力权重，而是直接使用其他层计算好的注意力权重。具体来说，可以选择一个或多个“源层”，这些层正常计算注意力权重，然后将这些权重复制到其他“目标层”。目标层直接使用这些共享的注意力权重进行后续计算，而无需重新计算。

关键创新：与以往的KV缓存方法不同，该方法直接共享注意力权重，而不是中间的KV缓存。这避免了存储大量中间结果的需要，进一步降低了内存占用。此外，该方法利用了注意力分布的各向同性特性，使得共享注意力权重成为可能，而不会显著降低模型性能。

关键设计：论文中需要确定哪些层作为源层，哪些层作为目标层。一种简单的策略是随机选择，或者根据经验选择。更高级的策略可以根据不同层之间的注意力分布相似度来动态选择。此外，还可以引入一些微调机制，使得模型在共享注意力权重后能够更好地适应新的任务。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在多个LLM模型上应用共享注意力机制后，可以在标准benchmark上保持最小的精度损失。具体来说，在某些模型上，精度损失小于1%，但计算量和KV缓存大小可以显著降低，例如降低10%-20%。这表明共享注意力机制可以在不牺牲太多性能的情况下，显著提高LLM的效率。

🎯 应用场景

该研究成果可应用于各种资源受限的场景，例如移动设备、嵌入式系统和边缘计算环境。通过降低LLM的计算和内存需求，可以使这些模型在这些平台上更高效地运行，从而实现智能助手、自然语言处理和机器翻译等应用。此外，该方法还可以用于加速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.