Can I read DHA: Learning Decoupled-Head Attention from Transformer Checkpoints via Adaptive Heads Fusion on EtoBox?

DHA: Learning Decoupled-Head Attention from Transformer Checkpoints via Adaptive Heads Fusion by Chen, Yilong; Zhang, Linhao; Shang, Junyuan; Zhang, Zhenyu; Liu, Tingwen; Wang, Shuohuan; Sun, Yu is a scholarly article available to read on EtoBox.

What is DHA: Learning Decoupled-Head Attention from Transformer Checkpoints via Adaptive Heads Fusion about?

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\% o

Author

Chen, Yilong; Zhang, Linhao; Shang, Junyuan; Zhang, Zhenyu; Liu, Tingwen; Wang, Shuohuan; Sun, Yu

Published

2024

Language

EN