MH-MoE: Multi-Head Mixture-of-Experts

TL;DR

This paper introduces a novel multi-head mixture-of-experts model that improves language model performance while maintaining computational efficiency and compatibility with low-bit LLMs.

Contribution

The paper presents a new MH-MoE implementation that preserves FLOPs and parameters, enhancing model quality and compatibility with 1-bit LLMs.

Findings

Improved language model quality over vanilla and fine-grained MoE.

Maintains FLOPs and parameter parity with sparse MoE.

Compatible with 1-bit LLMs like BitNet.

Abstract

Multi-Head Mixture-of-Experts (MH-MoE) demonstrates superior performance by using the multi-head mechanism to collectively attend to information from various representation spaces within different experts. In this paper, we present a novel implementation of MH-MoE that maintains both FLOPs and parameter parity with sparse Mixture of Experts models. Experimental results on language models show that the new implementation yields quality improvements over both vanilla MoE and fine-grained MoE models. Additionally, our experiments demonstrate that MH-MoE is compatible with 1-bit Large Language Models (LLMs) such as BitNet.