Post-Training Sparse Attention with Double Sparsity

📄 arXiv: 2408.07092v2 📥 PDF

作者: Shuo Yang, Ying Sheng, Joseph E. Gonzalez, Ion Stoica, Lianmin Zheng

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

发布日期: 2024-08-11 (更新: 2024-08-18)

🔗 代码/项目: GITHUB

💡 一句话要点

提出双重稀疏注意力，通过后训练稀疏化加速大语言模型推理。

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

关键词: 大语言模型 推理加速 稀疏注意力 后训练量化 KV缓存 token稀疏性 通道稀疏性

📋 核心要点

  1. 大语言模型推理速度慢、内存占用高，关键瓶颈在于过多的KV缓存访问。
  2. 提出双重稀疏注意力，结合token和通道稀疏性，离线校准通道稀疏模式以高效识别重要token。
  3. 实验表明，该方法在多种任务和模型上实现了显著的加速和内存减少，精度损失极小。

📝 摘要（中文）

本文提出了一种名为“双重稀疏”的后训练稀疏注意力技术，旨在通过减少键值(KV)缓存访问来缓解大语言模型推理过程中的瓶颈。双重稀疏结合了token稀疏性（仅使用重要token计算自注意力）和通道稀疏性（使用重要特征通道识别重要token）。核心思想是通道稀疏模式相对静态，允许离线校准以提高运行时效率，从而准确高效地识别重要token。该方法可以与卸载技术结合，显著减少内存使用。实验结果表明，双重稀疏在Wiki-2困惑度、键值检索和长上下文基准测试等任务中，对Llama-2-7B、Llama-2-70B和Mixtral-8x7B等模型实现了1/16的token和通道稀疏性，且对精度影响极小。在GPU上，注意力操作加速高达14.1倍，端到端推理性能提升1.9倍。结合卸载技术，在256K序列长度下，解码速度比最先进的解决方案加速16.3倍。代码已公开。

🔬 方法详解

问题定义：大语言模型推理过程中，需要频繁访问KV缓存，导致速度慢、内存占用高。现有方法难以在保证精度的前提下，有效减少KV缓存访问量，尤其是在长序列场景下，问题更加突出。

核心思路：利用token稀疏性和通道稀疏性的互补优势，设计双重稀疏策略。通道稀疏性模式相对稳定，可以通过离线校准学习，用于指导token稀疏性的选择，从而在运行时高效地选择重要token，减少KV缓存访问。

技术框架：该方法属于后训练量化/稀疏化技术，无需重新训练模型。主要包含两个阶段：离线校准阶段和在线推理阶段。离线校准阶段，通过少量数据确定通道稀疏模式。在线推理阶段，利用校准后的通道稀疏模式选择重要token，进行稀疏注意力计算。同时，可以结合KV缓存卸载技术进一步减少内存占用。

关键创新：核心创新在于结合了token稀疏性和通道稀疏性，并利用通道稀疏性的稳定性进行离线校准。与传统的token稀疏方法相比，该方法能够更准确地识别重要token，从而在保证精度的前提下实现更高的稀疏率。与单纯的通道稀疏方法相比，token稀疏性的引入可以进一步减少计算量。

关键设计：通道稀疏模式通过离线校准数据集获得，具体方法未知（论文未详细描述）。Token选择策略基于校准后的通道稀疏模式，选择具有较高激活值的token。稀疏比例（token和通道）是关键参数，需要根据具体任务和模型进行调整。损失函数未知，但目标是最小化稀疏化带来的精度损失。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，双重稀疏在Llama-2-7B、Llama-2-70B和Mixtral-8x7B等模型上实现了1/16的token和通道稀疏性，且对精度影响极小。在GPU上，注意力操作加速高达14.1倍，端到端推理性能提升1.9倍。结合卸载技术，在256K序列长度下，解码速度比最先进的解决方案加速16.3倍。

🎯 应用场景

该研究成果可应用于各种需要加速大语言模型推理的场景，例如：移动设备上的本地部署、资源受限的边缘计算环境、以及需要处理长序列的应用程序。通过减少内存占用和提高推理速度，可以降低部署成本，提升用户体验，并扩展大语言模型在实际应用中的范围。

📄 摘要（原文）

The inference process for large language models is slow and memory-intensive, with one of the most critical bottlenecks being excessive Key-Value (KV) cache accesses. This paper introduces "Double Sparsity," a novel post-training sparse attention technique designed to alleviate this bottleneck by reducing KV cache access. Double Sparsity combines token sparsity, which focuses on utilizing only the important tokens for computing self-attention, with channel sparsity, an approach that uses important feature channels for identifying important tokens. Our key insight is that the pattern of channel sparsity is relatively static, allowing us to use offline calibration to make it efficient at runtime, thereby enabling accurate and efficient identification of important tokens. Moreover, this method can be combined with offloading to achieve significant memory usage reduction. Experimental results demonstrate that Double Sparsity can achieve $\frac{1}{16}$ token and channel sparsity with minimal impact on accuracy across various tasks, including wiki-2 perplexity, key-value retrieval, and long context benchmarks with models including Llama-2-7B, Llama-2-70B, and Mixtral-8x7B. It brings up to a 14.1$\times$ acceleration in attention operations and a 1.9$\times$ improvement in end-to-end inference on GPUs. With offloading, it achieves a decoding speed acceleration of 16.3$\times$ compared to state-of-the-art solutions at a sequence length of 256K. Our code is publicly available at https://github.com/andy-yang-1/DoubleSparse.