MoH: Multi-Head Attention as Mixture-of-Head Attention
Authors: Peng Jin
Bo Zhu
Li Yuan
Shuicheng YAN
Venue: ICML 2025 poster (2025)
Keywords: Multi-Head Attention
Mixture of Experts
Foundation Models
Primary area: deep_learning->attention_mechanisms
TL;DR: We propose Mixture-of-Head attention (MoH) that outperforms multi-head attention even by using only 50%~90% of the attention heads.

Abstract: In this work, we upgrade the multi-head attention mechanism, the core of the Transformer model, to reduce computational costs while maintaining or surpassing the previous accuracy level. We show that multi-head attention can be expressed in the summation form. Drawing on the insight that not all attention heads hold equal significance, we propose Mixture-of-Head attention (MoH), a new architecture that treats attention heads as experts in the Mixture-of-Experts (MoE) mechanism. MoH has two significant advantages: First, MoH enables each token to select the appropriate attention heads, enhancing inference efficiency without compromising accuracy or increasing the number of parameters. Second, MoH replaces the standard summation in multi-head attention with a weighted summation, introducing flexibility to the attention mechanism and unlocking extra performance potential. Extensive experiments on ViT, DiT, and LLMs demonstrate that MoH outperforms multi-head attention by using only 50%$\sim$90% of the attention heads. Moreover, we demonstrate that pre-trained multi-head attention models, such as LLaMA3-8B, can be further continue-tuned into our MoH models. Notably, MoH-LLaMA3-8B achieves an average accuracy of 64.0% across 14 benchmarks, outperforming LLaMA3-8B by 2.4% by utilizing only 75% of the attention heads. We believe the proposed MoH is a promising alternative to multi-head attention and provides a strong foundation for developing advanced and efficient attention-based models.

MoH: Multi-Head Attention as Mixture-of-Heat Attention
Peng Jin 123
Bo Zhu
Li Yuan
1235 Shuicheng Yan
Abstract
https//iittabbcmm/Syyywwt//////
In this work, we upgrade the multi-head attention mechanism, the core of the Transformer model, to reduce computational costs while maintain- ing or surpassing the previous accuracy level. We show that multi-head attention can be ex- pressed in the summation form. Drawing on the insight that not all attention heads hold equal significance, we propose Hixture-of-Hean atten- tion (MoH), new architecture that treats atten- tion heads as experts in the Mixture-of-Experts (MoE) mechanism. MoH has two significant ad- vantages: First, MoH enables each token to se- lect the appropriate attention heads, enhancing inference efficiency without compromising ac- curacy or increasing the number of parameters. Second, MoH replaces the standard summation in multi-head attention with a weighted summa- tion, introducing flexibility to the attention mech- anism and unlocking extra performance potential. Extensive experiments on ViT, DiT, and LLMs demonstrate that MoH outperforms multi-head tention by using only 50%~~90% of the attention heads. Moreover, we demonstrate that pre-trained multi-head attention models, such as LLaMA33 8B, can be further continue-tuned into our MoH models. Notably, MoH--LaMA3-8B achieves an average accuracy of 64.0% across 14 bench- marks, outperforming ILLMM3-8B by 2.4% by utilizing only 75% of the attention heads. We believe the proposed MoH is promising alterna- tive to heal attention and provides a strong foundation for developing advanced and efficient attention- models. The code is available at
1. Introduction
Since attention is introduced and becomes fundamental component of Transformers ( (Vaswani et al., 2017), multi- head attention has been the standard architecture for natural language processing (Kenton & Toutanova, 2019) and com- puter vision tasks (Dosovitskiy al., 2021). It well known that using multiple heads can improve model accuracy. How- ever, not all attention heads hold equal significance. Some works have shown that many attention heads can be pruned without affecting accuracy. For example, Voita et al. (2019) introduces a method to quantify the usefulness each at- tention head and prune those that are redundant. Similarly, Michel al. (2019) challenges the necessity of multiple heads by examining the impact of extensive pruning across various settings. In computer vision, some works also iden- tify attention head redundancy. Bhattacharyya et al. (2023) reduces redundancy to boost performance, while Yun & Ro (2024) develop single-head attention for efficiency. These findings demonstrate that vanilla multi-head attention con- tains redundant attention heads.
Besides, in multi-head attention, each head operates in par- allel, and the final output is the sum of all heads (please refer to Section 3.1). Given that these attention heads op- erate independently and some may be redundant, we argue that is possible to build a dynamic attention-head routing mechanism. Such a mechanism would enable each token to adaptively select the appropriate attention heads, enhancing inference efficiency without compromising accuracy.
To this end, we introduce Mixture-of-Heen Heed attention (MoH), new architecture that integrates multi-head attention with the dixture-of--xperr (MoE) mechanism (Jacobs et al., 1991). Specifically, we propose to treat attention heads as ex- perts within the MoE framework. Similar MoE, MoH con- sists of multiple attention heads and router that activates the Top-K K heads for each token. Moreover, we replace the standard summation in multi- head attention with weighted summation. This design offers two significant advantages: First, MoH allows each token to select the most relevant attention heads, improving inference efficiency without sac- rificing accuracy or increasing the parameters. Second, by replacing the standard summation in multi-head attention
'School of Electronic and Computer Engineering, Shen- zhen Graduate School, Peking University, Shenzhen, China PPnggheng Laboratory, Shenzhen, China School of AI for Sci- ence, Shenzhen Graduate School, Peking University, Shenzhen, China "Skywork AI, Singapore 'Rabbitpre Intelligence, Shenzhen, China Natiooaal University of Singapore. Singapore. Correspon- dence to: Li Yuan <yuanli-cor pku.edu.cn2, Shuicheng Yan <shuicheng yan com..
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Moli: Multi-Head Attention as Mixture-of-Heet Head Attention
Head
...
Head..
He
Heal,
Head
...
Heal
Headh
Head
Head
Heedly
Shared Heacs
Heari
Attention
Attention
Routed Heads
Head
...
Head
Head,
Heal
...
Head
Headh
Headyy
Head
Heeely
Head
Input O0
Input
O0
Router
(a) Mullt--FeaddAttenniion
(b) Our proposed Mixtue Head
Figure 1. A high-level comparison between the multi-head attention and our proposed mixture-of- head attention. Subfigure (a) illustrates a standard multi-head attention layer with attention heads, while subfigure (b) demonstrates our proposed Mixture-of-Heat attention (MoH) architecture. It is important to note that MoH does not increase the number of attention heads, ensuring that the total parameter Moll is comparable to that of multi-head attention.
with weighted summation, MoH enhances the flexibility of the attention mechanism and increases the performance po- tential. Moreover, to efficiently capture common knowledge across different contexts, we designate a subset of attention heads as shared heads that remain always activated.
In addition to training from scratch, demonstrate that pre-trained multi-head attention models, such as LLaMA3-8B, can be further continue-tuned into our MoH models, greatly enhancing the applicability of the proposed MoH method.
We evaluate our proposed MoH across various pop- ular model frameworks, including Vision Transform- (ViT) (Dosovitskiy et al., 2021) for image classification, Diffusion models with Transformers (DiT) (Peebles & Xie, 2023) for class-conditional image generation, and Large Language Models (LLMs) (Brown et al., 2020; OpenAI, 2022; Ouyang et al., 2022). We show that MoH achieves competitive performance, or even outperforms multi-head attention with only 50%~90% of the attention heads. For ample, MoH-VTT-B achieves 84.9% Top-1 accuracy on the ImageNet-I IK (Deng al., 2009) classification bench- mark, surpassing well-tuned multi-head attention baselines with only 75% 750% of the attention heads.
Extensive experiments across various popular model frameworks, including ViT, DIT, and L.LMs, confirm that MoH is a promising alternative to vanilla multi- head attention, laying a solid foundation for developing advanced and efficient attention-based models.
2. Related Work
Multi-Head Attention. Transformers (Vaswani et al., 2017) have garnered significant interest and success in both natural language processing and computer vision. The suc- cess of transformers has been long attributed to the multi- head attention mechanism (Cordonnier et al., 2020). Multi- head attention mechanism is proposed by Vaswani et al. (2017) to enhance the representation power of an attention layer by allowing multiple attention heads to operate on dif- ferent now-dimensional projections of the input. The outputs from these heads are then concatenated to form the final result. Alternatively, by decomposing the output projection matrix by rows, multi-head attention can be expressed in summation form. In summation form, each head oper- ates in parallel, and the final output is the sum of all heads. Inspired by this observation, we propose MoH, a dynamic attention-head routing mechanism that allows each token to adaptively select the appropriate heads.
Furthermore, we demonstrate that pre-trained multi-head attention models, such as LLaMA3-8B (Dubey et al., 2024), can be further continue-tuned into our MoH models. Specif- ically, using only about 3% (400B tokens) of the origi- nal LLaMM3 pre-training data for continue-tuning MoH- LLaMA3-8B achieves an average accuracy of 64.0% across 14 benchmarks, outperforming LLaMA3-8B by 2.4% by utilizing only 75% of the attention heads. These results show that MoH is a promising alternative to vanilla multi- head attention, laying solid foundation for developing advanced and efficient attention- based models. The main contributions are summarized as follows:
Mixture-of-Experts Experss Models. The Mixture-of of -f-Experts (MoE) method (Du et al., 2022; Lewis et al., 2021; Rajbhan- dari et al., 2022; Roller et al., 2021; Zhou et al., 2022; Jin et al., 2025) is introduced to expand the capacity of deep neural networks without increasing computational costs. In this approach, only a subset of parameters, known as ex-
We propose dynamic attention-her head routing mech- anism that allows each token to adaptively select the appropriate attention heads, enhancing model perfor- mance and inference efficiency without increasing the number of parameters.
2

---

MolI: Multi-Head Attention as -ixtuaeeof--ead Attention
perts, is activated for each input. Shazeer et al. (2017) first introduces MoE. layer between LSTM layers. Switch Transformer (Fedus et al., 2022) further simplifies the gat- ing mechanism by selecting only the Top-1 expert per token. Gshard (Lepikhin et al., 2021) improves the Top-2 expert routing strategy. In contrast to MoE, which emphasizes efficient parameter scaling while maintaining manageable computational costs, our MoH focuses on reducing the acti- vation of redundant attention heads without increasing the number of parameters.
where Wa E Rduxde, WK E Rdinxde and Wy Rdin represent the projection matrices for the query, key, and value, respectively. In self-attention, the input to- kens are the same, i.e., X' = X, and it is common for the key and value dimensions to be equal, i.e., de = dk-
Concatenation Form. To enhance the representation power, Vaswani et al. (2017) proposes to allow multiple attention heads to operate on different low-dimensional pro- jections of the input tokens. Specifically, the multi-head attention mechanism computes 1 different low-dimensional projections of (Q, K, V). performs scaled dot-product at- tention for each head, concatenates the results, and applies a projection to the concatenated output. The concatenation form of the multi-head attention can be formulated as:
Attention Head Specialization and Efficiency. Many recent studies show that not all attention heads in Trans- formers are equally useful. Peng et al. (2020) proposes a mixture-of-heads approach, where only a few selected heads are used, yet the model performs just as well or even better. Csordás et al. (2024) pushes this further with SwitchHead, an MoE-style method that activates only small number of heads for each token, speeding up inference while keeping performance high. In long-contex( language models, this idea is even more clear. Wu et al. (2024) shows that a few special retrieval heads are mainly responsible for keeping facts consistent in long inputs. Fu et al. (2024) finds that keeping only the most useful heads in the KV cache can save memory. Xiao et al. (2024) proposes AuoAttention, which combines different types of heads to make long-contex. inference more efficient without losing quality. Similar pat- terns appear in vision models. Gandelsman et al. (2023) shows that CLIP's attention heads each focus on specific visual features, and this can be explained through related text prompts. Balasubramanian et al. (2024) finds that this kind of head specialization also exists in other vision models beyond CLIP. Basile et al. (2024) shows that using only few selected heads chosen by spectral methods can even beat the full model on zero-shot tasks.
MultiHead( X, x') Concat( (H', ,,H2, 3... H H' = Attention ( W X'W X'W),,
(2) where W E Rd Wk E Rdinxde/ and Wi Rd x /d/// represent the ith projection matrices for the query, key, and value, respectively. Wo E pdoxdout is the final output projection matrix.
Summation Form. The multi-head attention mechanism is typically represented in its concatenation form. How- ever, from another perspective, if we decompose Wo E Rd xdaut by rows, we can express multi-head attention in a summation form. Specifically, Wo can be divided into h matrices by rows, i.e., [W3, W W = Wo. where W E Rdd/hxduue Finally, the summation form of the multi-head attention can then be formulated as:
MultiHead(XXX) x')
Σ H'W:
(3)
The concatenation form can be viewed as a variant of the summation form, where the sum of the dimensions of all attention heads is exactly equal to the hidden size. As shown in Eq. 3, in standard multi-head attention, each attention head operates in parallel, and the final output is the sum of all attention heads. Since these attention heads func- tion independently, we can build dynamic attention-head routing mechanism allowing each token to adaptively se- lect the most relevant attention heads, improving inference efficiency without compromising accuracy.
3. Methodology
In this work, we aim reduce the activation of redundant attention heads without increasing the number of parameters. A high-level comparison between the vanilla multi-head attention and our proposed MoH is presented in Fig. 1.
3.1. Multi-Head Attention
We begin by reviewing the multi-head attention mechanism introduced by Vaswani et al. (2017). The multi-head atten- tion mechanism is based on scaled dot-product attention. Specifically, for T tokens X E x of din dimensions each and T' tokens X' = T' of din dimensions each, the scaled dot- dot-product attention is computed as follows:
3.2. Mixture-of-Head Attention
Recently, the Mixture-of-Expeee (MoE) method has emerged as popular approach for scaling the parameters of large language models (Jiang et al., 2024; Muennighoft et al., 2024). A MoE layer consists of multiple expert networks and a router that activates the Top-K experts. Generally, the number of activated experts K is significantly smaller than the total number of experts to ensure inference efficiency.
QK
Attention Q K, V) = =Softmxx
V.
VoT
Q K -''K, V = X'Wv,
3

---

MolI: Multi-Head Attention as Hixture-of-Head Attention
Table Comparisons current state-of-the-art methods on ImageNet-IK classification. Our MoH- VIT models, based on TransNeXt (Shi, 2024), are trained for 300 epochs using a resolution of 224x To ensure a fair comparison, we only replace the standard multi-head attention with our Mixture-of Head attention (MoH), keeping all other training parameters identical to Texts
#Params (M)
#Activated Heads (%)
Acc (%)
#Params (M)
#Activated Heads (%)
Acc (%)
Methods
Methods
DeïT-S (Tomvron et al., 2021)
22
100
79.8
DeïT-B (Tomsron et 2021)
86
100
81.8
T2T-ViT-19 (Yom at al., 2021)
39
100
81.9
TIT-V (Yoin at al., 2021)
100
82.3
Swin-S (Lia et 2021)
50
100
83.1
Swin-B (Lis el al., 2021)
88
100
83.5
PVTv2-B3 B3 (Wing at.l. 2022)
45
100
83.2
PVTv2-B3 (Wang al. 2022)
82
100
83,8
CoAtNet-1 1 etii 2021)
42
100
83.3
Focal-B (Yang al. 2021)
90
100
83.8
Focal-S (Yang 2021)
51
100
83.5
FocalNet-B (Yang et al., 20226)
89
100
83.9
FocalNet-S (Yang 20226)
50
100
83.5
CoArNet-2 (Dhi 202))
100
84.1
MViTv2-S (1. et 2022)
35
100
83.6
MViTv2-B (1. al., 2023)
100
84.4
UniFormer-B (Li al, 20236)
50
100
83.9
MOAT-2 (Yagg et al., 2022a)
73
100
84.7
CAFormer-S36 (Ys al, 2023)
39
100
84.5
iFormer-L (Si 2022)
87
100
84,8
TransNeXt-S (Shi, 2024)
50
100
84.7
TransNeXt-B B (Shi, 2024)
90
100
84.8
MoH-ViT-S
50
80
84.7
MoH-ViT-B
90
75
84.9
MoH-VIT-S
50
75
84.6
MoH--Vi-BB
90
50
84.7
Heads as Experts. Inspired by the great success of MoE, we propose Mixture-ou of- Head attention (MoH), which treats attention heads as experts. Specifically, MoH consists of h heads H = { !!.,.., } and a router that activates the Top-K attention heads. Formally, given input tokens and X', the output of MoH the weighted sum of outputs from the K selected attention heads:
R Txdia the routing score gi is defined as:
a Softmax W, ,aris if Sissi has if Head i is activated, otherwise,
9: =
Softmax W W.z:)i-hu
0,
where hy denotes the number of shared heads. W. & and W, (hhhh.)xdi, represent the projection matrices for the shared and routed heads, respectively. If lh, + 1 s h)),
MoH(X, Σ 9. H'W.
(4)
(W,zt)i-h, E Mop-K([(W,aa))) then the routed Head i is activated. The coefficients (r and a2 balance the contributions of the shared and routed heads, and are defined as:
where 9: represents the routing score. Ji is non-zero only when the in attention head is activated. This design pro- vides two key advantages: (i) On the one hand, MoH ables each token to select the most relevant attention heads, boosting inference efficiency while maintaining accuracy. (ii) On the other hand, in contrast to the standard summa- tion in multi-head attention, the weighted summation in MoH enhances the flexibility of the attention mechanism and unlocks performance potential.
[n, 02] Softmax Softmax(WW.ar))
(6)
where WA, 2 is the trainable projection matrix, and din is the hidden size of æt
Load Balance Loss Directly training an MoE layer often causes the majority of tokens to be routed to a small num- ber of experts, leaving the remaining experts insufficiently trained (Shazeer et al., 2017). To avoid the unbalanced load in the proposed MoH, following previous MoE meth- ods (Lepikhin et al., 2021; Wei al., 2024), apply load balance loss. Specifically, for the teh. input token E Rd.... in X E id'' the load balance loss La is formulated as:
Shared Heads. In attention mechanism, some attention
heads may capture common knowledge across different con- texts, such as grammatical rules in language. Inspired by Dai et al. (2024), we designate a subset of heads as shared heads that remain always activated. By consolidating com- mon knowledge within shared heads, we reduce redundancy among the other dynamically routed heads.
h Σ s=h,+1
1 1 Σ
T
C.
Pife P.
Softmax (W, W,
T
Two-Stage Routing. Moreover, to dynamically balance the weights between shared and routed heads, we propose a two-stage routing strategy. In this routing strategy, the routing scores are determined by both the score of each individual head and the score associated with the head type. Specifically, given the Lth input token at E dd in X E
f.
1 (Token æ, selects Head i),
T
(7)
where T denotes the number of tokens. 1(*) denotes the indicator function.

---

MoHI: Multi-Head Attention as Mixture-of-Iake -f-ae Attention
Table Comparisons to current state-of-the methods on the benchmarking class-conditional image generation on ImageNet- IK at 2 256 resolution. "I" denotes that higher is better. "I" denotes that lower is better. "cfg" denotes the classifier freee diffusion guidance scale. "400K" denotes the training budget is 400K training steps.
Methods
#Params (M)
#Activated Heads (%)
FID4
sFID),
IS*
Precision?
Recall
DiT-S/2 400K (Prebles 2023)
33
100
68.40
MolI-DiT-S/2 400K
33
90
67.25
12.15
20.52
0.37
0.58
MoH-DDT--/2 400K
33
75
69.42
12.85
19.96
0.36
0.55
DiT-B/2 400K (Peshles
130
100
43.47
MoH-DiT-B/2 400K
131
90
43,40
8.40
33.51
0.49
0.63
MolI-DiT-B/2 400K
131
75
43.61
8.48
33.43
0.49
0.62
DİT-L/2 400K (Peebles Xie. 2023)
458
100
23.33
MoHl-DiT-L/2 400K
459
90
23.17
6.16
58.92
0.61
0,63
MoH-DiT-L/2 400K
459
75
24.29
6.38
57.75
0.60
0.63
DiT-XL/2 7.000K iPeebles & Xic, 2023)
675
100
9.62
6.85
121.50
0.67
0.67
DiT-XL/2 7,000K
(cfg=1.25)
675
100
3.22
5.28
201.77
0.76
0,62
MoH-DiT-XL.2 2,000K
676
75
10.95
6.19
106.69
0.67
0.66
MolI-DiT-XL/2 2,000K
676
90
10.67
6.15
107.80
0.67
0.65
MoH-DiT-XL/2 7,000K
676
90
8.56
6.61
129.54
0.68
0.67
MoH-DiT-XLL/2 7,000K 1.2g
676
90
5.17
207.25
0.77
0.63
Total Training Objective. It is worth noting that the MoH is a general framework. Therefore, we evaluate our proposed MoH across various popular model frameworks, including Vision Transformers (ViT), Diffusion models with Transformers (DiT), and Large Language Models (LLMs). Depending on the specific task, we require the task-specific loss. Finally, the total training loss is the weighted sum of the task-specific loss Lrask and the load balance loss Lg:
ing (Szegedy et al., 2016) and DropPath (Huang et al., 2016) to regularize our models. We optimize our mod- els using AdamW optimizer (Loshchilov & Hutter, 2017) with a gradient clipping norm of 1.0 and weight decay of 0.05. The initial learning rate set to 1e-3, with 5-epoch warm-up starting at 1e-6. A cosine learning rate sched- uler (Loshchilov & Hutter, 2016) is employed to decay the learning rate. During training, images are randomly cropped to a size of 224 224. It is worth noting that we do not use Exponential Moving Average (EMA) weights.
-= Ctask + 3Cb,
(8)
where 3 is the trade-off syper-parameter mitigate the risk of routing collapse. By default, the weight 3 for the load balance loss is set to 0,01 for all tasks.
Results. As shown in Tab. 1, despite activating only a subset of attention heads, MoH-ViT achieves highly com- petitive performance compared to current fatte-of-the-an methods. For example, MoH-VTT-B achieves 84.9% Top- accuracy on the ImageNet- IK classification benchmark with just 75% the attention head. In contrast, the well- established ViT baseline, TransNeXt, attains a slightly lower accuracy of 84.8% while requiring 100% of the heads be activated. These results suggest that MoH is promising alternative multi-head attention for vision model design.
4. Experiments
4.1. VIT for Image Classification
Model Settings. For Vision Transformers (ViT) (Doso- vitskiy et al., 2021), MoH-ViT models are implemented based on the TransNeXt (Shi, 2024) framework and trained from scratch on the ImageNet-I1 1K dataset (Deng et al., 2009), which contains over 1.2 million images in 1,000 categories. To ensure a fair comparison, we only replace the standard multi-hes head attention with the proposed MoH, while keeping all other training parameters identical to TransNeXt.
4.2. DiT for Class-oodiitonal Image Generation
Model Settings. For Diffusion models with Transform- (DiT) (Peebles & Xie, 2023), only replace the stan- dard multi-head attention with our MoH in MoH-DiT mod- els, while keeping all other training parameters identical to DiT. We use the ImageNet-IK dataset for class-conditional image generation at a resolution of 256x256. To evalu- ate generation performance, we use Frechet Inception Dis- tance (FID) (Heusel al., 2017) assess overall sample quality, Precision and Recall (Kynkäänniemi al., 2019) to measure fidelity and diversity separately, and sFID (Nash
Training Details. Our MoH-ViT models are trained for 300 epochs using automatic mixed precision across 8 GPUs. We follow the training strategy of Tetn which in- cludes various data augmentation techniques, including Ran- dom Augmentation (Cubuk et al., 2020), Mixup (Zhang, 2017), CutMix (Yun et al., 2019), and Random Eras- ing (Zhong et al., 2020). We also apply Label Smooth-
5

---

MoHI: Multi-Head Attention as Mixturc-of- Attention
Table Comparisons between MoH-LLMSs and vanilla LLMs. "100B" denotes training budget of 100 billion tokens, while "200B"
denotes budget of 200 billion tokens. We observe that larger models, NoH-LLM-BB generally perform worse than smaller models, e.g., LMM-L, on TruthfulQA, consistent with the findings reported by (2022).
Methods
#Activated Heads (%)
SciQ
PIQA
WinoGrande
OpenbookQA
LogiQA
TruthfulQA
Average
LLM-S 100B
100
63.0
63.1
51.1
27.4
26.9
31.6
43.9
MoH-LLMM-S 100B
75
64.7
62.0
50.6
28.8
26.4
35.2
44.6
MoH-LLMM-S 1008
50
67.0
62.2
51.5
29.2
26.7
35.6
45.4
LLM-B
100
73.1
69.7
52.0
31.8
28.4
29.5
47.4
MoH-LLM--B 100B
74.7
69.2
52.8
30.0
28.1
32.2
47.8
MoH-LLM-B
75.2
67.0
52.0
29.0
26.9
32.8
47.2
LLLMBB 2008
100
73.1
70.3
53.3
32.4
29.0
29.5
47.9
MoH-LLMMB 2008
75
76.0
69.2
52.7
30.4
29.8
32.6
48.5
MoH-LLM-B 2008
50
75.6
66.9
53.5
29.4
26.7
32.7
47.5
Table Comparisons between MoH-LLaMA3-SB and LLaMA3-8B. Please refer to Tab. G in the Appendix for the performance of the model the end of the first stage of training.
Methods
#Activated Heads (%)
MMLU (5)
CEVAL
CMMLU (5)
GSM8K(8)
TruthfulQA
LLLMA3-8B (Duhay et 2024)
100
65.2
52.3
50.7
49.5
35.4
MoH-LLaMA3-8B
65.8
61.5
64.4
56.9
44.0
Methods
# #Activated Heads (%)
MMLU U (5)
CEVAL (5)
CMMLU (5)
GSM8K(8)
TruthfulQA
LLaMA3-8B MBBBBBBBBAAB (Dubey atal. 2024)
100
81.9
30.0
83.9
75.5
94.0
MoH-LLaMA3-8B
80.1
30.3
84.0
76.4
92.2
Methods
#Activated Heads (%)
MMLU (5)
CEVAL (5)
CMMLU (5)
GSM8K(8)
TruthfulQA
LLaMA3-8B 8B (Dbey ata 2024)
100
81.0
72.5
31.5
59.0
61.6
MoH-LLaaAA3-B
75
78.8
72.9
28.3
60.1
64.0
et al., 2021) as metric that better captures spatial relation- ships than FID. Moreover, we use Inception Score (IS) (Sal- imans et al., 2016) as another metric for fidelity.
tasks that require attention mechanisms to capture pixel- level fine- tne-grained relationships, leaving less redundancy in the attention heads compared to image classification tasks. These results suggest that MoH a promising alternative to multi-head attention for diffusion models.
Training Details. Following DiT, the final linear layer initialized with zeros, and all other layers follow stan- dard VIT weight initialization. We train all models using the AdamW optimizer (Loshchilov & Hutter, 2017) with constant learning rate of le-4, weight decay, and batch size of 256, applying horizontal flips for data augmen- tation. Following DiT, we employ the Exponential Mov- ing Average (EMA) of MoH-DiT weights during training with a decay rate of 0.9999, generating all images using the EMA model. We use an off-the-shelf pre-trained vari- ational autoencoder (Kingma, 2013) model from Stable Diffusion (Rombach et al., 2022). Following TransNeXt, our attention- head activation budget is unevenly distributed across layers, with fewer attention heads activated in the shallow layers and more in the deeper layers.
4.3. Training LLMs from Scratch
Model Settings. For training LLMs from scratch, we use Megatron (Shoeybi al., 2019), an open-source train- ing code, the training framework. Please refer to the Appendix for detailed hyper-parameter settings (Tab. C) of various MoH-LLMs. The evaluation is performed mul- tiple benchmarks using the Eleuther AI Language Model Evaluation Harness (Gao et al., 2024), unified framework for testing generative language models. Since the parame- ters are only about 0.2B for the smallest model, we select simple benchmarks the metric. Specifically, we report O-shot accuracy SciQ (Welbl al., 2017), PIQA (Bisk et al., 2020), WinoGrande (Sakaguchi et al., 2021), Open- book (Mihaylov et al., 2018), LogiQA (Liu et al., 2020), and TruthfulQA (Lin et al., 2022).
Results. As shown Tab. MoH-DiT models consis- tently outperform DiT models with 90% of heads activated. However, when only 75% of the heads are activated, MoH- DIT models perform worse than DiT models. This may be because image generation tasks are dense prediction
Training Details. We only use public datasets for train- ing, ensuring accessibility for academic research. Specifi-
6

---

Moli: Multi-Head Attention as Mixture-of-Heeet Attention
79.5
44.0
65.0
79.2
Accuracy
64.5
Accuracy
79.0
Acauray
43.0
64.0
78.5
63.5
78.2
41.0
MMLUGI 5
78.0
HellaSwag (10)
TruthfuIQA
63.0
77.8
35
30
20
20
Training Tokens (B)
Training Tokens (B)
Training Tokens (B)
Figure Performance evolution during continue-tuning. The MoHl model quickly recovers to over 95% of the performance of the original model within training budget of 10B tokens. Then, the performance gradually improves with the increase of the training tokens.
The image classification results are from MoH-ViT ViT-S, by utilizing 75% of the attention heads with a training budget of 100 epochs. The class-conditional image generation results come from Molt-DiT-S/2-400K, also by using 75% of the attention heads, with training budget of 400K training steps.
Table 5, Ablation study the impact of each component of the proposed MoH.
Shared
Two-Stage
Image Classification
Class-Conditioall Image Generation
Heads
Routing
Acc (t
FID
sFID
18t
Precision"
Recallf
75.6
71.97
13.58
19.06
0.35
0.55
78.3
69.54
12.80
19.67
0.36
0.55
78.6
69.42
12.85
19.96
0.36
0.55
cally, we sample from the RedPajama (Computer, 2023), Dolma (Soldaini et al., 2024), and Pile (Gao al., 2020) datasets according to different sampling probabili- ties. Please refer the Appendix for detailed sample ratios. Following previous works (Jin et al., 2025), we utilize the tokenizer from LLaMA2 (Touvron et al., 2023), which con- tains 65,536 vocabulary tokens.
routers: Integrating a randomly initialized router into the praitrdined model without compromising its original perfor- mance requires careful training techniques. To address this, propose a parameter-free router that determines routing scores using the l norm of the query of each attention head. (iii) Weighting attention heads: We observe that weighting the attention head outputs significantly alters the distribu- tion of the output of the attention layer, which necessitates large amount of training data to restore the original per- formance. To tackle this, we quantize the routing score and use the straight-through estimator (Bengio et al., 2013; Liu et al., 2022) to back-propagate the gradients through the sparsity function. Specifically, given the input token æ, we employ quantizer for activation routing scores, with its forward pass formulated as:
Results. As shown in Tab. 3, despite activating only a subset of attention heads, MoH-LLMs achieve highly com- petitive performance compared our baseline models. For example, MoH-LLM-S achieves average accuracy of 45.4% with just 50% of the attention heads activated. In contrast, the baseline model reaches a slightly lower accu- racy of 43.9% with 100% of the attention heads activated. These results suggest that MoH is a promising alternative to vanilla multi-head attention for training L.MMs from scratch. Surprisingly, we find that for MoH-LLM-S, activating only 50% of the attention heads outperforms activating 75%. We consider it may be because when both the model and dataset are small, activating fewer heads effectively regularizes the model. However, as the amount of data increases, activating more heads offers a higher potential for performance.
g = 1 (Tokeen æ selects Head i),
where 1(x) denotes the indicator function. 9 represents the quantized routing score. We then adopt a straight-through estimator, which assigns the incoming gradients to thresh- old operation to be the outgoing gradients:
dc ac
(10)
0g. 0g.
4.4. Continue-Tuning LLaMA3-8B
where Si denotes the real-valued routing score. This approx- imation function significantly mitigates the issue of gradient vanishing (Wang et al., 2024). Similar to training LLMs from scratch, we also use Megatron (Shoeybi et al., 2019), an open-source training code, as the training framework.
Model Settings. To significantly enhance the applica- bility of the proposed MoH method, we also attempt to further continue-tune pre-trained multi-head attention mod- els, such as LLaMA3-SB, into MoH models. However, this presents three challenges. (i) Determining the shared attention heads: We simply select the first 16 attention heads of each layer as shared heads. (ii) Adding head
Training Details. We find that if there is a discrepancy between the continue-training data and the original train- ing data distribution of the model, the performance of the
7

---

MolI: Multi-Head Attention as cixture-of-Head Attention
Table Ablation study on the impact the shared heads ratio among activated heads. All results are from MoH- ViT-S. by using 75% of the heads with training budget of 100 epochs.
Ratio of Shared Heads
13.9%
27.6%
31.3%
35.9%
37.5%
40.5%
46.8%
60.4%
74.0%
Accuracy (%)
78.6
78.5
78.4
78.4
78.5
78.6
78.4
78.6
78.4
Table Comparisons about inference time. We convert the Q. K, and V features into sparse matrices using the mask generated by the router and replace the dense matrix multiplication in the attention mechanism with sparse matrix multiplication. To eliminate the impact of underlying operator optimizations, we replaced matrix multiplications with sparse matrix multiplication when testing for speed.
Methods
#Head Num
#Head Dim
#Sequence Length
#Activated Heads (%)
Time (ms)
Multi-Head Attention
32
64
256
100
0.360
MoH (Ours)
32
64
256
90
0,352
MoH (Ours)
64
256
75
0.321
MoH (Ours)
32
64
256
50
0.225
Multi-Head Attention
64
512
100
1.376
MoH (Ours)
32
64
512
90
1.351
MoH (Ours)
32
64
512
75
1.180
MoH (Ours)
32
64
512
50
0.863
model may fluctuate wildly at the beginning of the train- ing process. Since we are unable to have access to the raw training data of LLaMA3, we address these potential performance fluctuations by dividing the training process into two stages. In the first stage, we continue-tune the original LLaMA3-8B model using 300B tokens to adapt the model our dataset. In the second stage, we continue- tune this adapted model into our proposed MoH model with 100B tokens. We utilize the in-evaluation-harnes. package to assess performance on comprehensive suite of down- stream tasks: (i) Following Pythia (Biderman et al., 2023), we report O-shot accuracy on LAMBADA (Paperno et al., 2016), LogiQA (Liu et al., 2020), PIQA (Bisk et al., 2020), SeiQ (Welbl et al., 2017), and WinoGrande (Sakaguchi et 2021). (ii) We report the accuracy of Chinese tasks, including 5-shot CEVAL (Huang et al., 2023) and 5-shot CMMLU (Li et al., 2023a). (iii) We report the accuracy of tasks from the Open LLM Leaderboard (Beeching et al., 2023), including 10-shot HellaSwag (Zellers et al., 2019), 25-shot ARC Challenge (ARC-C) (Clark et al., 2018), and 5-shot MMLU (Hendrycks et al., 2021). (iv) We re- port the exact match score for 32-shot Natural Questions (NQ) (Kwiatkowski et al., 2019) and the accuracy for 32- shot BoolQ (Clark et al., 2019). (v) We report the exact match score for 8-shot GSMSK (Cobbe et al., 2021) to evaluate the math ability. (vi) Moreover, we report O-shot accuracy on TruthfulQA (Lin et al., 2022) to assess the ability to generate truthful answers.
14 benchmarks, outperforming LLaMA3-8B by 2.4% by utilizing only 75% of the attention heads. These results demonstrate that pre-trained multi-head attention models can be further continue-tunde into our MoH models, signifi- cantly enhancing the applicability of the MoH method.
4.5. Ablative Analysis
Effect of Each Component of the Proposed MoH. To explore the impact of each component of our MoH method, we provide the ablation results Tab. "Shared Heads" refers to subset of attention heads that are always activated. "Two-Stage Routing" represents the dynamic coefficient that balances the weights between shared and routed heads over the routing score, as described in Eq. 5 and Eq. 6. As shown in Tab. 5, shared heads significantly improve model performance by effectively capturing common knowledge, allowing the routed heads to focus more on domain- -spccficc information. Moreover, two-stage routing further enhances model performance by dynamically balancing the weights between shared and routed heads. Our full model achieves the best performance, demonstrating that both components significantly benefit the attention mechanism.
Effect of the Shared Heads Ratio among Activated Heads. In Tab. 6, we provide the ablation study on the shared heads ratio among activated heads. We find that model performance remains relatively consistent across wide range of shared heads ratios (from 13.9% to 74.0%), These results indicate that the performance of the model is stable as long as the shared heads ratio is not extreme. From another perspective, shared heads can be viewed as a form of Soft MoE (Puigcerver et al., 2024). Based on the findings from the Soft MoE paper (Puigcerver et al.. 2024). recommend using higher ratio of shared heads among
Results. As shown in Fig. 2, MoH-LLLMA3--8 8B quickly recovers to over 95% of the performance of the origi- nal model within a training budget of 10B tokens. After continue-tuning with 100B tokens, as shown in Tab. 4, MoH- LLaMA3-8B achieves an average accuracy of 64.0% across
8

---

Moli: Multi-Head Attention as Mixturc-off-eeal Attention
MoH-ViT-B
MoH--DITXLL/2
MoH-LLM-B
0.8
Desk
0.40
Desk
0.80
LogiQA
Dersity
0.7
Density
0.38
0.39
0.73
0.70
0.83
0.37
0.80
0.36
0.59
11
17
25
0.35
0.50
Routed Hoads ID
Routed Heads ID
Routed Heads ID
0.8
Goldfish
Goldfish
0.80
PIQA
Dansity
0.39
0.38
0.75
0.70
0.6
0.65
0.37
0.5
0.36
0.35
0.50
11
17
23
25
Routed Heads ID
Routed Headss 1D
Routed Heads ID
Density
0.7
0.8
cream
Dersity
0.38
0.39
0.40
Ice cream
0.73
0.80
0.70
WinoGrande
0.6
0.65
0.37
0.50
0.36
0.58
0.50
11 Routed
Heads ID
Routed Hiads ID
Routed Heads
Figure 3. Visualization of the head load distribution the final MoH layer. For ViT and DIT, we present the head load distributions for the categories "Desk", "Goldfish", and "Ice cream".. For LLM, we display the head distributions for the tasks "LogiQA", "PIQA", and "WinoGrande". MiT-V MoH-DiT-XL/2, and MoH-LLM-B activate 75%, 90%, and 75% of the attention heads, respectively. "Density" denotes the ratio of the number of head activations to the total number of tokens.
the activated heads (greater than 40%).
and two-stage routing to enhance the standard MoE method. More importantly, we show that pre-trained multi-head at- tention models can be further continue-ture into our MoH models, greatly improving the applicability of the proposed MoH method. In contrast, MoA directly combines multi- head attention with MoE. Due to the adoption of shared keys and values, MoA must be trained from scratch, which limits its applicability. Finally, in terms of model frame- works, our MoH validated across various popular model frameworks and tasks, including ViT, DIT, and decoder-only LLMs, while MoA is only validated for language tasks.
5. Discussion
The Efficiency of Our Proposed MoH. To explore if our method performs better with longer sequences, we increase the input sequence length. For rows I to 4 of Tab. 7, the input length is 256. For rows 5 to 8, it is 512. As shown in Tab. 7, although dynamic routing introduces additional computational overhead, MoH still outperforms standard heati-head attention mechanisms. Furthermore, as the input sequence gets longer, the advantage of MoH grows.
Visualization of the Head Load Distribution. As shown in Fig. 3, we observe significant variation in attention head assignments across different categories and task topics, in- dicating that the MoH model adapts to diverse tasks by employing distinct head assignment patterns. This charac- teristic of MoH allows different attention heads to focus on different types of tasks, making parameter utilization more efficient than multi- attention. For additional visualiza- tions of MoH-LLaMA3-8B and a detailed analysis of the head load distribution, please refer to Appendix D.
6. Conclusion
In this work, we introduce MoH, promising alternative to multi-head attention. MoH enables each token to adaptively select the appropriate attention heads, improving both model performance and inference efficiency without increasing the number of parameters. Extensive experiments across var- ious popular model frameworks, including ViT, DiT. and LLMs, demonstrate that MoH outperforms multi-head at- tention, even when using only 50%~~90% of the attention heads. This work represents a promising step toward ad- vanced and efficient attention- -ased models, which may be helpful to both the research and industrial communities.
The Difference between MoH and MoA. We clarify the differences between MoH and MoA (Zhang et al., 2022) from the following three asepctt First, in terms of moti- vation, the goal of MoH is improve the efficiency and performance of the attention mechanism without increasing the number of parameters. In contrast, MoA shares the moti- vation of MoF, which is to expand model parameters while keeping inference costs low. Therefore, the model settings of MoH are more stringent than those of MoA. Second, in terms of methodology, our MoH introduces shared heads
A Acknowledgements
This work was supported in part by the Natural Science Foundation of China (No. 62202014, 62332002, 62425101, 62088102), and NUS Start-up Grant A-0010106-00-00, Be- sides, this work was performed when Peng Jin was an Intern at Skywork AI.

---

MolI: Multi-Head Attention as Hixture-of-Head Attention
Impact Statement
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L.,, Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110 1416,, 2021.
This work is an important step toward creating more ad- vanced and efficient attention-basee models, which could benefit both the rescarch and industrial communities. Effi- cient attention models will not only lower the training costs for researchers but also greatly reduce the expenses involved in deploying and using large models.
Computer, T. Redpajama: open dataset for training large language models, 2023. URL httpps://gtttttt com/
ama -Data.
Cordonnier, J.-B., Loukas, A., and Jaggi, M. Multi-head attention: Collaborate instead of concatenate. arXiv preprint arXiv:2006. 16362, 2020.
References
Balasubramanian, S., Basu, S., and Feizi, S. Decomposing and interpreting image representations via text in vits beyond clip. arXiv preprint arXivv2406.01583, 2024.
Csordás, R., Piçkos, P., Irie, K., and Schmidhuber, J. Switch- head: Accelerating transformers with mixture-of-experts attention. In NeurlPS, pp. 74411-74438, 2024.
Basile, L., Maiorca, V., Bortolussi, L., Rodolà, E., and Lo- catello, F. Residual transformer alignment with spectral decomposition. arXiv preprint 2t0::24/1.0024.. 2024.
Cubuk, E. D., Zoph, B., Shlens, J., and Le, Q. V. Ran- daugment: Practical automated data augmentation with a reduced search space. In CVPRW, pp. 702-703, 2020.
Beeching, E., Fourrier, C., Habib, N., Han, S., Lam- bert, N., Rajani, N., Sanseviero, O., Tunstall, L., and Wolf, T. Open I1m leaderboard (2023- 2024). https: / h/wgggiinggg nggaac. co/ spaces/ 1aa--11m-eeaeeette leaderboard, 2023.
Dai, D., Deng, C.,, Zhao, C., Xu, R., Gao, H., Chen, D., Li, J., Zeng, W., Yu, X., Wu, Y., et al. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experr lan- guage models. arXiv preprint wrXiv:2401.06066 2024.
old/open. 11m
Dai, Z., Liu, H., Le, Q. V., and Tan, M. Coatnet: Marrying convolution and attention for all data sizes. In NeurlPS, pp. 9975-3977, 2021.
Bengio, Y., Léonard, N., and Courville, A. Estimating or propagating gradients through stochastic neurons for con- ditional computation. arXiv preprint arXiv: 130..3432, 2013.
Deng. J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. CVPR, pp. 248-255,110 2009.
Bhattacharyya, M., Chattopadhyay, S., and Nag, S. Decatt: Efficient vision transformers with decorrelated attention heads. In CVPRW, pp. 4695-4699, 2023.
Dosovitskiy, A., Beyer, L.., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2021.
Biderman, S., Schoelkopf, H., Anthony, Q. G., Bradley, H., O'Brien, K., Hallahan, E., Khan, M. A., Purohit, S., Prashanth, U. S., Raff, E., et al. Pythia: suite for analyzing large language models across training and scaling. In ICML, pp. 2397-2430, 2023.
Du, N., Huang, Y.. Dai, M., Tong, S., Lepikhin, D., Xu, Y., Krikun, M., Zhou, Y., Yu, A. W., Firat, O., et al. Glam: Efficient scaling of language models with mixture- of-experts. In ICMI., pp. 5547-5569, 2022.
Bisk, Y., Zellers, R., Gao, J., Choi, Y., et al. Piqa: Reasoning about physical commonsense natural language. AAAL pp. 7432-7439, 2020.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle,
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. In NewrrPS, pp. 1877-1901, 2020.
A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21 2024.
Fedus, W. Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and ef- ficient sparsity. Journal of Machine Learning Research, 23(120): 2022.
Clark, C., Lee, K., Chang, M.-W., Kwiatkowski, T., Collins, M., and Toutanova, K. Boolq: Exploring the surprising difficulty of natural yes/no questions. In NAACL, 2019.
Clark, P., Cowhey, I., Etzioni, O., Khot, T.. Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try are, the ai2 reasoning challenge. arXiv preprint arXiv:1803 05457, 2018.
Fu, Y., Cai, Z., Asi, A., Xiong, W., Dong, Y., and Xiao, W. Not all heads matter: A head-level kv cache compression method with integrated retrieval and reasoning. arXiv preprint arXiv:2410. 14.188998 2024.
10

---

MolI: Multi-Head Attention as Mixture-of-Ike -fedd Attention
Gandelsman, Y., Efros, A. A., and Steinhardt, J. Interpreting clip's image representation via text-based decomposition. arXiv preprint arXiv:2310.00555
Jin, P., Zhu, B., Yuan, L.,, and Yan, S. Moe++: Accelerat- ing mixture-of-experts methods with zero-computation experts. In ICLR, 2025.
2023.
Gao, L., Biderman, S., Black, S., Golding, L.. Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027 2020.
Kenton, J. D. M.-W. C. and Toutanova, L. K. Bert: Pre- training of deep bidirectional transformers for language understanding. In NAACL, volume I, pp. 2, 2019.
Kingma, D. P. Auto-encoding variational bayes. arXiv preprint arXiv: 14, 2013.
Gao, I.,, Tow, J., Abbasi, B., Biderman, S., Black, S., DiPofi, A., Foster, C., Golding, L., Hsu, J., Le Noac'h, A., Li, H., McDonell, K., Muennighoff, N., Ociepa, C., Phang, J., Reynolds, L.., Schoelkopf, H., Skowron, A., Sutawika, L., Tang, E., Thite, A., Wang, B., Wang, K., and Zou, A. A framework for few-shot language model evaluation, 07 2024. URL pettps:/2eenoe org/ records/ 12608602.
Kwiatkowski, T., Palomaki, J., Redfield, O., Collins, M.. Parikh, A., Alberti, C.,, Epstein, D., Polosukhin, I., Devlin, J., Lee, K., et al. Natural questions: a benchmark for ques- tion answering research. Transactions of the Association for Computational Linguistics, 7:453 466, 2019.
Kynkäänniemi, T., Karras, T., Laine, S., Lehtinen, J., and Aila, T. Improved precision and recall metric for assess- ing generative models. In NeurIPS, 2019.
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding. In ICLR, 2021.
Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang. Y., Krikun, M., Shazeer, N., and Chen, Gshard: Scaling giant models with conditional computation and automatic sharding. In ICLR, 2021.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, Gans trained by a two time-scale update rule converge to a local nash equilibrium. In NeurIPS, 2017.
Lewis, M., Bhosale, S., Dettmers, T., Goyal, N., and Zettle- moyer, L. Base layers: Simplifying training of large, sparse models. ICML, pp. 6265-6274, 2021.
Huang, G., Sun, Y., Liu, Z., Sedra, D., and Weinberger, K. Q. Deep networks with stochastic depth. In ECCV, pp. 646-661, 2016.
Li, H., Zhang, Y., Koto, F., Yang, Y., Zhao, H., Gong, Y., Duan, N., and Baldwin, T. Cmmlu: Measuring mas- sive multitask language understanding in chinese. arXiv preprint arXiv:2306. 09212, 2023a.
Huang, Y., Bai, Y., Zhu, Z., Zhang, J.. Zhang, J.. Su, T., Liu, J., Lv, C., Zhang, Y., Fu, Y., et al. C-eval: A multi-level multi-discipling chinese evaluation suite for foundation models. In NeurlPS, 2023.
Li, K., Wang, Y., Zhang, J., Gao, P., Song, G., Liu, Y., Li, H., and Qiao, Y. Uniformer: Unifying convolution and self-attention for visual recognition. TPAMI, 45(10): 12581-12600, 2023b.
Jacobs, R. A., Jordan, M. I,, Nowlan, S. J., and Hinton, G. E. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Li, Y., Wu, C.-Y., Fan, H., Mangalam, K., Xiong, B., Malik, J., and Feichtenhofer, C. Mvitv2: Improved multiscale vision transformers for classification and detection. In CVPR, pp. 4804-8811 4814, 2022.
Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., Casas, D. I., Hanna, E. B., Bressand, F., et al. Mixtral of experts. arXiv 2024.
preprint ar arXiv:2401.0400
Jin, P., Huang, J., Xiong, P., Tian, S., Liu, C., Ji, X., Yuan, L., and Chen, J. Video-text game players: Hierarchi- cal banzhaf interaction for cross- cross-modal representation learning. In CVPR, pp. 2482,, 2023.
Lin, B., Zhu, B., Ye. Y., Ning, M., Jin, P., and Yuan, L. Video-llava: Learning united visual representa- tion by alignment before projection. arXiv preprint arXiv:2311.10111 2023.
Jin, P., Li, H., Yuan, I.., Yan, S., and Chen, J. Hierarchical banzhaf interaction for general video-language represen- tation learning. TPAMI, 2024a.
Lin, B., Tang, Z., Ye, Y., Cui, J., Zhu, B., Jin, P., Zhang, J., Ning, M., and Yuan, I.. Moe-Ilava: Mixture of ex- perts for large vision- -angun-a models. arXiv preprint arXiv:241..0 15947, 2024.
Jin, P., Takanobu, R., Zhang, W., Cao, X., and Yuan, L. Chat-univi: Unified visual representation empowers large language models with image and video understanding. In CVPR, pp. 13700-13710, 2024b.
Lin, S., Hilton, J., and Evans, O. TruthfulQA: Measuring how models mimic human falschoods. In ACL, pp. 3214- 3252, 2022.

---

MolI: Multi-Head Attention as Hixture-of-Head Attention
Liu, D., Zhang, R., Qiu, L.., Huang, S., Lin, W., Zhao, S., Geng, S., Lin, Z., Jin, P., Zhang, K., et al. Sphinx-x: Scaling data and parameters for a family of multi-modal large language models. In ICML, 2024.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In ICCV, pp. 4195-4205, 2023.
Peng, H., Schwartz, R., Li, D., and Smith, N. A. mixture of h - 1 heads better than h heads. arXiv preprint arXiv:2005.06557 2020.
Liu, J., Cui, L., Liu, H., Huang, D., Wang. Y., and Zhang, Y. Logiqa: A challenge dataset for machine reading comprehension with logical reasoning. In IJCAL, pp- 3622-3628, 2020.
Puigcerver, J., Ruiz, C. R., Mustafa, B., and Houlsby, N. From sparse to soft mixtures of experts. In ICLR, 2024.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. Exploring the limits of transfer learning with unified text-to-text transformer. The Journal of Machine Learning Research, 21()):5485-551, 2020.
Liu, Z., Lin, Y.. Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, pp. 10012- 10022, 2021.
Liu, Z., Cheng, K.TT., Huang. D., Xing, E. P., and Shen, Z. Nonuniform-to cowtunfoom-touunffooo to-ufforr quantization: Towards accurate quantization via generalized straight-thoough estimation. In CVPR, pp. 4942-4952, 2022.
Rajbhandari, S., Li, C., Yao, Z., Zhang, M., Aminabadi, Y., Awan, A. A., Rasley, J., and He, Y. Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale. In ICML, pp. 18332- 18346, 2022.
Loshchilov, I. and Hutter, F. Sgdr: Stochastic gra- dient descent with warm restarts. arXiv preprint arXiv:1608.03983, 2016.
Roller, S., Sukhbaatar, S., Weston, J., al. Hash layers for large sparse models. In NeurlPS, pp. 17555-17566. 2021.
Loshchilov, I. and Hutter, F. Decoupled weight decay regu- larization. arXiv preprint arXiv: 171.051711.05001. 2017.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In CVPR, pp.
Michel, P., Levy, O., and Neubig, G. Are sixteen heads really better than one? In NeurlPS, pp. 14014-14024, 2019.
2022.
Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge scale. Communications of the ACM, 64(9):99-106, 2021.
Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can a suit of armor conduct electricity? a new dataset for open book question answering. In EMNLP, 2018
Salimans, T., Goodfellow, I., Zaremba, W., Cheung, V., Radford, A., and Chen, X. Improved techniques for training gans. In NeurIPS, 2016.
Muennighoff, N., Soldaini, L., Groeneveld, D., Lo. K., Mor- rison, J., Min, S., Shi, W., Walsh, P., Tafjord, O., Lambert, N., et al. Olmoe: Open mixture-of-expee language models. arXiv preprint arXiv:2409.02200 2024.
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, J. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv: 2017.
Nash, C., Menick, J., Dieleman, S., and Battaglia, P. W. Generating images with sparse representations. arXiv preprint arXiv:2103.03841, 2021.
Shi, D. Transnext: Robust foveal visual perception for vision transformers. In CVPR, pp. 17773-17783, 2024.
OpenAl. Introducing chatgpt. CORR, 2022. URL https: //openai com
Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-Im: Training multi- billion parameter language models using model paral- lelism. arXiv preprint arXiv: 1900.0009.000 2019.
Ouyang, L.,, Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. In NeurlPS, pp. 27730-27744, 2022.
Si, C., Yu, W., Zhou, P., Zhou, Y., Wang. X., and Yan, Inception transformer. In NeurIPS, pp. 23495-23509, 2022.
Paperno, D., Kruszewski, G., Lazaridou, A., Pham, Q N., Bernardi, R., Pezzelle, S., Baroni, M., Boleda, G., and Fernández, R. The lambada dataset: Word prediction requiring a broad discourse context. In ACL, pp. 1525- 1534, 2016.
Soldaini, L., Kinney, R., Bhagia, A., Schwenk, D., Atkinson, D., Authur, R., Bogin, B., Chandu, K., Dumas, J., Elazar, Y., Hofmann, V., Jha, A. H., Kumar, S., Lucy, L., Lyu, X.,
12

---

MoHI: Multi-Head Attention as Hixture-of-Heed Attention
Lambert, N., Magnusson, I., Morrison, J., Muennighoff,
Wu, W., Wang, Y., Xiao, G., Peng, H., and Fu, Y. Re- trieval head mechanistically explains long-context factu- ality, arXiv preprint arXiv:2404.15571 2024.
N., Naik, A., Nam, C.,, Peters, M. E., Ravichander, A.,
Richardson, K., Shen, Z., Strubell, E., Subramani, N., Tafjord, O.. Walsh, P., Zettlemoyer, L., Smith, N. A., Hajishirzi, H., Beltagy, I., Groeneveld, D., Dodge, J., and Lo, K. Dolma: An Open Corpus of Three Trillion Tokens for Language Model Pretraining Research. arXiv preprint, 2024. URL https: //rrivi org/ abs/ 2402. 00159.
Xiao, G., Tang, J., Zuo, J., Guo, J., Yang. S., Tang, H., Fu, Y., and Han, S. Duoattention: Efficient long-context Ilm inference with retrieval and streaming heads. arXiv preprint arXiv:2400. 10819, 2024.
Yang, C., Qiao, S., Yu, Q.. Yuan, X., Zhu, Y., Yuille, A., Adam, H., and Chen, L.-C. Moat: Alternating mobile convolution and attention brings strong vision models. In ICLR, 2022a.
Szegedy, C., Vanhoucke, V., Ioffe, S., Shlens, J., and Wojna, Z. Rethinking the inception architecture for computer vision. In CVPR, pp. 2818 2826, 2016.
Yang. J., Li, C.,, Zhang, P., Dai, X., Xiao, B., Yuan, L.,, and Gao, J. Focal self-attention for local-global interactions in vision transformers. arXiv preprint arXiv:2107.00 00641, 2021.
Touvron, H.. Cord. M., Douze, M., Massa, F., Sablayrolles, A., and Jégou, H. Training data-efficient image trans- formers & distillation through attention. In ICML, pp. 10347-10357, 2021.
Yang, J., Li, C., Dai, X., and Gao, J. Focal modulation networks. In NeurlPS, pp. 4203-4217, 2022b.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint 2xx0:2307.09288, 2023.
Yu, W., Si, C., Zhou, P., Luo, M., Zhou, Y., Feng, J., Yan, S., and Wang, X. Metaformer baselines for vision. TPAMI, 2023.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. NeurlPS, 2017.
Yuan, L., Chen, Y., Wang. T., Yu, W., Shi, Y., Jiang, Z.-H., Tay, F. E., Feng, J., and Yan, S. Tokens-to-tokes vit: Training vision transformers from scratch on imagenet. In ICCV, pp. 558-567, 2021.
Voita, E., Talbot, D., Moiseev, F., Sennrich, R., and Titov, I. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In ACL, pp. 5797-5808, 2019.
Yun, S. and Ro, Y. Shvit: Single-head vision transformer with memory efficient macro design. In CVPR, pp. 5756- 5767, 2024.
Wan, Z. Wu, Z., Liu, C., Huang, J., Zhu, Z., Jin, P., Wang, L., and Yuan, L. Look-m: Look-once optimization in kv cache for efficient multimodal long-context inference. arXiv preprint ar arXiv:246..1 18139, 2024.
Yun, S., Han, D., Oh, S. Chun, S., Choe, J., and Yoo, Y. Cutmix: Regularization strategy to train strong classifiers with localizable features. In ICCV, pp. 6023 2019.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can machine really finish your sentence? In ACL, pp. 4791-4800, 2019.
Wang, H., Ma. S., Wang, R., and Wei, F. Q-sparse: All large language models can be fully preprint ar XXV:2407..10999 2024.
arXiv
Zhang, H. mixup: Beyond empirical risk minimization.
Wang, W., Xie, E., Li, X., Fan, D.-P., Song, K., Liang, D., Lu, T., Luo, P., and Shao, L. Pvt v2: Improved baselines with pyramid vision transformer. Computational Visual Media, 8(3):415-45 424 2022.
arXiv preprint arXiv:1710.09412, 2017.
Zhang, X., Shen, Y., Huang, Z., Zhou, J., Rong, W., and
Xiong, Z. Mixture of attention heads: Selecting attention heads per token. arXiv preprint arXiv:2210.0514: 2022.
Wei, T., Zhu, B., Zhao, L., Cheng, C., Li, B., Lü, W., Cheng, P., Zhang. J., Zhang, X., Zeng, L., et al. Skywork-moe' A deep dive into training techniques for mixture-of-exxprts language models. arXiv preprint 0rx.v:2406.066633, 2024.
Zhong, 7., Zheng, L.,, Kang, G., Li, S., and Yang. Y. Ran- dom erasing data augmentation. In AAAI, pp. 13001- 13008, 2020.
Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., Dai, A. M., Le, Q. V., Laudon, J., et al. -fxture-of-experts with expert choice routing, In NeurIPS, pp. 7103-7114, 2022.
Welbl, J., Liu, N. F. and Gardner, M. Crowdsoure- ing multiple choice science questions. arXiv preprint arXiv: /Tti://77.062000 2017.
13

---

MoHI: Multi-Head Attention as Mixture-of- -fede Attention
This appendix provides additional discussions (Appendix A), implementation details (Appendix B), several additional experiments (Appendix C), more qualitative analysis (Appendix D), and details of quantitative evaluations for LLMs (Appendix E).
Abstract.
A. Additional Discussions
A.I. Why is MoH Superior to Vanilla Multi-Head Attention?
We demonstrate that MoH is superior to vanilla multi-head attention from both theoretical and experimental perspectives.
Specifically, MoH not only improves efficiency and model performance but also helps different attention heads to specialize better compared multi-head attention.
From the theoretical perspective, in standard multi-head attention, all heads use the same data, which can cause them to learn similar features. Many studies have pointed out that there are redundant heads in multi-head attention. Given a minibatch of data D, the gradient of each attention head in multi-head attention can be written E (
In contrast, in MoH, routed heads are trained only on smaller subsets of data specifically assigned to them. In MoH's routing mechanism, the data is divided into h 1--h, subsets {D1, D2. DR-h, J. with each subset corresponding to a routed head. Besides, the routing score for each attention head acts as an adaptive adjustment to the learning rate, enabling the attention heads in MoH to specialize more effectively. Given a minibatch of data D and the router G(*), the gradient of each routed head in MoH can be written as [G(2) ( The gradient of each shared head in MoH can be written as
(Ge (c(=) As shown in Tab. A, the routing mechanism and adaptive weights in MoH enable attention heads to specialize more effectively compared to standard multi-head attention.
Table A. Comparisons between the multi-head attention and our proposed mixture-of-hees attention.
Methods
#Head Type
#Dat
#Weight (learning rate)
#Gradient
Multi-H Helli-Head Attention
D
Excol (eta))
MoH
routed head
DEED
G(æ);
E,ED, [G(x); [Gt); dc(()))
MoH
shared head
D
G(x);
82(2)
From the experimental perspective, we calculated the similarity of attention patterns and output features of different attention heads (include routed heads and shared heads). As shown Tab, B, the similarity of attention patterns and output features among attention heads in MoH lower than in standard multi head attention, indicating reduced redundancy and greater differentiation among the attention heads in MoH.
Table B. The similarity of attention patterns and output features among attention heads. Given pair of attention score matrices A and A', we calculate the similarity of attention patterns as 1 B[| A' [1] Since attention scores form a probability distribution for each query, the similarity is always between 0 to 1.
Similarity of Attention Patterns
Cosine Similarity of Output Features
Methods
VIT
LLM
VIT
LLM
Multi-Head Attention
0,5159 0.3978
0.4795 0.4333
0.0411 0.0165
0.2550
MoH
0.2042
A.2. Limitations and Future Work
In this section, we delineate the limitations of our work and outline avenues for future research.
Heterogeneous Attention Heads. We find that different attention heads operate in parallel within the attention mechanism, suggesting that different heads can have varying hidden sizes. Future work could explore the use of heterogeneous attention heads based on our MoH framework.
Lower Activation Rate. Currently, MoH outperforms multi-head attention by utilizing only 50%~99%% of the attention
14

---

MoHI: Multi-Head Attention as Mixture-of- -fede Attention
heads. However, this is still a relatively high proportion. Future work could aim further optimize MoH, reducing head activation to less than 50%.
Multimodal Inputs. Effectively processing information from multiple modalities in the attention mechanism remains an open question. Recent work (Wan et al.. 2024) has shown that visual and textual tokens exhibit distinct attention patterns in multi-head attention. Future work could explore the attention patterns of MoH with different modal inputs, for example within multimodal large language models (Jin et al., 2024b; Lin et al., 2023; 2024; Liu et al., 2024; Jin et al., 2023; 2024a).
More Downstream Tasks. We evaluate our proposed MoH across various popular model frameworks, including ViT for image classification, DiT for class-conditional image generation, and LLMs for language tasks. Future work explore the application of MoH in more downstream tasks, such as audio tasks and multimodal tasks.
More Parameters. Due to computational constraints, the maximum number of MoH model parameters in our experiments limited to 8B (MoH-LLLMA3--BB) However, our MoH method is highly generalizable and can be scaled to larger models in future research.
B. Implementation Details
B.1. VIT for Image Classification
Training Details. Our MoH-ViT models trained for 300 epochs using automatic mixed precision across 8 GPUs. We follow the training strategy of TransNeXt, which includes various data augmentation techniques, including Random Augmentation (Cubuk et al., 2020), Mixup (Zhang, 2017), CutMix (Yun et al., 2019), and Random Erasing (Zhong et al., 2020). We also apply Label Smoothing (Szegedy et al., 2016) and DropPath (Huang et al., 2016) to regularize our models. We optimize our models using A optimizer (Loshchilov & Hutter, 2017) with gradient clipping norm of 1.0 and a weight decay of 0.05. The initial learning rate set to 1e-3, with a 5-epoch warm-up starting at 1e-6. cosine learning rate scheduler (Loshchilov & Hutter, 2016) employed decay the learning rate. During training, images are randomly cropped to size of 224x224. It is worth noting that we do not use Exponential Moving Average (EMA) weights.
B.2. DIT for Class-Conditional Image Generation
Training Details. Following DiT, the final linear layer initialized with zeros, and all other layers follow standard ViT weight initialization. We train all models using the AdamW optimizer (Loshchilov & Hutter, 2017) with constant learning rate of le-4, no weight decay, and batch size of 256, applying horizontal flips for data augmentation. Following DiT, we employ the Exponential Moving Average (EMA) of MoH-DiT weights during training with a decay rate of 0.9999, generating all images using the EMA model. We use an off- thf-tf pre-trained variational autoencoder (Kingma, 2013) model from Stable Diffusion (Rombach et al., 2022). Following TransNeXt, our attention head activation budget is unevenly distributed across layers, with fewer attention heads activated the shallow layers and more the deeper layers.
B.3. Training LLMs from Scratch
Model Settings. For training LLMs from scratch, we use Megatron (Shoeybi et al., 2019), an open-source training code, as the training framework. The detailed hyper- sypem-parameter settings of various MoH-LLMSs are shown in Tab. C.
Table C. Sizes and architectures of MoII-LLMs and LLMss "MoH-LLM-BB" has more parameters than "LLM-B" due to the additional parameters introduced by the router network.
Methods
#Params
#Layers
#Hidden Size
#Intermediate Size
#Heads
#Head Dim
LLM-S
186
12
768
2048
12
64
MoH-LLM-S
186
LLM-B
881
24
1536
4096
16
96
MoH-LLM-B
882
Data Details.
Consistent with previous works, we use the tokenizer of LLaMA2, which contains 65,536 vocabulary tokens. is worth noting that MoH-LLM trained exclusively on public datasets, making it accessible for academic research settings. Tab. D shows the detailed sample ratios of different open-source datasets. Specifically, we sample from
15

---

MoH:: Multi-Heaa Attention as Mixture-of- Head Attention
the following datasets according different sampling probabilities:
The RedPajama (Computer, 2023) includes training data from seven domains: CommonCrawl, C4, Github, Wikipedia, Books, ArXiv, and StackExchange.
The Dolma (Soldaini et al., 2024), a large and diverse open English text corpus, contains 3 trillion tokens sampled from seven sources, including web pages from Common Crawl, code from The Stack, curated web data from C4 (Raffel et al., 2020), social media conversations from Reddit, academic papers from PeS2o, public domain books from Project Gutenberg, and comprehensive content from Wikipedia and Wikibooks.
The Pile (Gao et al., 2020), an open-source English text corpus for training large language models, includes 22 diverse, publicly available datasets such as Wikipedia, NIH ExPorter, ArXiv, Books3, BookCorpus2, OpenSubtitles, YoutubeSubtitles, and Enron Emails.
Table D. Sampling ratio of different open-source datasets for MoH--LMMs. MoH-LLM trained exclusively on public datasets, making it accessible for academic research settings.
Sampling Ratio
Redpajama Books
4.24%
Redpajama Wikipedia
3.50%
Redpajama ArXiv
4.37%
Redpajama StackExchange
3.19%
Redpajama C4
10.94%
Dolma
61.28%
Pile
12.48%
Training Hyper-Parameterrs Tab. E shows the detailed training typper-arametees of MoH-LLMs. Specifically, all MoH-LLMs are trained with the AdamW optimizer (Loshchilov & Hutter, 2017), using batch size of 4 million tokens with a sequence length of 2048. The final learning rate is set to 10% the maximum. During training, a weight decay of 0.1 and gradient clipping of 1.0 are applied. For LLM-S and MoH-LLM-SS, the maximum learning rate is set to 3e-4. For LLM-B and MoH-LLLM-B, the maximum learning rate is set to 5e-4.
Table E. Training hyper-parameters of MoH-LLMs.
MoH-LLM-S 100B
MoH-LLM-B 100B
MoH-LLMM-B 200B
(LLM-S 1008)
(LLM-BB 1008)
(LLM--B 200B)
Training budget
100B
100B
200B
Maximum learning
3e-4
5e-4
5c-4
Final learning rate
3e-5
5c-5
5c-5
LR warmup init
1e-7
le-7
le-7
LR warmup iters
2000
500
500
Sequence length
2048
2048
2048
Batch size (tokens)
4M
4M
4M
3 for La
0.01
0.01
0.01
Tensor parallel
1
1
Pipeline parallel
1
B.4. Poatinue-Tuning LLaMA3-8B
Training Hyper-Paramen Paper-Parrmmters. Tab. F shows the detailed training syper-parametent of MoH-LLaMA3-SB. We find that if there is discrepancy between the contimue-training data and the original training data distribution of the model, the performance of the model may fluctuate wildly at the beginning of the training process. Since we do not have access to the raw training data of LLaMA3, we address these potential performance fluctuations by dividing the training process into two stages. the first stage, we continue- tune the original LLaMA3-89 8B model using 300B tokens to adapt it to our dataset. In addition, during the first stage, to enhance the Chinese ability of the model, we expand the vocabulary size. Specifically, we
16

---

MoH:: Multi-Head Attention as Hixture-of-Head Attention
increase the original LLaMA3-8B vocabulary size from 128,256 160,896. In the second stage, we continue-tune this adapted model into our proposed MoH model with 100B tokens. During the first stage, the maximum learning rate is set 6c-5, and the final learning rate is 6e-6, In the second stage, the maximum learning rate is set to 2e-5, and the final learning rate is le-6. For both stages, we employ the Adam W optimizer (Loshchilov & Hutter, 2017), with a batch size of 16 million tokens with a sequence length of 8192. During training, we use weight decay of 0. and gradient clipping of 1.0.
Table Training hyper-parameters of MoH-LLaMA3-SB. We divide the training process into two stages. In the first stage, we continuc-tune the LLaMA3-8B model using 300B tokens. In the second stage, we continuc-tun this adapted model into our proposed MoH model with 100B tokens.
The First Stage
The Second Stage
Training budget
300B
100B
Maximum learning rate
6e-5
2c-5
Final learning rate
6e-6
1c-6
LR warmup iters
50
50
Sequence length
8192
8192
Batch size (tokens)
16M
16M
3 for Ls
0.01
Tensor parallel
2
]
Pipeline parallel
1
Table G. Comparisons between HoM3-LB and LLaMA3-8B-stagel. Motl-LLaMA3-BB outperforms LLaMA3-8B-stagel by
utilizing only 75% of the attention heads.
#Activated Heads (%)
Methods
MMLU (5)
CMMLU (5)
NQ (32)
GSMBK(8)
TruthfulQA
LLaMA3-8B-stage stagel
100
66.2
66.0
28.1
58.6
41.9
MoH--LLMMA3-8
75
65.8
64.4
28.3
56.9
44.0
#Activated Heads (%)
Methods
HellaSwag (10)
LogiQA
BoolQ (32)
LAMBADA
SciQ
LLaMA3-8B stagel
100
79.4
30.4
85.1
75.8
92.2
MoH-LLaMA3-8B
75
80.1
30.3
84.0
76.4
92.2
#Activated
Methods
PIQA
WinoGrande
ARC-E
ARC-C (25)
Average
Heads (%)
LLaMA3-8B-stan stagel
100
79.1
73.0
70.9
59.6
64.7
MoH-LLaMA3-88
75
78.8
72.9
72.5
60.1
64.8
C. Additional Experiments
Comparison between MoH-LLaMA3-8B and L.AMA3--B--stgeel We divide the training process into two stages. Tab. G shows the comparison between MoH-LLaMA3-SB and the model at the end of the first training stage (LLaMA3- 8B-stagel), 1). As shown in Tab. G, MoH-LLaMA3-SB -B quickly recovers the performance of LLaMA3-8B- -Bstgge within training budget of 100B tokens. Notably, in English language tasks, MoH-LLLMMA3-8! surpasses ---tt-----tt while using only 75% of the attention heads. However, for Chinese language and math tasks, the recovery performance of the MoH model is not as strong as for English. For example, MoH--LLMMA3-BB achieves an accuracy 64.4% on CMMLU, compared 66.0% for LLaMA3-8B-stagee capabilities primarily established during the first training stage. Since the first training stage uses only 300B tokens, significantly less than the 15T tokens in LLaMA3-8B's pre-training, the model's abilities in these areas are not fully stable. In the second training stage, after switching to the MoH model, the model experiences more significant forgetting in Chinese and math tasks. Overall, as shown in Tab, G. MoH-LLaMA3-SB achieves an average accuracy of 64.8% across 14 benchmarks, outperforming LLaMA3-88 -taggell by utilizing only 75% of the attention heads.
We attribute this the fact that the model's Chinese and mathematical
Effect of the Activated Head Ratio. shown Tab. H. activating more attention heads generally leads to improved
17

---

MolI: Multi-Head Attention as Hixture-of-Heeal Attention
MoH-VVT-B
MoH--DIT-XL/2
0.8
Baskethall
Baskethall
Density"
0.29
0.38
0.8
0.5
0.36
13
15
19
21
25
27
0.55
2
Routod Heeds ID
Routed Heads ID
0.8
Boolshop
Boolshop
Density
0.7
Density
0.38
0.37
15
25
0.25
Routed Heads
Roued Heads ID
0.40
0.8
Cart
Cart
Density
0.7
Density
0.30
0.47
0.4
0.35
Routed Heads
Roued Heads ID
Husky
0.40
Husky
Dessit
0.7
Density
0.39
0.6
0.37
0.36
0.4
15
19
23
25
Routed Heads
Routed Heads ID)
0.8
J eem
Jeem
0.39
Density
Density
0.38
0.6
0.26
3
11
15
17
20
23
23
25
27
0.35
2
3
Routed Heads ID
Routed Heads ID
Figure A. Additional visualization of the head load distribution in the final MoH layer. MoH- NOT-- activates 75% of the attention heads. MoH-DiT-XI/2 activates 90% of the attention heads.
model performance. These results are intuitive, as activating more attention heads equates to utilizing more parameters and performing additional computations on the input.
Table H. Ablation study on impact the activated head ratio. All results are from MoH-VVT--S by using training budget 100 epochs.
Activated Heads
50%
55%
60%
65%
70%
75%
80%
Accuracy (%)
78.32
78.38
78.44
78.50
78.42
78.58
78.78
D. Additional Qualitative Analysis
Additional Visualization of the Head Load Distribution. We provide additional visualization of the head load distribution in Fig. A. As illustrated in both Fig. 3 and Fig. A, there is notable variation in attention head assignments across different categories and task topics. This suggests that the MoH model adapts to wide range of tasks by utilizing distinct head assignment patterns. This ability enables MoH to allocate attention heads more effectively to specific task types, leading to more efficient parameter utilization compared to standard multi-heat attention.
Additional Visualization the Head Load Distribution in MoH-LLaMA3-8B. We provide additional visualization of the head load distribution in Fig. B. As shown Fig. B. MoH-LLLMMA3-8 exhibits similar characteristics to MoH-LLMMs trained from scratch, with significant variation in attention head assignments across different categories and task topics. This indicates that continue-tunin enables the model to adopt different head assignment patterns quickly. These results demonstrate that pre-trained multi-head attention models can be effectively continue-tuned into MoH models, significantly broadening the applicability of the proposed MoH approach.
Additional Visualization of the Head Routing Score Distribution. We provide additional visualization of the head routing score distribution in Fig. C, Fig. D, and Fig. E. As illustrated in Fig. C, Fig. D. and Fig. E, these head routing scores also vary across categories and task types. This dynamic weighting mechanism allows MoH adjust the importance of
18

---

MoHI: Multi-Head Attention as Mixture-of-Heead Attention
0.8
ARC Challenge
0.8
MMUU
Density
0.7
Denelty
15
13
15
Routed Heads ID
Routed Heads IT
0.8
ARC Easy
0.8
NQ
Dansity
0.7
Density
9
11
15
13
15
Routed Heads ID
Routed Heads ID
0.8
CEVAL
0.8
PIQA
Density
0.7
Densitr
15
13
15
Routed Heads ID
Routed Heads ID
0.8
LAMBADA
0.8
SiiQ
Donsiy
0.7
Dansity
0.6
15
13
15
Rouid Headss ID
Roded Heads
0.8
LogiQA
0.5
WinGrande
Dendty
0.7
Density
0.7
0.6
0.5
15
15
Routed Hoads ID
Roued Heads ID
Figure B. Additional visualization of the head load distribution in Mott-LLaMA3-8B.
each head in response to different task requirements, further enhancing its flexibility and performance. Besides, we find that the routing scores of shared heads change more across categories than those of routing headers. We consider this because routed heads adapt different categories by adjusting their activation, while shared heads remain activated all the time. Therefore, shared heads primarily rely on changes in routing scores to adapt to different categories.
Images Generated from the Proposed MoH-DiT-XI/2 Model. Fig. F shows samples generated by our class-conditional MoH-DiT-XL/2 model. These results demonstrate the ability of HoHH-DiTXL/2 to generate semantically correct content with accurate spatial relationships.
E. Details of Quantitative Evaluations for LLMs
We conduct comparative comparisons of MoH-LLM (MoH--LLaMA3-8B) against vanilla LLMs (LLaMA3-8B). The evaluation is performed on multiple key benchmarks using the Eleuther AI Language Model Evaluation Harness (Gao et al., 2024), a unified framework for testing generative language models across a wide range of tasks. The benchmarks used for evaluation include:
ARC (Clark et al., 2018) multiple-choice question-anss answering resource featuring questions from science exams for grades to 9. It is divided into partitions: Easy and Challenge, with the latter containing more difficult questions that necessitate reasoning. Most questions offer four answer choices, while less than 1% feature either three five choices. Additionally, ARC includes supporting knowledge base with 14.3 million unstructured text passages. We report O-shot accuracy on ARC Easy and 25-shot accuracy on ARC Challenge.
LAMBADA (Paperno et al., 2016) is an open-ended cloze task consisting of approximately 10,000 passages from BooksCor- pus, where the objective is to predict missing target word in the last sentence of each passage. The missing word is always the last word of the final sentence, with options provided. We report 0-shot accuracy on LAMBADA.
19

---

MoHI: Multi-Head Attention as Mixture-of- -fad Attention
Routed Heads
Shared Heads
Basketball
Reuting Score
Basketball
0.15
Routing
Routed Heads ID
Shared Hexds ID
Bookshop
Rouling Soore
Bookshop
0.5
0.08
0.00
Routed Heads ID
Shevrll Hieee IT
Jean
Routingt Score
Jen
8.00
0.65
0.co
10
21
Routed Heads ID
Shared Heads ID
Goldfish
Rouding Score
Goldfish
0.30
Roubnng
0.
0.00
Routed Heads ID
Shared Heads ID
Ice cream
Rouding Soore
Ice cream
0.11
0.
0.00
Routed Heads
ID
Shard Heads ID
Figure C. Additional visualization the head routing score distribution in MoH-ViT-B. MoH-ViT-B activates 75% of the attention heads.
LogiQA (Liu et al., 2020) comprises 8,678 question-ami instances that encompass various types of deductive reasoning. The dataset serves as benchmark for reexamining logical Al within the context of deep learning in NLP. We report 0-shot accuracy on LogiQA.
PIQA (Bisk et al., 2020) is a dataset designed for commonsense reasoning, aimed at evaluating the physical knowledge of current models. We report 0-shot accuracy on PIQA.
SeiQ (Welbl et al., 2017) includes 13,679 crowdsourced science exam questions covering subjects such as Physics, Chemistry, and Biology. Each question is presented in multiple-choice format with four answer options, and for most questions, an additional paragraph provides supporting evidence for the correct answer. We report O-shot accuracy on SciQ.
WinoGrande (Sakaguchi et al., 2021) is a large-scale dataset comprising 44,000 problems, inspired by the original WSC design but enhanced to increase both its scale and difficulty. We report O-shot accuracy on WinoGrande.
al., 2019) is challenging dataset designed to evaluate commonsense natural language inference, which proves difficult for state-of-the-art models but poses significant challenge for humans. We report the accuracy for the 10-shot HellaSwag.
HellaSwag (Zellers
MMLU (Hendrycks et al., 2021) is benchmark designed to assess models' knowledge acquired during pretraining, making it more challenging and human-lii like in evaluation. It covers 57 subjects across STEM, humanities, social sciences, and more, ranging from elementary to advanced professional levels. The benchmark tests both world knowledge and problem-solving skills, with subjects spanning traditional areas like math and history to specialized fields such as law and ethics, offering a comprehensive tool for identifying model blind spots. We report the accuracy for the 5-shot MMLU.
Natural Questions (NQ) (Kwiatkowski al., 2019) is question- answering dataset based on real, anonymized Google queries. Annotators label long and short answers (or null if no answer is found) from Wikipedia pages in the top search results. The dataset includes 307,373 training examples, 7,830 development examples, and 7,842 test examples with 5-way annotations. We report the exact match score for 32-shot Natural Questions to measure the factual knowledge in the model.
BoolQ (Clark et al., 2019) is dnestion-answerre dataset consisting of 15,942 yes/no questions. These questions are naturally occurring, and generated in unprompted and unconstrained contexts. Each example provided as triplet of
20

---

MoHI: Multi-Head Attention as Mixture-of- -ead Attention
Routed Heads
Shared Heads
Routing Routing Store
Basketball
Routing Stur
Basketball
0.04
0.33
0.02
0.01
0.00
0.00
Routed Heads
Shared Heads IDD
0.05
0.04
Bookshop
Routing Samre
0.06
0.07
Booksbop
0.05
0.04
0.03
0.1
0.82
0.01
0.00
0.00
5
33
15
Routed Heads ID)
Shared Hisdss ID
Routing Store
Jean
Routiry Stire
Jean
0.33
0.02
0.91
0.00
0.00
Routed Heads
Shared Heads ID
Routing Soore
0.0
cs
Goldfish
Routing SSore
0.06
0.05
0.07
Goldfish
0.04
0.00
01
0.01
0.02
0.00
0.00
13
15
Routod Heads
Shared Heads ID
Rouiing Soort
0.05
cream
0.06
0.07
Ice cream
0.05
0.04
0.02
Routed Heads
Shared Heads ID
Figure D. Additional visualization of the head routing score distribution in MoH-DiT-XL/2. MoH-DiT-XL DLL- activates 90% of the attention heads.
(question, passage, and answer), with the page title optionally included as additional context. We report the accuracy for the 32-shot BoolQ.
(Mihaylov et al., 2018) is a question-answerrngg dataset designed assess understanding of elementary-leve science, similar open-book exams. It contains 5,957 multiple-choice questions based "book" of 1,326 science facts. The dataset requires not only knowledge of these facts but also the application of broad common knowledge. includes mappings from each question to the core fact it targets and additional common knowledge facts. The dataset also provides scores of human accuracy and clarity, well as crowd-sourced data for further analysis. We report O-shot accuracy on OpenbookQA.
OpenbookQA
TruthfulQA (Lin et al., 2022) is a benchmark designed to evaluate the truthfulness of a language model's responses. consists of 817 questions across 38 categories, such as health, law, finance, and politics. The questions are crafted reflect common false beliefs misconceptions that might lead humans to answer inaccurately. We report O-shot accuracy TruthfulQA.
GSM8K (Cobbe et al., 2021) is a dataset containing 8.5K high-quality, linguistically diverse grade school math word problems. It is divided into 7.5K training problems and 1K test problems. Each problem requires 2 to 8 steps to solve, typically involving sequence of elementary calculations using basic arithmetic operations. capable middle school student should be able to solve all the problems, making the dataset suitable evaluating multi-step mathematical reasoning. report the exact match score for 8-shot GSM8K.
CEVAL (Huang al., 2023) is comprehensive Chinese evaluation suite designed to assess the advanced knowledge and reasoning abilities of LLMs in Chinese context. It includes multiple-choice questions across four difficulty levels (middle school, high school, college, and professional) and spans 52 diverse disciplines. We report the accuracy for the 5-shot CEVAL.
CMMLU (Li et al., 2023a) is a comprehensive Chinese benchmark designed to evaluate the knowledge and reasoning abilities of LLMs across various subjects, including natural sciences, social sciences, engineering, and humanities. We
21

---

MoHI: Multi-Heaa Attention as coxture-of-Head Attention
Routed Heads
Shared Heads
Routing Score
0.05
LogiQA
Routing Soore
0.25
LogiQA
0.02
0.10
0.03
0.05
0.00
0.00
Routed Heads
Shardd Heads ID
Routing Soore
0.05
PIQA
Routing Soore
0.25
PIQA
0.04
0.20
0.03
0.15
0.10
0.00
0.00
11
Routed Heads
Shared Heads ID
Routing Score
WinoGrande
Routing Score
WnnGrande
0.03
0.15
0.02
0.10
0.01
0.05
0.00
0.00
Routed Heads ID
Shared Heads ID
Figure E. Additional visualization of the head routing score distribution MoH-LLL-BB, MoH-LLM BB activate 75% the attention
heads.
report the accuracy for the 5-shot CMMLU,
22

---

MolI: Multi-Heat Head Attention as Mixture-of- Head Attention
Figure F. Images generated from the proposed MoH--DT-XLL2 model. We show samples generated from class-condition MoH-DiT-XIL22 model trained on ImageNet at 256 resolution. MoH-Dii-XXL-2 activates 90% of the attention heads.
23


Reviews:
[{"review_id": "42uCTgfYqx", "reviewer": "Reviewer_ekLt", "summary": "This paper aims to enhance the efficiency of multi-head self-attention by integrating mixture of experts into the attention. The authors propose Mixture of Head attention, which selectively activates subsets of attention heads for each token and gets a weighted sum of these selected heads to get the final output. The approach demonstrates improved efficiency and performance across some tasks, including image classification, image generation, and both fine-tuning and training large language models from scratch.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thoughtful comments and for recognizing that \"The experiments in the paper are well-conducted.\" Below, we address your questions in detail.\n\n**Q1:** One weakness of the method is its relatively high activation rate.\n\n**A1:** To demonstrate the robustness of our method, we only replace multi-head attention with MoH in various structures while keeping the original training parameters unchanged. Therefore, the results presented in the manuscript may not be optimal. Our latest experimental results indicate that, with tuning, our method can achieve even higher performance with a lower activation rate. Besides, we are developing deep learning-based methods to automatically determine the optimal activation ratio. We believe that our proposed MoH is promising and can be further optimized for even better performance.\n\n|     | # Activated Heads (%) | **Image Classification (Acc)** |\n|:---:|:---------------------:|:-----------:|\n| Multi-Head Attention |          100          |     84.8%     |\n| MoH |          50          |    **85.0%**     |\n\n\n**Q2:** Lack of comparison with mixture of feedforward experts.\n\n**A2:** Thanks for your insightful advice. We explain the difference between MoH and MoE from the following three aspects:\n* Attention and FFNs are the core components of Transformers. While MoE applies sparse activation at the FFN level, MoH introduces sparsity at the attention level. MoH not only extends the scope of MoE but also offers a more effective approach to reducing Transformer computation. **This is particularly significant because, as the input sequence length increases, FFN computation grows linearly, while attention computation scales quadratically.** Consequently, MoH has greater potential to alleviate the computational burden of Transformers.\n* MoH presents greater technical challenges than MoE. Unlike the MoE upcycling technique, which copies the FFN to increase the model size, **our MoH prunes the original model to reduce the activation parameters, making it more challenging.**\n* **MoH naturally leverages the multi-head structure in Transformers, while MoE requires additional FFN replication.** From this perspective, MoH matches the original Transformer design better.\n\n**Q3:** Lack of integration with mixture of feedforward experts.\n\n**A3:** Thanks for your valuable suggestion. Due to the time constraints of the rebuttal, we designed a 28M-sized small model as a baseline, with a training budget of 100 epochs on ImageNet-1K classification. As shown in the table below, MoH can be combined with MoE, where MoE enhances the model's performance by replicating the FFN, while MoH optimizes computational efficiency by the dynamic activation of attention heads.\n\n|     | #  Params | # Activated Heads (%) | **Image Classification (Acc)** |\n|:---:|:---------------------:|:-----------:|:-----------:|\n| Multi-Head Attention |  28M  |          100          |     77.0%   |\n| MoH |  28M  |          75          |     77.2%   |\n| MoE Top-1/2E |   45M     |         100          |     **78.1%**    |\n| MoH & MoE Top-1/2E  |    45M        |        75          |    **78.1%**     |\n\nWe sincerely thank you for your constructive comments. We will add the above important discussions in the final manuscript and highlight them. Thanks again for taking the time and effort on our paper.", "claims_and_evidence": "The paper claims to reduce computational costs while maintaining comparable or favorable performance. The experiments confirm this to a certain extent. However, one weakness of the method is its relatively high activation rate.", "methods_and_evaluation_criteria": "The proposed method is validated on well-known tasks, including image classification and generation, as well as training and fine-tuning LLMs.", "theoretical_claims": "The paper does not introduce new theoretical contributions, so there are no proofs to verify.", "experimental_designs_or_analyses": "The experiments in the paper are well-conducted, and the efficiency of the proposed method is confirmed across various tasks and datasets. Additionally, the experimental details are clearly reported.  \n\nHowever, one major weakness of the paper is the lack of discussion on the standard MoE applied to feedforward layers:  \n\n- *Lack of comparison with mixture of feedforward experts:* Traditional MoE is typically applied to feedforward layers and has demonstrated significant improvements in both efficiency with a highly sparse activation rate and performance. Furthermore, training feedforward MoE is generally simpler compared to the proposed MoH.  \n\n- *Lack of integration with mixture of feedforward experts:* Would incorporating MoE in both feedforward and attention layers result in even better results? It would be great to explore this possibility and see if it provides further efficiency gains and performance improvements.", "relation_to_broader_scientific_literature": "This introduces a new type of mixture of experts, which I find both important and interesting.", "essential_references_not_discussed": "I did not identify any essential or significant related works that are missing.", "other_strengths_and_weaknesses": "Please refer to the discussion above.", "other_comments_or_suggestions": "Please refer to the discussion above.", "questions_for_authors": "Please refer to the discussion above.", "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": "L0viIQ4EdM", "reviewer": "Reviewer_6acz", "summary": "The paper introduces Mixture-of-Head Attention (MoH) as an enhancement to the multi-head attention (MHA) mechanism in Transformer models, aiming to reduce computational costs while maintaining or improving model accuracy. The key insight is that not all attention heads contribute equally, and some can be pruned or dynamically selected without significantly affecting performance. Inspired by Mixture-of-Experts (MoE) models, MoH treats attention heads as experts and introduces a router that selects the most relevant heads for each token. This allows MoH to activate only a subset of attention heads dynamically, improving efficiency without increasing the number of model parameters. Additionally, MoH replaces the standard summation of multi-head attention with a weighted summation, enhancing flexibility.\n\nThe authors validate MoH across multiple architectures, including Vision Transformers (ViT), Diffusion Transformers (DiT), and Large Language Models (LLMs). Results show that MoH can match or outperform standard MHA while using only 50%–90% of the attention heads. Notably, MoH-LLaMA3-8B achieves a 2.4% accuracy improvement over LLaMA3-8B using only 75% of the attention heads.\n\n---\n\n## Update Post-rebuttal\n\nI am happy with the authors response and would like to see the paper accepted.\n\n---", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the constructive comments, and for noting that our method provides \"a new perspective on structured sparsity in Transformers\" and that \"the significance is high.\" We address the questions as below.\n\n**Q1:** If MoH is inspired by Mixture-of-Experts, then why does it manually designate some heads as \"always active\" instead of letting the router handle all head selection dynamically?\n\n**A1:** We explain this problem in three aspects: \n* **Load balance loss (also called MoE loss) pushes experts to focus on specific areas rather than general knowledge.** As shown below, load balance loss ensures all experts are chosen equally and become specialized. However, this goes against the idea that essential experts should be selected more often and should learn broader, general knowledge. **Due to the load balance loss, even though some heads are critical enough to remain active at all times, the routing function cannot naturally prioritize them.**\n\n$L_b = \\sum_{i=h_s+1}^{h} P_i f_i,$\n\n$ P_i = \\frac{1}{T} \\sum_{t=1}^{T} \\text{Softmax}(W_{r} x_{t})_{i-h_s}, $\n\n$f_i = \\frac{1}{T} \\sum_{t=1}^{T} {1}(\\text{Token } {x}_t \\text{ selects Head } i)$.\n\n* **From the perspective of training stability, keeping some heads active at all times helps maintain a stable gradient.** If all heads are selected freely, the gradients and loss spikes can increase significantly, reducing training efficiency.\n\n* In recent MoE work [1], some experts are also selected as shared experts to extract general knowledge. Besides, in attention mechanisms, some heads may capture common knowledge across different contexts, such as grammatical rules in language. Inspired by this idea, we designate a subset of heads as shared heads that remain always activated. **We also compare the gradients and training data distribution of shared heads and routed heads in Appendix Table A to further demonstrate that shared heads play a key role in capturing general knowledge.**\n\n[1] Dai, Damai, et al. \"Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models.\" arXiv preprint arXiv:2401.06066 (2024).\n\n**Q2:** Claims about computational efficiency are not supported by FLOPs/memory benchmarks.\n\n**A2:** Thanks for your valuable suggestion. We have provided the comparison of the efficiency between MoH and multi-head attention mechanisms in Table 7, and we have added FLOPs and memory as suggested. Specifically, MoH surpasses standard multi-head attention mechanisms, with its advantage becoming more pronounced as the input sequence length increases.\n\nWe find that MoH slightly reduces GPU memory usage, though the difference is not significant. This is because GPU memory is primarily used to store model parameters, gradients, and the KV cache. Since MoH only optimizes attention computation, it does not substantially reduce the GPU memory of these three components.\n\n|                      | # Head Num | # Head Dim | # Sequence Length | #Activated Heads (%) | Time (ms) | FLOPs | Memory |\n|:--------------------:|:----------:|:----------:|:-----------------:|:--------------------:|:---------:|:---------:|:---------:|\n| Multi-Head Attention |     32     |     64     |        256        |         100          |   0.360   | 23M | 998M |\n|         MoH          |     32     |     64     |        256        |          90          |   0.352   | 22M | 980M |\n|         MoH          |     32     |     64     |        256        |          75          |   0.321   | 18M | 978M |\n|       **MoH**        |     32     |     64     |        256        |          50          | **0.225** | **13M** | **978M** |\n| Multi-Head Attention |     32     |     64     |        512        |         100          |   1.376   | 83M | 3356M |\n|         MoH          |     32     |     64     |        512        |          90          |   1.351   | 77M | 3354M |\n|         MoH          |     32     |     64     |        512        |          75          |   1.180   | 65M | 3328M |\n|       **MoH**        |     32     |     64     |        512        |          50          | **0.863** | **45M** | **3302M** |\n\n**Q3:**  No comparisons are made against alternative dynamic attention mechanisms (e.g., MoE-based attention, sparse routing).\n\n**A3:** Thanks for your insightful advice. As shown in the table below, MoH outperforms both sparse attention and MoE-based attention. These results also demonstrate the benefits of shared heads from an experimental perspective.\n\n|     | # Activated Heads (%) | **Image Classification (Acc)** |\n|:---:|:---------------------:|:-----------:|\n| Sparse attention S |          -         |    78.4%     |\n| MoE-based attention S |          75          |    75.6%     |\n| MoH S |          75          |    **78.6%**     |\n\nWe sincerely thank you for your constructive comments. We will add the above important discussions in the final manuscript and highlight them. Thanks again for taking the time and effort on our paper.", "claims_and_evidence": "The paper presents MoH (Mixture-of-Head Attention) as an alternative to standard Multi-Head Attention (MHA) and claims that it improves efficiency and accuracy across multiple architectures (ViTs, DiTs, and LLMs). While the empirical results generally support these claims, some claims require further validation. For instance, the authors claim that MoH improves efficiency without increasing parameter count, but didn't provide FLOP counts or memory benchmarks to confirm efficiency gains. Moreover, the authors mentioned that MoH's routing strategy is optimal for balancing shared and routed heads, but didn't compare  this strategy against say dynamic attention strategies such as sparse attention.", "methods_and_evaluation_criteria": "MoH seems to be well-suited for improving Transformer efficiency, and the evaluation covers a diverse range of architectures (ViTs, DiTs, and LLMs) using 14 benchmark datasets (e.g., MMLU, CEVAL, GSM8K, TruthfulQA), making the results broadly applicable. However, the paper lacks efficiency metrics (e.g., FLOPs, memory usage, latency) despite claiming computational improvements. Additionally, MoH’s manual selection of shared vs. routed heads is not compared to fully dynamic routing methods (e.g., MoE-based attention), leaving some doubts about whether the hybrid design is necessary. Addressing these gaps could strengthen the empirical validation of MoH.", "theoretical_claims": "The core theoretical claims about selective activation and weighted summation are well-supported, but the manual head selection strategy and efficiency gains lack mathematical justification.", "experimental_designs_or_analyses": "The experimental design provides a broad evaluation of MoH across ViTs, DiTs, and LLMs, using many benchmark datasets (e.g., MMLU, CEVAL, GSM8K, TruthfulQA), which strengthens the generalization claims. The experimental setup is strong in terms of dataset diversity, but the lack of efficiency benchmarks and alternative routing comparisons in MoH’s empirical claims.", "relation_to_broader_scientific_literature": "The paper builds on prior work in Multi-Head Attention and Mixture-of-Experts models by introducing Mixture-of-Head Attention, which selectively activates attention heads to improve efficiency. The concept of routing-based selection is inspired by MoE-based Transformers, but MoH applies it at the attention head level instead of full feedforward layers, making it a more lightweight alternative. The idea of structured sparsity in Transformers aligns with research on sparse attention mechanisms and adaptive token selection methods, though MoH introduces a unique hybrid model where some heads are manually fixed while others are routed dynamically. However, the paper does not compare MoH against fully dynamic routing mechanisms like MoE-based attention, which raises questions about whether its manual head selection strategy is necessary.", "essential_references_not_discussed": "The paper provides theoretical justification for MoH, focusing on the routing mechanism and its impact on efficiency and expressivity. The derivations related to selective head activation and weighted summation appear logically consistent, following standard MoE formulations. However, the manual selection of shared heads is not rigorously justified—there is no proof explaining why certain heads must always be active, rather than letting the routing function learn to retain essential heads automatically. Additionally, the paper does not compare MoH's routing formulation against fully dynamic MoE-based attention models, leaving its mathematical superiority unverified.\n\n**If MoH is inspired by Mixture-of-Experts, then why does it manually designate some heads as \"always active\" instead of letting the router handle all head selection dynamically? I think this partially undermines the motivation for using a routing mechanism in the first place.**", "other_strengths_and_weaknesses": "The paper presents a refinement of Multi-Head Attention by introducing Mixture-of-Head Attention, which selectively activates heads to improve efficiency without increasing parameter count. This hybrid approach balances shared and dynamically routed heads, offering a new perspective on structured sparsity in Transformers. The significance is high, as MoH generalizes across ViTs, DiTs, and LLMs, making it applicable to a wide range of architectures. However, clarity is hindered by the lack of justification for manually fixing some heads as always active, which contradicts the motivation for routing-based selection. Additionally, claims about computational efficiency are not supported by FLOP/memory benchmarks, and no comparisons are made against alternative dynamic attention mechanisms (e.g., MoE-based attention, sparse routing). Strengthening these aspects would solidify MoH’s contributions and further validate its impact.", "other_comments_or_suggestions": "None.", "questions_for_authors": "If MoH is inspired by Mixture-of-Experts, then why does it manually designate some heads as \"always active\" instead of letting the router handle all head selection dynamically?\n\nIf certain heads are critical enough to be always active, why doesn’t the routing function naturally prioritize them?", "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": "O7Lg0hWfNc", "reviewer": "Reviewer_gWSK", "summary": "The paper introduces MoH (Mixture-of-Head Attention), a novel perspective on multi-head attention that formulates each head as an expert in a Mixture of Experts (MoE) framework. It employs a two-stage routing mechanism—comprising shared and non-shared experts—to reduce computational costs while enhancing accuracy. The authors validate their approach through experiments on both training-from-scratch and fine-tuning settings across vision tasks (such as image classification and class-conditional image generation) and language tasks. The results demonstrate that MoH outperforms or matches vanilla multi-head attention while activating fewer parameters.\n\n---\n\n## Update after rebuttal\nAfter reading the rebuttal, I appreciate that the authors thoroughly attempted to address and clarify all of my questions. However, my primary concern remains the incremental nature of the current version of the paper.\n\nFirst of all, the idea of using shared and non-shared heads originates from DeepSeek-MoE. While the authors claim to extend this idea by analyzing the differences between shared and routed heads through feature rank analysis, I view this more as an experimental insight rather than a fundamentally novel contribution.\n\nMoreover, the core idea of the paper is based on the equivalence between the summation and concatenation forms of multi-head attention. Building on this, the authors propose treating each head as an expert and applying a mixture-of-experts framework. This idea is relatively straightforward, and many other components are drawn from existing literature. In response to my concerns, the authors explained their method for integrating MoE into multi-head attention in the presence of  dynamic KV cache. Specifically, they compute the key and value for all heads, even those not selected by a given token. While this is a thoughtful engineering solution, I do not find it substantial enough to constitute the primary contribution of the work. Since the core idea of the paper is simple and the effectiveness is demonstrated through empirical results, I think the paper requires stronger mathematical justification to support the effectiveness of MoE, which could enhance the overall contribution.\n\nNevertheless, I appreciate the new perspective introduced in the paper, especially the reinterpretation of multi-head attention in summation form, which to my knowledge has not been explored in previous works. Along with the improved empirical results, I believe this work sheds light on potential advancements in multi-head attention using MoE and paves the way for more theoretical works in this field.\n\nTherefore, I have increased my score from 2 to 3 and leaned towards acceptance after considering the rebuttal.\n\n---", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thoughtful comments and your recognition that our method \"provides a promising suggestion to advance attention-based models.\" Below, we provide detailed responses to your questions.\n\n**Q1:** While prior work supports this claim when heads are combined through concatenation, the proposed method instead uses summation.\n\n**A1:** First, in the forward pass, the summation and concatenation forms give the same results, i.e., $ \\textrm{MultiHead}({X},{X}')=\\sum_{i=1}^{h} {H}^{i}{W}_O^{i}=\\textrm{Concat}({H}^{1}, {H}^{2}, ..., {H}^{h}){W}_O$. \n\nSecond, the gradients are also identical, i.e., $ \\frac{ \\partial  \\textrm{MultiHead}({X},{X}') }{\\partial H_{i} }={W}_O^{i}$. Therefore, the summation and concatenation forms are mathematically equivalent.\n\n**Q2:**  Theoretical guarantee for the effectiveness of the proposed methods.\n\n**A2:** Thanks for your valuable suggestion. **As mentioned by  Reviewer xCL1, we put the theoretical guarantee that MoH is superior to multi-head attention in Appendix A.** Specifically, we proved that MoH not only improves efficiency and model performance but also helps different attention heads to specialize better compared to multi-head attention, from both theoretical and experimental perspectives.\n\n**Q3:** The methodology is incremental.\n\n**A3:** **The motivation for our work is to upgrade the multi-head attention mechanism, the core of the Transformer model, to reduce computational costs while maintaining or surpassing the previous accuracy level, rather than making improvements to the MoE.** We propose combining multi-head attention and MoE, so we adopt the MoE technique:\n\n* For the two-stage routing, as shown in our response (A1) to Reviewer 6acz, the reason for our two-stage routing is to capture general knowledge. Besides, we compare the gradients and training data distribution of shared heads and routed heads in Appendix Table A to further demonstrate that shared heads play a key role in capturing general knowledge.\n \n* For the load-balancing loss, since we decompose multi-head attention into a summation form, which is similar to MoE structure, we directly adopt the auxiliary loss used in MoE.\n\n**Q4:** The lack of discussion of the new experts defined.\n\n**A4:** As suggested, we have defined the expert as $H^i W^i_O$ to avoid possible misunderstanding.\n\n**Q5:** Further theoretical discussion regarding the effectiveness, the convergence rate and the optimization scheme.\n\n**A5:** Mathematically, we prove that the summation and concatenation forms are equivalent. Besides, we show in Appendix A that the gradient per head in the MoH differs from the gradient of multi-head attention by only a single weight. **Finally, we replace multi-head attention with MoH in various structures while keeping the original training parameters unchanged. The experimental results demonstrate that MoH enhances the performance of multi-head attention, providing experimental evidence of its effectiveness.**\n\n**Q6:** The abuse of notation $W$.\n\n**A6:** As suggested, we have replaced the notation $W$ in the router with $E$, referring to them as expert embeddings.\n\n**Q7:** Concerns about the claim of the author that the new methods can be fine-tuned from pretrained models.\n\n**A7:** We explain this problem in three aspects:\n* We simply select the first 16 attention heads of each layer as shared heads. **Even if the structure is not optimal, the experimental results show that MoH-LLaMA3-8B has a significant advantage over LLaMA3-8B. This result shows the robustness of our method.**\n* Unlike the MoE upcycling technique, which copies the FFN to increase the model size, **our MoH prunes the original model to reduce the activation parameters**, making it more challenging.\n* **To the best of our knowledge, our work is the first to attempt to reduce the computational amount of the attention mechanism without degrading the model performance by continue-tuning on a pre-trained model.**\n\n**Q8:** Acknowledge the base code they adapted.\n\n**A8:** Thanks for your advice. We will acknowledge the contributors in our official code, and release our trained models.\n\n**Q9:** Existing works are not thoroughly discussed.\n\n**A9:** As suggested, we have expanded and improved the discussion of related works.\n\n**Q10:** What are the motivations to define $\\alpha_1$ and $\\alpha_2$ as in Eq. 6 rather than considering them as hyperparameters?\n\n**A10:** We choose to predict $\\alpha_1$ and $\\alpha_2$ based on the input so that different tokens can dynamically combine general knowledge from the shared heads and specialized knowledge from the routed heads.\n\n**Q11:** Explain the parameter-free router mentioned in Section 4.4 (line 319)?\n\n**A11:** We use the $l_2$ norm (which measures the magnitude of a vector in Euclidean space) of each head to represent its importance. We then normalize the $l_2$ norm of all heads using SoftMax. This simple router achieves results comparable to learnable routers.", "claims_and_evidence": "1. The motivation for proposing MoH as a dynamic-head routing mechanism is clear. It builds on existing literature showing that redundant heads in traditional multi-head attention can be pruned to reduce computational cost while maintaining accuracy. However, while prior work supports this claim when heads are combined through concatenation, the proposed method instead uses summation. This raises concerns about whether the claim from the literature still holds in this new form.  \n\n2. I cannot find any theoretical guarantee for the effectiveness of the proposed methods, particularly when the heads are combined using summation via MoE gate. Even the appendix does not provide a clear justification or proof.  \n\n3. The methodology is incremental. I will provide detailed comments in the following section.\n\n4. Existing works that explore the mixture of heads in attention mechanism are not thoroughly discussed in the related works section.", "methods_and_evaluation_criteria": "While I appreciate the authors' effort in developing a general framework that integrates Mixture of Experts (MoE) into attention layers in Transformers, the paper appears to be an incremental extension of existing work rather than presenting a novel contribution. Additionally, the writing and the formulation of the proposed methods present several issues.\n\n1. Most of the components of the proposed methods are directly adapted from previous works, making the paper more of a combination of existing approaches rather than a novel contribution.  \n\n- The core idea of two-stage routing involves introducing shared experts to capture global information, as presented in [1]. The original motivation of Dai et al. for this approach was to enhance expert specialization. However, this paper adapts it directly without discussing its suitability or potential additional benefits for the proposed mixture-of-head attention setting.  \n\n- The load-balancing loss is also directly adapted from previous works.  \n\n2. The summation form of combining heads is inconsistent with the authors' claim.\n- The author claims that in the proposed method, each head is divided by rows and then concatenated according to Eq. (3). Nevertheless, only $W_O$ is divided into smaller $W_O^{i}$, and $H_i$ remains undiscussed, meaning that $H_i$ is obtained as the vanilla multi-head attention, which is not divided by rows. I recommend the authors state clearly here that the expert defined is $H_i.W_O^{i}$ to avoid confusion that the experts are $W_O^{i}$ only.\n\n3. The lack of discussion of the new experts defined.\n- While traditional MoE considers experts as a FFN, the new formulation of each expert $i$ as $H_i.W_O^{i}$ needs further theoretical discussion regarding the effectiveness, the convergence rate and the optimization scheme.\n\n4. The abuse of notation $W$ in computing gating score makes the paper confusing.\n\n- As far as I understand, the paper implies $W_r$ and $W_s$ as the expert embeddings. However, the notation $W$  and its description as a *projection matrix* may be confusing, as it resembles the projection matrices used in a standard Transformer.\n\n---\nReferences\n\n[1] Dai, Damai, et al. \"Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models.\" arXiv preprint arXiv:2401.06066 (2024).", "theoretical_claims": "I find no theoretical guarantees or mathematical proofs in the paper.", "experimental_designs_or_analyses": "I have checked all experimental parts and found that the experimental design is appropriate, which includes a wide range of tasks from vision to language. However, I still have concerns about the claim of the author that the new methods can be fine-tuned from pretrained models, which is an advantage compared to [1]. \n\n- As far as I understand, the authors initialize the router weights and copy the pretrained weight for $W_O$, $W_Q^{i}$, $W_{K}^{i}$ to finetune it. However, with that scheme, the decision of the shared heads and the division of $W_O$ into smaller $W_i$ will affect the initialization of the models (eg. which pretrained head is better to share).\n- Additionally, if the proposed method can be fine-tuned by copying weights from pretrained models, then other approaches involving a mixture of heads, such as [1], can also leverage this technique by copying pretrained weights and initializing routers (e.g., sparse upcycling [2], [3]). This suggests that the empirical advantages of MoH are not unique since upcycling techniques could be applied to other methods as well, diminishing the claimed advantage of MoH.  \n\n---\nReferences\n\n1. Zhang, Xiaofeng, et al. \"Mixture of attention heads: Selecting attention heads per token.\" arXiv preprint arXiv:2210.05144 (2022).\n\n2. Komatsuzaki, Aran, et al. \"Sparse upcycling: Training mixture-of-experts from dense checkpoints.\" arXiv preprint arXiv:2212.05055 (2022).\n\n3. Zhang, Qizhen Irene, et al. \"Bam! just like that: Simple and efficient parameter upcycling for mixture of experts.\" Advances in Neural Information Processing Systems 37 (2024): 56304-56321.", "relation_to_broader_scientific_literature": "From my own perspective, this work provides a promising suggestion to advance attention-based models despite the weaknesses mentioned above.", "essential_references_not_discussed": "The authors should provide a more in-depth discussion of related works involving the mixture of heads in Section 2, along with a critical analysis of their weaknesses to better highlight the paper’s contribution. While I acknowledge the comparison with Mixture of attention heads (MoA) [1] in Section 5, the broader literature on this topic remains insufficiently discussed ([2], [3]).\n\n---\nReferences\n\n[1] Zhang, Xiaofeng, et al. \"Mixture of attention heads: Selecting attention heads per token.\" arXiv preprint arXiv:2210.05144 (2022).\n\n[2] Csordás, Róbert, et al. \"Switchhead: Accelerating transformers with mixture-of-experts attention.\" Advances in Neural Information Processing Systems 37 (2024): 74411-74438.\n\n[3] Hao Peng, Roy Schwartz, Dianqi Li, and Noah A. Smith. A mixture of $h - 1$ heads is better than $h$ heads. In Proc. Association for Computational Linguistics (ACL), pages 6566–6577, Virtual only, July 2020.", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "I suggest that the authors move the comparison between MoH and MoA to the related works and consider the limitations of Mixture of attention heads (MoA) as a motivation for proposing MoH.", "questions_for_authors": "See the weaknesses mentioned in comments about methodology and experimental results above. Furthermore, I have additional questions for the authors:\n\n1. What are the motivations to define $\\alpha_1$ and $\\alpha_2$ as in Eq. 6 rather than considering them as hyperparameters? \n\n2. Explain the *parameter-free router* mentioned in Section 4.4 (line 319)?", "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": "SK8pJO04Fg", "reviewer": "Reviewer_v1g7", "summary": "The paper proposes leveraging the Mixture-of-Experts (MoE) mechanism to upgrade the standard Multi-Head Attention into a novel Mixture-of-Heads (MoH) Attention. Specifically, MoH replaces the standard summation in multi-head attention with a weighted summation, where the weights are determined by a newly introduced Two-Stage Routing strategy. Experiments on Vision Transformers (ViT), Diffusion Transformers (DiT), and large language models (LLMs) are conducted to demonstrate the effectiveness of the proposed MoH technique.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thoughtful comments and for recognizing that our method is a \"practical contribution,\" \"can be fine-tuned from pre-trained multi-head attention models,\" and \"the results are promising and suggest broad applicability.\" Below, we address your questions in detail.\n\n**Q1:** There are inconsistencies between the theoretical formulation of MoH and its implementation in the LLaMA3-8B experiment.\n\n**A1:** We explain this problem in three aspects:\n* Experiments on ViT, DiT, and LLMs are conducted with models trained from scratch, giving us more flexibility to modify the architecture. **However, LLaMA3-8B is a pre-trained model trained on 15T tokens, while we had only 400B tokens (about 3% of the pre-training data) for continue-tuning. This limited our ability to alter the model structure significantly.** To address this, we adjusted our original formula to maximize weight reuse from the pre-trained model while preserving its output distribution.\n\n* **We have conducted experiments on LLMs trained from scratch, demonstrating the advantages of our method.** Since much of the current research focuses on fine-tuning open-source LLMs, we introduced additional experiments to continue-train LLaMA3-8B. **This contribution has been acknowledged by all other reviewers.**\n\n* To the best of our knowledge, our work is the first to attempt to reduce the computational amount of the attention mechanism without degrading the model performance by continue-tuning on a pre-trained model. We consider our proposed training techniques to be a valuable additional contribution.\n\n**Q2:** The paper does not provide a clear rationale for how the value of K in the Top-K selection is determined.\n\n**A2:** Thanks for your insightful advice. There are two ways to set this up: (1) manual configuration and (2) learning through learnable masks. In the manuscript, we show the results of manual configuration. For the learnable mask approach, we follow a method similar to [1]. Specifically, we introduce a mask module with binary values {0,1} applied to the heads. These masks are dynamically learned rather than statically assigned, allowing the model to determine which heads should be shared. The remaining heads are designated as routed heads, and we set K based on a predefined ratio. For example, if the ratio is 1/4 and there are 8 routed heads, then K is set to 2. As shown in the table below, our latest experimental results show that learning through learnable masks is generally better than manual configuration.\n\n|     | # Activated Heads (%) | **Image Classification (Acc)** |\n|:---:|:---------------------:|:-----------:|\n| Baseline |          100          |     84.8%     |\n| Manual Configuration |           75          |     84.9%     |\n| Manual Configuration |          50          |     84.7%    |\n| Learning Through Learnable Masks |          50          |    **85.0%**     |\n\n[1] Liu, Peiyu, et al. \"MOEfication by Experts as Masks.\"\n\n**Q3:** The relationship between the K value and the ratio of shared heads is not discussed.\n\n**A3:** Thanks for your valuable suggestion. It is worth noting that when using learnable masks to determine the number of shared heads, we only need to set the ratio of K among the routed heads. As shown in the table below, this ratio is a trade off parameter. If this ratio is too small, the number of activated heads will be insufficient, potentially degrading performance. Conversely, if this ratio is too large, it reduces the model’s sparsity, limiting efficiency improvements.\n\n|     | # The Ratio of K | **Image Classification (Acc)** |\n|:---:|:---------------------:|:-----------:|\n| Learning Through Learnable Masks |          1/8          |    84.8%     |\n| Learning Through Learnable Masks |          1/4         |    **85.0%**     |\n| Learning Through Learnable Masks |          1/2          |    84.9%     |\n\nWe sincerely thank you for your constructive comments. We will add the above important discussions in the final manuscript and highlight them. Thanks again for taking the time and effort on our paper.", "claims_and_evidence": "Part of.\n\nThere are inconsistencies between the theoretical formulation of MoH and its implementation in the LLaMA3-8B experiment. Specifically, the router used in the experiments differs from the one described in Eqs. (5) and (6); moreover, the weighting mechanism also deviates from the theoretical presentation. These discrepancies raise concerns about the alignment between the proposed theory and its practical application, and they should be addressed to ensure the validity of the results.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "NA.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "The key contributions of the paper are related to building a general transfrormer-based network architecture.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths:\n\n- The proposed MoH technique is simple and easy to integrate into existing architectures, making it a practical contribution to the field.\n\n - The idea of using a weighted summation to enhance multi-head attention is intuitively appealing and expected to improve performance.\n\n- The experiments span a wide range of tasks, including classification, generation, and LLMs. The results are promising and suggest broad applicability.\n\n- The paper is well-written and easy to follow.\n\nWeaknesses:\n\n- There are inconsistencies between the theoretical formulation of MoH and its implementation in the LLaMA3-8B experiment. Specifically, the router used in the experiments differs from the one described in Eqs. (5) and (6); moreover, the weighting mechanism also deviates from the theoretical presentation.\nThese discrepancies raise concerns about the alignment between the proposed theory and its practical application, and they should be addressed to ensure the validity of the results.\n\n- The paper does not provide a clear rationale for how the value of K in the Top-K selection is determined. This is a critical parameter that directly impacts the behavior of the MoH mechanism. \nBesides, the relationship between the K value and the ratio of shared heads is not discussed. Understanding this relationship is essential for interpreting the results and optimizing the method for different tasks.", "other_comments_or_suggestions": "No.", "questions_for_authors": "Please see the above comments.", "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": "prkpUi3x1X", "reviewer": "Reviewer_6df6", "summary": "The paper proposes Mixture-of-Head (MoH), a replacement for the standard attention mechanism, in which attention heads can be adaptively switched on and off and reweighted for each token. This proposal is motivated by the already studied redundancy/specialization of attention heads, and the authors show that MoH can maintain or even improve the performance of a variety of transformer networks across different tasks while activating just a fraction of the attention parameters.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thoughtful comments and for recognizing that \"the idea of adding a MoE-like routing in the attention mechanism is fascinating and worth investigating, especially for efficiency purposes.\" Below, we address your questions in detail.\n\n**Q1:** In most cases, performance enhancements provided by MoH are marginal.\n\n**A1:**  We explain this problem in two aspects:\n* The motivation for our work is to upgrade the multi-head attention mechanism to reduce computational costs while maintaining or surpassing the previous accuracy level. Therefore, **our method activates fewer parameters than multi-head attention**, and our method performs better if it activates the same level of parameters.\n* To demonstrate the robustness of our method, **we only replace multi-head attention with MoH in various structures while keeping the original training parameters unchanged**. Our latest experimental results indicate that, with tuning, our method can achieve even higher performance.\n\n|     | # Activated Heads (%) | **Image Classification (Acc)** |\n|:--:|:--:|:--:|\n| Multi-Head Attention |          100          |     84.8%     |\n| MoH |          50          |    **85.0%**     |\n\n**Q2:**  The accessibility of the method could be improved by rephrasing the paragraph on two-step routing.\n\n**A2:**  Thanks for your insightful advice. As suggested, we have rewritten the paragraph on two-step routing to make it easier for readers to understand.\n\n**Q3:** The difference between the setting of ViT and DiT w.r.t. Llama.\n\n**A3:** We explain our experimental setup in detail:\n* **ViT for Image Classification**. We trained a MoH model **from scratch** based on TransNeXt. To ensure a fair comparison, **we only replace the standard multi-head attention with the proposed MoH, while keeping all other training parameters identical to TransNeXt**. \n* **DiT for Class-Conditional Image Generation**. We trained a MoH model **from scratch** based on DiT by replacing the standard multi-head attention with the proposed MoH. **We also keep all other training parameters identical to DiT**. \n* **Training LLMs from Scratch**. We trained the LLMs **from scratch, maintaining the multi-head attention baseline with exactly the same training parameters as the MoH model.**\n* **Continue-Tuning LLaMA3-8B**. To significantly enhance the applicability of the proposed MoH method, we attempt to further continue-tune pre-trained multi-head attention models, such as LLaMA3-8B, into MoH models.\n\n**Q4:** What strategy is used to choose shared heads?\n\n**A4:** All MoH models contain shared heads. There are two ways to choose shared heads: (1) manual configuration and (2) learning through learnable masks. In the manuscript, we show the results of manual configuration. For the learnable mask approach, we follow a method similar to [1]. Specifically, we introduce a mask module with binary values {0,1} applied to the heads. These masks are dynamically learned rather than statically assigned, allowing the model to determine which heads should be shared. The remaining heads are designated as routed heads, and we set K based on a predefined ratio. For example, if the ratio is 1/4 and there are 8 routed heads, then K is set to 2. As shown in the table below, our latest experimental results show that learning through learnable masks is generally better than manual configuration.\n\n|     | # Activated Heads (%) | **Image Classification (Acc)** |\n|:--:|:--:|:--:|\n| Baseline |          100          |     84.8%     |\n| Manual Configuration |          50          |     84.7%     |\n| Learning Through Learnable Masks |          50          |    **85.0%**     |\n\n[1] Liu, Peiyu, et al. \"MOEfication by Experts as Masks.\"\n\n**Q5:** Essential references not discussed.\n\n**A5:**  Thanks for your advice. We have expanded [1,2,3] and improved the discussion of related works.\n\n**Q6:** The dynamic routing works at the single token level.\n \n**A6:** We input the sentence \"Give me a short introduction to large language model.\" into MoH-LLaMA3-8B. We find the tokens forming the phrase share most of the activated heads. For example, all three tokens in \"large language model\" activate heads {6,7,8,9,11,12}. For other tokens, no clear pattern emerged in the activation of routed heads. Notably, besides routed heads, the shared heads create a stable semantic interaction between all tokens.\n\n|  Token   | IDs of the Activated Heads | \n|:--:|:--:|\n| Give |          0,3,7,8,9,12,14,15          |   \n| me |          2,7,8,10,11,13,14,15           |    \n|  a |          0,1,4,5,6,8,9,12         |    \n| short |       1,4,5,6,9,11,12,15             |    \n| introduction |          0,2,3,6,7,8,9,11          |    \n| to |          1,2,7,8,10,12,13,14          |    \n| **large** |        1,**6,7,8,9,11,12**,15            |    \n| **language** |       **6,7,8,9**,10,**11,12**,15            |    \n| **model** |          0,**6,7,8,9**,10,**11,12**          |   \n| . |          0,2,3,8,9,10,13          |", "claims_and_evidence": "Claims are generally well supported by evidence. My only concern in this direction is that, in most cases, performance enhancements provided by MoH are marginal or even absent compared with the standard multi-head attention mechanism, which is probably insufficient to sustain the claim that MoH consistently enhances model performance.", "methods_and_evaluation_criteria": "Yes, methods and evaluations are reasonable. The accessibility of the method could be improved by rephrasing the paragraph on two-step routing, as the roles of the learnable projection matrices are not immediately clear.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Experimental design is generally sound, though unclear in some parts. Specifically, the difference between the setting of ViT and DiT w.r.t. Llama could be made clearer. From my understanding, the MoH module is trained on a pre-trained transformer in both cases. What is then the peculiarity of continual tuning in the case of Llama? Moreover, in the cases of ViT and DiT, it is not clear whether shared heads are assigned and what strategy is used to choose them.", "relation_to_broader_scientific_literature": "N/A (there is some information regarding this in other sections of the form).", "essential_references_not_discussed": "I would point the authors to the following related works that might be worth discussing in the manuscript:\n\n- Interpreting CLIP's Image Representation via Text-Based Decomposition; Gandelsman et al., ICLR 2024. In this paper, the authors find that the attention heads of CLIP tend to **specialize** in specific input attributes.\n\n- Decomposing and Interpreting Image Representations via Text in ViTs Beyond CLIP; Balasubramanian et al., NeurIPS 2024. This paper generalizes such findings to non-contrastive Vision Transformers;\n\n- ResiDual Transformer Alignment with Spectral Decomposition; Basile et al., 2024. In this paper, the specialization property of CLIP heads is used to show that **few heads can outperform the whole model** on zero-shot classification tasks.", "other_strengths_and_weaknesses": "As an additional strength, I think the idea of adding a MoE-like routing in the attention mechanism is fascinating and worth investigating, especially for efficiency purposes (but also interpretability!).\n\nOn the weaknesses side, they are scattered around aspect-specific parts of the review (e.g., Experimental Designs).", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. It's unclear to me how the dynamic routing works at the single token level. This is a core contribution to the paper (as stated repeatedly, including in the abstract), so I would like to better understand it and see it more clearly explained in the paper, where I feel most of the focus is at head level. What exactly happens when a token $t$ in sentence $S$ is routed/activated in head $h$? Does it interact only with the other ones routed to $h$?", "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": "IME9Zg63Oy", "reviewer": "Reviewer_xCL1", "summary": "This paper proposes Mixture-of-head attention (MoH) to replace standard multi-head attention(MHA) in transformers. The key idea is treat each attention head as an expert in mixture-of-expert framework. The experiments demonstrate that MoH can be applied to vision transformers(ViT) for image classification, DiT for diffusion-based image generation, and LLM. The results shows MoH can achieve competitive perforamnce with only 50%-90% of the heads/experts to be active.\n\nStrengths:\n\n  The idea of applying the mixture of experts on attention head is intuitive and effective.\n\n  MoH can achieve competitive performance with only 50%-90% heads/experts.\n\n  Detailed ablation study and well-organized paper structure.\n\nWeaknesses:\n\n  The idea is not entirely new. It shares a similar idea with MoA, plus implementation optimization of shared heads and two-stage routing.\n\n  No experiments compared with MoA.\n\nGiven these strengths and weaknesses, I am leaning towards weak accept. I am willing to raise my scores accordingly if the authors eventually address these concerns.\n\n## update after rebuttal\nThe authors solved my concern toward MOA comparison, therefore I raised my scores to 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thoughtful comments and for recognizing our work as \"intuitive and effective,\" acknowledging that \"MoH can achieve competitive performance,\" and highlighting our \"detailed ablation study and well-organized paper structure.\" Below, we address your questions in detail.\n\n**Q1:** The idea is not entirely new. It shares a similar idea with MoA.\n\n**A1:** We explain this problem in three aspects:\n* In terms of motivation, **MoH aims to make the attention mechanism more efficient and effective without adding extra parameters.** In contrast, MoA, like MoE, focuses on increasing model size while keeping inference costs low. Therefore, the model settings of MoH are more stringent than those of MoA. \n* In terms of methodology, **we maintain the original structure of multi-head attention as much as possible, allowing MoH to seamlessly replace standard multi-head attention across different tasks without extra tuning.** In contrast, MoA introduces shared keys and values to add heads without increasing the KV cache, which disrupts the original multi-head attention design.\n* In terms of flexibility, **we show that pre-trained multi-head attention models can be further continue-tuned into our MoH models, making MoH highly practical.** In contrast, MoA integrates multi-head attention with MoE but relies on shared keys and values, requiring training from scratch, which reduces its flexibility.\n\n**Q2:** No experiments compared with MoA.\n\n**A2:** Thanks for your insightful advice. As suggested, we compare MoA and MoH on the translation task and additionally provide results for image classification. As shown in the table below, MoH outperforms MoA, primarily for two reasons: (1) MoA's use of shared keys and values across heads limits its expressiveness; (2) MoH's shared heads and two-stage routing improve the model's ability to capture general knowledge (please refer to our response (A1) to Reviewer 6acz).\n\n|     | WMT14EnDe (BLEU) | **Image Classification (Acc)** |\n|:---:|:---------------------:|:-----------:|\n| MoA |          28.3          |    75.4%    |\n| MoH |          **29.0**         |    **78.6%**     |\n\n**Q3:** How to find a good balance between specialization and generalization.\n\n**A3:** There are two ways to choose shared heads and routed heads: **(1) manual configuration** and **(2) learning through learnable masks**. In the manuscript, we show the results of manual configuration. For the learnable mask approach, we follow a method similar to [1]. Specifically, we introduce a mask module with binary values {0,1} applied to the heads. These masks are dynamically learned rather than statically assigned, allowing the model to determine which heads should be shared. The remaining heads are designated as routed heads, and we set K based on a predefined ratio. For example, if the ratio is 1/4 and there are 8 routed heads, then K is set to 2. As shown in the table below, our latest experimental results show that learning through learnable masks is generally better than manual configuration.\n\n|     | # Activated Heads (%) | **Image Classification (Acc)** |\n|:---:|:---------------------:|:-----------:|\n| Baseline |          100          |     84.8%     |\n| Manual Configuration |           75          |     84.9%     |\n| Manual Configuration |          50          |     84.7%     |\n| Learning Through Learnable Masks |          50          |    **85.0%**     |\n\n[1] Liu, Peiyu, et al. \"MOEfication by Experts as Masks.\"\n\n**Q4:** Table 5 presents the Ablation study on the impact of each component of the proposed MoH, is the first row equal to MoA?\n\n**A4:** The first row in Table 5 follows a structure similar to MoA, but without shared keys and values or the use of additional z-loss.\n\n**Q5:** Providing state-of-the-art mechanism explanation presented in 2024.\n\n**A5:**  Thanks for your valuable suggestion. We have added your additional references to the Introduction to better explain that not all attention heads hold equal significance.\n\n**Q6:** TYPO: Sec 4.3 \"Please refer to the Appendix for detailed hyper-parameter settings (Tab. C)\" link to wrong table of Table 3.\n\n**A6:** Thank you for your thorough review. This issue may be a bug in LaTeX, and we will work on fixing it.\n\n**Q7:** Does the implementation of MoH affect GPU memory usage? if so, it increases of decreases the total memory usage? In that case, can we support larger model or small model by using MoH?\n\n**A7:** As shown in our response (A2) to Reviewer 6acz, MoH slightly reduces GPU memory usage, though the difference is not significant. This is because GPU memory is primarily used to store model parameters, gradients, and the KV cache. Since MoH only optimizes attention computation, it does not substantially reduce the GPU memory of these three components. Although MoH doesn't allow training larger models, it can make training and inference faster.", "claims_and_evidence": "1. Claim:  not all attention heads hold equal significance\n\nThe evidence is supported with Voita et al. (2019) and Michel et al. (2019). I suggest providing state-of-the-art mechanism explanation presented on 2024, like [1], [2] or [3]. \n\n2. Claim:  MoH outperforms multi-head attention by using only 50%∼90% of the attention heads\n\nEmpirical evaluations on ViT, DiT, and LLMs consistently show that MoH performs better.", "methods_and_evaluation_criteria": "MoH replace multiple head attention by MOH module, adding a router as MOE architecture to sellect top-K heads, and keep some heads always open. It makes sense for transformers.", "theoretical_claims": "Appendix A provides the theoretical claim that MoH is superior to vanilla multi-head attention. Given the 8-page limit, it is understandable to include it in the appendix.\n\nThe claim seems to tell a story that reduced redundancy and greater differentiation are better choices and lead to better model architecture. However, if it is true, complete MoH should be used, and all attention heads should be routed to moe-based design. Instead, this paper deploys a mixed technology in which parts of attention heads are specialized, and parts of attention heads are shared. In this case, specialization is not always good, finding a good balance between specialization and generalization is more important, how to find such a balance is one open question to answer.", "experimental_designs_or_analyses": "No experiments compared with MoA. Although section 5 discusses the difference between MoH and MoA, but not experiment results support this claim point.\n\nAs MoA is only validated for language tasks, can authors provide a comparison between MoA and MoH for language tasks?\n\nTable 5 presents the Ablation study on the impact of each component of the proposed MoH, is the first row equal to MoA? But it results from image classification, not from language tasks.", "relation_to_broader_scientific_literature": "Attention Head Pruning: Prior research of [1,2,3] shows that many heads can be removed without noticeable harm, MoH extends this insight by routing attention experts with performance gain.", "essential_references_not_discussed": "[1] Wu, Wenhao, et al. \"Retrieval head mechanistically explains long-context factuality.\" arXiv preprint arXiv:2404.15574 (2024).\n[2] Fu, Yu, et al. \"Not all heads matter: A head-level KV cache compression method with integrated retrieval and reasoning.\" arXiv preprint arXiv:2410.19258 (2024).\n[3] Xiao, Guangxuan, et al. \"Duoattention: Efficient long-context llm inference with retrieval and streaming heads.\" arXiv preprint arXiv:2410.10819 (2024).\nThese papers [1,2,3] propose the same insights that not all attention heads hold equal significance.", "other_strengths_and_weaknesses": "Strengths: Extensive Experimentation on vision classification, diffusion-based generation, large language modeling.\nStrengths: able to do both pretrain and finetune.", "other_comments_or_suggestions": "TYPO: Sec 4.3 \"Please refer to the Appendix for detailed hyper-parameter settings (Tab. C)\" link to wrong table of Table 3.", "questions_for_authors": "Does the implementation of MOH affect GPU memory usage? if so, it increases of decreases the total memory usage? In that case, can we support larger model or small model by using MOH?", "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: In the paper, the authors propose Mixture-of-Head attention (MoH) that treats attention heads as experts in mixture of experts. The proposed method allows each token to select the appropriate attention heads and yields extra flexibility to the attention mechanism to potentially improve the performance of Transformer.

After the rebuttal, the paper still received mixed reviews from the reviewers. On the positive side, most of the reviewers agree that the proposed method is sound and the experiments are well-conducted (e.g., they consist of a broad evaluation across ViTs, DiTs, and LLMs, using many benchmark datasets, including MMLU, CEVAL, GSM8K, and TruthfulQA). On the negative side, the method is not novel and quite incremental given the literature on the mixture of attention heads and Deepseek-MoE. Furthermore, the effectiveness of the MoH is purely demonstrated from the empirical results. It is important to have mathematical justification for the effectiveness of the MoH. 

After considering both the strengths and weaknesses of the paper, I think that the contributions of the proposed framework MoH are still sufficient and of interest to ICML. Therefore, I recommend accepting it in its current form. However, I encourage the authors to address the reviewers’ suggestions and integrate their feedback into the camera-ready version of their paper.
Decision: Accept (poster)