DHA: Learning Decoupled-Head Attention from Transformer Checkpoints via Adaptive Heads Fusion
Authors: Yilong Chen
Linhao Zhang
Junyuan Shang
Zhenyu Zhang
Tingwen Liu
Shuohuan Wang
Yu Sun
Venue: NeurIPS 2024 poster (2024)
Keywords: Large Language Models
Multi-Head Attention
Pre-training Acceleration
Efficient Inference
Model Fusion
Primary area: natural_language_processing

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
Yilong Chen Linhao Zhang Junyuan Shang3; Zhenyu Zhang",
Tingwen Liu Lit Shuohuan Wang?, Yu Sun
Institute of Information Engineering, Chinese Academy of Sciences
School of Cyber Security, University of Chinese Academy Sciences Baidu Inc.
(chenyilong, Liutingwen)@iie.ee
(zhanglinhao, shangjunyuan, zhangzhenyu07, vangshuohuan, sunya02)@baiduu.com
Abstract
Large language models (LLMs) with billions of parameters demonstrate impres- sive 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 necessitate substantial continued pre-training costs to restore performance. Based on the analysis of attention redundancy, we design a Decoupled-HHee (DHA) mechanism. DHA adaptively configures group sharing for key heads and value heads across various layers, achieving better balance between performance and efficiency. Inspired by the observation of clustering similar heads, 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 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 's pre- pre-training budgets 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.
Attention
Introduction
Transformer- based large language models (LLMss) shine in various natural language tasks due to their powerful understanding and generation capabilities [ 2 3 Multi-Head Attention (MHA) widely used 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 hillion-parameter model with 32 heads and 32 layers, input batch size of 4, and sequence length of 32k results in 64GB of KV cache, which is 4.7x larger than the model weights.
To reduce computational and memory overhead during inference, a widely used approach involves adapting the MHA model more efficient structure through the reuse of parameters across multiple heads
4 5 6 such as Multi-Query Attention (MQA) 4 and Strouped-uuey Attention (GQA) [5] These methods utilize portion of the original training computation which avoid information loss
denotes equal contribution. denotes corresponding author.
denotes project lead.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Multi-Head Attention (MHA)
Qrouped-Query Attention (GQA)
Decoupled-Head Attention (DHA)
Layern
Layer m+
Layern n
Layer
Layer
GQA Initialization
DHA Initialization
1/x
Fusion
New Heads
New Heads
Group size
Group size
Heads
Grouped
InsFusion Heads
Pooling
Head Dependence
Learning-Based
Figure I: Upper: Overview of Decoupled-hor hecadpled-head method. Multi- Heald-Headdd attention (MHA) has equal query, key and value hcads. Grouped-Query attention (GQA) instead shares single key and value heads for each group of query heads. Decoupled-Head attention (DHA) shares key heads and value heads for different groups of query heads in different layers. Lower: GQA Initialization: Heads are mean pooled into a single head; DHA Initialization: DIAA search head grouping and progressively fuse heads to maintain parameter functions.
vraining-inferene inference inconsistencies, common issue in pruning-based 7 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.
due training-
9 10 11
works.
Thus, this work, we seek to address the following question:
How can we construct more efficient model while keeping costs as low as possible?
With the limited understanding of parameter characteristics in modern LLMs, first perform empirical analysis from the perspectives 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 decoupled distribution, meaning that there is significant variation in the distribution of head-cluster similarities across layers, key heads and value heads, illustrated Fig. 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 significant performance drop (see Sec. 3.1). Based on this observation, we believe that selectively fusing corresponding heads in clusters can construct more efficient architecture with minimum loss.
In this paper, we propose Decoupled-Head Attention (DHA), 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 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 [4 [5), 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 perforinfo initial points for DHA from MHA checkpoints, requiring only minimal amount of Continued Pre-training to restore performance.

---

WK Head Similarity
Wi Head Similarity
Wk Head Similarity
Wv Head Similarity
3.8
0.0
10
0.6
0.7
-.0
25
89
0.6
69
22
0.5
0.6
(a) Head Parameter Similarity in Oth Layer
(b) Head Parameter Similarity 21st Layer
Figure 2: Visualization of the similarity between heads within the MHA of LLaMA2-7B model at the Oth layer (a) and the 21st layer (b). Details in Appendix E.I Key heads and value heads exhibit decoupled distributions.
To verify the effeetiveness, we construct DHA on models of different sizes, such as LLaMA2-7B [3). Sheared-LLaMMA-2777 & -133 (16) with the heads budget ratio set at 50% and 25%. With 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% MHA pre-training budget. Meanwhile, DHA is capable of reducing KV Cache by up 75% compared to MHA with minimal accuracy trade-off (maximum of 5.6%). Compared to GQA, DHA achieves a 5 x training acceleration, 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 xp) E RPxxaam denote the input prompts of hidden states of Transformer layer, where p stands for the number of tokens and dyoode for the hidden state dimension.
performs the attention with H different heads. For h-th head, different weight matrices Wh w Wh Remas Xd: are used project the input sequence into query, key, value vector, where de represents head dim. Denote softmax funcion as a, we have:
Multi-Head Attention (MHA) MHA [17
MIHA Concat (headl...... heady ) Wo, where head,
(xW) (XW2)
XW
(1)
Vdk
Ultimately MHA combines heads' outputs through the output projection Wo E IR Cantaal dante
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:
xW!(h), where wath) E EK/V oth Vdla
Wajo
head, ( (XWahhi)
(2)
K//V o(h).
Here, g/V refers to MHA key/value heads parameters within the ( -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 =
9(4)
Due to mean pooling for initialization, GQA results in loss of capability when converting from MHA, necessitating expensive pre-training training to recover. We aim identify better initialization and more refined head mapping relationships to achieve superior performance with reduced training costs.
2 MHA consists of heads stores a x H x x dmotel dimension KV cache for accelerating inference. PGQA and MQA consist of H + G heads in total and store a 2 G p dmolel dimension KV cache.
3

---

3
Observation
To study the inherent characteristics of head parameters in MHA, we use Centered Kernel Align- ment [8] to calculate the heads' similarity within each layer's W., W,, Based on the average heads' similarity, we define the redundancy of each MHA layer. For details, please refer to Appendix B.
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 Wx and W., 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. S.2. As shown in Fig. 3a the loss remains unchanged as the proportion of Head decreases and only increases when four heads parameters' ratios approach an 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 sim-
ilar head clusters varies between differ- ent layers. As illustrated in Fig [bb] the Oth layer of MHA shows few similar head clusters, while the 21st layer ex- hibits many. Within the same layer, value heads exhibit more clusters and higher similarity compared to key heads, indicat- ing 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, redun- dancy significantly exceeding that of W
2.4
Head
Head 1
Head
Los
1.0
Rodundan
0.5
0.5
20
0.4
0.2
Vilue Heads
Query Heads
Heeds
100
200
300
400
10
20
30
Steps
Layer
(a) Loss in 4-Head Fusion
(b) 32-Layer Redundancy
Figure 3: (a) Model loss with heads proportions in linear fusion. (b) Layer Redundancy of the query, key, value head parameter matrices in the LLaMA2-7B model MHA.
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, shows higher redundancy than W.. 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.
decouple and allocate more heads budget to critical key components, while
4 Method
In this section, propose 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 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.

---

Dependence Search
In-Group Head Fusion
Post-Process
Fusion Initiation
Fusion Process
KV
W
W
W
W
We
W
W,
W
W.
W
W
W
Wi
W
W
Mas
Micc (i2.
Saa od
dirg
W
W
2
W
Wo
W2
W
W
W
W
W
W
Wo
W2
Wi
W3
W,.W;
W. W.
4422
(PNy m2
Frucus Groggummx
Σ Σ (wey -mps)"
W
minge (XC10"H0)) Alspare
(2)( ALcuass
IW W1
Figure 4: An illustration DHA. First, we reconstruct the a single head forward as a linear combination of multiple heads' forward with proportions w. grouping heads with similar functions based on multi-step optimization. Next, we initialize and optimize the fusion operators. <> indicates the optimization narrows the distance between proportions Finally, we fuse heads within groups and continued pre-training DHA model.
Definition Defined model with layers and Hº heads in a layer. the numbers of Key heads and Value heads in the (-th layer are denoted as H1, .HY. We define the many- many-to-one mapping functions d (h, 1) and d (h, () representing key and value head corresponding to the h-th query head in 1-th DHA layer. The computation formalized as follows:
heads.1 = X (XW (unny
Xw
X (h.1)
Vde.
DHA shares key and value head in multi query heads' computation based independent mapping functions at different layers" 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 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, 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 model with Layer number L and Head number H 1r.H [W1,... ....... W.,, where W, E DxD denotes the weight of layer / with input and output dimen- sion D. In the initialization, our goal is transfer knowledge from MHA model MHA L.H,
DHA model D where H1 > H2 By learning a fusion operation that minimizes the functional difference between MHA and DHA model, the goal can be formalized as
Ll
arg min En-D [Cim (2: + XL (x: (E MHA 9.M
TIFA
(4)
Where M is the fusion operator, D is the training dataset, Lim is the training loss function, Ltusion 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 (a V (h,)) represents key, value heads group corresponding the h-th query head in 1-th layer, g {g* .g VV the group size. By introducing variables Wh. { E M represents the proportion of j-th key, value head
'DHA consists of / 0 + ΣL H Et. H. heads in total.
5

---

involved in the h-th query head forward within group. For each group, head have forward pass as:
htads,1 = X (0. OADTT
Xwd (h,)) where W w vds.
(b.))
who
W
k/v
(5)
where whj will be initialized to Kronecker delta function, which equals 1 if and only if h- and equals O otherwise. Under this initialization setting, the forward computation DHA is completely equivalent to that of MHA, see Fig.
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 e/v is many-to-many mapping, denoted by the function d 1), This indicates that in the forward process of each query, the key head or value head can be expressed as different linear combinations of MHA heads. According to Eq. 3 we aim to achieve many-to-on mapping that single fused key head or value head are shared across multiple query heads in DHA, denoted by the function dK/V (h. 7). Thus, we design fusion loss Liusion to optimize the initial mapping functions to converge to single mapping function, i.e., dK/ 1h + d (h, 1), Vh E K /V, Specifically, we define the optimization objective as minimizing the difference between the mapping functions of different query heads and / within the /-th layer and n-th group:
2 Why
Euead (h, =
Σ
W k/U
Σ
W 6/0
Σ (WNs wh's) k/0.5)
W
(6)
where g KNV represents the number of heads within group. Since Wi/h! can be regarded an orthogonal scalar, and thus we only need to optimize fusion variables so have:
1.
Liusion ΣΣΣ E Cuenap (h,
h'), subject to Lhewa (h, h²)
> (wns whr))
9 h=ljm
[=1 h= R'mb4+
Where N represents the number of groups, = Hr, The fusion loss can measured as the mean squared error loss 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 corre- sponding 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 []4 15] for incremental architectural transformations. Define t as the target loss, b as the base decay factor, s as the current global step, as the total number of steps in the warm-up phase, the overall training optimization is adversarial game:
max min EneD [Aim (2c; (a: M(MMI))))) + > Amax (Crusion t, 1.0]]] where max (0. (1.6((11 -) (8)
Our Augmented Lagrangian approach enforces the constraint Lussion 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 Liusion for each layer and Lhead for all heads. Based on the Lhead. we perform head grouping intending to minimize the average loss heads within each group and maximize the average loss of heads between groups and groups. Based Ltusion we use dynamic programming

---

algorithm to allocate more head budget to layers with higher loss within total budget. It allows us to fuse the most similar heads 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 Crusion that marks this DHA fusion process decreases. When Lfusion is less than le-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 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 all MHA layers. Notably, expanded the dimension of each head's fusion coefficient w from to the head's dimension du, 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.
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-Exchanp This dataset comprises validation with 2 million tokens, training set containing 4 billion tokens and an additional pre-training set totaling 50 billion tokens.
Training. Our experimental framework utilizes the Sheared- -LaMA codebase [16] implemented on the Composer package [20], and is executed on 8 NVIDIA A 100 GPUs (80GB). The models 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 le-4 for language modeling loss, and 1e-2 for Lagrangian multipliers and fusion operators respectively.
Budget. DHA models were trained for 1000 steps (O.2B token budget) during the fusion phases. For the continued pre- pre-training, trained both baseline models and DHA for up 50000 steps (50B token budget). To evaluate the training acceleration capability of DHA, we evaluate its performance under two budget scenarios. First, set 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.
2 to evaluate our models. For common sense and reading comprehension tasks, we report 0-shot accuracy results for SciQ [22). PIQA WinoGrande (Wino.) (24]. ARC Easy(ARC-E.) 25), and HellaSwag (HellaS.) (26). alongside
Evaluation. We employed the Im-evaluation--anne
23
25-shot accuracy for ARC Challenge (ARC-C.) 27 In the assessments continued QA and text un- derstanding, we report O-shot accuracy for LogiQA 28 32-shot BoolQ 22), 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 10,000 instruction-responss pairs from the ShT dataset and evaluate on another 1,000 instructions, using GPT-4 for response evaluator (33 The win rate of our model relative the baseline reported. For detailed information, refer Appendix C.I
https://sharogpt.com
7

---

Table 1: Comprehensive assessment of model's fundamental capabilities, in which DHA models demonstrate competitive performance while requiring significantly fewer training resources. Models with use MHA.
Commonsense & Comprehension
Continued
LM
Model
Budget SeiQ
PIQA
Wino. ARC-E ARC-C HellaS. LogiQA BoolQ LAMB. Average
LLaMA2-7B*
2T
94.1
78.1
69.1
76.3
49.7
58.9
25.7
80.8
74.1
67.4
DHA-7B-50%
50B
93.4
78.5
69.1
73.8
45.9
58.6
22.5
79.1
71.1
65.8
DHA-7B -5%
50B
92.4
78.5
68.6
72.9
43.9
57.6
22.4
76.7
70.2
64.8
GQA-7B-500 50%
IB
90.7
76.8
66.5
71.3
41.9
53.6
22.4
70.5
67.0
62.3
DHA-78-50%
18
90.8
76.5
66.7
71.3
44,6
55.1
22,4
74.8
67.2
63.3
GQA-7B-25%
1B
86.5
74.3
59.1
67.6
37.5
49.2
24.1
65.8
58.3
58.0
DHA-7B-2555
90.0
75.2
63.8
70.4
39.3
52.2
21.1
72.3
62.9
60.8
S.-LLaMA-2.77
2T
91.2
76.1
64.9
67.3
38,8
52.2
22.1
74,4
68.3
61.7
CQA-2.7B--33 50%
1B
86.7
74.8
59.0
64.0
34,2
48.2
23.8
64.9
60.3
57.3
DHA-2.7B-500
1B
86.8
75.1
59.5
64.6
35.1
48.7
22.4
66.4
61.7
57.8
GQA-2.7B-2555
82.0
72.8
54.9
58.4
31.0
42.9
21.7
58.5
49.6
52.4
DHA 2.7.-2.BB222 25%
85.6
74.1
57.6
61.5
32.4
45.9
21.7
63.1
56.9
55.4
S.-LLaMA-1.35 1.33
2T
87.0
73.6
58.2
60.9
29.5
45.4
21.8
65.5
61.3
55.9
IQA-1.3B-500
1B
84.3
72.3
55.8
57.5
28.2
41.8
20.7
62.9
52.9
$2.9
DHA-1.3B-50%5
1B
84.5
72.0
55.2
58.1
28.7
42.6
21.5
63.7
55.4
53.5
GQA-1.3B-25%%
76.6
70.0
52.9
51.9
23.5
37.6
21.0
59.9
41.0
48.3
DHA-1.30 1.3B-25%
82.8
54.0
55.4
25.8
21.5
57.6
48.6
50.8
Table 2: Ablation Results of DHA wo. Linear Heads Fusion and Adaptive Transformation. Experiments conducted with LLaMA2-7B with 25% heads budget and 0.5B IB training budget on O-shot Evaluation.
Models
SciQ
PiQA
Wino.
ARC-E.
ARC-C.
LogiQA
LAMB.
Average
Diff
DHA-7B-25% (0.5B)
88.6
75.9
61.3
68.2
36.1
23.8
63.2
59.6
w.o. Linear Heads Fusion
83.4
73.7
57.3
63.6
29.4
22.0
$1.9
54.5
-5.1
w.o. Adaptvie Transformation
87.9
74.1
60.1
69.4
34.7
19.5
62.1
58.3
-1.3
DHA-7B-25% (1B)
75.2
37.5
21.1
60.1
w.o. Lincar Heads Fusion
87.5
74.5
60.7
67.3
32.8
21.7
58.3
57.5
-2.6
Adaptvie Transformation
89.5
69.1
36.3
21.6
62.4
59.5
DHA-7B-25% (SB)
91.7
76.8
64.4
70.9
42.8
21.8
68.4
62.4
GQA-7B-25% (5B)
91.5
76.6
63.9
70.5
42.3
22.1
62.1
0.3
Baselines. We selected the LLaMA2-7B model and Sheard-LLLaMA-27BBB3t (S.-LLaMA- 2.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 (
5.2 Experimental Results
Foundational Capabilities. Tab. shows the foundational capabilities of DHA and GQA models at 50% and 25% compression rates (c.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 IB 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 '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
'LLaMA2 pre-trained 2T data; Sheared -LaMed pruned LLaMA on 50B RedPajama data.

---

G0A-78-251
Average Acc (96)
70
GQA
LM Loss
DHA: 78-288-
80
LM LoSs
DHA-40ba
DHA-240ba
DHA-400ba
DHA 1154-78-10 25%
6QA-18-25%
0.02
0.05
0.10
0:15
0:20
0.2 0.5
1.0
1.5
2.0
0.1
0.2 0.3 0.4 0.5
Fusion Training Budget (B)
Continued Pretraining Budget (B)
Continued Pretraining Budget
Figure LM Loss with Fusion Training (B) between GQA-7B- 25% and DHA-7B-25%.
Figure Task Average Accuracy (%) with CT (B) of DHA-7B-25% and GQA-7B-25%.
Figure 7: LM Loss with CT (B) between GQA-7B-25% and Cold- Start DHA after X Step Scarch.
1.65
Key heads
01
Value heads
10
6.9
697
011
0.75
5
7
11
13
15
17
19
21
23
25
27
31
Layer
Figure 9: Similarity of value heads in the 7th layer of LLaMA2-7B MHA (Left) the transformed DHA-7B 7B-25% DHA (Right).
Figure 8: Allocation of key-value head budgets with 32 layers in DHA-7B-25% after 240 step Search.
becomes more remarkable with reduced training budgets. It indicates that DHA effectively retains knowledge of larger models, significantly reducing pre-trate pre-training costs.
Better Initialization. We examined whether DHA offers a better initialization point than GQA by pre-training both DHA and GQA models the original RedPajama dataset. Fig. shows that the initial loss of the GQA model is high and decreases slowly. In contrast, the DHA model starting from MHA exhibits minor increase in LM loss as fusion progresses, maintaining a consistently lower loss. DHA converges with just O.1B data, demonstrating a 5 training speedup compared to GQA. Fig. Irpoort the average downstream task accuracy of DHA and GQA during continued pre-training. DHA achieves comparable performance to GQA's at IB 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 efficient architecture, whereas GQA loses information during initialization.
5.3 Analysis
Ablation Study. We report the effects of ablating Linear Heads Fusion and Adaptive Transformation in T When training with less data (O.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 IB, DHA's efficient architecture after Adaptive Transformation plays 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 computa- tion to the fusion phase aids in better retention of informa- tion 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 initial- ization point. Training with just O. 1B data is sufficient to achieve good starting point, and increasing fusion budget will not affect the final performance. This experiment also
Table 3: Data budget allocation to fusion and continued mee-trainnng(CT) and 0-shot Task Average Accuracy (%) D.A-1.3B.
Fusion
CT
Tokens
Avg.Acc
Tokens
Avg.Acc
0.05B
33.74
4,95B
59.08
0.10B
38.32
4,90B
0.15B
48.26
4.85B
59.46
0.20B
52.54
4,80B
59.16

---

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 the fusion stage to achieve 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. g 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 lower final loss. This indicates that DHA's architecture can accelerate training and achieve better performance, even without Linear Heads Fusion method.
For interpretability analysis, we visualized the parameter characteristics of the post-fusion DHA model in Fig. | (detials in Appendix [..2). and compared them with those prior to fusion. The DHA parameter distribution shows consistency with MHA's 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.
Parameter Characteristics in DHA.
6 Related Work
Advanced Multi-Head Attention. Some efforts have been converting the traditional Multi-Head Attention (MHA) [17] to Multi-Query Attention (MQA) 44, 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 holding transformations to expand the width by duplicating neurons, and uses a unitary layer implementation to expand the depth. LiGO (37) proposes 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 MHA checkpoint 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 Yingi 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.2021 YFB3 100600) and the Youth Innovation Promotion Association of CAS (Grant No. 202 153).
10

---

References
[1] Anthropic. Introducing claude. 2023.
[2] OpenAL. Gpt-4 technical report. ArXiv, page abs/2303.08774, 2023.
[3] 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, Xiao- qing 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.
[4] Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.021530 2019.
[5] 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 (XXi:23305.13244 2023.
[6] 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.
[7] Saurabh Agarwal, Bilge Acun, Basil Homer, Mostafa Elhoushi, Yejin Lee, Shivaram Venkatara- man, Dimitris Papailiopoulos, and Carole-Jean Wu. Chai: Clustered head attention for efficient IIm inference, 2024.
[8] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H __ Heavy-hitter oracle for efficient generative inference of large language models. arXiv preprint arXiv:2306. 14048, 2023.
[9] 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 Ilm kv cache compression at test time. arXiv preprint arXiv: 17118, 2023.
[10] 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 1rXiv:2310.01801, 2023.
[11] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309. 17453, 2023.
[12] 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.
[13] 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.
[14] Alexander M. Rush, David Sontag, Michael Collins, and Tommi Jaakkola. On dual decom- position 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.
[15] Ziheng Wang, Jeremy Wohlwend, and Tao Lei. Structured Pruning of Large Language Models. In Proceedings of the 2020 Conference Empirical Methods in Natural Language Processing (EMNLIP), 2020.
[16] Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared ILaMA: Accelerating Language Model Pre- Pre-training via Structured Pruning. October 2023.
11

---

[17]
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[18]
Simon Korbblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited, 2019.
[19]
TogetherAI. Redpajama: An open source recipe to reproduce llama training dataset, 2023.
[20] The Mosaic ML Team. composer. https:// http::///gthhb..... com http:/oosaiimmmmm wwwpm/////co///////////// 2021.
[21] 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.
[22] Johannes Welbl, Nelson 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.
[23] Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about physical commonsense in natural language. In The Thirty-Fouthh AAAI Conference on Artificial Intelligence, AAA! 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pages 7432-7439. AAAI Press, 2020.
[24] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. In The Thirty-Fourty Fourth AAAI Conference on Artificial Intelligence. AAAI 2020, The -hirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, FAAI 2020, New York, NY, USA, February 7-12, 2020, pages 8732-8740. AAAI Press, 2020.
[25] 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 1rXiv:1803005455 2018.
[26] 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, ACI. 2019, Florence, Italy, July 28- August 2, 2019, Volume I: Long Papers, pages 4791-4801 Association for Computational Linguistics, 2019.
[27] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick. and Oyvind Tafford. Think you have solved question answering? try arc, the A12 reasoning challenge. CoRR, abs/1803.05455 2018.
[28] 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.08/ 024, 2020.
[29] Christopher Clark, Kenton Lee, Ming-Wees Wei Chang, Tom Kwiatkowski. Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint 19Xiv:1905.1000 2019.
[30] 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 broad discourse context. arXiv preprint 1606. 06000, 2016.
[31] 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.
[32] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An cestructioo-follovviii llaaa model, 2023.
12

---

[33] 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 preprini ar3iv:2305. 14387, 2023.
[34] 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.
[35] Lemeng Wu, Dilin Wang, and Qiang Liu. Splitting steepest descent for growing neural architectures. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' 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.
[36] Tianqi Chen, lan Goodfellow, and Jonathon Shlens. Net2net: Accelerating learning via knowl- edge transfer. arXiv preprint arXiv:1511. 1.0641, 2015.
[37] 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.
[38] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arXiv preprint arXiv:2009.06733 2020.
[39] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 195009994,,009,, 2019.
[40]
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.
[41] Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 2020.
[42] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprinl 25 2020.
[43] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G. Carbonell, Quoc V. Le, and Ruslan Salakhut- dinov. Transformer-xl: Attentive language models beyond fixed-length context. CoRR, Abs/1991,00860, 2019.
[44] Siyu Ding, Junyuan Shang, Shuohuan Wang, Yu Sun, Hao Tian, Hua Wu, and Haifeng Wang. Ernie-doc: A retrospective long- documenttn modeling transformer. arXiv preprint arXiv:2012. 15688, 2020.
[45] Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. Recurrent memory transformer. Advances in Neural Information Processing Systems, 35:1 15:1--0911191, 2022.
[46] Alexis Chevalier, Alexander Wettig, Anirudh Ajith, and Danqi Chen. Adapting language models to compress contexts. arXiv preprint arXiv:2305. 14788, 2023.
[47] Yilong Chen, Junyuan Shang, Zhengyu Zhang, Jiawei Sheng, Tingwen Liu, Shuohuan Wang, Yu Sun, Hua Wu, and Haifeng Wang. Mixture of hidden- didden-dimenssons transformer, 2024.
[48] Matanel Oren, Michael Hassid, Yossi Adi, and Roy Schwartz. Transformers are multi-state mas. arXiv preprint arXiv:2401.06100 2024.
[49] Yilong Chen, Guoxia Wang. Junyuan Shang, Shiyao Cui, Zhenyu Zhang, Tingwen Liu, Shuo- huan 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 Computa- tional Linguistics ( (Volume 1: Long Papers), pages 7913-7926, Bangkok, Thailand, August 2024. Association for Computational Linguistics.
[50] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm.int80: 8-bit matrix multiplication for transformers at scale, 2022.
[51] Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers, 2023.
13

---

[52] Robert M. Gray and David L. Neuhoff. Quantization. IEEE transactions on information theory, 44(6):2325-2383, 1998.
[53] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network, 2015.
[54] Jianping Gou, Baosheng Yu, Stephen J. Maybank, and Dacheng Tao. Knowledge distillation: A survey, International Journal Computer Vision, (296))1789-1819 March 2021.
[55] 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.
[56] Zichang Liu, Jue Wang. Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shri- vastava, Ce Zhang, Yuandong Tian. Christopher Re, and Beidi Chen. Deja vu: Contextual sparsity for efficient LLMs 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-22177. PMLR, 23-29 Jul 2023.
[57] Elias Frantar and Dan Alistarh. SparseGPT: Massive Language Models Can Be Accurately
Pruned One-Shot, March 2023.
[58] Minjia Zhang and Yuxiong He. Accelerating training of transformer- language models with progressive layer dropping. Advances Neural Information Processing Systems, 33: 1- 14023, 2020.
[59] 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.
[60] 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.
[61] [2401.02415] LLaMA Pro: Progressive LLaMA with Block Expansion.
[62]
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 I: Long Pupers), pages 8005-8019, Bangkok, Thailand, August 2024. Association for Computational Linguistics.
14

---

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
45 46, 47] to address the self-attention operation which scales quadratically with the sequence length. For instance, Sparse Transformer [39 uses 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
44
4 maintain memory bank of past KV cache to process the long text in segments. However, the above methods either result in loss of model performance require retraining the model, which is unaffordable for the high computational resources of L.MMs. 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 [ 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 maintain consistent performance for training and inference. Pruning [I5] [6]. quantization [50] 511 [52] and distillation [53, [5 55] can reduce the number of model key and value headers, parameter dimensions, and activation reduce memory bandwidth overhead during model inference. Deja Vu (56) and CHAI ( prune pruning redundant heads through clustering methods for efficient inference. the LLM era, this leads 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.
10 g 48 49
Model Compression. Our approach is dedicated to obtaining a high- gegh-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 loss of model accuracy and the inability to freely build argee-sppeifii models. CRash [I3] and LayerDrop [58 [59] methods discard ineffective layers during training, which do not allow for target-speeific structuring and come with a large performance loss. Pruning [SS] minimizes the impact on performance by cutting out redundant neurons that pver-parameterize the model. In the L.M 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 [6] 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 specifie 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 L.Mss. Our work on Decoupled-Heed Attention (DHA) represents an advancement in optimizing the efficiency of Large Language Models (LLMs). By addressing the
15

---

substantial computational and memory costs associated with the widely used Multi- Head Attention (MHA), DHA enhances the applicability of LLMs various domains. The introduction 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 L.MMs more feasible and cost-effective. This efficiency allows for the broader accessibility of advanced L.Mss, 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 Al 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 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 models of any size future work.
A.2.3 Ethical Consideration
In study, utilize publicly available data and techniques 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 more inclusive and automated future for Al.
B More Implementation Details
B.1 Head Similarity and MHA Redundancy
Centered Kernel Alignment (CKA) 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 kernel function to map the original data into nigher-dimensional space, where the alignment of their central tendencies can be more easily measured. The CKA value ranges from 0 to 1, where O indicates no similarity and indicates identical representations.
The mathematical formulation of CKA, when using a linear kernel, is given by the following equation:
X
CKA(X,Y) = VIXTXX YTTIE
Here, X and Y are matrices whose columns are the vectors of the representations to be compared, ||- ||r denotes the Frobenius norm, and X" and Y 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:
1. Compute the similarity between heads: For each pair of heads within a given layer, calculate the similarity using the CKA formula.
2. Compute the average similarity: Average the similarity scores of pairs of heads to define the redundancy of the layer.
16

---

B.1.1 Compute Similarity Between Heads
Consider layer with H heads, where the parameters of each head represented by the matrices W; (e.g., Wa., Wgz...., WyH for query weights). For each pair of heads and j. compute the CKA similarity using the following formula:
W 11 | ww.. [Wj will?
CKA(W& Wj) =
B.1.2 Compute Redundancy
Calculate the similarity for all pairs of heads and then compute the average similarity:
H-1 H
Redundancy = H(H 1) 1=1 j=it1
Σ Σ
CKA(W,, W,)
The coefficient - 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 higher level of redundancy.
B.2 Implementation in LLaMA2 Model
Our method can theoretically transform MHA architecture in any transformer model efficient DHA architecture. Using LLaMA models as case studies, we implemented DHA transformation with various compression rates on all MHA layers. Only 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 MHA layer. 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 noise in fusion. As Shown in Fig 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. 2x, 4x, 8x), the algorithm first performs multiple STEPs constrained only by the Lim constraints propagation, computing the Ctusion but not optimizing the lincar fusion operator based on it. Based on the Ltusion between head and head as 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, on the compression rate as the fusion target, and optimize their Liusson 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.
select multiple groups with the smallest loss based
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 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 small training budget and the recovery phase with continuous pre-training. Before Fusion phase, we define the total number of compressed headers budget then C' is split into compression rates 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 Lrusion 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 X, the Crusion that marks this DHA fusion process decreases.
17

---

When Ltusion is less than 1e-3 we terminate the fusion algorithm and enter the post post-procsssing phase. The fusion weights within each group are computed by averaging the weights corresponding 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 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 imput key and value tensors are first reshaped and grouped according to the number of key and value heads, respectively. Given the batch size (bsz), num- ber of heads (num heads), key length (k (e_ee), and head dimension (head_dim), the key tensor reshaped into keys grouped of shape [bsz, num sey_hs num heads num _ee_ss k_len, head dim]. Similarly, the value tensor is reshaped into num value heads, num heads num value _eadss k len, head _im]_ These grouped tensors num_heads
of shape [bsz,
are then expanded by repeating them along the group size dimension, resulting in leys_expanded and values
values_eepandees expanded. Correspondingly, the weight tensors weights_k and weights_ _ are reshaped to match the expanded dimensions and are then multiplied element-wise with the expanded key and value tensors.
Algorithm 1 Attention Module Initialization
Require: K
key tensor value tensor b weighted key tensor weighted value tensor
Require: V
Ensure: K"
Ensure: V
1: b & batch size
2: H +- number of heads
3: La <- key length
4: <- head dimension
5: K2 <- eeff.num_ke__ite heads
6: V + self.num_value_heee
7: K
sel..keygrroup...
group_size 9: K self.weights. k
8: Vg f value wwwaaaallllluuseeeeeee
10: V <- weightss
11: K (iview(b, Ko H K,,K, Li, D)
12: V + Vb,,e((, Vs, H/V., Vo, D)
13: Ke Ke..reeet__iterkkeee interleave( dim = 1)
14: V. + Vg.repeat interleave (V,, dim = 1)
15: K t (iew((1 H, Kg. D)
16: V. + V Vi.iie((1, H, Vg 1, D)
17: WK & K. x K
18: Wy
Ve V
19: K' < WK.sum dim = 2)
20: V <- W Wy.sum (am 2)
21: return K", V
B.2.2 Attention Forward Pass
During the forward pass, the reshaping and expansion of the key and value tensors are performed in similar manner as in the initialization process with parameters specific the DHA fusion phase. The key tensor is reshaped into reys_groupee of shape [bsz, dha_ peemmup_grou_naee // dha srruuuppgrrp__ee [bsz, dha, warmup wurmup_group group num, num heads // dha warmup_growt group_ num, k len, head dim]. These
num _eads k_ k_len, head dim] and the value tensor into values grouuped of shape
18

---

grouped tensors are then expanded by repeating them according to the dha warmup_graee sizup The weights weights k and weights_ v are reshaped and expanded align with the dimensions of the expanded key and value tensors. Flement wise multiplication is performed between the expanded tensors and their corresponding weights, and the resulting weighted tensors are summed along the appropriate dimension.
Algorithm 2 Attention Forward Pass
Require: K
key tensor
Require:
b value tensor
Ensure: K'
> weighted key tensor weighted value tensor
Ensure: V'
I: b 4 batch size
2: + number of heads
3: L& f- key length
4: - head dimension
5: Gg seff.dha_warmup__eg
6: G. &- self,dha_ seff.dha_warrmp_ggonn
7: K self. self.weights_k
8: V self. weightsi
9: Kg K.view(b, Ggs H/Gqs Lks D)
10: V <- Vliew(b, Gg. H/Gg, Lg, , D)
dim = 1) 12: Ve & Vo.repeat interleave G., dim = 1)
11: Ke gg.repeat_inteee
interleave crepeat_interreave((
Kw Kg.view view(1, H. G., D)
VN V V'.view(1, H, G.,
15: K & K. x K
Wy Ve V
17: K' + Wk.sum dim = 2)
V' t Wy.sum 2)
19: return K", V
B.2.3 DHA Loss Calculation
The calculation the loss function 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 an exponential decay term and a linear decay term, ensuring 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 Mana is calculated as:
Maha max 0, (baseelobal x 1.0
global_step
dha @ata_warmup_scr step
MSE Loss are defined in Eq.
7 The adaptive DHA loss Laha is then:
T
Σ max(mse- Maha 0.0)
The overall loss L the adaptive DHA loss:
L Laha + Lim
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.
19

---

Algorithm 3 Adaptive DHA Loss Calculation
Require:
> blocks
Require: mie
Mean Squared Error tensor > loss dha diversity
Ensure: L
X self.a mse
2: - saff.mse_scale
3: I <-<0u 0.0
4: global_step +- 1000
5: Sha__armaap _teeppspppppppe 200
6: base 0.999
7: for each b E B do
8:
Li 0.0
loss_ foss_dha_die lihh__diverity_lle layer D attn atye_layer
9:
A - b.altn
10:
W& + Aweights k
11:
We + A. A.weights_
12:
G. + A.dha sacti__wamapp_aee
13:
G. < A.dha wwrrappparaaaeeo
14:
HE + A.num_kee foim__ee_reee bey_volue_iee
15:
+ A.num heads
D A.head dim
17:
We + reshape( Wk, HKo - G, D)
18:
W. reshape(W., (RR: G, D]
19:
for each r E W do
for each o Wa after do
21:
Li + L, + MSE Loss(r, o)
end for end for
23:
for each r E W. do
25:
for each o E W, after do
26:
L + La + MSE_Loss(n 0)
27:
end for
28:
end for
29: N
H x G. Dx2x (G-11)
L
sxla
N
31:
L+ L + Li
end for
33: L +
(en(B)
34: L
L X
35: expoment +- bascalohal_ast step
36: linear Sineardeecyy +- 1.0 tho_warnupp_ce step
global step
37: margin + max(0, exponent x lincar _dccrddd__d)))_
adaptive loss t max(mse margin, 0.0)
39: L 1.1. L adoplive_lesss
return
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 "ccorette_ccre"" function, which sums the scores from the matrix for each group, considering each connection twice and dividing two.
The algorithm starts with a high temperature (T=100) and gradually cools down (T_min=0.001) using 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.
20

---

If the new score is higher, or if randomly generated number is less than the exponential of delta divided by the temperature, the new grouping accepted. This allows the algorithm 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, use the MSE computed by the head and the head scores, and compute the matrix of scores between the head and the head for head clustering after forward.
Algorithm 4 Head Grouping Optimization on Fusion Loss
Require: M
score matrix >> number of groups final groups and their score
Require: G
Ensure: best_ _os__grouss,.eo
beet_escssreeeeeeeeeeeeeeeesseee score
1: function CALCULATE groups)
2: score 0
3: for each group groups do for each E group do for each group do score < score + M[]]]
4:
5:
6:
end for 8: end for
7:
9: end for
10: return score/2
B each connection counted twice
11: end function
12: function SIMULATED AMNNELLNGGN G)
13:
P - length(M)
number of points number of points per group
14:
N P/G
points array(range(P))
16:
nuuffee(points'
17:
groups poinds.reshapeeG....
18:
purrent_score < CALCULLAT_SCCRRF, groups)
T 100.0
initial temperature b minimum temperature cooling rate
20:
Tmin < 0.001
0.9
22:
while T Tmin do i.j + random integers [0, G)
24:
if i * then
a random integers in [0, NN
26:
new new_groups <- groups.copy(
temp <- new- Brtee_grruus______
28:
new new_groups à [a] new [ro__groopsssl
29:
new new_groups__ 5 f temp
30:
new_score + MALLCULATE_SCORE(M, new_grruuss
31:
+- n-nee__soom - currend_score
32:
if A Oor exp(A/T) > random() then
33:
groups + new groups
34:
current score new_
35:
end if end if
36:
37:
T- T x end while
> cooling down
38:
39: return groups, current_score
40: end function
41: best_groups, best_score + SIMULATED_ANNNEALNNGGMM G)
21

---

B.4 Layer Allocation Based on Fusion Loss
This algorithm efficiently allocates resources different layers based on their respective losses to optimize system performance. Initially, it assigns 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 the remaining allocation. allocates 16s to layers with the highest weights until reaching a predetermined limit. Next, redistributes the remaining allocation to layers with the highest loss-to-an aos-to-allocation ratios, assigning resources in multiples of 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 an efficient distribution of resources across the system. The total search process for the LLaMA2 model requires 42 minutes.
Algorithm 5 Layer Allocation Based on Losses
Require: L
losses for each layer possible allocations
Require: - [4, 8, 16]
Require: T 256
total allocation
Ensure: alloc t- [a1. an anl
b final allocations for each layer
1: length
2: alloc + [4] x n
initial allocation
3: + ΣT.
A weights proportional to losses D remaining allocation initial number of 16's allocate maximum number of 16's that can be allocated
4: R R+ R+T-E T Ealloc
5: k 1
6: M16 +R//16
7: t t-min(k,,u) M16)
8: for 1 do
9: idr t- argmax(W))
10: alloclidz] +- alloc idr) + 12
11: Wiid]] + 0
b prevent reallocation
end for 13: R + -allaoo
14: while 0 do
15:
if' R 8 then
idæ t argmax( attor alloc[idx] t alloc(idx) + 4
17:
R. R-4
19:
else if R 2 4 then
idx t- argmax( allop
21:
alloc[idx] t- alloc(adz) +4
22:
R
23:
end if
end while
25:
return alloc
22

---

B.5
Training Details
The hyperparameters used in our experiments are presented 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 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.
Table 4: Training hyper-parametery
Fusion
Contined Pre-training
Training budget
0.2B
5B
Learning rate of w, À
0.05
Learning Rate of
0.0001
0.0001
LR warmup ratio Batch size (tokens)
10%
3%
262K
1M
Evaluation interval m (steps)
40
40
Steps
800
5, 000
# GPUs
8
8
C Extended Experiments
C.1 Instruction Tuning Evaluation.
Instruction Tuning Evaluation. To assess our models' capabilities in downstream application after instruct tuning [3]] we fine-tune both DHA and the baseline models on 10,000 instruction- response pairs drawn from the initial round of multi-turn chat histories in the ShareGPT datasef" 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. I0 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.
DHA-2.7B-25% vs. GQA-2.7B-25%
DHA-1.38-25% vs. GQA--138-25%
15.75%
72.75%
27.25%
20
60
80
100
20
40
60
80
100
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 and 4 we demonstrated that DHA is 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 convert MHA to 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 MHA which inefficient and bestly.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
https://eeharegpt.com
23

---

Table 5: Comparison of log(PPL) between DHA and GQA with NACL.
Method
log(PPL)
GQA-7b-25%
2.89
DHA-7b-25%
2.84
GQA-7b-25% (NACL 25%)
3,01
DHA-7b-25% (NACL 25%)
2.93
Table 6: Comparison of Avg ACC and PPL between different methods at 7B-25% (5B).
Method
Avg ACC
PPL
DHA-7B-251 (5B)
62.4
7.29
GQA-7B-2544 (5B)
60.3
7.54
GQA CCKKA-Groupingg-77-2555 (5B)
60.4
7.51
and then perform mean- peoaannoooing initialization within each cluster. We compare this approach with the Vanilla GQA and DHA.
Tab. shows that (OA(CKA-Grouping)
(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
where we show the weight variation diagram. In the fusion process of heads 0-3, head initially constitutes 100% 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 slight increase in loss, but significantly.
How Merging Weights Change. Refer to Fig. Ta
DHA's 's Compatibility on GQA Model. DHA is primarily designed for models based on the Transformer Decoder architecture and can be adapted all models with this architecture. We chose LLaMA [6]] the experimental baseline because it is a classic model using the decoder architecture in LLMss. 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 most decoder-only models. GQA [5] an efficient variant of MHA, which optimizes the inference process through head grouping and sharing. Due to its simplicity and efficiency, GQA widely used. DHA can 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 minute, GQA can be losslessly converted into MHA by simply replicating the GQA' KV heads. Then, MHA architecture.
perform the DHA transformation the
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. 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 hetween 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 622 This paper represents an early exploration of applying parameter fusion methods within model parameters. The inter-layer fusion approach 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.
25

---

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 3x4 grid layout, with each row corresponding specific layer and each column corresponding to type of attention mechanism. Note: Layer numbers start from 1.
Liger WC Hold uiaalit
LAME www.eddeiiaeee..
100
100 too
Nod
rea 1entum
N0
0/ res enilut
Lieer Winddds
Wi/oodd
Undariy
Hedd
Figure 11: Visualization of query, key, value head parameters similarity from layer to layer lb in LLaMA2-7B.
26

---

Lyet Paie snaluh
Layer
11
10 Nornwedin wiy
08 unsturo
Layer 20 moaders w/y
900
Loc
Laye
2
27
Lyyer Heat oretriy
20 nonender w/y
Leyer aniieyy
Lyyer Wwneedd
Figure 12: Visualization of query, key, value head parameters similarity from layer layer 32 in I.aMM2-7B.
27

---

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. 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.
Lover Head pindlcity
Leven recroed sinderty
Lraen We Heat
Laynt WC srabrryy
Liget 1W reck-andonn
Layer EdDiiaddinion miartty
Layer NX Fasd snalery
Laer A/0 read unaartt
Layer Encesed candert
Uper EW aneirt
LovrAMo
Figure Visualization of query, key, value head parameters similarity from layer to layer 8 DHA-7B
28

---

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]
The abstract provides a concise summary of the key findings and experiment results. The introduction in Sec. outlines the research questions and objectives in paragraph 3,4 and contribution in paragraph 5.
Justification:
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 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 long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations 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 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails 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 fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant 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
29

---

Question: For each theoretical result, does the paper provide the full set of assumptions and 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. [ presents the conjecture in Sec. 3 and provides a detailed derivation of the form and opti-
mization process in Sec. 4 All assumptions made in the paper are thoroughly validated through experiments in Sec.
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 cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide 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
Does the paper fully disclose all the information needed to reproduce the main perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification:
The paper provides a detailed description of the experimental data setup and
hyperparameter settings in Sec. Additionally, in Sec. Dec. 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, 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 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 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 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 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 submis- sions 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
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
30

---

(d)
The authors should provide scripts to reproduce all experimental results for the new reproducible, they We recognize that reproducibility may be tricky in some cases, in which case authors are welcome 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 instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
The datasets, baseline methods, and models used in the paper are fully open- source and available on Hugging Face. The paper includes the key implementation steps and code in Sec. f and the Appendix. [B22] However, the complete code is still being organized and is under consideration for open sourcing.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
Bablic/guides/CCe http://eeuuiiiiiiooooooiiiioiiie
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for 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
/pips..ccppulic//guue wwwpc///guiesss/eeee /
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
proposed method and baselines. If only subset of experiments 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 supplemental material (appended the paper) 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 (c.g., data splits, hyper- parameters, 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. Additionally, in Sec. and in Appendix. B.2 sections 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 the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
31

---

Justification: All results are averaged over multiple tests, and we report the mean accuracy along with the standard deviation (ace norm) measure of error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi-
dence 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).
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 should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful 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 the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
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 well 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 htttp:://eeurppe http:///////oo.....///////////// wwtppublic/cciiccuuiiiieeaeceeee
Answer:[Yar Yes]
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.
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
32

---

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 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 direct path to any negative applications, the authors should point it out. For cxample, it is legitimate to point out that 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 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 unintentional) misuse of the technology.
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 system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
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 such risks.
Released models that have 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 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make 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.
33

---

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 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 released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 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 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 crowdsoureing 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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, 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]
34

---

Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending 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.
recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
35


Reviews:
[{"review_id": "uckc1B0mjf", "reviewer": "Reviewer_aEta", "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.", "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? \n2. Are there any experimental results available that demonstrate how the DHA model performs on larger datasets?", "limitations": "The  author has mentioned limitations in Appendix", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n2. There is a noticeable accuracy gap between DHA and MHA models on academic datasets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer aEta\n\nThank you for your appreciation to our work and insightful comments.\n\n**Q1: Comparison with GQA and MHA Training Data Requirements in Large Dataset.**\n\n**R1:** Under the same training budgets, DHA performs better than GQA.\n\n- **Steady improvement of DHA over GQA in limited training budget setting**. In Figure 6, we showed that in the setting of limited 5B training budget, DHA effectively and consistently outperforms GQA on downstreming tasks. At the end of the 5B training, the optimization is converging while the performance gap remains.\n\n- **Improvement remains with more training budgets**. We have supplemented our experiments with more data budgets, and the results are shown below. Our method consistently achieves better perplexity (PPL) than GQA. In the final version, we will include experiments with more training budgets.\n\n| Model\\\\ Training Tokens  | 5B  | 7.25B | 10B |\n|-------------------|-----|-------|-----|\n| GQA-7b-25\\% (PPL)  | 7.41| 7.15  | 6.72|\n| DHA-7b-25\\% (PPL)  | 7.2 | 7.03  | 6.61|\n\nThe superior performance of DHA over GQA is mainly due to two factors: \n1. DHA retains more important information through linear fusion; \n2. The DHA architecture is more efficient than GQA. By allocating more parameters to more important components, training directly with the DHA architecture also yields better performance than GQA. Refer to the [Anonymous Figure](https://mjj.today/i/jqaTqd).\n\n**Q2: Accuracy gap between DHA and MHA models on academic datasets.**\n\n**R2:** The performance gap between the results shown in the paper and MHA is primarily due to the following two reasons:\n\n- **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). In Table 1, we directly reported the performance of original MHA model (LLaMA2-7B) without continue-training as we found the performance decline of the MHA model when continue-training on the open-sourced pre-training data. Compared to S.-LLaMA-1.3B trained on the same dataset, our method only loses 3.5% when compressed by 50% with only 5B data budget. In other words, the improved pre-training data will close the gap between DHA and MHA. In the table below I provide a performance comparison of MHA after pre-training on the same data.\n\n| Model               | SciQ | PIQA | WinoGrande | ARC-E | ARC-C (25) | HellaSwag (10) | LogiQA | BoolQ (32) | LAMBADA | AVG  |\n|---------------------|------|------|------------|-------|-------------|----------------|--------|-------------|---------|------|\n| LLaMA2-7B           | 94.1 | 78.1 | 69.1       | 76.3  | 49.7        | 58.9           | 25.7   | 80.8        | 74.1    | 67.4 |\n| LLaMA2-7B-CT(5B)    | 93.7 | 77.9 | 68.8       | 76.1  | 49.2        | 58.3           | 23.7   | 79.2        | 73.8    | 66.7 |\n| GQA-7B-50% (5B)     | 90.8 | 76.2 | 64.8       | 69.2  | 41.6        | 52.8           | 22.5   | 70.5        | 65.2    | 61.5 |\n| DHA-7B-50% (5B)     | 93.8 | 76.5 | 66.6       | 72.0  | 43.8        | 55.3           | 21.4   | 75.7        | 67.2    | 63.6 |\n| GQA-7B-25% (5B)     | 89.5 | 75.7 | 61.9       | 68.0  | 38.9        | 50.9           | 23.8   | 67.1        | 64.1    | 60.0 |\n| DHA-7B-25% (5B)     | 91.7 | 76.8 | 64.4       | 70.9  | 42.8        | 54.2           | 21.8   | 74.6        | 68.4    | 62.8 |\n\n- **In the setting of limited training budget**. In Table 1, we presented DHA's experimental results which reflect the early training phase with about only 5B tokens, not the performance of Sheared-LLaMA after 50B tokens. More training budget can be allocated to DHA to further improve its performance.\n\n- **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.\n\n**Q3: Clarify why an increase in the fusion phase token budget leads to a decrease in average accuracy post-CT.**\n\n**R3:** Table 3 shows how the budget is allocated to fusion and CT throughout the training phase. Since DHA fusion is very efficient and fast, allocating only 0.1B of the fusion budget can provide a good starting point, after which pre-training can quickly restore the model's capabilities. If we allocate more (0.2B) to the fusion budget, we can achieve an even better starting point, but within the total constraint of a 5B budget, we have less budget (4.8B) for CT. When the starting point is sufficiently good, increasing the CT budget is a better choice. Overall, continuously allocating the fusion phase token budget will not degrade performance but will raise overall training costs.\n\n---\n(Xia et al., 2024) Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning. ICLR 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9hAB1jG73A", "reviewer": "Reviewer_FLzu", "summary": "The 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.", "questions": "Please answer my questions listed in \"Weaknesses\"", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer FLzu\n\nThank you for your appreciation to the novelty of our work and the thoughtful reviews. \n\n**Q1:DHA's Compatibility with GQA**\n\n**R1:** Thank you for your suggestions! Here, we provide two feasible methods to convert GQA to DHA.\n\n- **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 MHA architecture.\n\n- **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.\n\n**Q2: The performance of DHA in the end of the training process.**\n\n**R2:** Under the same training budgets, DHA performs better than GQA.\n\n- **Steady improvement of DHA over GQA in limited training budget setting**. In Figure 6, we showed that in the setting of limited 5B training budget, DHA effectively and consistently outperforms GQA on downstreming tasks. At the end of the 5B training, the optimization is converging while the performance gap remains.\n\n- **Improvement remains with more training budgets**. We have supplemented our experiments with more data budgets, and the results are shown below. Our method consistently achieves better perplexity (PPL) than GQA. In the final version, we will include experiments with more training budgets.\n\n| Model\\\\ Training Tokens  | 5B  | 7.25B | 10B |\n|-------------------|-----|-------|-----|\n| GQA-7b-25\\% (PPL)  | 7.41| 7.15  | 6.72|\n| DHA-7b-25\\% (PPL)  | 7.2 | 7.03  | 6.61|\n\nThe superior performance of DHA over GQA is mainly due to two factors: 1. DHA retains more important information through linear fusion; 2. The DHA architecture is more efficient than GQA. By allocating more parameters to more important components, training directly with the DHA architecture also yields better performance than GQA. Refer to the [Anonymous Figure](https://mjj.today/i/jqaTqd).\n\n\n**Q3: Accuracy loss after transformation.**\n\n**R3:** The performance gap between the results shown in the paper and MHA is primarily due to the following two reasons:\n\n- **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). In Table 1, we directly reported the performance of original MHA model without continue-training as we found the performance decline of the MHA model when continue-training on the open-sourced pre-training data. In other words, the improved pre-training data will close the gap between DHA and MHA.  In the table below I provide a performance comparison of MHA after pre-training on the same data.\n\n| Model               | SciQ | PIQA | WinoGrande | ARC-E | ARC-C (25) | HellaSwag (10) | LogiQA | BoolQ (32) | LAMBADA | AVG  |\n|---------------------|------|------|------------|-------|-------------|----------------|--------|-------------|---------|------|\n| LLaMA2-7B           | 94.1 | 78.1 | 69.1       | 76.3  | 49.7        | 58.9           | 25.7   | 80.8        | 74.1    | 67.4 |\n| LLaMA2-7B-CT(5B)    | 93.7 | 77.9 | 68.8       | 76.1  | 49.2        | 58.3           | 23.7   | 79.2        | 73.8    | 66.7 |\n| GQA-7B-50% (5B)     | 90.8 | 76.2 | 64.8       | 69.2  | 41.6        | 52.8           | 22.5   | 70.5        | 65.2    | 61.5 |\n| DHA-7B-50% (5B)     | 93.8 | 76.5 | 66.6       | 72.0  | 43.8        | 55.3           | 21.4   | 75.7        | 67.2    | 63.6 |\n| GQA-7B-25% (5B)     | 89.5 | 75.7 | 61.9       | 68.0  | 38.9        | 50.9           | 23.8   | 67.1        | 64.1    | 60.0 |\n| DHA-7B-25% (5B)     | 91.7 | 76.8 | 64.4       | 70.9  | 42.8        | 54.2           | 21.8   | 74.6        | 68.4    | 62.8 |\n\n- **In the setting of limited training budget**. In Table 1, we presented DHA's experimental results which reflect the early training phase with about 5B tokens, not the performance of Sheared-LLaMA after 50B tokens. More training budget can be allocated to DHA to further improve its performance.\n\n- **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.\n\n**Q4: DHA combines with existing KV cache compression techniques.**\n\n**R4:** DHA is an efficient GQA architecture, so it has similarly good compatibility. In Sections 2 and 4.1, we demonstrated the relationship between DHA and GQA. We tested the compatibility of the DHA model with the KVCache eviction method NACL (Chen et al., 2024). \"NACL 25\\%\" indicates retaining only 25\\% of the KVCache. The experiment results are shown in the table below. DHA and GQA exhibit equally good compatibility with KV cache compression techniques. We will include more experiments in the final version.\n\n| Method                           | log(PPL) |\n|-------------------------------|----------|\n| GQA-7b-25\\%                    | 2.89     |\n| DHA-7b-25\\%                    | 2.84     |\n| GQA-7b-25\\% (NACL 25\\%)         | 3.01     |\n| DHA-7b-25\\% (NACL 25\\%)         | 2.93     |\n\n\n---\n(Xia et al., 2024) Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning. ICLR 2024.\n\n(Chen et al., 2024) NACL: A General and Effective KV Cache Eviction Framework for LLMs at Inference Time. ACL 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vmWmcE3qIu", "reviewer": "Reviewer_Txec", "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.", "questions": "See the weaknesses.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "-The paper demonstrates a remarkable balance between performance and efficiency, showing that DHA can achieve high performance with significantly reduced computational resources.\n\n-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?\n\nThe 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.\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer Txec\n\nThank you for your appreciation to the novelty of our work and the thoughtful reviews. \n\n**Q1:The effectiveness of DHA on a wider range of model architecture**\n\n**R1:** Thanks for your suggestions! 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 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. Thank you for your suggestion. Due to time constraints, we will include experiments with other models in the final version.\n\n**Q2: Complex to implement and reproduce for researchers without access to similar computational resource.**\n\n**R2:** DHA is primarily used in the pre-training phase of very large models and can save a considerable amount of training costs compared to the widely used GQA. For researchers without access to similar computational resources, we can consider the following methods:\n\n- Conduct experiments with the DHA architecture on smaller models before applying it to large-scale models.\n- The efficient architecture of DHA can serve as a reference, and adopting the DHA configuration in model design can directly accelerate training speed. Refer to the [Anonymous Figure](https://mjj.today/i/jqaTqd), when not constructed based on MHA, both DHA and GQA start training from random initial points, and DHA demonstrates a faster training speed than GQA.\n- Models trained with DHA can be combined with other optimization methods.\n\n\n**Q3: Exploring non-linear fusion techniques.**\n\n**R3:** Thank you for your valuable suggestions! To our best knowledge, DHA appears to be the initial attempt to explore learning linear fusion within models for fine-grained LLM's parameter compression, and we believe it holds significant value in the current academic field. We are continuing to explore this in our subsequent research, including using nonlinear fusion methods for more refined integration.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RHoZr67xrU", "reviewer": "Reviewer_wgtQ", "summary": "The 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.", "questions": "* 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.", "limitations": "The paper adequately addresses the limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* 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.", "weaknesses": "* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Response to Reviewer wgtQ\n\nThank you for your thorough review and positive comments on our work.\n\n**Q1: Inter-layer grouping of heads or only intra-layer grouping?**\n\n**R1:** **Only intra-layer grouping and fusion** is conducted in DHA. Figure 1 meant to illustrate the *decoupled*-heads where the number of key and value heads can be different among layers. We apologize for any misunderstandings and will improve the Figure 1 and intro. part for better clarification. \n\nThe DHA method employs parameter fusion within each layer for two reasons: \n\n- **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.\n\n- **More complex optimization for inter-layer fusion**. The optimization process between layers is very complex and, as you mentioned, requires memory operations for cross-layer calls, which inherently increases the inference cost. \n\n- **Promising future work by introducing inter-layer fusion**. This paper represents an early exploration of applying parameter fusion methods within model parameters. The inter-layer fusion approach you suggested is indeed a valuable direction for future exploration. \n\n**Q2: GQA method's initialization in Experiments.**\n\n**R2:** It's a common and effective approach to convert MHA to GQA using mean pooling instead of training from scratch. The author of GQA (Ainslie et al., 2023) 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 MHA which is inefficient and costly. \n\nIn response to the new baseline issue, we provided the experimental result as follows:\n\n- **CKA-Grouping, then mean pooling**. We have considered grouping the heads with high similarity using CKA, then performed GQA's mean pooling initialization as a more solid and strong baseline. However, this approach only achieved comparable performance to the original implementation in GQA (Ainslie et al., 2023) paper. See table below:\n\n| Method           | DHA-7B-25\\% (5B) | GQA-7B-25\\% (5B) | GQA（CKA-Grouping）-7B-25\\% (5B) |\n|----------------|-----------------|-----------------|-------------------------|\n| Avg ACC      | 62.4            | 60.3            | 60.4                    |\n| PPL          | 7.29            | 7.54            | 7.51                    |\n\n\n% - 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.\n\n\n**Q3: DHA's Training overhead compared with GQA.**\n\n**R3:** As mentioned in our response to Question 2, both DHA and GQA are constructed based on MHA. The difference is that DHA employs a training method, whereas GQA uses pooling.\n\n- **Construction based on MHA**. DHA achieved comparable performance to GQA using only 0.2B tokens, whereas GQA requires 12 times the amount of training data to reach the same level of loss as DHA. (see line 253). This is because DHA effectively retains information through its fusion method.\n\n- **Training from scratch**. We supplement the experiment in the setting of training from scratch, see the [Anonymous Figure](https://mjj.today/i/jqaTqd). DHA demonstrates a faster training speed than GQA. This is sufficient to prove that the DHA architecture is more efficient in training.\n\n**Q4: DHA's conversion process requires task-oriented training using real data?**\n\n**R4:** You are correct that DHA requires data to learn the fusible relationships between heads, but it needs only a small portion (0.25\\%) of the total training data. DHA is primarily used in the pre-training phase of very large models and can save a considerable amount of training cost compared to the widely used GQA.\n\n\n**Q5: Time in DHA's Search Step.**\n\n**R5:** Using the experimental setup described in the paper, the search process for the LLaMA2 model requires 42 minutes.\n\n**Q6: Providing a demonstration of how merging weights change.**\n\n**R6:** Thank you for your suggestions. You can refer to Figure 3(a), 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. We will include more relevant information in the final version‘s appendix.\n\n**Q7: Visualizes the weight similarity of CKA.**\n\n**R7:** We apologize for any misunderstanding. The parameter similarities defined in Figures 2, 3b, and 9 are calculated using CKA, as noted in lines 93-94 and Appendix B.1. Based on this consistent measurement method, we can clearly see that DHA is inspired by similarity redundancy and ultimately achieves a similar effect through merging. We will provide more detailed explanations in the final version.\n\n\n**Q8: Typos**\n\n**R8:** Thank you for your careful proofreading! We will address this in the final version.\n\n---\n\n(Ainslie et al., 2023) GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This submission received mixed reviews before the rebuttal phase. Reviewers are in general treating the idea in this paper as novel and interesting, but are not satisfied by various aspects, in particular, experimental validations. During the rebuttal phase, the authors answers most of these concerns. Further discussions after the rebuttal and author-reviewer discussion phases also cleared some remaining concerns from the reviewers.

In the end, 3 reviewers upgraded their scores for this submission, and all 4 reviwers are suggesting or inclined to accept this submission.
Decision: Accept (poster)