MoH: Multi-Head Attention as Mixture-of-Head Attention

## Abstract

Abstract In this work, we upgrade the multi-head attention mechanism, the core of the Transformer model, to improve efficiency 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% ∼ similar-to \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.

## 1 Introduction

Since attention is introduced and becomes a fundamental component of Transformers (Vaswani et al., ), multi-head attention has been the standard architecture for natural language processing (Kenton & Toutanova, ) and computer vision tasks (Dosovitskiy et al., ). It is well known that using multiple heads can improve model accuracy. However, 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. () introduces a method to quantify the usefulness of each attention head and prune those that are redundant. Similarly, Michel et al. () challenges the necessity of multiple heads by examining the impact of extensive pruning across various settings. These findings demonstrate that vanilla multi-head attention contains redundant attention heads.

Besides, in multi-head attention, each attention head operates in parallel, and the final output is the sum of all attention heads (please refer to Section 3.1). Given that these attention heads operate independently and some may be redundant, we argue that it 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-Head attention (MoH), a new architecture that integrates multi-head attention with the Mixture-of-Experts (MoE) mechanism (Jacobs et al., ; Jin et al., ). Specifically, we propose to treat attention heads as experts within the MoE framework. Similar to MoE, MoH consists of multiple attention heads and a router that activates the Top-K heads for each token. Moreover, we replace the standard summation in multi-head attention with a weighted summation. This design offers two significant advantages: First, MoH allows each token to select the most relevant attention heads, improving inference efficiency without sacrificing accuracy or increasing the parameters. Second, by replacing the standard summation in multi-head attention with a weighted summation, MoH enhances the flexibility of the attention mechanism and increases the performance potential. Moreover, to efficiently capture common knowledge across different contexts, we designate a subset of attention heads as shared heads that remain always activated.

We evaluate our proposed MoH across various popular model frameworks, including Vision Transformers (ViT) (Dosovitskiy et al., ) for image classification, Diffusion models with Transformers (DiT) (Peebles & Xie, ) for class-conditional image generation, and Large Language Models (LLMs) (Brown et al., ; OpenAI, ; Ouyang et al., ) for language tasks. We show that MoH achieves competitive performance, or even outperforms multi-head attention with only 50%$\sim$90% of the attention heads. For example, MoH-ViT-B achieves 84.9%/84.7% Top-1 accuracy on the ImageNet-1K (Deng et al., ) classification benchmark, surpassing well-tuned multi-head attention baselines with only 75%/50% of the attention heads.

Furthermore, we demonstrate that pre-trained multi-head attention models, such as LLaMA3-8B (Dubey et al., ), can be further continue-tuned into our MoH models. Specifically, using only about 3% (400B tokens) of the original LLaMA3 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 a solid foundation for developing advanced and efficient attention-based models. The main contributions are summarized as follows:

- •
We propose a dynamic attention-head routing mechanism that allows each token to adaptively select the appropriate attention heads, enhancing model performance and inference efficiency without increasing the number of parameters.
- •
In addition to training from scratch, we 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.
- •
A wide range of experiments across various popular model frameworks, including ViT, DiT, and LLMs, 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., ) have garnered significant interest and success in both natural language processing and computer vision. The success of transformers has been long attributed to the multi-head attention mechanism (Cordonnier et al., ). Multi-head attention mechanism is proposed by Vaswani et al. () to enhance the representation power of an attention layer by allowing multiple attention heads to operate on different low-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 a summation form. In summation form, each head operates 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.

Mixture-of-Experts Models.
The Mixture-of-Experts (MoE) method (Du et al., ; Lewis et al., ; Rajbhandari et al., ; Roller et al., ; Zhou et al., ; Jin et al., ) is introduced to expand the capacity of deep neural networks without increasing computational costs. In this approach, only a subset of parameters, known as experts, is activated for each input. Shazeer et al. () first introduces an MoE layer between LSTM layers. Switch Transformer (Fedus et al., ) further simplifies the gating mechanism by selecting only the Top-1 expert per token. Gshard (Lepikhin et al., ) improves the Top-2 expert routing strategy. In contrast to MoE, which emphasizes efficient parameter scaling while maintaining manageable computational costs, the proposed MoH focuses on reducing the activation of redundant attention heads without increasing the number of parameters.

Figure: 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 $h$ attention heads, while subfigure (b) demonstrates the Mixture-of-Head attention (MoH) architecture. It is important to note that MoH does not increase the number of attention heads, ensuring that the total parameter for MoH is comparable to that of the multi-head attention.
Refer to caption: x1.png

## 3 Methodology

In this work, we aim to 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 Mixture-of-Head attention (MoH) is presented in Fig. 1.

### 3.1 Multi-Head Attention

We begin by reviewing the standard multi-head attention mechanism introduced by Vaswani et al. (). The multi-head attention mechanism is based on scaled dot-product attention. Specifically, for $T$ tokens $\bm{X}\in\mathbb{R}^{T\times d_{in}}$ of $d_{in}$ dimensions each and $T^{\prime}$ tokens $\bm{X}^{\prime}\in\mathbb{R}^{T^{\prime}\times d_{in}}$ of $d_{in}$ dimensions each, the scaled dot-product attention is computed as follows:

$$ $\displaystyle\textrm{Attention}(\bm{Q},\bm{K}$ $\displaystyle,\bm{V})=\textrm{Softmax}\Big{(}\frac{\bm{Q}\bm{K}^{\top}}{\sqrt{ d_{k}}}\Big{)}\bm{V},$ (1) $\displaystyle\bm{Q}=\bm{X}\bm{W}_{Q},$ $\displaystyle\bm{K}=\bm{X}^{\prime}\bm{W}_{K},\ \bm{V}=\bm{X}^{\prime}\bm{W}_{ V},$ $$

where $\bm{W}_{Q}\in\mathbb{R}^{d_{in}\times d_{k}}$, $\bm{W}_{K}\in\mathbb{R}^{d_{in}\times d_{k}}$, and $\bm{W}_{V}\in\mathbb{R}^{d_{in}\times d_{v}}$ represent the projection matrices for the query, key, and value, respectively. In self-attention, the input tokens are the same, i.e., $\bm{X}^{\prime}=\bm{X}$, and it is common for the key and value dimensions to be equal, i.e., $d_{v}=d_{k}$.

Concatenation Form.
To enhance the representation power of the attention layer, Vaswani et al. () proposes to allow multiple attention heads to operate on different low-dimensional projections of the input tokens. Specifically, the multi-head attention mechanism computes $h$ different low-dimensional projections of $(\bm{Q},\bm{K},\bm{V})$, performs scaled dot-product attention for each head, concatenates the results, and applies a final projection to the concatenated output. The concatenation form of the multi-head attention can be formulated as:

$$ $\displaystyle\textrm{MultiHead}($ $\displaystyle\bm{X},\bm{X}^{\prime})=\textrm{Concat}(\bm{H}^{1},\bm{H}^{2},... ,\bm{H}^{h})\bm{W}_{O},$ (2) $\displaystyle\bm{H}^{i}=$ $\displaystyle\textrm{Attention}(\bm{X}\bm{W}_{Q}^{i},\bm{X}^{\prime}\bm{W}_{K} ^{i},\bm{X}^{\prime}\bm{W}_{V}^{i}),$ $$

where $\bm{W}_{Q}^{i}\in\mathbb{R}^{d_{in}\times d_{k}/h}$, $\bm{W}_{K}^{i}\in\mathbb{R}^{d_{in}\times d_{k}/h}$, and $\bm{W}_{V}^{i}\in\mathbb{R}^{d_{in}\times d_{v}/h}$ represent the $i_{th}$ projection matrices for the query, key, and value, respectively. $\bm{W}_{O}\in\mathbb{R}^{d_{v}\times d_{out}}$ is the final projection matrix.

Summation Form.
The multi-head attention mechanism is typically represented in its concatenation form. However, from another perspective, if we decompose $\bm{W}_{O}\in\mathbb{R}^{d_{v}\times d_{out}}$ by rows, we can express multi-head attention in a summation form. Specifically, $\bm{W}_{O}$ can be divided into $h$ matrices by rows, i.e., $[\bm{W}_{O}^{1},\bm{W}_{O}^{2},...,\bm{W}_{O}^{h}]=\bm{W}_{O}$, where $\bm{W}_{O}^{i}\in\mathbb{R}^{d_{v}/h\times d_{out}}$. Finally, the summation form of the multi-head attention can then be formulated as:

$$ $\textrm{MultiHead}(\bm{X},\bm{X}^{\prime})=\sum_{i=1}^{h}\bm{H}^{i}\bm{W}_{O}^ {i}.$ (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 function independently, we can build a dynamic attention-head routing mechanism allowing each token to adaptively select the most relevant attention heads, improving inference efficiency without compromising accuracy.

### 3.2 Mixture-of-Head Attention

Recently, the Mixture-of-Experts (MoE) method has emerged as a popular approach for scaling the parameters of large language models (Jiang et al., ). A typical 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.

Heads as Experts.
Inspired by the great success of MoE, we propose Mixture-of-Head attention (MoH), which treats attention heads as experts. Specifically, MoH consists of $h$ heads $\bm{H}=\{H^{1},H^{2},...,H^{h}\}$ and a router that activates the Top-K heads. Formally, given input tokens $\bm{X}$ and $\bm{X}^{\prime}$, the output of MoH is the weighted sum of outputs from the $K$ selected heads:

$$ $\textrm{MoH}(\bm{X},\bm{X}^{\prime})=\sum_{i=1}^{h}g_{i}\bm{H}^{i}\bm{W}_{O}^{ i},$ (4) $$

where $g_{i}$ represents the routing score. $g_{i}$ is non-zero only when the $i_{th}$ attention head is activated. This design provides two key advantages: On the one hand, MoH enables each token to select the most relevant attention heads, boosting inference efficiency while maintaining accuracy. On the other hand, in contrast to the standard summation in multi-head attention, the weighted summation in MoH enhances the flexibility of the attention mechanism and unlocks performance potential.

Shared Heads.
In attention mechanism, some attention heads may capture common knowledge across different contexts, such as grammatical rules in language. Inspired by Dai et al. (), we designate a subset of heads as shared heads that remain always activated. By consolidating common knowledge within shared heads, we reduce redundancy among the other dynamically routed heads.

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 $t_{th}$ input token $\bm{x}_{t}\in\mathbb{R}^{d_{in}}$ in $\bm{X}\in\mathbb{R}^{T\times d_{in}}$, the routing score $g_{i}$ is defined as:

$$ $\quad g_{i}=\begin{cases}\alpha_{1}\text{Softmax}(\bm{W}_{s}\bm{x}_{t})_{i},& \text{if}\ 1\leq i\leq h_{s},\\ \alpha_{2}\text{Softmax}(\bm{W}_{r}\bm{x}_{t})_{i},&\text{if}\ (\bm{W}_{r}\bm{ x}_{t})_{i}\in\text{Top-K}\big{(}\{(\bm{W}_{r}\bm{x}_{t})_{i}|h_{s}+1\leq i \leq h\}\big{)},\\ 0,&\text{otherwise},\end{cases}$ (5) $$

where $h_{s}$ denotes the number of shared heads. $\bm{W}_{s}\in\mathbb{R}^{h_{s}\times d_{in}}$ and $\bm{W}_{r}\in\mathbb{R}^{(h-h_{s})\times d_{in}}$ represent the projection matrices for the shared and routed heads, respectively. The coefficients $\alpha_{1}$ and $\alpha_{2}$ balance the contributions of the shared and routed heads, and are defined as:

$$ $[\alpha_{1},\alpha_{2}]=\text{Softmax}(\bm{W}_{h}\bm{x}_{t}),$ (6) $$

where $\bm{W}_{h}\in\mathbb{R}^{2\times d_{in}}$ is the trainable projection matrix, and $d_{in}$ is the hidden size of $\bm{x}_{t}$.

Load Balance Loss
Directly training an MoE layer often causes the majority of tokens to be routed to a small number of experts, leaving the remaining experts insufficiently trained (Shazeer et al., ). To avoid the unbalanced load in the proposed MoH, following previous MoE methods (Lepikhin et al., ; Wei et al., ), we apply a load balance loss. Specifically, for the $t_{th}$ input token $\bm{x}_{t}\in\mathbb{R}^{d_{in}}$ in $\bm{X}\in\mathbb{R}^{T\times d_{in}}$, the load balance loss $\mathcal{L}_{b}$ is formulated as:

$$ $\mathcal{L}_{b}=\sum_{i=h_{s}+1}^{h}f_{i}P_{i},\ f_{i}=\frac{1}{T}\sum_{t=1}^{ T}\mathds{1}(\text{Token $\bm{x}_{t}$ selects Head $i$}),\ P_{i}=\frac{1}{T} \sum_{t=1}^{T}\text{Softmax}(\bm{W}_{r}\bm{x}_{t})_{i},$ (7) $$

where $T$ denotes the number of tokens. $\mathds{1}(*)$ denotes the indicator function.

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 $\mathcal{L}_{task}$ and the load balance loss $\mathcal{L}_{b}$:

$$ $\mathcal{L}=\mathcal{L}_{task}+\beta\mathcal{L}_{b},$ (8) $$

where $\beta$ is the trade-off hyper-parameter to mitigate the risk of routing collapse. By default, the weight $\beta$ for the load balance loss is set to 0.01 for all tasks.

**Table 1: Comparisons to current state-of-the-art methods on ImageNet-1K classification. All models are trained exclusively on the ImageNet-1K training set. Our MoH-ViT models, based on TransNeXt (Shi, ), are trained for 300 epochs using a resolution of 224$\times$224. 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 TransNeXt.**
| Methods | #Params | #Activated | Acc |
| --- | --- | --- | --- |
| (M) | Heads (%) | (%) |  |
| DeiT-S (Touvron et al., ) | 22 | 100 | 79.8 |
| T2T-ViT-19 (Yuan et al., ) | 39 | 100 | 81.9 |
| Swin-S (Liu et al., ) | 50 | 100 | 83.1 |
| PVTv2-B3 (Wang et al., ) | 45 | 100 | 83.2 |
| CoAtNet-1 (Dai et al., ) | 42 | 100 | 83.3 |
| Focal-S (Yang et al., ) | 51 | 100 | 83.5 |
| FocalNet-S (Yang et al., ) | 50 | 100 | 83.5 |
| MViTv2-S (Li et al., ) | 35 | 100 | 83.6 |
| UniFormer-B (Li et al., ) | 50 | 100 | 83.9 |
| CAFormer-S36 (Yu et al., ) | 39 | 100 | 84.5 |
| TransNeXt-S (Shi, ) | 50 | 100 | 84.7 |
| MoH-ViT-S | 50 | 80 | 84.7 |
| MoH-ViT-S | 50 | 75 | 84.6 |

## 4 Experiments

### 4.1 ViT for Image Classification

Model Settings.    For Vision Transformers (ViT) (Dosovitskiy et al., ), our MoH-ViT models are implemented based on the TransNeXt (Shi, ) framework and trained from scratch on the ImageNet-1K dataset (Deng et al., ), which contains over 1.2 million images in 1,000 categories. 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.

Training Details.    Our MoH-ViT models are 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., ), Mixup (Zhang, ), CutMix (Yun et al., ), and Random Erasing (Zhong et al., ). We also apply Label Smoothing (Szegedy et al., ) and DropPath (Huang et al., ) to regularize our models. We optimize our models using AdamW optimizer (Loshchilov & Hutter, ) with a gradient clipping norm of 1.0 and a weight decay of 0.05. The initial learning rate is set to 1e-3, with a 5-epoch warm-up starting at 1e-6. A cosine learning rate scheduler (Loshchilov & Hutter, ) is employed to decay the learning rate. During training, images are randomly cropped to a size of 224$\times$224. It is worth noting that we do not use Exponential Moving Average (EMA) weights.

Results.    As shown in Tab. 1, despite activating only a subset of attention heads, MoH-ViT achieves highly competitive performance compared to current state-of-the-art methods. For example, MoH-ViT-B achieves 84.9% Top-1 accuracy on the ImageNet-1K classification benchmark with just 75% of 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 to be activated. Tab. 1 demonstrates that MoH-ViT outperforms other models with fewer activated attention heads. This suggests that MoH is a promising alternative to vanilla multi-head attention for vision model design, offering the potential for competitive performance with more efficient attention head usage.

**Table 2: Comparisons to DiT on the benchmarking of class-conditional image generation on ImageNet-1K at 256$\times$256 resolution. To ensure a fair comparison, we only replace the standard multi-head attention with the MoH in MoH-DiT models, while keeping all other training parameters identical to DiT. “400K” denotes the training budget is 400K training steps.**
| Methods | #Params | #Activated | FID$\downarrow$ | sFID$\downarrow$ | IS$\uparrow$ | Precision$\uparrow$ | Recall$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- |
| (M) | Heads (%) |  |  |  |  |  |  |
| DiT-S/2 400K (Peebles & Xie, ) | 33 | 100 | 68.40 | - | - | - | - |
| MoH-DiT-S/2 400K | 33 | 90 | 67.25 | 12.15 | 20.52 | 0.37 | 0.58 |
| MoH-DiT-S/2 400K | 33 | 75 | 69.42 | 12.85 | 19.96 | 0.36 | 0.55 |
| DiT-B/2 400K (Peebles & Xie, ) | 130 | 100 | 43.47 | - | - | - | - |
| MoH-DiT-B/2 400K | 131 | 90 | 43.40 | 8.40 | 33.51 | 0.49 | 0.63 |
| MoH-DiT-B/2 400K | 131 | 75 | 43.61 | 8.48 | 33.43 | 0.49 | 0.62 |
| DiT-L/2 400K (Peebles & Xie, ) | 458 | 100 | 23.33 | - | - | - | - |
| MoH-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 |

**Table 3: Comparisons to current state-of-the-art methods on the benchmarking of class-conditional image generation on ImageNet-1K at 256$\times$256 resolution. “$\uparrow$” denotes that higher is better. “$\downarrow$” denotes that lower is better. “cfg” denotes the classifier-free diffusion guidance scale. We extend the training budget of our MoH-DiT-XL/2 to 7,000K training steps, aligning it with DiT-XL/2.**
| Methods | #Activated | FID$\downarrow$ | sFID$\downarrow$ | IS$\uparrow$ | Precision$\uparrow$ | Recall$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- |
| Heads (%) |  |  |  |  |  |  |
| ADM-G, ADM-U (Dhariwal & Nichol, ) | - | 3.94 | 6.14 | 215.84 | 0.83 | 0.53 |
| CDM (Ho et al., ) | - | 4.88 | - | 158.71 | - | - |
| LDM-8 (Rombach et al., ) | - | 15.51 | - | 79.03 | 0.65 | 0.63 |
| LDM-4 (Rombach et al., ) | - | 10.56 | - | 103.49 | 0.71 | 0.62 |
| LDM-4-G (cfg=1.25) | - | 3.95 | - | 178.22 | 0.81 | 0.55 |
| DiT-XL/2 7,000K (Peebles & Xie, ) | 100 | 9.62 | 6.85 | 121.50 | 0.67 | 0.67 |
| DiT-XL/2 7,000K (cfg=1.25) | 100 | 3.22 | 5.28 | 201.77 | 0.76 | 0.62 |
| MoH-DiT-XL/2 2,000K | 75 | 10.95 | 6.19 | 106.69 | 0.67 | 0.66 |
| MoH-DiT-XL/2 2,000K | 90 | 10.67 | 6.15 | 107.80 | 0.67 | 0.65 |
| MoH-DiT-XL/2 7,000K | 90 | 8.56 | 6.61 | 129.54 | 0.68 | 0.67 |
| MoH-DiT-XL/2 7,000K (cfg=1.25) | 90 | 2.94 | 5.17 | 207.25 | 0.77 | 0.63 |

### 4.2 DiT for Class-Conditional Image Generation

Model Settings.    For Diffusion models with Transformers (DiT) (Peebles & Xie, ), we only replace the standard multi-head attention with our MoH in MoH-DiT models, while keeping all other training parameters identical to DiT. We use the ImageNet-1K dataset (Deng et al., ) for class-conditional image generation at a resolution of 256$\times$256.

Training Details.    Following DiT, the final linear layer is initialized with zeros, and all other layers follow standard ViT weight initialization. We train all models using the AdamW optimizer (Loshchilov & Hutter, ) with a constant learning rate of 1e-4, no weight decay, and a 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-the-shelf pre-trained variational autoencoder (Kingma, ) model from Stable Diffusion (Rombach et al., ). 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.

Evaluation Benchmarks.
To evaluate generation performance, we use Frechet Inception Distance (FID) (Heusel et al., ) to assess overall sample quality, Precision and Recall (Kynkäänniemi et al., ) to measure fidelity and diversity separately, and sFID (Nash et al., ) as a metric that better captures spatial relationships than FID. Moreover, we use Inception Score (IS) (Salimans et al., ) as another metric for fidelity.

Results.    To conduct comparative evaluations of our proposed MoH-DiT models against vanilla DiT models, we start with Small models and expand to XLarge models. As shown in Tab. 2, MoH-DiT models consistently outperform vanilla DiT models with 90% of attention heads activated. However, when only 75% of the attention heads are activated, MoH-DiT models perform worse than DiT models with 100% of attention heads activated. This may be because image generation tasks are dense prediction tasks that require attention mechanisms to capture pixel-level fine-grained relationships, leaving less redundancy in the attention heads compared to image classification tasks. Moreover, we extend the training budget of our MoH-DiT-XL/2 to 7,000K training steps, aligning it with DiT-XL/2. As shown in Tab. 3, despite activating 90% attention heads, MoH-DiT-XL/2 achieves highly competitive performance compared to current state-of-the-art methods. These results suggest that MoH is a promising alternative to multi-head attention for diffusion models.

**Table 4: Comparisons between MoH-LLMs and vanilla LLMs. “100B” denotes a training budget of 100 billion tokens, while “200B” denotes a budget of 200 billion tokens. We observe that larger models, e.g., MoH-LLM-B, generally perform worse than smaller models, e.g., MoH-LLM-S, on TruthfulQA, consistent with the findings reported by Lin et al. ().**
| Methods | #Activated | Language Tasks | Avg. |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Heads (%) | SciQ | PIQA | WinoGrande | OpenbookQA | LogiQA | TruthfulQA |  |  |
| LLM-S 100B | 100 | 63.0 | 63.1 | 51.1 | 27.4 | 26.9 | 31.6 | 43.9 |
| MoH-LLM-S 100B | 75 | 64.7 | 62.0 | 50.6 | 28.8 | 26.4 | 35.2 | 44.6 |
| MoH-LLM-S 100B | 50 | 67.0 | 62.2 | 51.5 | 29.2 | 26.7 | 35.6 | 45.4 |
| LLM-B 100B | 100 | 73.1 | 69.7 | 52.0 | 31.8 | 28.4 | 29.5 | 47.4 |
| MoH-LLM-B 100B | 75 | 74.7 | 69.2 | 52.8 | 30.0 | 28.1 | 32.2 | 47.8 |
| MoH-LLM-B 100B | 50 | 75.2 | 67.0 | 52.0 | 29.0 | 26.9 | 32.8 | 47.2 |
| LLM-B 200B | 100 | 73.1 | 70.3 | 53.3 | 32.4 | 29.0 | 29.5 | 47.9 |
| MoH-LLM-B 200B | 75 | 76.0 | 69.2 | 52.7 | 30.4 | 29.8 | 32.6 | 48.5 |
| MoH-LLM-B 200B | 50 | 75.6 | 66.9 | 53.5 | 29.4 | 26.7 | 32.7 | 47.5 |

### 4.3 Training LLMs from Scratch

Model Settings.    For training LLMs from scratch, we use Megatron (Shoeybi et al., ), an open-source training code, as the training framework. Please refer to the Appendix for detailed hyper-parameter settings (Tab. A) of various MoH-LLMs. All models are trained with the AdamW optimizer (Loshchilov & Hutter, ), using a batch size of 4 million tokens with a sequence length of 2048. The final learning rate is set to 10% of the maximum. During training, a weight decay of 0.1 and gradient clipping of 1.0 are applied. For LLM-S and MoH-LLM-S, the maximum learning rate is set to 3e-4. For LLM-B and MoH-LLM-B, the maximum learning rate is set to 5e-4.

Training Details.    We only use public datasets for training, ensuring accessibility for academic research. Specifically, we sample from the RedPajama (Computer, ), Dolma (Soldaini et al., ), and Pile (Gao et al., ) datasets according to different sampling probabilities. Please refer to the Appendix for detailed sample ratios (Tab. B). Following previous works, we utilize the tokenizer from LLaMA2 (Touvron et al., ), which contains 65,536 vocabulary tokens.

Evaluation Benchmarks.
The evaluation is performed on multiple benchmarks using the Eleuther AI Language Model Evaluation Harness (Gao et al., ), a unified framework for testing generative language models. Since the parameters are only about 0.2B for the smallest model, we select 6 simple benchmarks as the metric. Specifically, we report 0-shot accuracy on SciQ (Welbl et al., ), PIQA (Bisk et al., ), WinoGrande (Sakaguchi et al., ), OpenbookQA (Mihaylov et al., ), LogiQA (Liu et al., ), and TruthfulQA (Lin et al., ).

Results.    As shown in Tab. 4, despite activating only a subset of attention heads, MoH-LLMs achieve highly competitive performance compared to our baseline models. For example, MoH-LLM-S achieves an average accuracy of 45.4% with just 50% of the attention heads activated. In contrast, the baseline model reaches a slightly lower accuracy 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 LLMs 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.

**Table 5: Comparisons between MoH-LLaMA3-8B and LLaMA3-8B. Please refer to the Appendix for the performance of the model at the end of the first stage of training.**
| Methods | #Activated | MMLU (5) | CEVAL (5) | CMMLU (5) | GSM8K(8) | TruthfulQA |
| --- | --- | --- | --- | --- | --- | --- |
| Heads (%) |  |  |  |  |  |  |
| LLaMA3-8B (Dubey et al., ) | 100 | 65.2 | 52.3 | 50.7 | 49.5 | 35.4 |
| MoH-LLaMA3-8B | 75 | 65.8 | 61.5 | 64.4 | 56.9 | 44.0 |
| Methods | #Activated | HellaSwag (10) | LogiQA | BoolQ (32) | LAMBADA | SciQ |
| Heads (%) |  |  |  |  |  |  |
| LLaMA3-8B (Dubey et al., ) | 100 | 81.9 | 30.0 | 83.9 | 75.5 | 94.0 |
| MoH-LLaMA3-8B | 75 | 80.1 | 30.3 | 84.0 | 76.4 | 92.2 |
| Methods | #Activated | PIQA | WinoGrande | NQ (32) | ARC-C (25) | Average |
| Heads (%) |  |  |  |  |  |  |
| LLaMA3-8B (Dubey et al., ) | 100 | 81.0 | 72.5 | 31.5 | 59.0 | 61.6 |
| MoH-LLaMA3-8B | 75 | 78.8 | 72.9 | 28.3 | 60.1 | 64.0 |

Figure: Figure 2: Performance evolution during continue-tuning. The MoH model quickly recovers to over 95% of the performance of the original model within a training budget of 10B tokens. Then, the performance gradually improves with the increase of the training tokens.
Refer to caption: x2.png

### 4.4 Continue-Tuning LLaMA3-8B

Model Settings.    To significantly enhance the applicability of the proposed MoH method, we also attempt to further continue-tune pre-trained multi-head attention models, such as LLaMA3-8B, 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 routers: Integrating a randomly initialized router into the pre-trained model without compromising its original performance requires careful training techniques. To address this, we propose a parameter-free router that determines routing scores using the $\ell_{2}$ norm of the query of each attention head. (iii) Weighting attention heads: We observe that weighting the attention head outputs significantly alters the distribution of the output of the attention layer, which necessitates a large amount of training data to restore the original performance. To tackle this, we quantize the routing score and use the straight-through estimator (Bengio et al., ; Liu et al., ) to back-propagate the gradients through the sparsity function. Specifically, given the input token $\bm{x}$, we employ a quantizer for activation routing scores, with its forward pass formulated as:

$$ $g_{i}^{q}=\mathds{1}(\text{Token $\bm{x}$ selects Head $i$}),$ (9) $$

where $\mathds{1}(*)$ denotes the indicator function. $g_{i}^{q}$ represents the quantized routing score. We then adopt a straight-through estimator, which assigns the incoming gradients to a threshold operation to be the outgoing gradients, which is formulated as:

$$ $\frac{\partial\mathcal{L}}{\partial g_{i}^{q}}=\frac{\partial\mathcal{L}}{ \partial g_{i}},$ (10) $$

where $g_{i}$ denotes the real-valued routing score. This simple approximation function significantly mitigates the issue of gradient vanishing (Wang et al., ). Similar to training LLMs from scratch, we also use Megatron (Shoeybi et al., ), an open-source training code, as the training framework.

Training Details.    We find that if there is a discrepancy between the continue-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 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 to our dataset. 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 to 6e-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 1e-6. For both stages, we employ the AdamW optimizer (Loshchilov & Hutter, ), with a batch size of 16 million tokens with a sequence length of 8192. During training, we use a weight decay of 0.1 and gradient clipping of 1.0.

Evaluation Benchmarks.
We use the Eleuther AI Language Model Evaluation Harness (Gao et al., ) to evaluate models on multiple key benchmarks. Specifically, we utilize the lm-evaluation-harness package to assess performance on a comprehensive suite of downstream tasks: (i) Following Pythia (Biderman et al., ), we report 0-shot accuracy on LAMBADA (Paperno et al., ), LogiQA (Liu et al., ), PIQA (Bisk et al., ), SciQ (Welbl et al., ), and WinoGrande (Sakaguchi et al., ). (ii) We report the accuracy of Chinese tasks, including 5-shot CEVAL (Huang et al., ) and 5-shot CMMLU (Li et al., ). (iii) We report the accuracy of tasks from the Open LLM Leaderboard (Beeching et al., ), including 10-shot HellaSwag (Zellers et al., ), 25-shot ARC Challenge (ARC-C) (Clark et al., ), and 5-shot MMLU (Hendrycks et al., ). (iv) We report the exact match score for 32-shot Natural Questions (NQ) (Kwiatkowski et al., ) and the accuracy for 32-shot BoolQ (Clark et al., ). (v) We report the exact match score for 8-shot GSM8K (Cobbe et al., ) to evaluate the math ability. (vi) Moreover, we report 0-shot accuracy on TruthfulQA (Lin et al., ) to assess the ability to generate truthful answers.

Results.    As shown in Fig. 2, MoH-LLaMA3-8B quickly recovers to over 95% of the performance of the original model within a training budget of 10B tokens. After continue-tuning with 100B tokens, as shown in Tab. 5, 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 demonstrate that pre-trained multi-head attention models can be further continue-tuned into our MoH models, significantly enhancing the applicability of the MoH method.

**Table 6: Ablation study on the impact of each component of the proposed MoH. The image classification results are from MoH-ViT-S, by utilizing 75% of the attention heads with a training budget of 100 epochs. The class-conditional image generation results come from MoH-DiT-S/2-400K, also by using 75% of the attention heads, with a training budget of 400K training steps.**
| Shared | Two-Stage | Image Classification | Class-Conditional Image Generation |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Heads | Routing | Acc (%)$\uparrow$ | FID$\downarrow$ | sFID$\downarrow$ | IS$\uparrow$ | Precision$\uparrow$ | Recall$\uparrow$ |
|  |  | 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 |

**Table 7: Ablation study on the impact of the shared heads ratio among activated heads. All results are from MoH-ViT-S, by using 75% of the heads with a 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 |

### 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 in Tab. 6. “Shared Heads” refers to a 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. 6, shared heads significantly improve model performance by effectively capturing common knowledge, allowing the routed heads to focus more on domain-specific 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. 7, we provide the ablation study on the shared heads ratio among activated heads. We find that model performance remains relatively consistent across a 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., ). Based on the findings from the Soft MoE paper (Puigcerver et al., ), we recommend using a higher ratio of shared heads among the activated heads (greater than 40%).

Figure: Figure 3: Visualization of the head load distribution in 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”. MoH-ViT-B, MoH-DiT-XL/2, and MoH-LLM-B activate 75%, 90%, and 75% of the attention heads, respectively.
Refer to caption: x3.png

## 5 Discussion

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, indicating that the MoH model adapts to diverse tasks by employing distinct head assignment patterns. This characteristic of MoH allows different attention heads to focus on different types of tasks, making parameter utilization more efficient than multi-head attention. For additional visualizations of MoH-LLaMA3-8B and a detailed analysis of the head load distribution, please refer to Appendix D.

The Difference between MoH and MoA.    We clarify the differences between MoH and MoA (Zhang et al., ) from the following three aspects. First, in terms of motivation, the goal of MoH is to improve the efficiency and performance of the attention mechanism without increasing the number of parameters. In contrast, MoA shares the motivation of MoE, 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 and two-stage routing to enhance the standard MoE method. More importantly, we show that pre-trained multi-head attention models can be further continue-tuned 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 frameworks, our MoH is validated across various popular model frameworks and tasks, including ViT, DiT, and decoder-only LLMs, while MoA is only validated on the encoder-decoder architecture for language tasks.

## 6 Conclusion

In this paper, we introduce MoH, a 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 various popular model frameworks, including ViT, DiT, and LLMs, demonstrate that MoH outperforms multi-head attention, even when using only 50%$\sim$90% of the attention heads. More encouragingly, we show that pre-trained multi-head attention models, such as LLaMA3-8B, can be further continue-tuned into our MoH models, significantly enhancing the applicability of the proposed MoH method. This work represents a promising step toward advanced and efficient attention-based models, which may be meaningful and helpful to both the research and industrial communities.

## Appendix A Appendix

Abstract.    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).

Code.    We have attached the code to the supplementary material. In this code, we also provide the evaluation process of the proposed method. We promise to release a more detailed and clean code version upon publication.

## Appendix A Additional Discussions

### A.1 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%$\sim$90% of the attention heads. However, this is still a relatively high proportion. Future work could aim to 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., ) 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., ; Lin et al., ; ; Liu et al., ; Jin et al., ).

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 can explore the application of MoH in more downstream tasks, such as audio tasks and multimodal tasks (Jin et al., ; ).

More Parameters.
Due to computational constraints, the maximum number of MoH model parameters in our experiments is limited to 8B (MoH-LLaMA3-8B). However, our MoH method is highly generalizable and can be scaled to larger models in future research.

## Appendix B Implementation Details

### B.1 Training LLMs from Scratch

Model Settings.    For training LLMs from scratch, we use Megatron (Shoeybi et al., ), an open-source training code, as the training framework. The detailed hyper-parameter settings of various MoH-LLMs are shown in Tab. A.

**Table A: Sizes and architectures of MoH-LLMs and LLMs. “MoH-LLM-B” 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. It is worth noting that MoH-LLM is trained exclusively on public datasets, making it accessible for academic research settings. Tab. B shows the detailed sample ratios of different open-source datasets. Specifically, we sample from the following datasets according to different sampling probabilities:

- •
The RedPajama (Computer, ) includes training data from seven domains: CommonCrawl, C4, Github, Wikipedia, Books, ArXiv, and StackExchange.
- •
The Dolma (Soldaini et al., ), 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., ), 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., ), 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 B: Sampling ratio of different open-source datasets for MoH-LLMs. MoH-LLM is 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-Parameters.
Tab. C shows the detailed training hyper-parameters of MoH-LLMs. Specifically, all MoH-LLMs are trained with the AdamW optimizer (Loshchilov & Hutter, ), using a batch size of 4 million tokens with a sequence length of 2048. The final learning rate is set to 10% of the maximum. During training, a weight decay of 0.1 and gradient clipping of 1.0 are applied. For LLM-S and MoH-LLM-S, the maximum learning rate is set to 3e-4. For LLM-B and MoH-LLM-B, the maximum learning rate is set to 5e-4.

**Table C: Training hyper-parameters of MoH-LLMs.**
|  | MoH-LLM-S 100B | MoH-LLM-B 100B | MoH-LLM-B 200B |
| --- | --- | --- | --- |
|  | (LLM-S 100B) | (LLM-B 100B) | (LLM-B 200B) |
| Training budget | 100B | 100B | 200B |
| Maximum learning rate | 3e-4 | 5e-4 | 5e-4 |
| Final learning rate | 3e-5 | 5e-5 | 5e-5 |
| LR warmup init | 1e-7 | 1e-7 | 1e-7 |
| LR warmup iters | 2000 | 500 | 500 |
| Sequence length | 2048 | 2048 | 2048 |
| Batch size (tokens) | 4M | 4M | 4M |
| $\beta$ for $\mathcal{L}_{b}$ | 0.01 | 0.01 | 0.01 |
| Tensor parallel | 1 | 1 | 1 |
| Pipeline parallel | 1 | 1 | 1 |

### B.2 Continue-Tuning LLaMA3-8B

Training Hyper-Parameters.
Tab. D shows the detailed training hyper-parameters of MoH-LLaMA3-8B. We find that if there is a discrepancy between the continue-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. In the first stage, we continue-tune the original LLaMA3-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 increase the original LLaMA3-8B vocabulary size from 128,256 to 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 to 6e-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 1e-6. For both stages, we employ the AdamW optimizer (Loshchilov & Hutter, ), with a batch size of 16 million tokens with a sequence length of 8192. During training, we use a weight decay of 0.1 and gradient clipping of 1.0.

**Table D: Training hyper-parameters of MoH-LLaMA3-8B. We divide the training process into two stages. In the first stage, we continue-tune the LLaMA3-8B model using 300B tokens. In the second stage, we continue-tune 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 | 2e-5 |
| Final learning rate | 6e-6 | 1e-6 |
| LR warmup iters | 50 | 50 |
| Sequence length | 8192 | 8192 |
| Batch size (tokens) | 16M | 16M |
| $\beta$ for $\mathcal{L}_{b}$ | - | 0.01 |
| Tensor parallel | 2 | 1 |
| Pipeline parallel | 1 | 8 |

**Table E: Comparisons between MoH-LLaMA3-8B and LLaMA3-8B-stage1. MoH-LLaMA3-8B outperforms LLaMA3-8B-stage1 by utilizing only 75% of the attention heads.**
| Methods | #Activated | MMLU (5) | CMMLU (5) | NQ (32) | GSM8K(8) | TruthfulQA |
| --- | --- | --- | --- | --- | --- | --- |
| Heads (%) |  |  |  |  |  |  |
| LLaMA3-8B-stage1 | 100 | 66.2 | 66.0 | 28.1 | 58.6 | 41.9 |
| MoH-LLaMA3-8B | 75 | 65.8 | 64.4 | 28.3 | 56.9 | 44.0 |
| Methods | #Activated | HellaSwag (10) | LogiQA | BoolQ (32) | LAMBADA | SciQ |
| Heads (%) |  |  |  |  |  |  |
| LLaMA3-8B-stage1 | 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 |
| Methods | #Activated | PIQA | WinoGrande | ARC-E | ARC-C (25) | Average |
| Heads (%) |  |  |  |  |  |  |
| LLaMA3-8B-stage1 | 100 | 79.1 | 73.0 | 70.9 | 59.6 | 64.7 |
| MoH-LLaMA3-8B | 75 | 78.8 | 72.9 | 72.5 | 60.1 | 64.8 |

## Appendix C Additional Experiments

Comparison between MoH-LLaMA3-8B and LLaMA3-8B-stage1.
We divide the training process into two stages. Tab. E shows the comparison between MoH-LLaMA3-8B and the model at the end of the first training stage (LLaMA3-8B-stage1). As shown in Tab. E, MoH-LLaMA3-8B quickly recovers the performance of LLaMA3-8B-stage1 within a training budget of 100B tokens. Notably, in English language tasks, MoH-LLaMA3-8B surpasses LLaMA3-8B-stage1 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-LLaMA3-8B achieves an accuracy of 64.4 on CMMLU, compared to 66.0 for LLaMA3-8B-stage1. We attribute this to the fact that the model’s Chinese and mathematical capabilities are 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 mathematical tasks. Overall, as shown in Tab. E, MoH-LLaMA3-8B achieves an average accuracy of 64.8% across 14 benchmarks, outperforming LLaMA3-8B-stage1 by utilizing only 75% of the attention heads.

Figure: Figure A: Additional visualization of the head load distribution in the final MoH layer. For MoH-ViT-B and MoH-DiT-XL/2, we present the head load distributions for the categories “Basketball”, “Bookshop”, “Cart”, “Husky”, and “Jean”. MoH-ViT-B activates 75% of the attention heads. MoH-DiT-XL/2 activates 90% of the attention heads.
Refer to caption: x4.png

## Appendix 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 a 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-head attention.

Additional Visualization of the Head Load Distribution in MoH-LLaMA3-8B.
We provide additional visualization of the head load distribution in Fig. B. As shown in Fig. B, MoH-LLaMA3-8B exhibits similar characteristics to MoH-LLMs trained from scratch, with significant variation in attention head assignments across different categories and task topics. This indicates that continue-tuning 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.

Images Generated from the Proposed MoH-DiT-XL/2 Model.
Fig. C shows samples generated by our class-conditional MoH-DiT-XL/2 model. These results demonstrate the ability of MoH-DiT-XL/2 to generate semantically correct content with accurate spatial relationships.

## Appendix 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(^4^44[https://github.com/EleutherAI/lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness)) (Gao et al., ), a unified framework for testing generative language models across a wide range of tasks. The benchmarks used for evaluation include:

Figure: Figure B: Additional visualization of the head load distribution in MoH-LLaMA3-8B.
Refer to caption: x5.png

ARC (Clark et al., ) is a multiple-choice question-answering resource featuring questions from science exams for grades 3 to 9. It is divided into two 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 or five choices. Additionally, ARC includes a supporting knowledge base with 14.3 million unstructured text passages. We report 0-shot accuracy on ARC Easy and 25-shot accuracy on ARC Challenge.

LAMBADA (Paperno et al., ) is an open-ended cloze task consisting of approximately 10,000 passages from BooksCorpus, where the objective is to predict a missing target word in the last sentence of each passage. The missing word is always the last word of the final sentence, with no options provided. We report 0-shot accuracy on LAMBADA.

LogiQA (Liu et al., ) comprises 8,678 question-and-answer instances that encompass various types of deductive reasoning. The dataset serves as a benchmark for reexamining logical AI within the context of deep learning in NLP. We report 0-shot accuracy on LogiQA.

PIQA (Bisk et al., ) is a dataset designed for commonsense reasoning, aimed at evaluating the physical knowledge of current models. We report 0-shot accuracy on PIQA.

Figure: Figure C: Images generated from the proposed MoH-DiT-XL/2 model. We show samples generated from our class-conditional MoH-DiT-XL/2 model trained on ImageNet at 256$\times$256 resolution. MoH-DiT-XL/2 activates 90% of the attention heads.
Refer to caption: x6.png

SciQ (Welbl et al., ) includes 13,679 crowdsourced science exam questions covering subjects such as Physics, Chemistry, and Biology. Each question is presented in a multiple-choice format with four answer options, and for most questions, an additional paragraph provides supporting evidence for the correct answer. We report 0-shot accuracy on SciQ.

WinoGrande (Sakaguchi et al., ) 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 0-shot accuracy on WinoGrande.

HellaSwag (Zellers et al., ) is a challenging dataset designed to evaluate commonsense natural language inference, which proves difficult for state-of-the-art models but poses no significant challenge for humans. We report the accuracy for the 10-shot HellaSwag.

MMLU (Hendrycks et al., ) is a benchmark designed to assess models’ knowledge acquired during pretraining, making it more challenging and human-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 et al., ) is a 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 5 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., ) is a question-answering dataset consisting of 15,942 yes/no questions. These questions are naturally occurring, and generated in unprompted and unconstrained contexts. Each example is provided as a triplet of (question, passage, and answer), with the page title optionally included as additional context. We report the accuracy for the 32-shot BoolQ.

OpenbookQA (Mihaylov et al., ) is a question-answering dataset designed to assess understanding of elementary-level science, similar to open-book exams. It contains 5,957 multiple-choice questions based on a “book” of 1,326 core science facts. The dataset requires not only knowledge of these facts but also the application of broad common knowledge. It 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, as well as crowd-sourced data for further analysis. We report 0-shot accuracy on OpenbookQA.

TruthfulQA (Lin et al., ) is a benchmark designed to evaluate the truthfulness of a language model’s responses. It consists of 817 questions across 38 categories, such as health, law, finance, and politics. The questions are crafted to reflect common false beliefs or misconceptions that might lead humans to answer inaccurately. We report 0-shot accuracy on TruthfulQA.

GSM8K (Cobbe et al., ) 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 a sequence of elementary calculations using basic arithmetic operations. A capable middle school student should be able to solve all the problems, making the dataset suitable for evaluating multi-step mathematical reasoning. We report the exact match score for 8-shot GSM8K.

CEVAL (Huang et al., ) is a comprehensive Chinese evaluation suite designed to assess the advanced knowledge and reasoning abilities of LLMs in a 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., ) 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 report the accuracy for the 5-shot CMMLU.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Please see weaknesses.
soundness: 3
strengths: MoH can achieve competitive performance while using fewer attention heads.By introducing shared heads and a two-stage routing mechanism, MoH enhances the standard Mixture-of-Experts (MoE) method, enabling the model to capture shared knowledge more effectively across different contexts.

MoH can be fine-tuned from pre-trained multi-head attention models, such as LLaMA3-8B, significantly enhancing the applicability of the model.

The method has been validated across various popular model frameworks, including Vision Transformers (ViT), Diffusion Models (DiT), and Large Language Models (LLMs), demonstrating superior performance in both image classification and language tasks.
summary: In the field of deep learning, multi-head attention mechanism has always been a core component of Transformer models, achieving great success in natural language processing and computer vision tasks. However, research has shown that not all attention heads are equally important, and many attention heads can be pruned without affecting model accuracy. Based on this insight, this paper proposes a new architecture called Mixture of Head Attention (MoH) aimed at improving the efficiency of attention mechanisms while maintaining or surpassing previous accuracy levels.
weaknesses: It is suggested to provide more evidence about the diversity within the selected heads. Visualizations and statistics of the distribution may provide more insights.

What is the effect of MoH on multi-task joint learning. More discussions or experiments are welcomed.

What is the `density' in Figure.3. Is it a weight used to select whether to activate？

The discussion section indicate that MoA only suitable for encoder-decoder architecture.  It requires more evidence and explanation.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: The author claims that the inference efficiency is improved by MoH. However, there is limited ground for this claim in the experiment. Is there any more evidence to support this claim? If further experiments are difficult, even a theoretical interpretation should be presented.
soundness: 3
strengths: - The author conducts comprehensive verification experiments to assess the performance of MoH, demonstrating results that are equal to or surpass previous methods. 
- MoH can significantly reduce the head's resources, which can tackle the most important problem of heavy MHSA operations.
summary: The authors introduce an approach that mitigates redundancy among attention heads through the use of the MoE, which adaptively selects attention heads according to input tokens. This method enhances inference efficiency by employing only those heads that are crucial for feature extraction during the inference process. The MoH demonstrates enhanced performance, even when utilizing a limited number of heads, as evidenced by comprehensive validation experiments.


I appreciate the authors' response. The response provided to my comment appears to be satisfactory, so I keep my score.
weaknesses: - The author performed ablation studies utilizing different ratios to determine the optimal configuration of shared heads or activated heads; however, this approach is heuristic. The ablation study concerning the ratio of shared heads presented in Table 7 indicates that identifying the optimal head ratio shows significant challenges.
- Given that the primary focus of this paper is to enhance inference efficiency through the reduction of multi-head ratios, it is essential to conduct an experiment that compares this approach with prior methods aimed at decreasing multi-head ratios.
- The author asserts that the shared head acquires common knowledge in Line [180-183], yet this paper is limited to providing evidence that the shared head genuinely learns common knowledge.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: See Weaknesses.
soundness: 3
strengths: 1. The idea of applying the mixture-of-experts paradigm to attention heads is novel.
2. MoH shows clear effectiveness for reducing computational overhead by activating fewer attention heads without sacrificing accuracy.
3. The paper presents a wide range of experiments across different model types, demonstrating the effectiveness of MoH. The ability to fine-tune pre-trained multi-head attention models like LLaMA adds practical value to the method.
summary: The paper introduces Mixture-of-Head Attention (MoH) to the multi-head attention mechanism in Transformer models, incorporating a routing mechanism that activates the most relevant attention heads for each token. Extensive experiments across diverse model architectures demonstrate that MoH achieves comparable or better performance with fewer attention heads than traditional multi-head attention.
weaknesses: 1. The contribution is incremental. Replacing the summation of heads with a weighted sum and using expert selection are not entirely new ideas in machine learning, and their application here may not be sufficiently ground-breaking to warrant significant attention without a stronger theoretical basis.
2. The ablation studies are limited in scope and fail to deeply explore the design choices behind MoH. For example, there is little discussion on the impact of different numbers of activated heads beyond the experiments shown. The use of shared heads is also not well-motivated, and the reported improvements may be marginally due to tuning specific hyperparameters.
3. The paper could benefit from evaluations on more diverse and challenging tasks, such as object detection and instance segmentation, in line with prior research on ViT designs.
4. MoH introduces additional complexity with its routing mechanism. The added complexity is not fully justified by the performance gains, especially given that the gains appear marginal in some cases (e.g., DiT models).

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: Please see weakness section
soundness: 3
strengths: 1. Without altering the number of parameters, this work treats standard multi-head attention as Mixture-of-Head attention, which enhances the flexibility of the attention mechanism and shows improved performance.
2. To consolidate common knowledge, it introduces shared heads along with a corresponding routing mechanism. The ablation study results have validated the effectiveness of these designs.
3. Popular tasks including ViT for image classification, DiT for image generation, and LLM for language generation, demonstrate superior performance. Furthermore, the proposed MoH attention can continue-tune pre-trained standard multi-head LLaMA3-8B, significantly enhancing its applicability.
summary: This paper observes that each attention head in multi-head attention operates in parallel. By formulating the multi-head attention in summation form, it builds a dynamic mixture-of-head attention without increasing the number of parameters. Meanwhile, it introduces shared heads and a two-stage routing mechanism to enhance the standard MoH method. Extensive experiments across popular model frameworks, including ViT, DiT, LLMs and continue-tuning LLMs demonstrate strong performance and applicability.
weaknesses: 1. This work claims enhanced inference efficiency multiple times. However, there is a lack of experimental evidence to support this claim. Upon reviewing the implementation provided in the Supplementary Material, it appears that you simply mask the useless heads after obtaining the whole multi-head results, which may not genuinely improve inference efficiency. Additionally, the process of dynamically routing each token to the appropriate heads could potentially increase inference costs.
2. The cited works that show multi-head attention contains redundant attention heads primarily focus on natural language processing, it would be better incorporate additional studies from the field of computer vision to provide a more comprehensive perspective.
3. In Equation 5 on line 190, the dimension of $W_r x_t $ is $h-h_s$, so for indices where $h_s+1<i \leq h$, $i$ in $(W_r x_t)_i$ should be 
${i-h_s}$.
4. In Table 5 on line 383, does the baseline LLaMA3-8B refer to the model after continue-tuning with standard multi-head attention, consistent with the configuration of MoH models, or does it represent the starting point for continue-tuning? If the baseline is the starting point, it would be better to add a baseline that reflects the results after continue-tuning with multi-head attention, as continue-tuning is likely to improve performace.

Meta-review: This paper investigates neural architecture design and proposes a new architecture called a Mixture of Head Attention aimed at improving the efficiency of attention mechanisms while maintaining accuracy. The manuscript was reviewed by four experts in the field. The recommendations are ("5: marginally below the acceptance threshold", 3 x "6: marginally above the acceptance threshold"). The reviewers raised several concerns regarding the paper, e.g., incremental technical contribution, unconvincing experimental evaluations, etc. Considering the reviewers' concerns, we regret that the paper cannot be recommended for acceptance at this time. The authors are encouraged to consider the reviewers' comments when revising the paper for submission elsewhere.

Decision: Reject