CHAI: Clustered Head Attention for Efficient LLM Inference

📄 arXiv: 2403.08058v2 📥 PDF

作者: Saurabh Agarwal, Bilge Acun, Basil Hosmer, Mostafa Elhoushi, Yejin Lee, Shivaram Venkataraman, Dimitris Papailiopoulos, Carole-Jean Wu

分类: cs.LG, cs.CL

发布日期: 2024-03-12 (更新: 2024-04-27)

💡 一句话要点

提出聚类头注意力机制以提高大语言模型推理效率

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

关键词: 聚类头注意力 大型语言模型 推理效率 自注意力机制 内存优化 计算资源 模型压缩 机器学习

📋 核心要点

  1. 现有大型语言模型在推理时的计算和内存需求极高，导致资源消耗大，效率低下。
  2. 提出的聚类头注意力（CHAI）通过结合相关性高的头部，优化自注意力计算，降低内存和计算需求。
  3. 实验结果显示，CHAI在不微调的情况下，内存需求减少21.4%，推理速度提升1.73倍，且准确率偏差小于3.2%。

📝 摘要（中文）

大型语言模型（LLMs）以数百亿参数改变了机器学习领域。然而，这些模型在推理时的计算和内存需求极高，单个请求可能需要多个GPU和数十GB内存。多头注意力是LLMs的关键组成部分，占据了超过50%的内存和计算需求。基于头部之间的冗余性，本文提出了聚类头注意力（CHAI），通过在运行时结合相关性高的头部进行自注意力计算，从而减少内存和计算需求。实验结果表明，CHAI在不需要微调的情况下，能够将K,V缓存的内存需求减少最多21.4%，推理延迟提高至1.73倍，同时在三种不同模型（OPT-66B、LLAMA-7B、LLAMA-33B）及五个评估数据集上，准确率的偏差最大为3.2%。

🔬 方法详解

问题定义：本文旨在解决大型语言模型在推理时的高计算和内存需求问题，现有多头注意力机制存在头部冗余，导致资源浪费。

核心思路：通过聚类相关性高的头部进行自注意力计算，CHAI有效减少了内存和计算需求，优化了推理效率。

技术框架：CHAI的整体架构包括头部聚类模块和自注意力计算模块，首先对头部进行聚类，然后在推理时使用聚类结果进行高效计算。

关键创新：CHAI的主要创新在于动态聚类相关性高的头部，显著降低了冗余计算，与传统多头注意力机制相比，提升了推理效率。

关键设计：在设计中，CHAI采用了动态聚类算法，确保在推理时选择最优的头部组合，同时保持模型的准确性，损失函数和网络结构经过优化以适应新的计算流程。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，CHAI在不进行微调的情况下，能够将K,V缓存的内存需求减少最多21.4%，推理延迟提升至1.73倍，且在三种不同模型和五个评估数据集上的准确率偏差最大仅为3.2%。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和机器翻译等，能够显著提高大型语言模型的推理效率，降低资源消耗，适用于资源受限的环境。未来，CHAI可能推动更多高效模型的开发，促进AI技术的广泛应用。

📄 摘要（原文）

Large Language Models (LLMs) with hundreds of billions of parameters have transformed the field of machine learning. However, serving these models at inference time is both compute and memory intensive, where a single request can require multiple GPUs and tens of Gigabytes of memory. Multi-Head Attention is one of the key components of LLMs, which can account for over 50% of LLMs memory and compute requirement. We observe that there is a high amount of redundancy across heads on which tokens they pay attention to. Based on this insight, we propose Clustered Head Attention (CHAI). CHAI combines heads with a high amount of correlation for self-attention at runtime, thus reducing both memory and compute. In our experiments, we show that CHAI is able to reduce the memory requirements for storing K,V cache by up to 21.4% and inference time latency by up to 1.73x without any fine-tuning required. CHAI achieves this with a maximum 3.2% deviation in accuracy across 3 different models (i.e. OPT-66B, LLAMA-7B, LLAMA-33B) and 5 different evaluation datasets.