Title: DHA: Learning Decoupled-Head Attention from Transformer Checkpoints via Adaptive Heads Fusion
Venue: NeurIPS 2024
Accepted: yes
Keywords: Large Language Models

Multi-Head Attention

Pre-training Acceleration

Efficient Inference

Model Fusion
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 96.1\% 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 5\% relative improvement under 0.05\% pre-training budget.

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

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 only 2.5% of the original model's pre-training budgets to achieve 96.1% of performance while saving 75% of KV cache. Compared to Group-Query Attention (GQA), DHA achieves a 5 × training acceleration, a maximum of 13.93% performance improvement under 0.01% pre-training budget, and 5% relative improvement under 0.05% pre-training budget.

1 Introduction

Transformer-based large language models (LLMs) shine in various natural language tasks due to their powerful understanding and generation capabilities [1, 2, 3]. Multi-Head Attention (MHA) is widely used in LLMs, with the number of heads increasing as the model size grows. However, MHA inference overhead increases linearly with the expansion of the context and model sizes, due to the surprisingly large memory consumption of the KV Cache mechanism. For instance, a 7 billion-parameter model with 32 heads and 32 layers, an input batch size of 4, and a sequence length of 32k results in 64GB of KV cache, which is 4 . 7 × larger than the model weights.

To reduce computational and memory overhead during inference, a widely used approach involves adapting the MHA model to a more efficient structure through the reuse of parameters across multiple heads [4, 5, 6] , such as Multi-Query Attention (MQA) [4] and Grouped-Query Attention (GQA) [5]. These methods utilize a portion of the original training computation which avoid information loss

Layer n+1

……

Layer n+1

……

Layer n+1

…..…

V

K

Q

V

K

Q

Layer n

……

Layer n

……

Layer n

……

V

K

Q

Multi-Head Attention (MHA)

Grouped-Query Attention (GQA)

Decoupled-Head Attention (DHA)

Head Dependence Search

Learning-Based Head Fusion

Fusion

Fusion

Group size

𝑔

Heads

New Heads

Group size

𝑔

1/g

Heads

New Heads

GQA Initialization

DHA Initialization

In-Fusion Heads

due to training-inference inconsistencies, a common issue in pruning-based [7, 8, 9, 10, 11] works. However, the training computation is prohibitively expensive for recovering the model's performance, due to the information loss in the parameters when creating the initial point.

Thus, in this work, we seek to address the following question:

How can we construct a more efficient model while keeping costs as low as possible ?

With the limited understanding of parameter characteristics in modern LLMs, we first perform an empirical analysis from the perspectives of heads' parameter similarity. We observe that there are some head-clusters with high internal similarity in MHA checkpoints. Similar head clusters imply enormous redundancy in MHA, which coincides with the sparsity found in previous studies [12, 13]. In particular, the clusters of key heads and value heads across different layers show a decoupled distribution, meaning that there is a significant variation in the distribution of head-cluster similarities across layers, key heads and value heads, as illustrated in Fig. 2a,2b. Intuitively, we can prune redundant heads based on the above characteristics. Nonetheless, each head has its unique role, and thus no heads should be arbitrarily discarded. Furthermore, we find that linear fusion based on multiple similar heads can reconstruct the original head functionality without causing a significant performance drop (see Sec. 3.1). Based on this observation, we believe that selectively fusing corresponding heads in clusters can construct a more efficient architecture with minimum loss.

In this paper, we propose Decoupled-Head Attention (DHA) , an efficient attention architecture developed through the Adaptive Head Fusion of checkpoints' parameters. Recalling the decoupled heads parameter characteristics, DHA allocates different numbers of key heads and value heads at different layers to balance model efficiency and performance. The MHA checkpoint can be rapidly transformed into DHA with three stages: Search , Fusion , and Continued Pre-training (CT) . During the Search stage, we group similar functional heads together and determine reasonable allocations of key heads and value heads for each layer. Specifically, we reconfigure the original key and value head into multiple linear combinations of heads within the same layer. Thus, we can allocate the heads based on the loss after replacement. In the Fusion stage, we perform linear fusion on similar heads, ensuring the preservation of original functionality. Leveraging the Augmented Lagrangian approach [14, 15], the Fusion operator initializes from MHA and explores possible head combinations in the early training, followed by refined intra-group head fusion in the later. Based on well-trained operators on unlabeled data, we can rapidly obtain high-performing initial points for DHA from MHA checkpoints, requiring only a minimal amount of Continued Pre-training to restore performance.

0

3

6

9

12

15

18

21

24

27

30

0

3

6

30 27 24 21 18 15 12 9

Wk Head Similarity

0

3

6

9

12

15

18

21

24

27

30

0

3

6

30 27 24 21 18 15 12 9

Wv Head Similarity

0.2

0.4

0.6

0.8

1.0

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0

3

6

9

12

15

18

21

24

27

30

0

3

6

30 27 24 21 18 15 12 9

Wk Head Similarity

0

3

6

9

12

15

18

21

24

27

30

0

3

6

30 27 24 21 18 15 12 9

Wv Head Similarity

0.6

0.7

0.8

0.9

1.0

0.9

0.9

0.9

0.9

1.0

1.0

1.0

To verify the effectiveness, we construct DHA on models of different sizes, such as LLaMA2-7B [3], Sheared-LLaMA-2.7B & -1.3B [16] with the heads budget ratio set at 50% and 25%. With a modest fusion training of just 0.2 billion tokens, DHA learns sufficiently competent initial points. As the continued pretraining progresses, DHA continuously outperforms GQA narrowing the gap with MHA on 9 representative downstream tasks. DHA only requires 0.25% of MHA pre-training budget. Meanwhile, DHA is capable of reducing KV Cache by up to 75% compared to MHA with minimal accuracy trade-off (maximum of 5.6%). Compared to GQA, DHA achieves a 5 × training acceleration, a maximum 13.93% performance improvement under 0.01% pre-training budget, and 5% relative improvement under 0.05% pre-training budget. Overall, DHA exhibits great performance and efficiency, which can be quickly adapted to various existing MHA Transformer models.

2 Background

Let X = ( x 1 , . . . , x p ) ∈ R p × d model denote the input prompts of hidden states of a Transformer layer, where p stands for the number of tokens and d model for the hidden state dimension.

Multi-Head Attention (MHA) MHA[17] performs the attention with H different heads. For h -th head, different weight matrices W h q , W h k , W h v ∈ R d model × d k are used to project the input sequence into query, key, value vector, where d k represents head dim. Denote softmax funcion as σ , we have:

Ultimately MHA combines heads' outputs through the output projection W O ∈ R d model × d model 2 .

Grouped-Query Attention (GQA) & Multi-Query Attention (MQA) To accelerate inference, MQA[4] and GQA [5] have been proposed based on the idea of reusing head parameter weights. In these variants, H different query heads are divided into G groups, where the heads within the same group share the same key heads and value heads parameter matrices. Given the mapping relationship from the h -th query head to a GQA key and value heads using the many-to-one function g ( h ) , we define the h -th head forward pass as:

Here, K / V g ( h ) init refers to MHA key/value heads parameters within the g ( h ) init -th group during GQA initialization. When transitioning from an MHA checkpoint, GQA uses the mean pooling method for heads within the group. MQA is a special case of GQA where G = 1 3 .

Due to mean pooling for initialization, GQA results in loss of capability when converting from MHA, necessitating expensive pre-training to recover. We aim to identify better initialization and more refined head mapping relationships to achieve superior performance with reduced training costs.

3 Observation

To study the inherent characteristics of head parameters in MHA, we use Centered Kernel Alignment [18] to calculate the heads' similarity within each layer's W k , W v . Based on the average heads' similarity, we define the redundancy of each MHA layer. For details, please refer to Appendix B.1.

3.1 Head clusters in MHA

Observation From Fig. 2a and Fig. 2b, we observe that clusters form spontaneously among heads, with high similarity within clusters and low similarity between clusters. It indicates that heads among different clusters may have distinct functionalities, processing linguistic features in various aspects.

Analysis Given the numerous similar head clusters in W k and W v , we identified the opportunity to linearly fuse functionally similar heads within clusters while retaining each head's parameterized knowledge. We conducted an empirical study, transforming the parameters of Head 0 in MHA into a linear fusion of the parameters from Heads 0, 1, 2, and 3. We share the fusion head across four query heads and progressively optimize the fusion ratio under the LmLoss. For details, please refer to Sec. 4.2. As shown in Fig. 3a, the loss remains unchanged as the proportion of Head 0 decreases and only increases when four heads parameters' ratios approach an even distribution. It suggests that fusing similar parameters can reduce the number of heads without significant information loss.

3.2 Variability across Layers and KV pairs

Observation The distribution of similar head clusters varies between different layers. As illustrated in Fig.2a, 2b, the 0th layer of MHA shows few similar head clusters, while the 21st layer exhibits many. Within the same layer, value heads exhibit more clusters and higher similarity compared to key heads, indicating a divergence between the two. Fig. 3b shows that the redundancy is lower in the initial and final layers, and higher in the middle layers. Moreover, W v redundancy significantly exceeding that of W k

.

Analysis Inspired by layer and key-value head variability, we propose allocating more heads to layers with lower redundancy to enhance learning and expression. Since W v shows higher redundancy than W k , we can decouple and allocate more heads budget to critical key components, while compressing redundant value heads at a higher compression rate. Finer grouping and sharing based on the parameters function may contribute to compression rates and performance improvements.

4 Method

In this section, we propose a more efficient Decoupled Head Attention (DHA) architecture and its construction process. We define DHA in Sec. 4.1 and Adaptive Head Fusion algorithm in Sec. 4.2. Then we demonstrate the adaptive construction based on the MHA checkpoint, which can be divided into: Search , Fusion , and Continued Pre-training (Discussed in in Sec. 4.3). Finally, we introduce practical application of our DHA architecture on the LLaMA2 model in Sec. 4.3.

4.1 Decoupled-Head Attention (DHA)

We present a more efficient attention architecture called Decoupled-Head Attention (DHA). Based on observed significant functional differences among different layers' key value heads, DHA adaptively allocates more heads to critical components, thus enhancing overall model efficiency and performance.

0

100

200

300

400

Steps

2.0

2.2

2.4

Loss

Loss

0.0

0.2

0.4

0.6

0.8

1.0

Head Ratios

Head 0

Head 1

Head 2

Head 3

(a) Loss in 4-Head Fusion

0

10

20

30

Layer

0.2

0.4

0.6

0.8

Redundancy

Key Heads

Value Heads

(b) 32-Layer Redundancy

𝐖

!

𝐖

"

𝐖

#

𝐖

$

𝐖

!

𝐖

"

𝐖

#

𝐖

$

𝐖

!

𝐖

"

𝐖

#

𝐖

$

ℒ	 =

!

min

",ℳ

ℒ

%&

ℳ 𝚯

'()

+𝜆ℒ

*+,-./

ℒ

!"#$%&

= 𝐺𝑟𝑜𝑢𝑝_𝑛𝑢𝑚×

1

8

$

'

()*

$

'

(

!

)(+*

$

'

,)*

𝜔

(,

-𝜔

(

!

,

'

In-Group Head Fusion

𝜔

-,-

𝜔

',-

𝐖

!

𝐖

"

𝐖

#

𝐖

$

𝐖

!

𝐖

"

𝐖

#

𝐖

$

𝜔

-,'

= 𝜔

',-

= 0

𝜔

-,-

= 𝜔

','

= 1

𝜔

-,-

𝜔

','

𝜔

*,*

𝜔

/,/

𝜔

*,/

= 𝜔

/,*

= 0

𝜔

*,*

= 𝜔

/,/

= 1

𝜔

-,'

𝜔

','

Fusion Initiation

𝜔

*,*

𝜔

/,*

𝜔

*,/

𝜔

/,/

𝜔

-,-

𝜔

','

𝜔

*,*

𝜔

/,/

Fusion Process

𝜔

-,'

𝜔

',-

𝜔

*,/

𝜔

/,*

𝐖

!

𝐖

#

𝐖

"

𝐖

$

Q

K/V

Dependence Search

𝐖

!

𝐖

#

𝐖

"

𝐖

$

ℒ	 =

min

",ℳ

ℒ

%&

ℳ 𝚯

'()

+𝜆ℒ

*+,-./

ℒ

!"#$%&

=

1

64

$

0

()*

$

0

(

!

)(+*

$

0

,)*

𝜔

(,

-𝜔

(

!

,

'

𝐖

!

𝐖

"

𝐖

#

𝐖

$

𝐖

!

%

𝜔

!,#

𝜔

$,%

𝜔

!,!

𝜔

$,$

Post-Process

𝐖

#

%

𝜔

-,-

= 𝜔

',-

𝜔

-,'

= 𝜔

','

𝜔

*,/

= 𝜔

/,/

𝜔

*,*

= 𝜔

/,*

𝚯

&/(

)*+

= 𝐖′

0

, ⋯ , 𝐖′

1

𝐖′

-

= 𝜔

-,-

𝐖

-

+𝜔

-,'

𝐖

'

𝐖′

*

= 𝜔

*,*

𝐖

*

+𝜔

*,/

𝐖

/

𝐖

-

, 𝐖

'

𝐖

*

, 𝐖

/

𝜔

-,'

𝜔

',-

𝜔

*,/

𝜔

/,*

Definition Defined model with L layers and H Q heads in a layer, the numbers of Key heads and Value heads in the l -th layer are denoted as H K l , H V l . We define the many-to-one mapping functions d K ( h, l ) and d V ( h, l ) representing key and value head corresponding to the h -th query head in l -th DHA layer. The computation be formalized as follows:

DHA shares a key and value head in multi query heads' computation based on independent mapping functions at different layers 4 . GQA can be considered a special case of DHA, where not only all layers share the same mapping functions, but the mapping functions for keys and values are identical.

4.2 Learning Efficient MHA Transformation via Linear Heads Fusion

Due to the high cost of building an efficient Attention mechanism in LLM from scratch, we construct DHA based on the existing MHA checkpoint using minimal computational budgets. Based on the head clustering phenomenon in MHA, we propose a linear fusion method for similar heads within clusters. By incrementally fusing head parameters, we compress the number of heads while retaining the original model's knowledge, significantly reducing training budgets and improving performance.

Goal Formally, we define a model with Layer number L and Head number H as Θ L,H = [ W 1 , · · · , W L ] , where W l ∈ R D × D denotes the weight of layer l with input and output dimension D . In the initialization, our goal is to transfer knowledge from a MHA model Θ MHA L,H 1 to a DHA model Θ DHA L,H 2 , where H 1 > H 2 . By learning a fusion operation that minimizes the functional difference between MHA and DHA model, the goal can be formalized as

Where M is the fusion operator, D is the training dataset, L lm is the training loss function, L fusion measures the transformation from MHA to DHA, and λ is the learnable scale factor.

Fusion Operator During DHA initialization, the fusion operator M constructs new heads based on the linear combinations of the original key and value heads within the group, and shares the new heads among the query heads' forward. Define each group K d K ( h,l ) , V d V ( h,l ) represents key, value heads group corresponding to the h -th query head in l -th layer, g = { g K , g V } as the group size. By introducing variables ω h = { ω hj } g j =1 , ω ∈ M represents the proportion of j -th key, value head

involved in the h -th query head forward within group. For each group, a head have forward pass as:

where ω hj will be initialized to Kronecker delta function, which equals 1 if and only if h = j , and equals 0 otherwise. Under this initialization setting, the forward computation of DHA is completely equivalent to that of MHA, see Fig. 4.

Optimization During the optimization phase, we design a fusion loss to optimize the initialized model towards DHA target architecture. Note that after initialization, the mapping of heads within the group W h,l q → W j k/v is a many-to-many mapping, denoted by the function d K/V init ( h, l ) . This indicates that in the forward process of each query, the key head or value head can be expressed as different linear combinations of g MHAheads. According to Eq. 3, we aim to achieve a many-to-one mapping that a single fused key head or value head are shared across multiple query heads in DHA, denoted by the function d K/V ( h, l ) . Thus, we design a fusion loss L fusion to optimize the initial mapping functions to converge to a single mapping function, i.e., d K/V init ( h, l ) → d K/V ( h, l ) , ∀ h ∈ K n / V n . Specifically, we define the optimization objective as minimizing the difference between the mapping functions of different query heads h and h ′ within the l -th layer and n -th group:

where g = g K/V represents the number of heads within a group. Since W j, MHA k/v,ij can be regarded as an orthogonal scalar, and thus we only need to optimize fusion variables ω , so we have:

Where N represents the number of groups, N = H 1 g . The fusion loss can be measured as the mean squared error loss of the head and head fusion variables within each group at each layer.

Augmented Lagrangian approach When the fusion loss is zero, the key and value heads corresponding to query heads within the group are optimized to share the same fusion variables. This allows the new DHA key-value head parameters to be effectively shared among the queries in the group. Given that it is challenging to optimize the loss to a very small value, we use an augmented Lagrangian approach [14, 15] for incremental architectural transformations. Define t as the target loss, b as the base decay factor, s as the current global step, k as the total number of steps in the warm-up phase, the overall training optimization is an adversarial game:

Our Augmented Lagrangian approach enforces the constraint L fusion ≤ t , where the Lagrange multiplier λ is updated during training. The update increases the loss unless the constraint is satisfied. Early in training, the model tolerates more significant discrepancies between head weights, promoting exploration. As training progresses, the margin shrinks, enforcing stricter adherence to minimizing discrepancies and refining head alignment within the group.

4.3 Adaptive DHA Transformation on LLaMA Model

Based on the observation of similar head clusters and key-value head parameter variability across layers, DHA employs the adaptive transformation. It allows DHA to search for and fuse similar heads while allocating different group sizes across layers. As shown in Fig. 4, the transformation can be divided into three stages: Search , Fusion and Continued Pre-training .

In the beginning, we initialize the DHA operators to the MHA model. Next, we perform 240 Search steps, calculating L fusion for each layer and L head for all heads. Based on the L head, we perform head grouping intending to minimize the average loss of heads within each group and maximize the average loss of heads between groups and groups. Based on L fusion, we use a dynamic programming

algorithm to allocate more head budget to layers with higher loss within a total budget. It allows us to fuse the most similar heads to minimize loss during the fusion process and selectively compress the model's most redundant components. For more details, see Apendix B.3, B.4.

During Fusion phase, we modified the forward propagation path of MHA in the form of DHA based on the layer head budget and head grouping obtained during the Search phase. Then we antagonistically optimize the fusion operator and update Lagrangian multipliers λ , the L fusion that marks this DHA fusion process decreases. When L fusion is less than 1e-3, we terminate the fusion algorithm and enter the Continued Pre-training phase.

During the Continued Pre-training phase, we fuse MHA head parameters based on averaged fusion weights to construct DHA initialization. DHA initialization can recover the performance with a small amount of restorative pre-training. For more information, please refer to Appendix B.2.

Our method can theoretically transform MHA architecture in any transformer model to efficient DHA architecture. Using LLaMA models as case studies, we implemented DHA transformation with various compression rates on all MHA layers. Notably, we expanded the dimension of each head's fusion coefficient ω from 1 to the head's dimension d k , allowing for finer-grained parameter fusion and better knowledge retention. Intuitively, we learn different fusion ratios for each dimension of the head. Only a very small number of additional parameters need to be introduced, DHA significantly accelerates training and improves performance.

5 Empirical Evaluation

5.1 Experimental Setup

Data. To train DHA operators and extend pre-training, we employ the RedPajama [19], which parallels the LLaMA training data across seven domains: CommonCrawl, C4, GitHub, Wikipedia, Books, ArXiv, and Stack-Exchange. This dataset comprises a validation set with 2 million tokens, a training set containing 4 billion tokens and an additional pre-training set totaling 50 billion tokens.

Training. Our experimental framework utilizes the Sheared-LLaMA codebase [16] implemented on the Composer package [20], and is executed on 8 NVIDIA A100 GPUs (80GB). The models are trained with a sequence length of 4096, employing a global batch size of 64 during the fusion phase and 256 during the continued pre-training phases. The learning rates were set at 1e-4 for language modeling loss, and 1e-2 for Lagrangian multipliers and fusion operators respectively.

Budget. DHA models were trained for 1000 steps (0.2B token budget) during the fusion phases. For the continued pre-training, we trained both baseline models and DHA for up to 50000 steps (50B token budget). To evaluate the training acceleration capability of DHA, we evaluate its performance under two budget scenarios. First, we set a budget of 1B tokens to compare the early-stage rapid convergence capabilities of DHA and GQA. Then, we set a budget of 50B tokens to further assess the performance of DHA over a more extended training period.

Evaluation. We employed the lm-evaluation-harness [21] to evaluate our models. For common sense and reading comprehension tasks, we report 0-shot accuracy results for SciQ [22], PIQA [23], WinoGrande (Wino.) [24], ARC Easy(ARC-E.) [25], and HellaSwag (HellaS.) [26], alongside 25-shot accuracy for ARC Challenge (ARC-C.) [27]. In the assessments of continued QA and text understanding, we report 0-shot accuracy for LogiQA [28], 32-shot BoolQ [29], and 0-shot LAMBADA [30]. All reported results were calculated with the mean and stderr of multiple experiments.

Instruction tuning evaluation. To assess our models' capabilities after instruct tuning [31, 32], we fine-tune both DHA and baseline models on 10,000 instruction-response pairs from the ShareGPT dataset 5 and evaluate on another 1,000 instructions, using GPT-4 for response evaluator [33]. The win rate of our model relative to the baseline is reported. For detailed information, refer to Appendix C.1.

Baselines. We selected the LLaMA2-7B model and Sheared-LLaMA-2.7B&1.3B (S.-LLaMA2.7B&1.3B) as the MHA baselines. For each scaled model's checkpoint, we constructed 25% and 50% compressed GQA and DHA models in 0.5B & 1B tokens (0.025% & 0.05% of pretrain budget 6 ).

5.2 Experimental Results

Foundational Capabilities. Tab. 1 shows the foundational capabilities of DHA and GQA models at 50% and 25% compression rates (e.g., 64 key value heads compress to 16) across different scales. DHA was obtained by transforming LLaMA using adaptive head fusion and then further pre-trained with 1B tokens. For comparison, we constructed GQA with the same compression rates and training budget. Experiments show that DHA can achieve efficient architecture with only 0.05% of the original model's pre-training cost without significant performance loss. Compared to GQA, DHA consistently achieved better performance across all model scales and pre-training cost settings. Under the same checkpoint and training budget settings, DHA demonstrates significant improvements at higher compression rates. For example, with LLaMA7B at a 25% compression rate, DHA achieved a 5% relative performance improvement over GQA. This showcases DHA's fusion algorithm's ability to efficiently retain knowledge at high compression rates and the advantage of DHA's decoupled architecture in adaptively compressing redundant components. Possibly due to the lack of relevant data, DHA performed on par with LogiQA. As shown in Fig. 6, DHA's performance advantage

0.02

0.05

0.10

0.15

0.20

Fusion Training Budget (B)

2

4

6

LM Loss

GQA-7B-25%

DHA-7B-25%

0.2

0.5

1.0

1.5

2.0

Continued Pretraining Budget (B)

40

50

60

70

Average Acc (%)

DHA-7B-25%

GQA-7B-25%

1

3

5

7

9

11 13 15 17 19 21 23 25 27 29 31

Layer

0

5

10

15

# Heads

Key heads

Value heads

0.1

0.2

0.3

0.4

0.5

Continued Pretraining Budget (B)

2

3

4

LM Loss

GQA

DHA-40ba

DHA-240ba

DHA-400ba

1,0

0.9

0.9

0.8

0.8

0.7

30

1.0

0.9

0.9

0.9

0.9

0.9

7

becomes more remarkable with reduced training budgets. It indicates that DHA effectively retains knowledge of larger models, significantly reducing pre-training costs.

Better Initialization. We examined whether DHA offers a better initialization point than GQA by pre-training both DHA and GQA models on the original RedPajama dataset. Fig. 5 shows that the initial loss of the GQA model is high and decreases slowly. In contrast, the DHA model starting from MHAexhibits a minor increase in LM loss as fusion progresses, maintaining a consistently lower loss. DHA converges with just 0.1B data, demonstrating a 5 × training speedup compared to GQA. Fig. 6 reports the average downstream task accuracy of DHA and GQA during continued pre-training. DHA achieves comparable performance to GQA's at 1B tokens with only 0.2B tokens, outperforming GQA's 0.2B token performance by 13.93%. This demonstrates DHA's effectiveness in retaining parameter information. Ultimately, DHA achieves a higher performance ceiling than GQA due to retaining information from the original model and its more efficient architecture, whereas GQA loses information during initialization.

5.3 Analysis

Ablation Study. Wereport the effects of ablating Linear Heads Fusion and Adaptive Transformation in Tab. 2. When training with less data (0.5B), ablating Linear Heads Fusion leads to significant performance degradation, indicating that this method preserves crucial knowledge in LLMs, greatly accelerates DHA model training, and enhances performance. Adaptive Transformation allocates parameters more efficiently during construction, thereby strengthening the model's capability and reducing training difficulty. When we allocate more training budget to 1B, DHA's efficient architecture after Adaptive Transformation plays a more significant role, enhancing the model's performance ceiling. When continuing to pre-train the DHA model, it demonstrated strong learning capabilities and sustained performance improvements, ultimately achieving 96.1% of the performance with just 2.5% of the original model's pre-training budget, while saving 75% of the KV cache.

Training Budget Allocation. Allocating more computation to the fusion phase aids in better retention of information within the checkpoint. Our experiments assessed the effects of budget allocations between the fusion and CT phases within a fixed budget of 2 billion tokens . Tab. 3 shows that increasing the fusion budget consistently from 0.05B to 0.2B improves model performance at the initialization point. Training with just 0.1B data is sufficient to achieve a good starting point, and increasing fusion budget will not affect the final performance. This experiment also

demonstrates the necessity and effectiveness of the fusion stage under low-resource conditions. When we have a larger training budget, we can allocate more resources to the fusion stage to achieve a better initialization point for DHA.

Heads Budget Allocation. We investigated how the model adaptively allocates decoupled head group sizes across different layers under global head budgets. As illustrated in Fig. 8, the head numbers of DHA layers decrease from higher to lower across layers. Deeper layers exhibit higher compression rates due to greater redundancy. However, the initial and crucial layers need more heads, suggesting they may have specialized functions. As shown in Fig. 7, we presented the LM loss for the cold-start training of DHA models initialized with parameter averaging under different DHA configurations obtained at various search steps. Despite using the same initialization method as GQA, DHA exhibits a faster loss decline and a lower final loss. This indicates that DHA's architecture can accelerate training and achieve better performance, even without Linear Heads Fusion method.

Parameter Characteristics in DHA. For interpretability analysis, we visualized the parameter characteristics of the post-fusion DHA model in Fig. 9 (detials in Appendix E.2), and compared them with those prior to fusion. The DHA parameter distribution shows consistency with MHA's. This indicates that DHA effectively aggregates multiple similar functional heads within clusters and new fused heads successfully reconstruct the functionalities of multiple origin heads in MHA. It is noteworthy that the significant reduction in the number of similar heads within the DHA architecture indicates that our method effectively reduces redundancy among the heads.

6 Related Work

Advanced Multi-Head Attention. Some efforts have been converting the traditional Multi-Head Attention (MHA) [17] to Multi-Query Attention (MQA) [4], Group-Query Attention (GQA) [5] or GQKVA [6]. These methods achieve a balance between performance and efficiency by reducing the number of head parameters through parameter reuse across grouped heads. DHA is inspired by these methods and has a much higher optimization rate and much less training overhead.

Efficient Pre-training Approaches. In recent years, the ability of incremental training to accelerate large-scale model training by studying how to obtain the optimal initialization point for training has thus attracted much attention [34, 35]. Net2Net [36] uses function-holding transformations to expand the width by duplicating neurons, and uses a unitary layer implementation to expand the depth. LiGO [37] proposes a learnable expansion method that can be used at the initial initialization point of a transformer. DHA is inspired by these methods, but we investigate how to learn to map the parameter matrix from large to small without losing the ability of the larger model itself. For additional related work, please refer to Appendix A.

7 Conclusion

In this paper, we propose an efficient attention architecture and a method for fast converting an MHAcheckpoint into an efficient structure. By grouping similar heads and performing controlled linear fusion, we develop an initial DHA architecture that decouples head components at various layers, reducing training overhead while maintaining performance. Experimental results show that our method preserves the knowledge of the original model, improving training acceleration, inference efficiency, and computational cost savings. This transformation paradigm offers research value and potential for broader application with minimal performance loss and reduced computational effort.

Acknowledgments

We would like to thank Yinqi Yang, Jiawei Sheng, Xinhua Zhang, Shicheng Wang, Chuanyu Tang and members of the IIE KDsec group for their valuable feedback and discussions. We are very grateful to Mengzhou Xia for providing the concise and effective ShearingLLaMA experimental code and for her assistance during the reproduction process. Work done during Yilong Chen's internship in Baidu Inc. This research is supported by the National Key Research and Development Program of China (grant No.2021YFB3100600) and the Youth Innovation Promotion Association of CAS (Grant No. 2021153).

References

Anthropic. Introducing claude. 2023.

OpenAI. Gpt-4 technical report. ArXiv , page abs/2303.08774, 2023.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open Foundation and Fine-Tuned Chat Models, July 2023.

Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150 , 2019.

Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245 , 2023.

Farnoosh Javadi, Walid Ahmed, Habib Hajimolahoseini, Foozhan Ataiefard, Mohammad Hassanpour, Saina Asani, Austin Wen, Omar Mohamed Awad, Kangling Liu, and Yang Liu. Gqkva: Efficient pre-training of transformers by grouping queries, keys, and values, 2023.

Saurabh Agarwal, Bilge Acun, Basil Homer, Mostafa Elhoushi, Yejin Lee, Shivaram Venkataraman, Dimitris Papailiopoulos, and Carole-Jean Wu. Chai: Clustered head attention for efficient llm inference, 2024.

Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H _ 2 o: Heavy-hitter oracle for efficient generative inference of large language models. arXiv preprint arXiv:2306.14048 , 2023.

Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time. arXiv preprint arXiv:2305.17118 , 2023.

Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for llms. arXiv preprint arXiv:2310.01801 , 2023.

Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453 , 2023.

Yujia Qin, Cheng Qian, Jing Yi, Weize Chen, Yankai Lin, Xu Han, Zhiyuan Liu, Maosong Sun, and Jie Zhou. Exploring Mode Connectivity for Pre-trained Language Models, October 2022.

Kaiyan Zhang, Ning Ding, Biqing Qi, Xuekai Zhu, Xinwei Long, and Bowen Zhou. CRaSh: Clustering, Removing, and Sharing Enhance Fine-tuning without Full Large Language Model, October 2023.

Alexander M. Rush, David Sontag, Michael Collins, and Tommi Jaakkola. On dual decomposition and linear programming relaxations for natural language processing. In Hang Li and Lluís Màrquez, editors, Proceedings of the 2010 Conference on Empirical Methods in Natural Language Processing , pages 1-11, Cambridge, MA, October 2010. Association for Computational Linguistics.

Ziheng Wang, Jeremy Wohlwend, and Tao Lei. Structured Pruning of Large Language Models. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP) , 2020.

Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning, October 2023.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited, 2019.

TogetherAI. Redpajama: An open source recipe to reproduce llama training dataset, 2023.

The Mosaic ML Team. composer. https://github.com/mosaicml/composer/ , 2021.

Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023.

Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. In Leon Derczynski, Wei Xu, Alan Ritter, and Tim Baldwin, editors, Proceedings of the 3rd Workshop on Noisy User-generated Text, NUT@EMNLP 2017, Copenhagen, Denmark, September 7, 2017 , pages 94-106. Association for Computational Linguistics, 2017.

Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about physical commonsense in natural language. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020 , pages 7432-7439. AAAI Press, 2020.

Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020 , pages 8732-8740. AAAI Press, 2020.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457 , 2018.

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Anna Korhonen, David R. Traum, and Lluís Màrquez, editors, Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers , pages 4791-4800. Association for Computational Linguistics, 2019.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the AI2 reasoning challenge. CoRR , abs/1803.05457, 2018.

Jian Liu, Leyang Cui, Hanmeng Liu, Dandan Huang, Yile Wang, and Yue Zhang. Logiqa: A challenge dataset for machine reading comprehension with logical reasoning. arXiv preprint arXiv:2007.08124 , 2020.

Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint arXiv:1905.10044 , 2019.

Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. arXiv preprint arXiv:1606.06031 , 2016.

Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems , 35, 2022.

Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-following llama model, 2023.

Yann Dubois, Xuechen Li, Rohan Taori, Tianyi Zhang, Ishaan Gulrajani, Jimmy Ba, Carlos Guestrin, Percy Liang, and Tatsunori B Hashimoto. Alpacafarm: A simulation framework for methods that learn from human feedback. arXiv preprint arXiv:2305.14387 , 2023.

Di Xie, Jiang Xiong, and Shiliang Pu. All you need is beyond a good init: Exploring better solution for training extremely deep convolutional neural networks with orthonormality and modulation. In 2017 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2017, Honolulu, HI, USA, July 21-26, 2017 , pages 5075-5084. IEEE Computer Society, 2017.

Lemeng Wu, Dilin Wang, and Qiang Liu. Splitting steepest descent for growing neural architectures. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d'AlchéBuc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada , pages 10655-10665, 2019.

Tianqi Chen, Ian Goodfellow, and Jonathon Shlens. Net2net: Accelerating learning via knowledge transfer. arXiv preprint arXiv:1511.05641 , 2015.

Peihao Wang, Rameswar Panda, Lucas Torroba Hennigen, Philip Greengard, Leonid Karlinsky, Rogerio Feris, David Daniel Cox, Zhangyang Wang, and Yoon Kim. Learning to Grow Pretrained Models for Efficient Transformer Training, March 2023.

Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arXiv preprint arXiv:2009.06732 , 2020.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 , 2019.

Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020 . OpenReview.net, 2020.

Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems , 2020.

Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150 , 2020.

Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G. Carbonell, Quoc V. Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. CoRR , abs/1901.02860, 2019.

Siyu Ding, Junyuan Shang, Shuohuan Wang, Yu Sun, Hao Tian, Hua Wu, and Haifeng Wang. Ernie-doc: A retrospective long-document modeling transformer. arXiv preprint arXiv:2012.15688 , 2020.

Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. Recurrent memory transformer. Advances in Neural Information Processing Systems , 35:11079-11091, 2022.

Alexis Chevalier, Alexander Wettig, Anirudh Ajith, and Danqi Chen. Adapting language models to compress contexts. arXiv preprint arXiv:2305.14788 , 2023.

Yilong Chen, Junyuan Shang, Zhengyu Zhang, Jiawei Sheng, Tingwen Liu, Shuohuan Wang, Yu Sun, Hua Wu, and Haifeng Wang. Mixture of hidden-dimensions transformer, 2024.

Matanel Oren, Michael Hassid, Yossi Adi, and Roy Schwartz. Transformers are multi-state rnns. arXiv preprint arXiv:2401.06104 , 2024.

Yilong Chen, Guoxia Wang, Junyuan Shang, Shiyao Cui, Zhenyu Zhang, Tingwen Liu, Shuohuan Wang, Yu Sun, Dianhai Yu, and Hua Wu. NACL: A general and effective KV cache eviction framework for LLM at inference time. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 7913-7926, Bangkok, Thailand, August 2024. Association for Computational Linguistics.

Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm.int8(): 8-bit matrix multiplication for transformers at scale, 2022.

Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers, 2023.

Robert M. Gray and David L. Neuhoff. Quantization. IEEE transactions on information theory , 44(6):2325-2383, 1998.

Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network, 2015.

Jianping Gou, Baosheng Yu, Stephen J. Maybank, and Dacheng Tao. Knowledge distillation: A survey. International Journal of Computer Vision , 129(6):1789-1819, March 2021.

Lianshang Cai, Linhao Zhang, Dehong Ma, Jun Fan, Daiting Shi, Yi Wu, Zhicong Cheng, Simiu Gu, and Dawei Yin. Pile: Pairwise iterative logits ensemble for multi-teacher labeled distillation. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing: Industry Track , pages 587-595, 2022.

Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, and Beidi Chen. Deja vu: Contextual sparsity for efficient LLMs at inference time. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 22137-22176. PMLR, 23-29 Jul 2023.

Elias Frantar and Dan Alistarh. SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot, March 2023.

Minjia Zhang and Yuxiong He. Accelerating training of transformer-based language models with progressive layer dropping. Advances in Neural Information Processing Systems , 33:1401114023, 2020.

Hassan Sajjad, Fahim Dalvi, Nadir Durrani, and Preslav Nakov. On the effect of dropping layers of pre-trained transformer models. Computer Speech & Language , 77:101429, 2023.

Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems , 35:16344-16359, 2022.

[2401.02415] LLaMA Pro: Progressive LLaMA with Block Expansion.

Yilong Chen, Junyuan Shang, Zhenyu Zhang, Shiyao Cui, Tingwen Liu, Shuohuan Wang, Yu Sun, and Hua Wu. LEMON: Reviving stronger and smaller LMs from larger LMs with linear parameter fusion. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 8005-8019, Bangkok, Thailand, August 2024. Association for Computational Linguistics.

Appendix

A Extended Related Works, Discussions, and Limitations

A.1 Extended Related Works

Efficient Transformers. Efficient Transformers[38] have been extensively explored [39, 40, 41, 42, 43, 44, 45, 46, 47] to address the self-attention operation which scales quadratically with the sequence length. For instance, Sparse Transformer [39] uses a dilated sliding window the reduces the attention complexity. Longformer [42] and Bigbird [41] reduced the complexity of self-attention by combining random, window and global attention. Recurrence Transformers [43] maintain a memory bank of past KV cache to process the long text in segments. However, the above methods either result in a loss of model performance or require retraining the model, which is unaffordable for the high computational resources of LLMs. DHA requires very little computation to transform checkpoints into an efficient architecture that balances performance and computational resources.

KV Cache Compression. KV Cache Compression methods emerged for reducing the prominent inference bottleneck caused by KV cache, particularly for long content input. A series of methods [8, 10, 9, 48, 49] explored the sparsity among Transformer's attention block, then evicted unnecessary tokens from KV Cache for efficient inference. However, these methods discard information from the context and use algorithms for inference that are inconsistent with the training phase, which can cause model performance degradation. DHA does not need to discard information from the context and is able to maintain consistent performance for training and inference. Pruning [15, 16], quantization [50, 51, 52] and distillation [53, 54, 55] can reduce the number of model key and value headers, parameter dimensions, and activation to reduce memory bandwidth overhead during model inference. Deja Vu [56] and CHAI [7] prune pruning redundant heads through clustering methods for efficient inference. In the LLM era, this leads to a significant reduction in neuron redundancy as models move from task-specific to generalized [57]. The application of these methods to LLMs is computationally expensive and leads to performance degradation at larger pruning magnitudes.

Model Compression. Our approach is dedicated to obtaining a high-performance lightweight language model, which is the same goal as the task of model compression. Quantization [52] reduces the numerical accuracy of model weights and activations, and speeds up training and inference, but results in a loss of model accuracy and the inability to freely build target-specific models. CRash [13] and LayerDrop [58, 59] methods discard ineffective layers during training, which do not allow for target-specific structuring and come with a large performance loss. Pruning [15] minimizes the impact on performance by cutting out redundant neurons that over-parameterize the model. In the LLM era, this leads to a significant reduction in neuron redundancy as models move from task-specific to generalized [57]. Pruning LLM leads to performance degradation at larger pruning magnitudes. LLMsheairng [16] uses the results of pruning as initialization for continuous pre-training of the model to recover performance, but this approach requires more data and computational overhead. We avoid the information loss caused, by learning the parameter fusion matrix of the model to reach a specific structure, thus obtaining better initialization points and reducing the overhead of continuous pre-training.

A.2 Broader Impact and Limitations

A.2.1 Broader Impact

In this paper, we observe the MHA head mechanism and report the phenomenon of modular clustering of heads in MHA. This paper innovatively proposes linearly fusible parameters within the model, and designs linear fusion operators and related experiments to verify the low-loss fusible nature of the parameters. This helps to advance parameter fusion theory and LLM interpretability studies, which provide a foundation and inspiration for future algorithmic advancements, encouraging further optimization and innovation in LLMs. Our work on Decoupled-Head Attention (DHA) represents an advancement in optimizing the efficiency of Large Language Models (LLMs). By addressing the

substantial computational and memory costs associated with the widely used Multi-Head Attention (MHA), DHA enhances the applicability of LLMs in various domains. The introduction of DHA not only achieves a remarkable balance between performance and efficiency but also significantly reduces the need for extensive pre-training, making the deployment of LLMs more feasible and cost-effective. This efficiency allows for the broader accessibility of advanced LLMs, democratizing technology and fostering innovation across industries. Furthermore, by requiring only 0.25% of the original model's pre-training budgets to achieve near-original performance while saving 75% of KV cache, DHA contributes to significant energy savings, aligning with sustainable and environmentally friendly AI practices. The enhanced performance and reduced training costs accelerate the development of AI applications, enhancing productivity in fields such as natural language processing, healthcare, and finance.

A.2.2 Limitation

There are two limitations to our current approach. Firstly, we have only utilized linear methods for parameter fusion in our model. Future research should explore nonlinear methods, as they may offer a better way to link different parameters and achieve optimal results. Secondly, due to computational resource constraints, we have only experimented with models of 7 billion, 3 billion, and 1.3 billion parameters. However, our method is scalable and can be extended to models of any size in future work.

A.2.3 Ethical Consideration

In our study, we utilize publicly available data and techniques to address privacy concerns. Our approach focuses on improving model parameter efficiency and reducing model size to develop robust, compact, and accessible models, thus promoting the open dissemination and democratization of NLP technologies. By implementing pre-training strategies, we aim to mitigate biases through comprehensive training on large datasets, contributing to ethical AI development that prioritizes transparency, efficiency, and bias reduction. Our work is dedicated to advancing accessible and efficient NLP technologies, fostering a more inclusive and automated future for AI.

B More Implementation Details

B.1 Head Similarity and MHA Redundancy

Centered Kernel Alignment (CKA) is a statistical measure used to quantify the similarity between two sets of data representations. Unlike traditional correlation measures, CKA is designed to be invariant to orthogonal transformations and scaling of the data.

To calculate the similarity between two sets of representations using CKA, we employ a kernel function to map the original data into a higher-dimensional space, where the alignment of their central tendencies can be more easily measured. The CKA value ranges from 0 to 1, where 0 indicates no similarity and 1 indicates identical representations.

The mathematical formulation of CKA, when using a linear kernel, is given by the following equation:

Here, X and Y are matrices whose columns are the vectors of the representations to be compared, ∥ · ∥ F denotes the Frobenius norm, and X T and Y T are the transposes of X and Y , respectively. To mathematically define the redundancy of each layer based on the average similarity between heads, we follow these steps:

Compute the similarity between heads: For each pair of heads within a given layer, calculate the similarity using the CKA formula.

Compute the average similarity: Average the similarity scores of all pairs of heads to define the redundancy of the layer.

B.1.1 Compute Similarity Between Heads

Consider a layer with H heads, where the parameters of each head are represented by the matrices W i (e.g., W q 1 , W q 2 , . . . , W qH for query weights). For each pair of heads i and j , compute the CKA similarity using the following formula:

B.1.2 Compute Redundancy

Calculate the similarity for all pairs of heads and then compute the average similarity:

The coefficient 2 H ( H -1) ensures that the average similarity is computed over all pairs of heads. This redundancy measure reflects the degree of similarity between the parameters of different heads within each layer. A higher redundancy indicates that the parameters of different heads are more similar, implying a higher level of redundancy.

B.2 Implementation in LLaMA2 Model

Our method can theoretically transform MHA architecture in any transformer model to efficient DHA architecture. Using LLaMA models as case studies, we implemented DHA transformation with various compression rates on all MHA layers. Only a very small number of additional parameters need to be introduced, DHA significantly accelerates training and improves performance.

DHA adaptively gives search heads and heads connectivity relationship with redundancy in each MHAlayer. Thus DHA assigns different group sizes at different layers and aggregates similar heads into one group to speed up fusion and reduce knowledge loss due to noise in fusion. As Shown in Fig 4, the transformation process of MHA to DHA can be divided into three stages.

In order to keep the performance of the DHA model at the fusion start consistent with the MHA model, we initialize the operators of the DHA model to the MHA model with the corresponding scaling factors of query-key, query-value set to 1, and the corresponding scaling factors within the rest of the groups set to 0. At the beginning of every fusion process (e.g. 2 × , 4 × , 8 × ), the algorithm first performs multiple STEPs constrained only by the L lm constraints to propagation, computing the L fusion but not optimizing the linear fusion operator based on it. Based on the L fusion between head and head as a measure of the distance between head and head we perform head clustering with the goal of minimizing the average loss of heads within each group and maximizing the average loss of heads between groups and groups. Afterwards, we select multiple groups with the smallest loss based on the compression rate as the fusion target, and optimize their L fusion for back propagation. This algorithm ensures that the most redundant components of the model are fused and compressed during each transformation, while components requiring more parameters retain their original properties.

Our approach is theoretically applicable to transforming parameters across various transformer model designs, focusing on preserving the knowledge within MHA parameters.

Using LLaMA models as a case study, we implement our DHA transformation on all MHA layer. The whole transformation process can be divided into two phases: the Fusion phase with a small training budget and the recovery phase with continuous pre-training. Before Fusion phase, we define the total number of compressed headers budget C then C is split into compression rates at different compression levels. During Fusion phase, we modified the forward propagation path of MHA in the form of DHA refer to Eq. 5 and optimize L fusion refer to Eq. 7. At the beginning, the fusion operators of each layer will be initialized making the DHA and the original MHA functionally equivalent. As we antagonistically optimize the fusion operator and upadte Lagrangian multipliers λ , the L fusion that marks this DHA fusion process decreases.

When L fusion is less than 1e-3 we terminate the fusion algorithm and enter the post-processing phase. The fusion weights within each group are computed by averaging the weights corresponding to each query-key and query-value within the group.We construct new DHA heads' parameters from the original MHA heads based on the fusion operator. After that, the fused model parameters can recover the performance and complete the transformation with a small amount of restorative pre-training.

We implemented the DHA algorithm with different compression ratios on models of different sizes. Experiments show that the DHA algorithm is adapted to models of various sizes. Only a very small number of additional parameters need to be introduced, and DHA preserves parameter knowledge in the model and improves performance.

B.2.1 Attention Module Initialization

In the module initialization process, the input key and value tensors are first reshaped and grouped according to the number of key and value heads, respectively. Given the batch size (bsz), number of heads (num_heads), key length (k_len), and head dimension (head_dim), the key tensor is reshaped into keys_grouped of shape [bsz, num_key_heads, num_heads // num_key_heads, k_len, head_dim]. Similarly, the value tensor is reshaped into values_grouped of shape [bsz, num_value_heads, num_heads // num_value_heads, k_len, head_dim]. These grouped tensors are then expanded by repeating them along the group size dimension, resulting in keys_expanded and values_expanded. Correspondingly, the weight tensors weights_k and weights_v are reshaped to match the expanded dimensions and are then multiplied element-wise with the expanded key and value tensors.

B.2.2 Attention Forward Pass

During the forward pass, the reshaping and expansion of the key and value tensors are performed in a similar manner as in the initialization process but with parameters specific to the DHA fusion phase. The key tensor is reshaped into keys_grouped of shape [bsz, dha_warmup_group_num, num_heads // dha_warmup_group_num, k_len, head_dim] and the value tensor into values_grouped of shape [bsz, dha_warmup_group_num, num_heads // dha_warmup_group_num, k_len, head_dim]. These

grouped tensors are then expanded by repeating them according to the dha_warmup_group_size. The weights weights_k and weights_v are reshaped and expanded to align with the dimensions of the expanded key and value tensors. Element-wise multiplication is performed between the expanded tensors and their corresponding weights, and the resulting weighted tensors are summed along the appropriate dimension.

B.2.3 DHA Loss Calculation

The calculation of the loss function in this model involves the adaptive DHA loss. This loss is computed based on the global step, warmup steps, and a base value. The DHA margin is calculated as the product of an exponential decay term and a linear decay term, ensuring it is non-negative. The adaptive DHA loss is derived by comparing the mean squared error (MSE) with the DHA margin and summing the positive differences.

Formally, the DHA margin M dha is calculated as:

MSE Loss are defined in Eq. 7. The adaptive DHA loss L dha is then:

The overall loss L is the adaptive DHA loss:

This combined loss function effectively utilizes the adaptive component to optimize the attention mechanism in the model. The calculation process ensures that the model adapts dynamically during training, reducing the loss progressively as the training steps increase.

B.3 Head Grouping Based on Fusion Loss

This algorithm uses simulated annealing to optimize group scores based on a given score matrix. It begins by defining the number of groups and distributing the points among them randomly. The initial score for these groups is calculated using the 'calculate_score' function, which sums the scores from the matrix for each group, considering each connection twice and dividing by two.

The algorithm starts with a high temperature (T=100) and gradually cools down (T_min=0.001) using a cooling rate (alpha=0.9). During each iteration, two random points from different groups are swapped, creating a new grouping. The score for this new grouping is calculated, and the difference in score (delta) is evaluated.

If the new score is higher, or if a randomly generated number is less than the exponential of delta divided by the temperature, the new grouping is accepted. This allows the algorithm to escape local optima. The temperature is then reduced according to the cooling rate. This process continues until the temperature reaches the minimum threshold. The algorithm returns the final group configuration and its corresponding score, which represents an optimized grouping based on the initial score matrix.

In practice, we use the MSE computed by the head and the head as scores, and compute the matrix of scores between the head and the head for head clustering after forward.

̸

B.4 Layer Allocation Based on Fusion Loss

This algorithm efficiently allocates resources to different layers based on their respective losses to optimize system performance. Initially, it assigns a minimum allocation to each layer. Then, it calculates weights for each layer based on their losses, prioritizing layers with higher losses. The algorithm determines the number of times 16 can be allocated based on the remaining allocation. It allocates 16s to layers with the highest weights until reaching a predetermined limit. Next, it redistributes the remaining allocation to layers with the highest loss-to-allocation ratios, assigning resources in multiples of 8 or 4. This process ensures that layers with higher losses receive more resources, optimizing the overall system performance. Finally, the algorithm returns the final allocation for each layer, resulting in an efficient distribution of resources across the system. The total search process for the LLaMA2 model requires 42 minutes.

B.5 Training Details

The hyperparameters used in our experiments are presented in Tab. 4. We employ fully sharded data parallel to efficiently train our models in parallel, and we utilize FlashAttention V1 [60] to accelerate the training process. A cosine learning rate scheduler is used, with the learning rate decaying to a minimum of 10% of the peak value. Preliminary experiments were conducted to determine the optimal peak learning rate for learning the fusion variables and Lagrange multipliers.

C Extended Experiments

C.1 Instruction Tuning Evaluation.

Instruction Tuning Evaluation. To assess our models' capabilities in downstream application after instruct tuning [31, 32], we fine-tune both DHA and the baseline models on 10,000 instructionresponse pairs drawn from the initial round of multi-turn chat histories in the ShareGPT dataset 7 . For evaluation, we select another 1,000 instructions from ShareGPT, generate responses using our fine-tuned models and other baseline models and employ GPT-4 as an evaluator to compare these responses [33]. We report the win rate of our model relative to the baseline model.

Instruction Tuning. As shown in Fig. 10, the tuned DHA model outperforms all GQA baselines of comparable scale . This demonstrates that the DHA model effectively retains the foundational capabilities of the MHA model and can be activated through instruction tuning to produce long, coherent, and informative responses.

0

20

40

60

80

100

84.25%

15.75%

DHA-2.7B-25% vs. GQA-2.7B-25%

0

20

40

60

80

100

72.75%

27.25%

DHA-1.3B-25% vs. GQA-1.3B-25%

Figure 10: In model scale of 7B, 3B, and 1.3B, DHA significantly outperforms GQA and achieves comparable performance with MHA after instruction tuning .

Combination with KV Cache Compression Techniques. In Sections 2 and 4, we demonstrated that DHA is a more efficient GQA architecture, so it has similarly good compatibility. We tested the compatibility of the DHA model with the KVCache eviction method NACL [49]. NACL 25% indicates retaining only 25% of the KVCache. The experiment results are shown in the Tab. 5. DHA and GQA exhibit equally good compatibility with KV cache compression techniques.

Compare with Advance GQA Initialization. It's a common and effective approach to convert MHAto GQA using mean pooling instead of training from scratch. The author of GQA tested several methods for the initialization of GQA and found it works best using simple mean pooling from MHA. Indeed, training GQA from scratch will cost trillions tokens budget to match the performance of MHAwhich is inefficient and costly.Inspired by the similarity of head parameters, we improved the initialization method of GQA: instead of direct grouping, we first cluster similar heads using CKA

and then perform mean-pooling initialization within each cluster. We compare this approach with the Vanilla GQA and DHA.

Tab. 6 shows that GQA(CKA-Grouping)-7B-25% (5B) achieved comparable performance to the original implementation in Vanilla GQA. We believe the reason for this is that the head grouping learned by DHA is based on the fusible nature between heads, which cannot be completely equated with CKA similarity. More importantly, DHA not only groups heads based on similarity but also learns the fusible parameters. This allows it to eliminate the influence of redundant parameters and retain more important information during the initialization process, which is not possible with mean initialization.

D Extend Analysis

How Merging Weights Change. Refer to Fig. 3a, where we show the weight variation diagram. In the fusion process of heads 0-3, head 0 initially constitutes 100% as the starting head of the MHA. As the fusion process progresses, the parameters of the important heads increase, and the proportions of all heads become more balanced. This indicates that the algorithm attempts to retain information from different heads by balancing the parameter proportions of each head. This process results in a slight increase in loss, but not significantly.

DHA's Compatibility on GQA Model. DHA is primarily designed for models based on the Transformer Decoder architecture and can be adapted to all models with this architecture. We chose LLaMA [61] as the experimental baseline because it is a classic model using the decoder architecture in LLMs. Other open-source LLM models differ from LLaMA only in certain details (such as activation functions and training methods), which do not affect DHA's training. Successfully applying DHA to LLaMA indicates that it can be used in most decoder-only models. GQA [5] is an efficient variant of MHA, which optimizes the inference process through head grouping and sharing. Due to its simplicity and efficiency, GQA is widely used. DHA can be similarly constructed based on GQA, requiring only minor adjustments to the construction process. Here, we provide two feasible methods to convert GQA to DHA.

Easiest method in less than 1 minute. GQA can be losslessly converted into MHA by simply replicating the GQA' KV heads. Then, we can perform the DHA transformation on the MHAarchitecture.

Minor modification by grouping KV. DHA only needs to group and fuse the Key and Value heads. When constructing DHA on GQA, we initially group the Key and Value, maintaining alignment with GQA functionality. During the training phase, the fused head parameters can replace the original GQA heads for sharing.

Inter-layer Grouping of Heads or Only Intra-layer Grouping? Only intra-layer grouping and fusion is conducted in DHA. Fig. 1 meant to illustrate the decoupled-heads where the number of key

and value heads can be different among layers. The DHA method employs parameter fusion within each layer for three reasons:

Higher redundancy of heads within layer for fusion. The heads within a layer exhibit high similarity and redundancy, which provides a good starting point for parameter fusion.

More complex optimization for inter-layer fusion. The optimization process between layers is very complex and requires memory operations for cross-layer calls, which inherently increases the inference cost.

Promising future work by introducing inter-layer fusion [62]. This paper represents an early exploration of applying parameter fusion methods within model parameters. The inter-layer fusion approach is indeed a valuable direction for future exploration.

Accuracy Loss after Transformation. The performance gap between the results shown in the paper and MHA is primarily due to the following two reasons:

The gap of pre-training data. The MHA model was not trained on the same data used for DHA. Since LLaMA's training data is not directly open-sourced, we used an experimental open-sourced pre-training data following Sheared-LLaMA (Xia et al., 2024). The improved pre-training data will close the gap between DHA and MHA.

Parameter size difference. Compared to MHA, DHA compresses 50% or 25% of attention heads, requires only 0.05% of pre-training data and achieves approximately 5% loss. The number of parameters of MHA is much larger than that of DHA, so performance loss is inevitable during conversion. Compared with GQA, a strong baseline with the same number of parameters, DHA has shown higher training efficiency and performance advantages. Due to the high efficiency of DHA, DHA can use more heads than MHA with the same number of parameters, and has the opportunity to achieve better performance.

E Extend Observation

E.1 Header parameter characteristics in MHA

We show more of our head similarity observations in the LLaMA2-7b model MHA. Each subfigure represents the similarity between heads within the same layer for three different types of attention mechanisms: WQ (query), WK (key), and WV (value). The matrices are arranged in a 3x4 grid layout, with each row corresponding to a specific layer and each column corresponding to a type of attention mechanism. Note: Layer numbers start from 1.

Layer 1 WQ Head similarity

Layer 1 WK Head similarity

0 5 9 6 21518112220

Layer 1 WV Head similarity

Layer 5 WQ Head similarity

888688

Layer 5 WK Head similarity

.

Layer 5 WV Head similrity

1.00

0.95

0.90

0.85

0.80

0.75

0369 2151 1242230

0.4

0 3 6 9 12 15 18 21242730

0 3 6 9 12 15 18 21 24 27 30

0.4

0 3 6 9 12151821242730

Layer 2 WQ Head similarity

Layer 2 WK Head similarity

1.0

Layer 2 WV Head similarity

Layer 6 WQ Head similarity

0  9 6215181122 0

Layer 6 WK Head similarity

2

Layer 6 WV Head similarity

1.000

0.975

0.950

0.925

0.900

0.875

0.2

0.2

0.850

0369251822420

0369215182242 3

0369215 182 2220

0 3 6 9 12 15 18 21242730

0.5

9 12 15 18 21 24 27 30

9 12 15 18 21 24 27 30

0.825

Layer 3 WQ Head similarity

-0.8

Layer 3 WK Head similarity

1.0

Layer 3 WV Head similarity

1.00

0.95

Layer 7 WQ Head similarity

Layer 7 WK Head similarity

Layer 7 WV Head similarity

0.95

1.00

0.6

0.90

0.85

0.90

.4

0.80

0.85

0.75

0.80

692151382242230

0.2

0.2

亿忆亿69

0.70

0 36 9 12151821242730

0 3 6 9 121518 21242730

0 3 6 9 12151821242730

0.75

Layer 4 WQ Head similarity

88.....

Layer 4 WK Head similarity

Layer 4 wV Head similarity

1.00

Layer 8 WQ Head similarity

98868

Layer 8 WK Head similarity

1.0

Layer 8 WV Head similarity

1.00

0.90

0.95

0.90

0.80

0.75

0.85

0.70

0.80

M

6の

2 15 18 2 2422 30

2 15 1 2 24 22 0

12 15 18 22 24 22 0

0 3 6 9 12 15 18 212427 30

0 3 6 9 1215 18 21242730

0 3 6 9 12 15 18212427 30

Layer 9 WQ Head similarity

0.9

Layer 9 WK Head similarity

0.9

0  9 6 215t 81 122206

Layer 9 WV Head similarity

0.96

0.98

1.00

Layer 13 WQ Head similarity

Layer 13 WK Head similarity

09

1.0

Layer 13 WV Head similarity

0.98

1.00

0.8

0.8

.94

0.92

0.8

0.96

0.94

0.7

-0.7

0.90

0.92

0 3 6 9 1215 18 21242730

0.6

0 36 9 12 15 18 2124 2730

0.6

0 3 6 9 1215 18 21242730

0.88

0.86

0 3 6 9 1215 18 21 242730

0.70

0.65

0 3 6 9 12 15 18 21242730

0.7

0.6

0 3 6 9 1215 18 212427 30

0.90

0.88

Layer 10 WQ Head similarity

1.00

Layer 10 WK Head similarity

0 6 9 6 21s181122 06

Layer 10 WV Head similarity

0.98

1.00

Layer 14 WQ Head similarity

Layer 14 WK Head similarity

Layer 14 WV Head similarity

S260

1.000

0.96

昌

0.950

0.94

0.85

0.925

0.92

0.90

0.80

0.900

S480

0.70

880

0.70

0.850

0 3 6 9 12 15 18 21 242730

0 3 6 9 12 15 18 2124 27 30

0.65

0 3 6 9 1215 18 212427 30

0.86

0 3 6 9 12151821242730

0.70

03 6 9 12 1518 21242730

0 3 6 9 12151821242730

0.825

Layer 11 WQ Head similarity

1.00

Layer 11 WK Head similarity

0.9

1.0

0 6 9 6 21 5811220

Layer 11 WV Head similarity

1.00

0.96

Layer 15 WQ Head similarity

100

0.95

0.90

Layer 15 WK Head similarity

1.00

Layer 15WVHead similrity

1.00

0.98

0.96

036

91215 18 21242730

F.80

0.75

0.70

0 3 6 9 12 15 1821242730

0.8

0.7

0 3 6 9 1215 18 212427 30

0.94

0.90

0.88

0 3 6 9 1215 18 21242730

0.80

0.70

0.65

0 3 6 9 12151821242730

0.70

0.65

0 3 6 9 1215 18 212427 30

0.94

0.92

0.90

0.88

Layer 12 WQ Head similarity

Layer 12 WK Head similarity

1.0

Layer 12 WV Head similarity

1.00

Layer 16 WQ Head similarity

1.00

Layer 16 WK Head similarity

Layer 16 WV Head similarity

1.00

0.9

0.9

S60

0.95

0.98

0.96

0.8

0.8

0.90

0.94

Z60

0.7

-0.85

0.90

0.6

0.6

0.80

0.65

0.70

0.65

0.88

0.86

0 3 6 9 1215 18 21242730

0 3 6 9 12 15 1821242730

0 3 6 9 1215 1821242730

0 3 6 9 1215 18 21 242730

0 3 6 9 1215 18 21242730

0 3 6 9 1215 1821242730

08s

Layer 17 WQ Head similarity

1.0

Layer 17 WK Head similarity

1.0

Layer 17 WV Head similarity

1.000

Layer 21 WQ Head similarity

1.00

Layer 21 WK Head similarity

Layer 21 WV Head similarity

1.00

.9

0.8

0.9

0.975

0.950

0.925

0.900

0.95

0.85

0.90

0.80

m

0.95

0.90

0.85

0.7

0.7

0.875

0.75

0 3 6 9 12151821242730

0.6

022512690

0.6

022690

0.850

0.70

品

0.65

0 3 6 9 1215 1821242730

0.80

Layer 18 WQ Head similarity

Layer 18 WK Head similarity

0.9

1.0

Layer 18 WV Head similarity

0.95

1.00

Layer 22 WQ Head similarity

0.95

0.90

1.00

。

Layer 22 WK Head similarity

0.9

1.0

Layer 22 WV Head similarity

0.975

1.000

0.950

0.8

0.90

0.85

0.8

0.925

0.85

0.80

0.900

0.7

0.80

0.75

0.7

0.875

.65

9°0

0.75

0.70

0.6

0.850

9 12 15 18 21 24 27 30

9 12 15 18 21 24 27 30

0 3 6 9 1215 18 212427 30

16 9 12 15 18 21 24 27 30

0 3 6 9 12 15 18 2124 2730

0 3 6 9 12 15 18 212427 30

Layer 19 WQ Head similarity

0 3 6 9 12 15 18 21 242730

三

0.90

0.95

0.80

0.85

0.75

0.70

Layer 19 WK Head similarity

0 3 6 9 12 1518 2124 2730

三

0.95

0.85

0.80

0.75

0.70

1.00

96

Layer 19 WV Head similarity

0 3 6 9 1215 18 212427 30

0.975

0.950

0.875

0.850

0.925

0.900

1.000

0.825

2

Layer 23 WQ Head similarity

0 3 6 9 1215 18 21 2427 30

0.95

0.90

0.85

0.80

0.75

0.70

1.00

0.65

m

Layer 23 WK Head similarity

0 3 6 9 12 15 18 2124 2730

0.80

1.00

0.65

0 3 6 9 1215 18 212427 30

Layer 23 WV Head similarity

0.975

0.950

0.925

1.000

S480

0.850

0.900

0

Layer 20 WQ Head similarity

0

0.85

0.70

Layer 20 WK Head similarity

F0.80

0.75

0.95

0.85

0.70

1.00

.90

。

Layer 20 WV Head similarity

0.975

1.000

0.950

0.925

0.900

0.875

0.850

Layer 24 WQ Head similarity

0.95

0.75

0.70

0.90

0.80

1.00

Q.85

Layer 24 WK Head similarity

0.95

0.75

0.70

0.85

0.80

1.00

0.65

Layer 24 WV Head similarity

0.95

0.90

0.85

1.00

08'0

9 12 15 18 21 2427 30

036

9 12 15 18 2124 27 30

0 3 6 9 1215 18212427 30

03

69 1215 18 21 2427 30

0 3 6 9 12 151821242730

9 12 15 18 212427 30

。

Layer 25 WQ Head similarity

1.00

Layer 25 WK Head similarity

0.9

Layer 25 WV Head similrity

0.95

1.00

Layer 29 WQ Head similarity

1.00o

0.95

0.90

96

m

Layer 29 WK Head similarity

Layer 29 WV Head similarity

1.00

0.95

0.8

0.90

0.85

2

0.90

0.7

0.85

0.80

0.75

8 17

0.85

0.70

0.6

0.80

0.70

0.80

0.65

0.75

0.65

0.75

0 3 6 9 12 15 18 2124 27 30

03 6 9 12 15 18 21242730

0 3 6 9 1215 18 21 2427 30

0 3 6 9 1215 18 212427 30

03 6 9 12 15 18 212427 30

0 3 6 9 1215 18212427 30

S 81

2

Layer 26 WQ Head similarity

0.95

0.90

0.85

1.00

7 24 21 18 15 12

Layer 26 WK Head similarity

0.8

0°7

E

6

Layer 26 WV Head similarity

0.95

1.00

0.90

Layer 30 WQ Head similarity

H.

a.80

0.85

1.00

9 6 2151 81

。

Layer 30 WK Head similarity

0.9

1.0

Layer 30 WV Head similarity

0.90

0.95

0.85

1.00

0.80

0.7

0.80

0 3 6 9 12 15 18 212427 30

0.75

0.70

品

03 6 9 12 15 18 21242730

0.7

0.6

36

9 12 15 18 21 2427 30

0.80

0.85

0 3 6 9 1215 18 21 2427 30

0.65

2

品

036

9 12 15 18 21 24 27 30

03 6 9 1215 18212427 30

0.75

Layer 27 WQ Head similarity

1.0

Layer 27 WK Head similarity

1.0

Layer 27 WV Head similarity

1.00

Layer 31 WQ Head similarity

1.00

9 6 21 51 81

。

Layer 31 WK Head similarity

1.00

。

Layer 31 WV Head similarity

1.00

0.9

0.9

0.95

0.95

0.8

0.8

2

0.90

0.85

0.90

0.7

0.7

0.85

0.75

0.85

a70

0.65

0 3 6 9 12 15 18 2124 27 30

0.6

0 3 6 9 12 1518 21 242730

36

9 12 15 18 21 24 27 30

0.80

036

9 12 15 18 21 24 27 30

0.65

0 3 6 9 12 15 18 21242730

0 3 6 9 12151821242730

0.80

Layer 28 WQ Head similarity

1.00

Layer 28 WK Head similarity

1.00

Layer 28 WV Head similarity

1.00

Layer 32 WQ Head similarity

Layer 32 WK Head similarity

Layer 32 WV Head similarity

1.00

m

0.95

0.95

0.95

0.95

0.90

0.90

0.90

0.90

0.85

0.85

0.80

0.80

0.80

0.75

0.75

0.75

0.75

0 3 6 9 12 15 18 2124 2730

0.70

0 3 6 9 12 1518 21 242730

0.70

0 3 6 9 12 15 18 21 24 27 30

0.70

0 3 6 9 1215 18 21242730

0 3 6 9 12 15 18 21 2427 30

0 3 6 9 1215 18212427 30

E.2 Header parameter characteristics in DHA

The DHA parameter distribution of shows consistency with MHA's. It indicates that DHA effectively aggregates multiple similar functional heads within clusters and new fused heads successfully reconstruct the functionalities of multiple origin heads in MHA. It is noteworthy that the significant reduction in the number of similar heads within the DHA architecture indicates that our method effectively reduces redundancy among the heads.

Layer 1 WQ Head similarity

0.8

1.0

0t  + s 9

Layer 1 WK Head similarity

0.8

1.0

Layer 1 WV Head similarity

Layer 5 WQ Head similarity

1.0

0.9

Layer 5 WK Head similarity

0.95

Layer 5 WV Head similarity

0.98

1.00

0.6

N

0.8

0.90

0.96

0.4

0.2

0.4

0.2

D

0.3

m

0.7

0.6

0.5

n

0.85

0.75

0.80

0.70

0.94

0.92

0.90

0123

45

6

7

01234567

01234567

01234567

Layer 2 WQ Head similarity

1.0

Layer 2 WK Head similarity

Layer 2 WV Head similarity

1.0

Layer 6 WQ Head similarity

1.0

Layer 6 WK Head similarity

1.00

。

Layer 6 WV Head similarity

1.00

.8

.9

0.9

0.98

.8

E

0.8

0.85

0.96

0.4

0.7

0.80

0.94

0.7

0.6

0.75

n

260

1

3

4

5

67

01234567

0.3

01234567

0.6

0i23

4567

0.5

012

4567

0.65

0.70

0123

4567

0.90

Layer 3 WQ Head similarity

Layer 3 WK Head similarity

。

Layer 3 WV Head similarity

0.98

1.00

Layer 7 WQ Head similarity

0.9

1.0

Layer 7 WK Head similarity

0.95

1.00

。

Layer 7 WV Head similarity

0.98

1.00

0.96

N

0.90

0.96

0.8

0.7

0.94

0.92

m

0.7

n

0.80

0.85

0.92

0.94

0123 4 56

7

0.3

N

0i

0.6

0i

2

9SE

7

0.90

0 12 3 4 5 6 7

0.5

0.6

0123456 7

0.75

0 123 45 6 7

0.90

0.88

。

Layer 4 WQ Head similarity

F0.6

1.0

0.7

0.5

Layer 4 WK Head similarity

-.8

1.0

0.7

Layer 4 WV Head similarity

1.00

0.98

0.96

0.94

0.90

0.92

0.88

0.86

m

Layer 8 WQ Head similarity

0.9

0.8

1.0

0.7

0.6

Layer 8 WK Head similarity

0.95

0.75

0.90

0.80

0.85

0.70

。

0

Layer 8 WV Head similarity

1.00

0.96

0.98

0.92

0.94

0.90

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract provides a concise summary of the key findings and experiment results. The introduction in Sec. 1 outlines the research questions and objectives in paragraph 3,4 and contribution in paragraph 5.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The paper discusses the limitations of the work performed by the authors in detail in Appendix Appendix. A.2, highlighting two specific limitations and the broader impact.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: This paper is mainly based on observation, making conjectures and methods and proving the effects through experiments. The paper defines the background in Sec. 2, presents the conjecture in Sec. 3, and provides a detailed derivation of the form and optimization process in Sec. 4. All assumptions made in the paper are thoroughly validated through experiments in Sec. 5.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The paper provides a detailed description of the experimental data setup and hyperparameter settings in Sec. 5. Additionally, in Sec. 4 and in Appendix. B.2 sections we thoroughly explain the derivation and implementation process, ensuring all necessary information for reproducing the main experimental results is disclosed.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: The datasets, baseline methods, and models used in the paper are fully opensource and available on Hugging Face. The paper includes the key implementation steps and code in Sec. 4 and the Appendix. B.2. However, the complete code is still being organized and is under consideration for open sourcing.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes] ,

Justification: The paper provides a detailed description of the experimental data setup and hyperparameter settings in Sec. 5. Additionally, in Sec. 4 and in Appendix. B.2 sections we thoroughly explain the derivation and implementation process.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: All results are averaged over multiple tests, and we report the mean accuracy along with the standard deviation (acc_norm) as a measure of error bars.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: In Sec. 5.1, we report the GPUs we used, the memory, and detailed training information. For more information you can refer to the Appendix B.2.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer:[Yes]

Justification: The discussion of the ethics and impact can be consulted in Appendix. A.2. We are open and transparent throughout the study and do not design for human subjects, privacy data bias, or other issues.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: The discussion of the broader impacts can be consulted in Appendix. A.2. Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper presents an improved approach based on the existing model architecture, but does not release any new models. The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: This article uses assets reasonably in compliance with the license, and the assets used are cited in the article.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a novel method to convert the Transformer attention mechanism from multi-head attention (MHA) to the proposed decoupled-head attention (DHA). DHA employs fewer unique KV heads than MHA or GQA by representing the heads as a linear combination of unique heads within the group, where each group contains similar heads. This conversion process is performed at the parameter level, so basically any MHA model can be converted to a DHA model. After conversion, continued pre-training helps mitigate performance drops. The core novelty of the proposed conversion mechanism comes from its head fusion strategy: first, it searches for similar head groups, optimizes the fusion weights to obtain many-to-one mapping, and then allocates a different number of heads to different layers within the resource budget. Experiments show that DHA requires considerably fewer training resources while achieving lower LM loss compared to GQA. Furthermore, DHA is beneficial for reducing the KV cache size.\n\nSTRENGTHS:\n* The paper identifies significant redundancies in the attention KV heads and proposes a practical solution to address the problem. Empirical evidence for the motivation is also provided.\n* The learning-based fusion optimization and adaptive budget allocation offer a systematic approach to finding optimal mappings, adding non-trivial value to the proposed method.\n* Related works are sufficiently addressed.\n\nWEAKNESSES:\n* It is unclear whether the paper uses inter-layer grouping of heads or only intra-layer grouping. The introduction and Figure 1 suggest that KV heads from different layers can also be merged, but this does not seem to be reflected in the method and the models used in the experiments.\n  * If the grouping is only performed in an intra-layer manner, the paper should be revised accordingly for clarity.\n  * If the grouping is performed in an inter-layer manner, the optimization process would become much more complex. Also, during inference, each layer would need to call a large number of parameters (essentially all weights of the corresponding group across layers), which would slow down the inference due to the memory-bound nature of LLMs.\n  * If the grouping is performed in an inter-layer manner, a visualization similar to Figure 2 across layers could strengthen the claim.\n* It is assumed that \"GQA\" in all experiments refers “MHA -> GQA” conversion using simple mean pooling. Is this a correct understanding? Is it a common approach to convert MHA to GQA instead of training a GQA model from scratch? Can we consider the uniform division of heads and apply mean pooling for GQA construction as a solid and strong baseline? For example, using CKA to identify similar heads before mean pooling could improve GQA’s performance.\n* It is difficult to say that DHA shows \"much less training overhead\" because it starts from MHA (line 294), especially compared to the models that begin with GQA from scratch.\n\nQUESTIONS:\n* Although the paper introduces the method as performing MHA to DHA conversion at the parameter level (checkpoint-level), the Lagrangian min-max training requires LM loss computation for training. This means the fusion step in the conversion process requires task-oriented training using real data. Is this correct?\n* How long does it take for the searching process? The paper mentions 100 search steps (line 197) for each layer.\n* Providing a demonstration of how merging weights change after each process step would help in understanding the method better.\n* Figure 2 visualizes the weight similarity instead of CKA, which makes sense because similar weights would generate similar outputs, and the method merges similar parameters. Then, is CKA-based analysis necessary (also there is no CKA visualization)?\n* (minor) Typo: Adaptvie --> Adaptive in the caption of Table 2.\n\nLIMITATIONS:\nThe paper adequately addresses the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

comments: SUMMARY:
This paper introduces a novel mechanism to optimize large language models (LLMs) by addressing the computational and memory costs associated with the Multi-Head Attention (MHA) mechanism. The authors propose Decoupled-Head Attention (DHA), which adaptively configures the sharing of key and value heads across layers, aiming to strike a balance between performance and efficiency. The transformation from MHA to DHA is achieved through a process of head parameter clustering and linear fusion, which retains the knowledge from the original model with minimal performance degradation.

STRENGTHS:
-The paper demonstrates a remarkable balance between performance and efficiency, showing that DHA can achieve high performance with significantly reduced computational resources.

-The proposed method is not limited to specific model size, indicating that it could be broadly applied to various existing MHA Transformer models.

WEAKNESSES:
The effectiveness of DHA on a wider range of model architecture remains unexplored. Could the method work well on models beyond LLAMA?

The process of transforming MHA to DHA involves several stages of training and fusion that might be complex to implement and reproduce for researchers without access to similar computational resources.

The paper relies solely on linear methods for parameter fusion. There might be room for improvement by exploring non-linear fusion techniques that could potentially offer better optimization.

QUESTIONS:
See the weaknesses.

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces Decoupled-Head Attention (DHA), a new efficient attention mechanism for large language models. DHA adaptively configures group sharing for key and value heads across layers, transforming Multi-Head Attention checkpoints through a three-stage process. It achieves 97.6% of the original performance while using only 0.25% of the pre-training budget and reducing KV cache by 75%. DHA outperforms other efficient attention methods in terms of training speed and performance under limited resources.\n\nSTRENGTHS:\n1. Novel approach: Introduces Decoupled-Head Attention (DHA), an innovative method to optimize attention mechanisms in large language models.\n2. Efficiency gains: Achieves significant reductions in computational and memory costs (75% KV cache reduction) with small performance loss.\n3. Rapid adaptation: Requires only 0.25% of the original pre-training budget to achieve 97.6% of performance.\n\nWEAKNESSES:\n1. While the proposed method seems to be effective for MHA, the authors did not demonstrate its compatibility with GQA, which is increasingly popular recently and starts to replace MHA in almost all state-of-the-art models.\n2. While efficient adaptation is important, it is more important to see whether the performance of DHA is similar to GQA/MQA in the end of the training process, when the entire model converges. If DHA just converges faster, I think its value is rather limited.\n3. While the accuracy loss is not as large as other KV compression methods, a ~5% performance drop for Llama2-7B is still too large. It's almost equivalent to the performance gap between Llama-2-7B and Sheared-Llama-2.7B. This indicates that the process is extremely lossy, and the resulting model might not be useful after DHA process.\n4. It is unclear how DHA combines with existing KV cache compression techniques, such as KV cache quantization.\n\nQUESTIONS:\nPlease answer my questions listed in \"Weaknesses\"\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

comments: SUMMARY:
The paper introduces a novel attention mechanism, Decoupled-Head Attention (DHA), designed to enhance the efficiency of large language models (LLMs) with minimal performance loss. DHA adaptively configures key and value heads across layers by leveraging insights from attention redundancy, leading to significant savings in computational and memory costs. The authors propose a method to transform a standard multi-head attention (MHA) model into DHA with a minimal pre-training budget, achieving near-original performance while reducing the key-value (KV) cache requirements substantially. The experiments demonstrate the effectiveness of DHA in maintaining performance with reduced resources.

STRENGTHS:
The paper proposes a novel Decoupled-Head Attention (DHA) mechanism that addresses the computational and memory costs associated with Multi-Head Attention (MHA) in large language models (LLMs). The idea of adaptively configuring group sharing for key and value heads is innovative. DHA initialization can recover the performance with a very small amount of restorative pre-training, compared to Group-Query Attention.

WEAKNESSES:
1. **Comparison with GQA and MHA Training Data Requirements:**  It would be beneficial for the authors to report the performance of the DHA model when continue pretrain with a larger dataset, compared to GQA and MHA models. 
2. There is a noticeable accuracy gap between DHA and MHA models on academic datasets.

QUESTIONS:
1. Could the authors clarify why an increase in the fusion phase token budget leads to a decrease in average accuracy post-CT as observed in Table 3? 
2. Are there any experimental results available that demonstrate how the DHA model performs on larger datasets?

LIMITATIONS:
The  author has mentioned limitations in Appendix
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13