@inproceedings{
chen2024dha,
title={{DHA}: Learning Decoupled-Head Attention from Transformer Checkpoints via Adaptive Heads Fusion},
author={Yilong Chen and Linhao Zhang and Junyuan Shang and Zhenyu Zhang and Tingwen Liu and Shuohuan Wang and Yu Sun},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=g92nu7knRq}
}
Large language models (LLMs) with billions of parameters demonstrate impressive performance. However, the widely used Multi-Head Attention (MHA) in LLMs incurs substantial computational and memory costs during inference. While some efforts have optimized attention mechanisms by pruning heads or sharing parameters among heads, these methods often lead to performance degradation or necessitate substantial continued pre-training costs to restore performance. Based on the analysis of attention redundancy, we design a Decoupled-Head Attention (DHA) mechanism. DHA adaptively configures group sharing for key heads and value heads across various layers, achieving a better balance between performance and efficiency. Inspired by the observation of clustering similar heads, we propose to progressively transform the MHA checkpoint into the DHA model through linear fusion of similar head parameters step by step, retaining the parametric knowledge of the MHA checkpoint. We construct DHA models by transforming various scales of MHA checkpoints given target head budgets. Our experiments show that DHA remarkably requires a mere 0.25\% of the original model's pre-training budgets to achieve 97.6\% of performance while saving 75\% of KV cache. Compared to Group-Query Attention (GQA), DHA achieves a 12$\times$ training acceleration, a maximum of 24.85\% performance improvement under 0.2B tokens budget, and finally 2.3\% overall performance improvement.
DHA: Learning Decoupled-Head Attention from Transformer Checkpoints via Adaptive Heads Fusion
Yilong Chen Linhao Zhang Junyuan Shang Zhenyu Zhang Tingwen Liu Shuohuan Wang Yu Sun