CHAI: Clustered Head Attention for Efficient LLM Inference

TL;DR

This paper introduces CHAI, a method that reduces memory and computation in large language models by clustering similar attention heads, achieving significant efficiency gains with minimal accuracy loss.

Contribution

The paper proposes Clustered Head Attention (CHAI), a novel approach that groups similar attention heads to improve inference efficiency in LLMs without fine-tuning.

Findings

Memory for K,V cache reduced by up to 21.4%

Inference latency decreased by up to 1.73x

Accuracy deviation limited to 3.2% across models

Abstract

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.