Align Attention Heads Before Merging Them: An Effective Way for Converting MHA to GQA (2025)

Abstract: Large language models (LLMs) have demonstrated exceptional performance across diverse natural language processing tasks. However, as the model size and the input sequence’s length increase, the linearly increasing key-value (KV) cache significantly degrades inference throughput. Therefore, grouped-query attention (GQA), as an alternative to multi-head attention (MHA), has been widely introduced into LLMs. In this work, we propose a cost-effective method for converting MHA into GQA with any compression ratio of KV heads. The key point of our method lies in the application of Procrustes analysis to the attention heads, which enhances the similarity among attention heads while preserving computational invariance, thereby improving the model’s post-training performance. Subsequently, we employ \mathit{L_0} regularization to prune redundant parameters. The model after pruning can be adapted to the standard GQA framework. Experimental results show that our strategy can compress up to 87.5% KV heads of LLaMA2-7B model and 75% KV heads of Sheared-LLaMA-1.3B with acceptable performance degradation. Our code is released at https://github.com/fpcsong/mha2gqa.

## introduction recently, large language models (llms) @xcite @xcite show remarkable performance on a variety of natural language processing tasks. however, since most llms are based on transformer architecture @xcite , the expansion of the sequence length during inference leads to a linear increase in the memory footprint of keyvalue (kv) cache, which substantially increases on-device memory consumption. a reduced kv cache footprint not only lowers inference costs but also facilitates the processing of longer sequences * corresponding author. and improves inference speed. therefore, reducing the size of the kv cache is a key issue for llms. many methods for kv cache compression have been proposed, including kv cache quantization @xcite @xcite , tokens dropping @xcite @xcite and so on. however, these approaches often introduce additional computational procedures, which are incompatible with general llm frameworks. another approach to kv cache compression is to directly change the attention architecture. multi-query attention (mqa) @xcite and grouped-query attention (gqa) @xcite reduce kv cache by allowing multiple attention heads to share a single kv head, offering a simple and effective solution for attention optimization. since gqa has better inference stability and performance, it has been widely used in llama 2 @xcite , llama 3 @xcite , qwen2 @xcite , mistral @xcite and other llms @xcite . multi-head latent attention (mla) @xcite further reduces kv cache through low-rank projection of the cached data. mla has been used in the deepseek model series @xcite . these efficient attention architectures achieve kv cache compression with better universality. in this study, we present a method that converts mha to gqa to compress kv cache. inspired by the idea of computational invariance in llms @xcite and procrustes analysis @xcite , we apply proper orthogonal transformations to the projection matrices in attention heads to simplify the conversion from mha to gqa: specifically, we regroup attention heads based on the similarity of their kv caches and use generalized procrustes analysis (wikipedia contributors, 2022) to maximize the similarity of attention heads within each group. this transformation pre-serves the model's output invariance. finally, l 0 regularization @xcite is applied to transfer original kv heads to new ones to get gqa. figure 1 illustrates the overall framework of our method. experimental results show that attention head transformation can significantly improve the performance of the pruned model. our contributions are as follows. in cofi @xcite , the l 0 method is applied directly to llms by introducing pruning masks with different granularities. they prune the hidden dimension, the intermediate dimension, the number of attention heads, and even an entire mha or ffn layer. the subsequent work sheared-llama @xcite incorporates previous methods and specifies the target structure so that the pruned model can be directly adapted to standard llm frameworks. ## method in this section, we will specifically describe our method. our method consists of two parts, transformation of attention heads and pruning training. transformation of attention heads represents employing procrustes analysis to align projection matrices in order to enhance the similarity between figure 1: an illustration of our method. 1. in each attention layer of llm, after the average similarity score is evaluated by procrustes analysis between every two value caches (or key caches). use simulated annealing algorithm to search the optimal grouping result that maximizes the similarity score among value caches (or key caches) in each layer. 2. after grouping, then fuse the orthogonal matrix obtained from generalized procrustes analysis into each projection matrix to enhance the similarity among attention heads in each group without changing the model. 3. during training, l 0 loss is used to gradually transfer original key-value projection matrices to newly added ones within each group (there are g groups in one attention layer). after pruning, original key-value projection matrices will be discarded, then we get a standard gqa model. ## conclusion in this paper, we propose a general method for converting an mha model into a gqa model with any compression ratio of kv heads. we find that applying orthogonal transformations to attention heads based on procrustes analysis can enhance the similarity between kv heads without changing the model, thereby improving its performance after pruning. furthermore, we introduce l 0 regularization during pruning training, which reduces the impact of directly eliminating parameters on the model. our method is applicable to all kv head pruning conditions.