Back

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

Peng JinBo ZhuLi YuanShuicheng Yan

article2025ICML64 citations

Proposes Mixture-of-Head attention, a parameter-neutral architecture that dynamically routes tokens to a subset of attention heads using weighted summation to cut computation while improving accuracy across vision models, diffusion models, and large language models.

Listen

Modern artificial intelligence models rely heavily on the Transformer architecture, where multi-head attention serves as the core computational mechanism. However, standard multi-head attention activates every attention head uniformly for every piece of data, despite growing evidence that many heads are redundant. This full activation introduces unnecessary computational overhead and raises inference costs when deploying large models at scale.

The article introduces Mixture-of-Head attention (MoH), a dynamic mechanism designed to reduce computational demands during inference while maintaining or exceeding baseline model accuracy, all without increasing total model parameters.

MoH treats individual attention heads as specialized experts within a dynamic routing framework. For each token, the system selects only a subset of the most relevant heads to activate, while a designated group of shared heads remains constantly active to retain common foundational knowledge. The architecture also incorporates a two-stage routing mechanism to weight head contributions and a load-balancing loss to keep heads evenly trained. The authors evaluated MoH across three mainstream AI domains: vision models for image classification (Vision Transformers), generative image diffusion models (Diffusion Transformers), and large language models (trained from scratch and continue-tuned on an existing 8-billion-parameter foundation model).

The evaluation produced several key findings. First, MoH models consistently matched or outperformed standard multi-head baselines while activating only 50% to 90% of total attention heads across vision, diffusion, and language domains. Second, existing foundation models can be converted into MoH models efficiently; continue-tuning a baseline 8-billion-parameter model with about 3% of its original pre-training budget produced a 2.4% average accuracy improvement across 14 benchmarks while using only 75% of heads. Third, dynamic routing reduced empirical inference latency, demonstrating greater time savings as sequence lengths expanded. Finally, ablation studies showed that always-active shared heads and two-stage routing are critical to preventing performance degradation.

These findings indicate that artificial intelligence architectures can achieve higher efficiency and specialized parameter use without parameter expansion or costly training from scratch. For organizations deploying generative and visual AI, adopting MoH offers a viable path to lower computational hardware costs, accelerate response latencies, and reduce energy expenses during inference. The successful conversion of pre-trained foundation models also provides an economical upgrade path that avoids multi-million-dollar re-training cycles.

Organizations evaluating this architecture should consider piloting MoH in high-throughput inference pipelines where latency and serving costs are critical. Implementers should maintain a relatively balanced proportion of shared heads (recommended at greater than 40% of activated heads) to ensure stability. Dense prediction tasks like image generation require conservative head reduction (e.g., activating around 90% of heads) compared to language or classification tasks (50% to 75%). Future work should explore extending head sparsity below 50%, scaling evaluations to larger parameter regimes beyond 8 billion, and testing performance on multimodal and audio architectures.

Confidence in these findings is high for classification and standard language benchmarks up to the 8-billion-parameter scale under controlled experimental conditions. However, decision-makers should note certain limitations: converting models on new data distributions showed slight performance drops in non-English and math tasks due to catastrophic forgetting during tuning. Real-world cost benefits will also depend on low-level software kernel support for sparse operations.

National University of SingaporePeking UniversityPeng Cheng LaboratoryRabbitpre IntelligenceSkywork AImixture-of-head attentionmulti-head attentiondynamic attention-head routing
arXiv: 2410.11842SkyworkAI/MoH

No sufficiently relevant recommendations were found.

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

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%~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. The code is available at https://github.com/SkyworkAI/MoH.

Keywords

mixture-of-head attentionmulti-head attentiondynamic attention-head routingattention-head activation budgethead specializationlong-context inferencemassive multitask language understanding

Table of Contents

  • 1 Introduction
  • 2 Related Work
  • 3 Methodology
  • 3.1 Multi-Head Attention
  • 3.2 Mixture-of-Head Attention
  • 4 Experiments
  • 4.1 ViT for Image Classification
  • 4.2 DiT for Class-Conditional Image Generation
  • 4.3 Training LLMs from Scratch
  • 4.4 Continue-Tuning LLaMA3-8B
  • 4.5 Ablative Analysis
  • 5 Discussion
  • 6 Conclusion
  • References
  • A Additional Discussions
  • A.1 Why is MoH Superior to Vanilla Multi-Head Attention?
  • A.2 Limitations and Future Work
  • B Implementation Details
  • B.1 ViT for Image Classification
  • B.2 DiT for Class-Conditional Image Generation
  • B.3 Training LLMs from Scratch
  • B.4 Continue-Tuning LLaMA3-8B
  • C Additional Experiments
  • D Additional Qualitative Analysis
  • E Details of Quantitative Evaluations for LLMs

Knowls

  1. Knowl 1 — Mixture-of-Head attention sparsifies and reweights the head sum

    model/method

    For a token xtx_t attending to context tokens X′X', standard multi-head attention can be written as a sum of projected head outputs. MoH uses the same heads and output projections but multiplies each head’s contribution by a token-dependent routing score:

    ytMHA=∑i=1hHi(xt,X′)WO(i),ytMoH=∑i=1hgt,iHi(xt,X′)WO(i).y_t^{\mathrm{MHA}}=\sum_{i=1}^{h} H_i(x_t,X')W_O^{(i)},\qquad y_t^{\mathrm{MoH}}=\sum_{i=1}^{h} g_{t,i}H_i(x_t,X')W_O^{(i)}.

    Here, hh is the total number of attention heads; Hi(xt,X′)H_i(x_t,X') is the scaled-dot-product attention output of head ii for query token xtx_t and context X′X'; WO(i)W_O^{(i)} is the output-projection block for that head; and gt,ig_{t,i} is its routing weight for token tt. An unselected head has gt,i=0g_{t,i}=0. In the basic sparse routing design, the router selects the top KK heads for each token. Thus, MoH replaces the fixed sum with a token-specific weighted sum, while retaining the existing attention heads rather than adding head experts. The router adds routing parameters, so the resulting model has a comparable—not necessarily identical—parameter count to standard multi-head attention.

  2. Knowl 2 — Shared heads and two-stage routing balance common and token-specific attention

    model/method

    MoH divides its hh attention heads into hsh_s shared heads, which remain active for every token, and h−hsh-h_s routed heads, of which the router selects the top KK per token. For token representation xt∈Rdinx_t\in\mathbb{R}^{d_{\mathrm{in}}}, the two-stage router first computes separate softmax distributions over the shared and routed head groups, then scales each group by a token-dependent coefficient:

    gt,i={αt,1 [Softmax⁡(Wsxt)]i,1≤i≤hs,αt,2 [Softmax⁡(Wrxt)]i−hs,i>hs and routed head i is selected,0,otherwise,[αt,1,αt,2]=Softmax⁡(Whxt).g_{t,i}=\begin{cases} \alpha_{t,1}\,[\operatorname{Softmax}(W_sx_t)]_i, & 1\le i\le h_s,\\ \alpha_{t,2}\,[\operatorname{Softmax}(W_rx_t)]_{i-h_s}, & i>h_s\text{ and routed head }i\text{ is selected},\\ 0, & \text{otherwise}, \end{cases} \qquad [\alpha_{t,1},\alpha_{t,2}]=\operatorname{Softmax}(W_hx_t).

    Here, Ws∈Rhs×dinW_s\in\mathbb{R}^{h_s\times d_{\mathrm{in}}} and Wr∈R(h−hs)×dinW_r\in\mathbb{R}^{(h-h_s)\times d_{\mathrm{in}}} produce shared-head and routed-head logits; routed head ii is selected when its corresponding component of WrxtW_rx_t is among the top KK routed logits. The trainable matrix Wh∈R2×dinW_h\in\mathbb{R}^{2\times d_{\mathrm{in}}} produces the two mixing coefficients. The shared heads provide an always-available path for information common across contexts, while the routed heads can specialize to each token; the two coefficients dynamically set the relative contribution of those groups.

  3. Knowl 3 — A query-norm router enables conversion of pretrained attention models

    model/method

    To turn a pretrained multi-head attention model into MoH without inserting a randomly initialized learned router, the authors use a parameter-free routing score based on each attention head’s query-vector norm. In the LLaMA3-8B conversion, the first 16 heads of each layer are designated shared heads; the other heads are routed. For token tt, let qt,i∈Rdkq_{t,i}\in\mathbb{R}^{d_k} be the query vector for head ii, where dkd_k is the query dimension. The router ranks routed heads using st,i=∥qt,i∥2s_{t,i}=\lVert q_{t,i}\rVert_2 and selects the required heads. Its discrete forward gate is

    gt,iq=1(token t selects head i),g^q_{t,i}=\mathbf{1}(\text{token }t\text{ selects head }i),

    where 1\mathbf{1} is the indicator function. Because this hard gate would obstruct ordinary back-propagation, training uses a straight-through estimator: if gt,ig_{t,i} is the corresponding real-valued routing score and LL is the training loss, the backward pass sets

    ∂L∂gt,iq=∂L∂gt,i.\frac{\partial L}{\partial g^q_{t,i}}=\frac{\partial L}{\partial g_{t,i}}.

    This lets the forward pass use sparse, binary head selection while passing the incoming gradient through the sparsity operation during training.

  4. Knowl 4 — A load-balance auxiliary loss discourages routing collapse

    equation

    During training, MoH adds an auxiliary loss over routed heads to discourage most tokens from selecting only a small subset of them. For a batch containing TT tokens xtx_t, total head count hh, shared-head count hsh_s, routed-head projection WrW_r, and routed head i∈{hs+1,…,h}i\in\{h_s+1,\ldots,h\}, define the mean router probability PiP_i and empirical selection frequency fif_i by

    Pi=1T∑t=1T[Softmax⁡(Wrxt)]i−hs,fi=1T∑t=1T1(token t selects head i),P_i=\frac{1}{T}\sum_{t=1}^{T}[\operatorname{Softmax}(W_rx_t)]_{i-h_s},\qquad f_i=\frac{1}{T}\sum_{t=1}^{T}\mathbf{1}(\text{token }t\text{ selects head }i),

    where 1\mathbf{1} is the indicator function. The load-balance loss and total training objective are

    Lb=∑i=hs+1hPifi,L=Ltask+βLb.L_b=\sum_{i=h_s+1}^{h}P_if_i,\qquad L=L_{\mathrm{task}}+\beta L_b.

    Here, LtaskL_{\mathrm{task}} is the task-specific loss and β\beta controls the auxiliary-loss contribution; the paper uses β=0.01\beta=0.01 by default across tasks. The auxiliary term is computed for routed heads, not the always-active shared heads.

  5. Knowl 5 — Continue-tuned MoH-LLaMA3-8B exceeds the original model’s reported average

    empirical result

    The authors converted LLaMA3-8B to MoH using a two-stage continuation procedure because the original pretraining data were unavailable. They first continued training the original model for 300B tokens to adapt it to the available data, then trained the MoH version for another 100B tokens, for 400B tokens in total—about 3% of the original 15T-token pretraining budget. During conversion, 75% of attention heads were activated. The first stage also expanded the vocabulary from 128,256 to 160,896 tokens.

    The final MoH-LLaMA3-8B achieved a reported average score of 64.0% across 14 benchmarks, outperforming the original LLaMA3-8B by 2.4% while using 75% of its attention heads. The benchmarks used mixed few-shot settings: 0-shot LAMBADA, LogiQA, PIQA, SciQ, WinoGrande, and TruthfulQA; 5-shot CEVAL, CMMLU, and MMLU; 10-shot HellaSwag; 25-shot ARC Challenge; 32-shot Natural Questions and BoolQ; and 8-shot GSM8K. The reported comparison on five of these benchmarks was:

    BenchmarkLLaMA3-8B (%)MoH-LLaMA3-8B (%)
    MMLU (5-shot)65.265.8
    CEVAL (5-shot)52.361.5
    CMMLU (5-shot)50.764.4
    GSM8K (8-shot)49.556.9
    TruthfulQA35.444.0

    The reported training curve also shows recovery to over 95% of the original model’s performance within 10B tokens of MoH training on the plotted MMLU, HellaSwag, and TruthfulQA evaluations.

  6. Knowl 6 — MoH-ViT matches or exceeds TransNeXt with fewer active heads

    empirical result

    MoH-ViT models replace only the standard attention layers in TransNeXt and are trained on ImageNet-1K for 300 epochs at 224×224 resolution, with the remaining training settings kept the same as TransNeXt. The comparison below reports parameter count, percentage of active attention heads, and ImageNet-1K Top-1 accuracy.

    ModelParameters (M)Active heads (%)Top-1 accuracy (%)
    TransNeXt-S5010084.7
    MoH-ViT-S508084.7
    MoH-ViT-S507584.6
    TransNeXt-B9010084.8
    MoH-ViT-B907584.9
    MoH-ViT-B905084.7

    Thus, MoH-ViT-B exceeds its full-head TransNeXt-B baseline by 0.1 percentage points at 75% head activation and matches it at 50%. MoH-ViT-S matches the full-head baseline at 80% activation and is 0.1 points lower at 75%.

  7. Knowl 7 — MoH language models trained from scratch retain or improve benchmark averages

    empirical result

    The authors trained decoder-only LLMs from scratch with Megatron using public text datasets and evaluated 0-shot accuracy on SciQ, PIQA, WinoGrande, OpenbookQA, LogiQA, and TruthfulQA. The table reports the average across those six benchmarks for two model sizes and training budgets; the small and large baselines have 186M and 881M parameters, respectively. MoH-LLM-B has about 1M additional router parameters relative to LLM-B.

    Model and training budgetActive heads (%)Average accuracy (%)
    LLM-S, 100B tokens10043.9
    MoH-LLM-S, 100B tokens7544.6
    MoH-LLM-S, 100B tokens5045.4
    LLM-B, 100B tokens10047.4
    MoH-LLM-B, 100B tokens7547.8
    MoH-LLM-B, 100B tokens5047.2
    LLM-B, 200B tokens10047.9
    MoH-LLM-B, 200B tokens7548.5
    MoH-LLM-B, 200B tokens5047.5

    MoH-LLM-S at 50% activation scores 1.5 points above its full-head baseline. For the larger model, 75% activation gives the best MoH average at both training budgets, exceeding the corresponding full-head baseline by 0.4 points at 100B tokens and 0.6 points at 200B tokens. The paper suggests that fewer active heads may regularize the smaller model, but presents this as a possible explanation rather than an established cause.

  8. Knowl 8 — MoH improves DiT FID at 90% head activation, while 75% can reduce quality

    empirical result

    MoH-DiT replaces the standard attention layers in DiT for class-conditional ImageNet-1K image generation at 256×256 resolution. Models use the same other training settings as DiT, including a pretrained variational autoencoder; the attention-head budget is distributed unevenly across layers, with fewer heads active in shallow layers and more in deeper layers. The table reports FID, for which lower is better, at the stated training budgets.

    ModelTraining stepsActive heads (%)FID ↓
    DiT-S/2400K10068.40
    MoH-DiT-S/2400K9067.25
    MoH-DiT-S/2400K7569.42
    DiT-B/2400K10043.47
    MoH-DiT-B/2400K9043.40
    MoH-DiT-B/2400K7543.61
    DiT-L/2400K10023.33
    MoH-DiT-L/2400K9023.17
    MoH-DiT-L/2400K7524.29
    DiT-XL/27,000K1009.62
    MoH-DiT-XL/22,000K7510.95
    MoH-DiT-XL/22,000K9010.67
    MoH-DiT-XL/27,000K908.56

    At 90% activation, MoH has lower FID than the listed full-head baseline at each scale with the same stated training budget. At 75%, the S, B, and L MoH models have higher FID than their full-head baselines. With classifier-free guidance scale 1.25 at 7,000K steps, the reported XL/2 FID is 3.22 for DiT and 2.94 for MoH-DiT at 90% head activation.

  9. Knowl 9 — Sparse-matrix timing shows lower attention latency as fewer heads are activated

    empirical result

    The authors compared multi-head attention and MoH with 32 heads of dimension 64 at sequence lengths 256 and 512. To reduce the effect of underlying operator optimizations, they converted query, key, and value features to sparse matrices using the router mask and used sparse matrix multiplication for all methods in the timing test. The reported times are milliseconds.

    Sequence lengthMethodActive heads (%)Time (ms)
    256Multi-head attention1000.360
    256MoH900.352
    256MoH750.321
    256MoH500.225
    512Multi-head attention1001.376
    512MoH901.351
    512MoH751.180
    512MoH500.863

    Under this sparse-multiplication test, lower activation reduces reported latency at both sequence lengths; the largest reduction is at 50% activation. These measurements demonstrate the timing benefit in the stated setup, not a hardware-independent guarantee for all implementations.

  10. Knowl 10 — Ablations identify benefits from shared heads and two-stage weighting

    empirical result

    An ablation compared MoH without shared heads or two-stage routing, MoH with shared heads only, and the full design with both components. Image-classification results use MoH-ViT-S at 75% active heads for 100 epochs; generation results use MoH-DiT-S/2 at 75% active heads for 400K steps. Higher accuracy, IS, precision, and recall are better; lower FID and sFID are better.

    Shared headsTwo-stage routingViT accuracy (%) ↑DiT FID ↓DiT sFID ↓DiT IS ↑Precision ↑Recall ↑
    NoNo75.671.9713.5819.060.350.55
    YesNo78.369.5412.8019.670.360.55
    YesYes78.669.4212.8519.960.360.55

    Adding shared heads improves ViT accuracy and all reported DiT metrics except unchanged recall. Adding two-stage routing on top gives a further 0.3-point accuracy increase and a 0.12 FID reduction; it also increases IS, while sFID rises slightly and precision and recall remain unchanged. The full configuration therefore does not improve every metric relative to the shared-head-only variant.

  11. Knowl 11 — MoH heads show lower inter-head similarity in vision and language tests

    empirical result

    The authors measured pairwise attention-pattern similarity and cosine similarity between output features of different heads for standard multi-head attention and MoH. For attention matrices AA and A′A', the pattern similarity is defined as 1−12E[∥A−A′∥1]1-\tfrac{1}{2}\mathbb{E}[\lVert A-A'\rVert_1]; since each attention row is a probability distribution, this score lies between 0 and 1. Lower pattern similarity and lower output-feature cosine similarity indicate less similarity between heads under these measurements.

    MethodAttention-pattern similarity, ViTAttention-pattern similarity, LLMOutput-feature cosine similarity, ViTOutput-feature cosine similarity, LLM
    Multi-head attention0.51590.47950.04110.2550
    MoH0.39780.43330.01650.2042

    MoH has lower values than standard multi-head attention in both model families and both measures. The paper interprets this pattern as evidence of reduced head redundancy and greater head differentiation; the reported comparisons are diagnostic measurements rather than a causal test of specialization.

Coverage note — The qualitative head-load visualizations and the paper’s future-work limitations (including activation below 50%, multimodal or audio tasks, heterogeneous heads, and scaling beyond 8B parameters) are not separate knowls: the plots provide diagnostic support for the extracted routing results, while the limitations are stated as open directions rather than demonstrated contributions.

References

  1. 1.Balasubramanian, S., Basu, S., and Feizi, S. Decomposing and interpreting image representations via text in vits beyond clip. arXiv preprint arXiv:2406.01583, 2024.
  2. 2.Basile, L., Maiorca, V., Bortolussi, L., Rodolà, E., and Locatello, F. Residual transformer alignment with spectral decomposition. arXiv preprint arXiv:2411.00246, 2024.
  3. 3.Beeching, E., Fourrier, C., Habib, N., Han, S., Lambert, N., Rajani, N., Sanseviero, O., Tunstall, L., and Wolf, T. Open llm leaderboard (2023-2024). https://huggingface.co/spaces/open-llm-leaderboard-old/open_llm_leaderboard, 2023.
  4. 4.Bengio, Y., Léonard, N., and Courville, A. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013.
  5. 5.Bhattacharyya, M., Chattopadhyay, S., and Nag, S. Decatt: Efficient vision transformers with decorrelated attention heads. In CVPRW, pp. 4695–4699, 2023.
  6. 6.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: A suite for analyzing large language models across training and scaling. In ICML, pp. 2397–2430, 2023.
  7. 7.Bisk, Y., Zellers, R., Gao, J., Choi, Y., et al. Piqa: Reasoning about physical commonsense in natural language. In AAAI, pp. 7432–7439, 2020.
  8. 8.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 NeurIPS, pp. 1877–1901, 2020.
  9. 9.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.
  10. 10.Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457, 2018.
  11. 11.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.14168, 2021.
  12. 12.Computer, T. Redpajama: an open dataset for training large language models, 2023. URL https://github.com/togethercomputer/RedPajama-Data.
  13. 13.Cordonnier, J.-B., Loukas, A., and Jaggi, M. Multi-head attention: Collaborate instead of concatenate. arXiv preprint arXiv:2006.16362, 2020.
  14. 14.Csordás, R., Pi˛ekos, P., Irie, K., and Schmidhuber, J. Switch-head: Accelerating transformers with mixture-of-experts attention. In NeurIPS, pp. 74411–74438, 2024.
  15. 15.Cubuk, E. D., Zoph, B., Shlens, J., and Le, Q. V. Randaugment: Practical automated data augmentation with a reduced search space. In CVPRW, pp. 702–703, 2020.
  16. 16.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-experts language models. arXiv preprint arXiv:2401.06066, 2024.
  17. 17.Dai, Z., Liu, H., Le, Q. V., and Tan, M. Coatnet: Marrying convolution and attention for all data sizes. In NeurIPS, pp. 3965–3977, 2021.
  18. 18.Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In CVPR, pp. 248–255, 2009.
  19. 19.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.
  20. 20.Du, N., Huang, Y., Dai, A. 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 ICML, pp. 5547–5569, 2022.
  21. 21.Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024.
  22. 22.Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1–39, 2022.
  23. 23.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.19258, 2024.
  24. 24.Gandelsman, Y., Efros, A. A., and Steinhardt, J. Interpreting clip’s image representation via text-based decomposition. arXiv preprint arXiv:2310.05916, 2023.
  25. 25.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.
  26. 26.Gao, L., 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 https://zenodo.org/records/12608602.
  27. 27.Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding. In ICLR, 2021.
  28. 28.Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to a local nash equilibrium. In NeurIPS, 2017.
  29. 29.Huang, G., Sun, Y., Liu, Z., Sedra, D., and Weinberger, K. Q. Deep networks with stochastic depth. In ECCV, pp. 646–661, 2016.
  30. 30.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-discipline chinese evaluation suite for foundation models. In NeurIPS, 2023.
  31. 31.Jacobs, R. A., Jordan, M. I., Nowlan, S. J., and Hinton, G. E. Adaptive mixtures of local experts. Neural computation, 3(1):79–91, 1991.
  32. 32.Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., Casas, D. d. l., Hanna, E. B., Bressand, F., et al. Mixtral of experts. arXiv preprint arXiv:2401.04088, 2024.
  33. 33.Jin, P., Huang, J., Xiong, P., Tian, S., Liu, C., Ji, X., Yuan, L., and Chen, J. Video-text as game players: Hierarchical banzhaf interaction for cross-modal representation learning. In CVPR, pp. 2472–2482, 2023.
  34. 34.Jin, P., Li, H., Yuan, L., Yan, S., and Chen, J. Hierarchical banzhaf interaction for general video-language representation learning. TPAMI, 2024a.
  35. 35.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.
  36. 36.Jin, P., Zhu, B., Yuan, L., and Yan, S. Moe++: Accelerating mixture-of-experts methods with zero-computation experts. In ICLR, 2025.
  37. 37.Kenton, J. D. M.-W. C. and Toutanova, L. K. Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL, volume 1, pp. 2, 2019.
  38. 38.Kingma, D. P. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
  39. 39.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 question answering research. Transactions of the Association for Computational Linguistics, 7:453–466, 2019.
  40. 40.Kynkäänniemi, T., Karras, T., Laine, S., Lehtinen, J., and Aila, T. Improved precision and recall metric for assessing generative models. In NeurIPS, 2019.
  41. 41.Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z. Gshard: Scaling giant models with conditional computation and automatic sharding. In ICLR, 2021.
  42. 42.Lewis, M., Bhosale, S., Dettmers, T., Goyal, N., and Zettlemoyer, L. Base layers: Simplifying training of large, sparse models. In ICML, pp. 6265–6274, 2021.
  43. 43.Li, H., Zhang, Y., Koto, F., Yang, Y., Zhao, H., Gong, Y., Duan, N., and Baldwin, T. Cmmlu: Measuring massive multitask language understanding in chinese. arXiv preprint arXiv:2306.09212, 2023a.
  44. 44.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.
  45. 45.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–4814, 2022.
  46. 46.Lin, B., Zhu, B., Ye, Y., Ning, M., Jin, P., and Yuan, L. Video-llava: Learning united visual representation by alignment before projection. arXiv preprint arXiv:2311.10122, 2023.
  47. 47.Lin, B., Tang, Z., Ye, Y., Cui, J., Zhu, B., Jin, P., Zhang, J., Ning, M., and Yuan, L. Moe-llava: Mixture of experts for large vision-language models. arXiv preprint arXiv:2401.15947, 2024.
  48. 48.Lin, S., Hilton, J., and Evans, O. TruthfulQA: Measuring how models mimic human falsehoods. In ACL, pp. 3214–3252, 2022.
  49. 49.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.
  50. 50.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 IJCAI, pp. 3622–3628, 2020.
  51. 51.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.
  52. 52.Liu, Z., Cheng, K.-T., Huang, D., Xing, E. P., and Shen, Z. Nonuniform-to-uniform quantization: Towards accurate quantization via generalized straight-through estimation. In CVPR, pp. 4942–4952, 2022.
  53. 53.Loshchilov, I. and Hutter, F. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983, 2016.
  54. 54.Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
  55. 55.Michel, P., Levy, O., and Neubig, G. Are sixteen heads really better than one? In NeurIPS, pp. 14014–14024, 2019.
  56. 56.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.
  57. 57.Muennighoff, N., Soldaini, L., Groeneveld, D., Lo, K., Morrison, J., Min, S., Shi, W., Walsh, P., Tafjord, O., Lambert, N., et al. Olmoe: Open mixture-of-experts language models. arXiv preprint arXiv:2409.02060, 2024.
  58. 58.Nash, C., Menick, J., Dieleman, S., and Battaglia, P. W. Generating images with sparse representations. arXiv preprint arXiv:2103.03841, 2021.
  59. 59.OpenAI. Introducing chatgpt. CoRR, 2022. URL https://openai.com/blog/chatgpt.
  60. 60.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 NeurIPS, pp. 27730–27744, 2022.
  61. 61.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.
  62. 62.Peebles, W. and Xie, S. Scalable diffusion models with transformers. In ICCV, pp. 4195–4205, 2023.
  63. 63.Peng, H., Schwartz, R., Li, D., and Smith, N. A. A mixture of h − 1 heads is better than h heads. arXiv preprint arXiv:2005.06537, 2020.
  64. 64.Puigcerver, J., Ruiz, C. R., Mustafa, B., and Houlsby, N. From sparse to soft mixtures of experts. In ICLR, 2024.
  65. 65.Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485–5551, 2020.
  66. 66.Rajbhandari, S., Li, C., Yao, Z., Zhang, M., Aminabadi, R. 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.
  67. 67.Roller, S., Sukhbaatar, S., Weston, J., et al. Hash layers for large sparse models. In NeurIPS, pp. 17555–17566, 2021.
  68. 68.Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In CVPR, pp. 10684–10695, 2022.
  69. 69.Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99–106, 2021.
  70. 70.Salimans, T., Goodfellow, I., Zaremba, W., Cheung, V., Radford, A., and Chen, X. Improved techniques for training gans. In NeurIPS, 2016.
  71. 71.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:1701.06538, 2017.
  72. 72.Shi, D. Transnext: Robust foveal visual perception for vision transformers. In CVPR, pp. 17773–17783, 2024.
  73. 73.Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019.
  74. 74.Si, C., Yu, W., Zhou, P., Zhou, Y., Wang, X., and Yan, S. Inception transformer. In NeurIPS, pp. 23495–23509, 2022.
  75. 75.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., Lambert, N., Magnusson, I., Morrison, J., Muennighoff, 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://arxiv.org/abs/2402.00159.
  76. 76.Szegedy, C., Vanhoucke, V., Ioffe, S., Shlens, J., and Wojna, Z. Rethinking the inception architecture for computer vision. In CVPR, pp. 2818–2826, 2016.
  77. 77.Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., and Jégou, H. Training data-efficient image transformers & distillation through attention. In ICML, pp. 10347–10357, 2021.
  78. 78.Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
  79. 79.Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. In NeurIPS, 2017.
  80. 80.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.
  81. 81.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 arXiv:2406.18139, 2024.
  82. 82.Wang, H., Ma, S., Wang, R., and Wei, F. Q-sparse: All large language models can be fully sparsely-activated. arXiv preprint arXiv:2407.10969, 2024.
  83. 83.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–424, 2022.
  84. 84.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-experts language models. arXiv preprint arXiv:2406.06563, 2024.
  85. 85.Welbl, J., Liu, N. F., and Gardner, M. Crowdsourcing multiple choice science questions. arXiv preprint arXiv:1707.06209, 2017.
  86. 86.Wu, W., Wang, Y., Xiao, G., Peng, H., and Fu, Y. Retrieval head mechanistically explains long-context factuality. arXiv preprint arXiv:2404.15574, 2024.
  87. 87.Xiao, G., Tang, J., Zuo, J., Guo, J., Yang, S., Tang, H., Fu, Y., and Han, S. Duoattention: Efficient long-context llm inference with retrieval and streaming heads. arXiv preprint arXiv:2410.10819, 2024.
  88. 88.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.
  89. 89.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.00641, 2021.
  90. 90.Yang, J., Li, C., Dai, X., and Gao, J. Focal modulation networks. In NeurIPS, pp. 4203–4217, 2022b.
  91. 91.Yu, W., Si, C., Zhou, P., Luo, M., Zhou, Y., Feng, J., Yan, S., and Wang, X. Metaformer baselines for vision. TPAMI, 2023.
  92. 92.Yuan, L., Chen, Y., Wang, T., Yu, W., Shi, Y., Jiang, Z.-H., Tay, F. E., Feng, J., and Yan, S. Tokens-to-token vit: Training vision transformers from scratch on imagenet. In ICCV, pp. 558–567, 2021.
  93. 93.Yun, S. and Ro, Y. Shvit: Single-head vision transformer with memory efficient macro design. In CVPR, pp. 5756–5767, 2024.
  94. 94.Yun, S., Han, D., Oh, S. J., Chun, S., Choe, J., and Yoo, Y. Cutmix: Regularization strategy to train strong classifiers with localizable features. In ICCV, pp. 6023–6032, 2019.
  95. 95.Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can a machine really finish your sentence? In ACL, pp. 4791–4800, 2019.
  96. 96.Zhang, H. mixup: Beyond empirical risk minimization. arXiv preprint arXiv:1710.09412, 2017.
  97. 97.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.05144, 2022.
  98. 98.Zhong, Z., Zheng, L., Kang, G., Li, S., and Yang, Y. Random erasing data augmentation. In AAAI, pp. 13001–13008, 2020.
  99. 99.Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., Dai, A. M., Le, Q. V., Laudon, J., et al. Mixture-of-experts with expert choice routing. In NeurIPS, pp. 7103–7114, 2022.

Citation

MLA
Jin, P., et al. “MoH: Multi-Head Attention as Mixture-of-Head Attention”. arXiv, 2024, http://arxiv.org/abs/2410.11842v3.
APA
Jin, P., Zhu, B., Yuan, L., & Yan, S. (2024). MoH: Multi-Head Attention as Mixture-of-Head Attention. arXiv. http://arxiv.org/abs/2410.11842v3
Chicago
Jin, P., B. Zhu, L. Yuan, and S. Yan. 2024. “MoH: Multi-Head Attention as Mixture-of-Head Attention”. arXiv. http://arxiv.org/abs/2410.11842v3.
Harvard
Jin, P. et al. (2024) “MoH: Multi-Head Attention as Mixture-of-Head Attention”, arXiv [Preprint]. Available at: http://arxiv.org/abs/2410.11842v3.
Vancouver
1. Jin P, Zhu B, Yuan L, Yan S (2024) MoH: Multi-Head Attention as Mixture-of-Head Attention. arXiv

BibTeX

@article{jin2024moh,
  title = {MoH: Multi-Head Attention as Mixture-of-Head Attention},
  author = {Jin, Peng and Zhu, Bo and Yuan, Li and Yan, Shuicheng},
  year = {2024},
  journal = {arXiv},
  url = {http://arxiv.org/abs/2410.11842v3},
  eprint = {2410.11842}
}
Metadata:arXiv

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF
License: https://creativecommons.org/licenses/by/4.0/