DHA: Learning Decoupled-Head Attention from Transformer Checkpoints via Adaptive Heads Fusion

TL;DR

This paper introduces DHA, a decoupled-head attention mechanism that adaptively shares parameters across heads, transforming existing transformer checkpoints to improve efficiency and performance with minimal additional training.

Contribution

The paper proposes DHA, a novel attention mechanism that reduces computational costs by adaptively sharing heads and transforming checkpoints, maintaining performance with less pre-training.

Findings

DHA achieves 97.6% of original performance with only 0.25% of pre-training budget.

DHA saves 75% of KV cache during inference.

DHA accelerates training by 5 times compared to GQA.

Abstract

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 5$\times$ training acceleration, a maximum of 13.93\% performance improvement under 0.01\% pre-training budget, and 4\% relative improvement under 0.05\% pre-training budget.