CHAI：面向高效大语言模型推理的聚类头注意力机制
机器学习
2024-04-30 v2 计算与语言
摘要
拥有数千亿参数的大语言模型 (LLM) 已改变了机器学习领域。然而，在推理时服务这些模型既计算密集又内存密集，单个请求可能需要多个 GPU 和数十 GB 的内存。多头注意力是 LLM 的关键组件之一，可能占 LLM 内存和计算需求的 50% 以上。我们观察到，不同注意力头在关注哪些词元上存在大量冗余。基于这一洞察，我们提出了聚类头注意力 (CHAI)。CHAI 在运行时将具有高度相关性的头组合起来进行自注意力计算，从而同时减少内存和计算量。在我们的实验中，我们展示了 CHAI 能够将存储 K,V 缓存的内存需求降低高达 21.4%，并将推理时间延迟降低高达 1.73 倍，且无需任何微调。CHAI 在 3 个不同模型（即 OPT-66B、LLAMA-7B、LLAMA-33B）和 5 个不同评估数据集上实现此效果时，精度偏差最大仅为 3.2%。
引用
@article{arxiv.2403.08058,
title = {CHAI: Clustered Head Attention for Efficient LLM Inference},
author = {Saurabh Agarwal and Bilge Acun and Basil Hosmer and Mostafa Elhoushi and Yejin Lee and Shivaram Venkataraman and Dimitris Papailiopoulos and Carole-Jean Wu},
journal= {arXiv preprint arXiv:2403.08058},
year = {2024}
}