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

🎯 Pitch

They made multi-head MoEs cost the same as standard sparse MoEs—achieving better perplexity by simply splitting tokens into sub-tokens before expert routing, with zero extra FLOPs or parameters. The method works so well it even boosts 1-bit LLMs.

1. Executive Summary

This paper proposes a new implementation of Multi-Head Mixture-of-Experts (MH-MoE) that maintains both FLOPs and parameter parity with standard sparse Mixture of Experts models—a constraint that the original MH-MoE formulation violated by incurring significantly higher computation. The authors evaluate their approach on language modeling using decoder-only Transformers trained on the RedPajama dataset, introducing a multi-head mechanism that splits input tokens into sub-tokens along the token dimension before routing them through expert FFNs (operationally: projecting input through a head layer, splitting into h sub-tokens, feeding each sub-token to a Mixture-of-Experts layer, then merging via a second linear projection). The proposed variants consistently outperform both vanilla SMoE and fine-grained SMoE baselines, with the 3-head MH-MoE achieving the lowest perplexity across all evaluated benchmarks—for instance, reaching 10.51 on RedPajama at 100K training steps versus 10.90 for standard SMoE and 10.74 for fine-grained SMoE—while using matched FLOPs and parameters. The paper further demonstrates that MH-MoE integrates effectively with 1-bit quantization (BitNet), establishing that the multi-head mechanism remains beneficial even under aggressive model compression.

2. Context and Motivation

The Core Problem: Multi-Head MoE's Practical Viability Was Compromised by FLOPs Inflation

The fundamental question this paper tackles is straightforward but practically crucial: can the multi-head mechanism simultaneously benefit both performance and computational efficiency in Mixture-of-Experts models? The original MH-MoE proposal by Wu et al. (2024) demonstrated superior performance over standard sparse MoE architectures by enabling experts to collectively attend to information from different representation spaces. However, this performance came with a significant, unaddressed cost: the original MH-MoE consumed substantially more FLOPs than a standard sparse MoE model with equivalent activated parameters.

To understand why this gap matters, we need to appreciate the central bargain of sparse MoE architectures. The promise of sparse MoE—established through a line of work stretching from Shazeer et al. (2017) through Mixtral (Jiang et al., 2024) and DeepSeekMoE (Dai et al., 2024)—is this: you can dramatically increase a model's total parameter count while keeping the per-token inference cost roughly constant. You achieve this by having a gating mechanism route each input token to only a small subset of available experts (typically 1 or 2), so while the full model may contain billions of parameters, each token only activates a fraction of them. This makes sparse MoE one of the primary strategies for building large language models that can be served economically at scale.

The original MH-MoE broke this bargain. As the paper's complexity analysis reveals (Section 2.2), when Wu et al. (2024) configured their model with 4 heads and a scaling factor $\beta = 63/64$ , the resulting scalar multiplications had a leading term of $67Bd^2$ compared to $16Bd^2$ for standard sparse MoE—a more than 4× increase in FLOPs for a model that, by design, activated the same number of parameters. This is not a minor inefficiency; it fundamentally undermines the economic and computational rationale for using the multi-head mechanism. A model that consumes 4× more inference FLOPs than a baseline MoE is not a drop-in replacement—it's a substantially more expensive system that requires different hardware provisioning, latency budgets, and cost structures.

Why This Gap Is Important: The FLOPs-Performance Frontier Defines What Gets Deployed

The significance of closing this FLOPs gap extends beyond a single architectural variant. Mixture-of-Experts models have become a default architectural choice for frontier LLM development precisely because they offer favorable scaling properties: for a given FLOPs budget, an MoE model can be trained with more total parameters than a dense model, yielding better performance (Clark et al., 2022). Any architectural innovation that wants to be adopted in this ecosystem must respect this FLOPs-performance tradeoff. If the multi-head mechanism—which showed genuine promise for improving expert utilization and representation diversity in Wu et al. (2024)—cannot be implemented at FLOPs parity, it will remain a theoretically interesting but practically irrelevant technique.

The paper's focus on FLOPs parity is therefore not merely an engineering optimization; it's a response to a specific failure mode in the prior work's practical applicability. By ensuring that MH-MoE operates at the same computational cost as standard sparse MoE, the authors are essentially asking: "If we remove the confounding factor of increased FLOPs, does the multi-head mechanism still provide meaningful gains?" This is a clean experimental question that the original paper could not answer, because its performance improvements were potentially attributable to simply using more computation rather than any inherent superiority of multi-head routing.

Prior Approaches and Their Limitations

Standard Sparse Mixture-of-Experts. The canonical approach, established by Shazeer et al. (2017) and refined through GShard (Lepikhin et al., 2020), GLaM (Du et al., 2021), and numerous subsequent works, replaces dense FFN layers with MoE layers containing multiple expert FFNs and a learned gating function. Each token is routed to the top-k experts (typically k=1 or k=2), enabling the model to scale total parameters without scaling per-token FLOPs proportionally. The key limitation of this architecture is single-representation routing: the gating function operates on the token's representation in a single, fixed embedding space. This means that routing decisions are made based on whatever features happen to be most salient in that particular representation, potentially missing complementary signals that exist in other representational subspaces.

Fine-Grained Mixture-of-Experts. A recognized variant of standard MoE reduces the size of individual experts (decreasing the intermediate dimension of each FFN) while increasing the total number of experts. This increases the combinatorial expressiveness of expert combinations—with more, smaller experts available, the model has finer control over which specialized computations are applied. The limitation is that fine-grained MoE does not fundamentally change how routing decisions are made; it simply increases the vocabulary of available experts while routing remains a single-representation operation. The gains from fine-graining therefore saturate: at some point, adding more experts with the same routing mechanism provides diminishing returns because the gating function, operating in a single representation space, cannot effectively discriminate among all the available experts.

The Original Multi-Head MoE (Wu et al., 2024). This work introduced two key innovations: (1) a head layer that projects the input through a learned linear transformation before splitting it into h sub-tokens, each residing in a $d/h$ -dimensional subspace, and (2) a merge layer that recombines the processed sub-tokens after they pass through the expert layer. The intuition, explicitly drawn from the multi-head attention mechanism (Vaswani et al., 2017), is that different representation subspaces encode different aspects of the input, and routing decisions should be made with awareness of all these subspaces simultaneously. Each sub-token is routed independently through the expert layer, meaning a single original token can be processed by different experts in different subspaces—something impossible in standard MoE.

The critical shortcoming of the original MH-MoE was its computational cost structure. The paper's analysis (Section 2.2) quantifies this precisely: the original implementation achieved approximately $67Bd^2$ leading-term scalar multiplications versus $16Bd^2$ for baseline MoE. This 4× inflation arose because the intermediate dimension of the expert FFNs was not adjusted downward to compensate for the added head and merge layers, resulting in a model that processed more FLOPs per token despite activating comparable parameters. In practical terms, this meant the original MH-MoE was not competing on a level playing field—its performance gains were potentially explainable by increased computation rather than architectural superiority.

Where Existing Approaches Fall Short Collectively. A pattern emerges across these prior approaches: they either modify expert granularity (fine-grained MoE), expert selection (various gating improvements), or representation spaces (original MH-MoE), but none simultaneously achieves three desirable properties: (1) multi-subspace routing that enables a single token to engage with different experts in different representational contexts, (2) FLOPs parity with standard sparse MoE, and (3) parameter parity with standard sparse MoE. The original MH-MoE achieves (1) and (3) but fails on (2). Standard MoE achieves (2) and (3) but lacks (1). Fine-grained MoE modifies the expert allocation within the constraints of (2) and (3) but does not introduce multi-subspace routing.

How This Paper Positions Itself

The paper frames its contribution as a re-implementation and re-analysis of MH-MoE, not a fundamentally new architecture. The authors explicitly credit Wu et al. (2024) for the multi-head mechanism and instead focus on the question: how should we adjust the expert configuration (intermediate dimension, number of experts, gating width) to achieve FLOPs parity while preserving the multi-head mechanism's benefits?

The answer comes through a derived set of equations (Section 2.3) that specify how to set the intermediate dimension $d_{\text{mhmoe}}$ and expert count $E_{\text{mhmoe}}$ as functions of the original MoE's configuration. The key insight is straightforward: the head and merge layers add computation ( $2Bd^2 - Bd$ scalar multiplications each, from Equation 6), so to maintain FLOPs parity, this cost must be subtracted from the expert computations. This is achieved by reducing each expert's intermediate dimension and compensating with more (smaller) experts to preserve total parameter count. Equation 8 captures this tradeoff:

$d_{\text{mhmoe}} = d_{\text{moe}} - \frac{d}{k}$

where $d_{\text{moe}}$ is the original expert intermediate dimension, $d$ is the model dimension, and $k$ is the number of activated experts. For a standard configuration with $d_{\text{moe}} = 4d$ and top-1 gating ( $k=1$ ), this yields $d_{\text{mhmoe}} = 3d$ —experts become narrower. To maintain parameter count, the number of experts must increase per Equation 9, leading to configurations like 40 experts for the 2-head variant and 96 experts for the 3-head variant (compared to 8 experts in the baseline).

This positioning is notably different from the original MH-MoE paper in one crucial respect: the authors are not claiming to invent multi-head routing for MoE. They are claiming to provide the first implementation that makes multi-head routing practically viable by ensuring FLOPs parity. This is an engineering and analysis contribution rather than a conceptual one, but it addresses the exact gap that prevented the original MH-MoE from being considered a legitimate alternative to standard MoE in production settings.

The paper also extends the validation of MH-MoE to the 1-bit quantization regime (Section 3.2), positioning the architecture as compatible with emerging trends in extremely compressed LLM deployment. This is forward-looking: as the field moves toward running LLMs on consumer hardware and edge devices, techniques like BitNet (Ma et al., 2024) become increasingly relevant, and the paper establishes early evidence that the multi-head mechanism's benefits survive aggressive quantization.

3. Technical Approach

3.1 Reader Orientation

This paper presents a re-engineered implementation of the Multi-Head Mixture-of-Experts (MH-MoE) layer—a drop-in replacement for standard MoE layers in Transformer architectures—that enables tokens to be processed through multiple representation subspaces simultaneously while consuming exactly the same computational budget (FLOPs) and model parameters as a conventional sparse MoE layer. The approach solves the practical viability problem left by the original MH-MoE proposal: by deriving closed-form expressions for how to shrink expert dimensions and expand expert counts to exactly offset the cost of the new head and merge projections, the authors transform a theoretically promising but computationally expensive architecture into one that competes on a genuinely level playing field with standard sparse and fine-grained MoE baselines.

3.2 Big-Picture Architecture (Diagram in Words)

The MH-MoE layer sits at the same position in the Transformer as a standard MoE layer (replacing dense FFN blocks at regular intervals) and processes tokens through four sequential stages:

Head Projection ( $W_{\text{head}}$ ): A learned linear transformation of shape $d \times d$ that projects each token's representation into a space suitable for multi-head splitting. This is conceptually analogous to the combined query/key/value projection in multi-head attention—it prepares the representation for decomposition into subspaces.
Token Splitting and Multi-Head Reshaping: The projected token $\hat{x} \in \mathbb{R}^d$ is split along the feature dimension into $h$ sub-tokens, each of dimension $d/h$ . These sub-tokens are arranged in parallel as if they were additional tokens in the batch, effectively increasing the number of routing decisions by a factor of $h$ . Each sub-token represents the original token as viewed through one of $h$ different representational lenses.
Expert Processing (per sub-token): Each sub-token is independently routed through a Mixture-of-Experts layer containing $E_{\text{mhmoe}}$ experts. This is where the multi-subspace routing occurs: different sub-tokens derived from the same original token can be sent to different experts, enabling a single token to benefit from multiple specialized computations simultaneously. Each expert is a standard FFN (using SwiGLU activation) with a reduced intermediate dimension $d_{\text{mhmoe}}$ —smaller than a standard MoE expert, as dictated by the FLOPs parity equations.
Concatenation and Merge Projection ( $W_{\text{merge}}$ ): After expert processing, the $h$ sub-token outputs (each of dimension $d/h$ ) are concatenated back into a single vector of dimension $d$ . This concatenated representation is then passed through a second learned linear projection $W_{\text{merge}} \in \mathbb{R}^{d \times d}$ that integrates information across the different subspace outputs, producing the final layer output.

The critical design constraint governing all component dimensions is that the total scalar multiplications across these four stages must equal the scalar multiplications in a standard MoE layer's expert computation alone. The configuration equations in Section 2.3 are precisely the solution to this constraint satisfaction problem.

3.3 Roadmap for the Deep Dive

This section unpacks the technical machinery in the following order, which mirrors both the architectural dataflow and the logical dependencies:

First, the complexity analysis framework that the paper uses to count FLOPs (Section 2.2 of the paper). Understanding how the authors measure and compare computational cost is foundational—every design decision flows from the constraint of FLOPs parity, so we need to establish what "FLOPs parity" means quantitatively before we can understand how it is achieved.
Second, the derivation of the parity-preserving configuration equations (Section 2.3). These equations—for the reduced intermediate dimension $d_{\text{mhmoe}}$ and the increased expert count $E_{\text{mhmoe}}$ —are the paper's primary technical contribution. We walk through the derivation from first principles, explaining why each term appears and what tradeoffs are encoded.
Third, the specific instantiation of the architecture, including the SwiGLU expert FFN, the gating mechanism choice, the experimental configurations (768-dim models, 12 layers, top-1/2/3 gating), and how the derived equations produce the concrete expert counts of 40 and 96 for the two experimental variants.
Fourth, the shared expert variant, which adds a residual-style MoE configuration where a single shared expert processes all tokens alongside the routed experts—an established technique from DeepSeekMoE that the paper incorporates to test MH-MoE's compatibility with broader MoE design patterns.
Fifth, the 1-bit quantization integration with BitNet, including what changes in the training procedure and why this compatibility test matters for deployment scenarios.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily an implementation and analysis paper whose core contribution is a set of derived equations that, given a standard sparse MoE configuration, specify exactly how to construct an MH-MoE variant with matched FLOPs and parameters. The supporting experiments validate that the resulting architecture provides genuine performance improvements that are not attributable to increased computation.

The FLOPs Accounting Framework

Why FLOPs counting matters here. Before modifying any architecture to achieve parity, you need an agreed-upon method for counting computational cost. The paper uses scalar multiplications as the atomic unit of FLOPs, which is standard practice in Transformer efficiency analysis. A scalar multiplication (multiplying two individual floating-point numbers) is the fundamental operation that composes all matrix multiplications and feedforward computations. Counting these gives a hardware-agnostic measure of computational work that is proportional to actual runtime on modern accelerators.

The authors analyze the MH-MoE layer by decomposing it into three additive sources of scalar multiplications, using $B$ for batch size (number of tokens), $d$ for the model's hidden dimension, $d_{\text{moe}}$ for the expert FFN's intermediate dimension, $h$ for the number of heads, and $k$ for the number of activated experts per token:

The head layer performs a matrix multiplication between an input of shape $B \times d$ and a weight matrix of shape $d \times d$ . The standard count for such a multiplication (without bias) is $2Bd^2 - Bd$ scalar multiplications. The $2Bd^2$ term comes from the fact that for each of the $Bd$ output elements, we need $d$ multiply-accumulate operations—but since we count only multiplications and treat additions separately, the multiplication count is $Bd \cdot d = Bd^2$ . The paper's formula of $2Bd^2 - Bd$ appears to count multiply-add pairs differently than the conventional approach, but the key point is that this term is linear in $B$ and quadratic in $d$ , and it does not appear in a standard MoE layer.

The activated experts (there are $k$ of them, since top- $k$ gating activates $k$ experts per token) each process their input through a two-layer FFN. For each expert, an input of shape $B \times (d/h)$ is first projected to dimension $d_{\text{moe}}$ via a weight matrix $W_1 \in \mathbb{R}^{(d/h) \times d_{\text{moe}}}$ , then projected back via $W_2 \in \mathbb{R}^{d_{\text{moe}} \times (d/h)}$ . The paper states the total multiplication count for the activated experts as $(4Bdd_{\text{moe}} - Bd - Bd_{\text{moe}}h) \cdot k$ . The leading term $4Bdd_{\text{moe}} \cdot k$ captures the dominant cost: for each of the $k$ experts, each of the $B$ tokens incurs $2 \cdot (d/h) \cdot d_{\text{moe}}$ multiplications in the up-projection and $2 \cdot d_{\text{moe}} \cdot (d/h)$ in the down-projection, which simplifies to $4d d_{\text{moe}} / h$ per token per expert. With $k$ experts and $B$ tokens, this becomes $4B d d_{\text{moe}} k / h$ . Note the division by $h$ : because each expert operates on sub-tokens of dimension $d/h$ rather than $d$ , the per-expert cost is reduced relative to a standard MoE expert operating on full-dimension tokens.

The merge layer performs another $d \times d$ matrix multiplication, identical in cost to the head layer: $2Bd^2 - Bd$ .

Summing these three components yields the total scalar multiplications for the MH-MoE layer:

$\text{Multiplications}_{\text{MH-MoE}} = \underbrace{2Bd^2 - Bd}_{\text{Head}} + \underbrace{(4Bdd_{\text{moe}} - Bd - Bd_{\text{moe}}h) \cdot k}_{\text{Activated Experts}} + \underbrace{2Bd^2 - Bd}_{\text{Merge}}$

What this equation reveals structurally. The head and merge layers each contribute approximately $2Bd^2$ of computational cost that does not exist in a standard MoE layer. For the total to match standard MoE, the expert computation term must be reduced by roughly $4Bd^2$ to compensate. Since the expert term scales with $d_{\text{moe}}$ , the intermediate dimension, the compensation must come from reducing $d_{\text{moe}}$ . This is the fundamental tradeoff: the multi-head mechanism adds fixed overhead (the projections), so experts must become narrower to keep total FLOPs constant. But narrower experts mean fewer parameters per expert, so to maintain total parameter count, you need more experts. The configuration equations in Section 2.3 formalize exactly how much narrower and how many more.

For comparison, the standard sparse MoE (without head/merge layers, using top- $k$ gating, experts operating on full $d$ -dimensional tokens with intermediate dimension $d_{\text{moe}}$ ) has scalar multiplications:

$\text{Multiplications}_{\text{SMoE}} = (4Bdd_{\text{moe}} - Bd_{\text{moe}} - Bd) \cdot k$

The leading term is $4Bdd_{\text{moe}} \cdot k$ . For the standard configuration used in the experiments ( $d_{\text{moe}} = 4d$ , $k=1$ ), this leading term becomes $16Bd^2$ —the baseline against which MH-MoE must match.

Why leading-term matching is sufficient. The paper explicitly states that they "only consider the leading term of the equation" when designing for parity. The lower-order terms (those linear in $B$ and $d$ rather than quadratic) represent edge effects—the subtraction of bias-like terms that arise from zero-index counting conventions or non-multiplicative operations. These are typically less than 1% of total FLOPs at the model scales discussed and are architecture-dependent in ways that would make exact matching unnecessarily complex. Matching leading terms ensures that the asymptotic computational scaling is identical, which is what matters for practical deployment at scale.

Derivation of the Parity-Preserving Configuration Equations

This subsection contains the paper's primary technical contribution. The goal is to answer: given a standard MoE configuration with expert intermediate dimension $d_{\text{moe}}$ , number of experts $E_{\text{moe}}$ , and top- $k$ gating, what should be the MH-MoE's expert intermediate dimension $d_{\text{mhmoe}}$ and number of experts $E_{\text{mhmoe}}$ such that both FLOPs and total parameters match?

Step 1: Matching FLOPs via intermediate dimension adjustment.

The intellectual move here is to isolate the added cost of the head and merge layers and offset it entirely by reducing the expert computation budget. From the complexity analysis, the MH-MoE has an additional leading term of $4Bd^2$ (from head + merge, each contributing $2Bd^2$ ) compared to standard MoE. The expert computation leading term is $4Bdd_{\text{mhmoe}}k$ for MH-MoE versus $4Bdd_{\text{moe}}k$ for standard MoE. Setting the total leading terms equal:

$4Bd^2 + 4Bdd_{\text{mhmoe}}k = 4Bdd_{\text{moe}}k$

Dividing through by the common factor $4Bd$ (valid since $B > 0$ and $d > 0$ ):

$d + d_{\text{mhmoe}}k = d_{\text{moe}}k$

Solving for $d_{\text{mhmoe}}$ :

$d_{\text{mhmoe}} = d_{\text{moe}} - \frac{d}{k}$

This is Equation 8 from the paper.

where $d_{\text{moe}}$ is the intermediate dimension of a standard MoE expert, $d$ is the model's hidden dimension, and $k$ is the number of activated experts per token (typically 1 for top-1 gating, 2 for top-2 gating).

What this equation computes: given a baseline MoE configuration, it tells you how much to shrink each expert's intermediate dimension to exactly offset the FLOPs cost of the head and merge projections. The subtracted term $d/k$ represents the per-expert dimension reduction needed: each unit of $d$ added to the head and merge layers must be compensated by reducing the expert's intermediate dimension proportionally to $1/k$ , because the expert computation is multiplied by $k$ (each token goes through $k$ experts).

Why this form: the subtraction is linear in $d$ and inversely proportional to $k$ . If you activate more experts per token ( $k$ is larger), each expert sees a smaller portion of the token's representation (each sub-token is $d/h$ -dimensional, but $k$ experts are activated per sub-token), so the per-expert cost is higher and the required per-expert dimension reduction is smaller. For $k=1$ (top-1 gating) and the standard $d_{\text{moe}} = 4d$ , this yields $d_{\text{mhmoe}} = 4d - d = 3d$ : experts become 25% narrower. For $k=2$ and $d_{\text{moe}} = 2048$ (as in the paper's 2-head variant with $d=768$ ), this yields $d_{\text{mhmoe}} = 2048 - 768/2 = 2048 - 384 = 1664$ .

A crucial subtlety: this FLOPs matching assumes the head and merge layers are present. If $d_{\text{moe}}$ is already small relative to $d$ , the required $d_{\text{mhmoe}}$ could become uncomfortably small or even negative, indicating that FLOPs parity is impossible with the given $k$ —you would need to increase $k$ or accept a smaller $d$ (model dimension) for the MH-MoE variant. The paper's experimental configurations operate safely within the regime where $d_{\text{mhmoe}}$ is positive and practically reasonable.

Step 2: Matching parameters via expert count adjustment.

Matching FLOPs guarantees that inference cost is identical, but the model still needs to have the same total number of trainable parameters to be considered a fair comparison. Shrinking experts makes each expert contain fewer parameters, so the number of experts must increase.

The parameter count for a standard MoE layer comes exclusively from the expert FFNs (ignoring the gating network, which is negligible). Each expert has two weight matrices: $W_1 \in \mathbb{R}^{d \times d_{\text{moe}}}$ and $W_2 \in \mathbb{R}^{d_{\text{moe}} \times d}$ . Using the SwiGLU activation (which requires three matrices— $W_{\text{gate}}$ , $W_{\text{up}}$ , and $W_{\text{down}}$ —but combines them into the effective parameter count), the total parameters for $E_{\text{moe}}$ experts in a standard MoE layer is approximately $2dd_{\text{moe}} \cdot E_{\text{moe}}$ (two matrices per expert, each with $d \cdot d_{\text{moe}}$ elements). The paper uses this simplified count:

$\text{Params}_{\text{SMoE}} = 2dd_{\text{moe}} \cdot E_{\text{moe}}$

For the MH-MoE layer, parameters come from three sources: the head layer ( $W_{\text{head}} \in \mathbb{R}^{d \times d}$ , contributing $d^2$ parameters), the merge layer ( $W_{\text{merge}} \in \mathbb{R}^{d \times d}$ , contributing another $d^2$ parameters), and the experts. Each MH-MoE expert operates on sub-tokens of dimension $d/h$ and has intermediate dimension $d_{\text{mhmoe}}$ , so its two weight matrices have shapes $(d/h) \times d_{\text{mhmoe}}$ and $d_{\text{mhmoe}} \times (d/h)$ . The per-expert parameter count is $2 \cdot (d/h) \cdot d_{\text{mhmoe}}$ . With $E_{\text{mhmoe}}$ experts:

$\text{Params}_{\text{MH-MoE}} = \underbrace{2d^2}_{\text{head + merge}} + \underbrace{2 \cdot \frac{d}{h} \cdot d_{\text{mhmoe}} \cdot E_{\text{mhmoe}}}_{\text{experts}}$

Setting these equal (Equation 9 from the paper):

$2dd_{\text{moe}} \cdot E_{\text{moe}} = 2d^2 + 2\frac{d}{h}d_{\text{mhmoe}} \cdot E_{\text{mhmoe}}$

What this equation does: given the standard MoE configuration (left side) and the already-computed $d_{\text{mhmoe}}$ from the FLOPs matching step, it solves for $E_{\text{mhmoe}}$ , the number of experts needed in the MH-MoE layer to achieve total parameter parity. The $2d^2$ term on the right represents the "parameter tax" paid by the head and merge layers—parameters that exist in MH-MoE but not in standard MoE. To compensate, the expert parameters must be smaller by exactly $2d^2$ , which is achieved by having many narrow experts rather than few wide ones.

Why this form encodes a genuine tradeoff: the head and merge layers consume parameters that could otherwise have been allocated to experts. The number of heads $h$ appears in the denominator of the expert parameter term: more heads means each sub-token is smaller (dimension $d/h$ ), so each expert has fewer parameters. To maintain total parameter count, you need even more experts. This creates a three-way tension: more heads enable richer multi-subspace routing (the supposed benefit), but they increase the FLOPs overhead not captured by the leading-term analysis (since the head and merge costs don't scale with $h$ ) and they reduce per-expert capacity (each expert processes a smaller slice of the representation), requiring more experts to compensate. The optimal $h$ balances these forces.

Worked example (3-head variant from Section 3): Starting from the standard SMoE configuration with $d_{\text{moe}} = 4d$ (in practice, 2048 for $d=768$ ), top-1 gating ( $k=1$ ), and $E_{\text{moe}} = 8$ . Step 1: compute $d_{\text{mhmoe}} = 4d - d/1 = 3d = 3 \times 768 = 2304$ . Wait—this doesn't match the paper's numbers. Let's re-derive carefully using the actual experimental values from Section 3.

The paper's standard SMoE uses $d=768$ , $d_{\text{moe}} = 2048$ (which is $2048/768 \approx 2.67d$ , not $4d$ as in the simplified derivation), $E_{\text{moe}} = 8$ , and top-1 gating. For the 3-head MH-MoE variant, the paper sets $d_{\text{mhmoe}} = 512$ , uses top-3 gating, and $E_{\text{mhmoe}} = 96$ .

Let's verify the FLOPs matching with Equation 8. With $d_{\text{moe}}=2048$ , $d=768$ , and $k=3$ : $d_{\text{mhmoe}} = 2048 - 768/3 = 2048 - 256 = 1792$ . But the paper uses $d_{\text{mhmoe}} = 512$ . This discrepancy tells us that the paper's Equation 8 is a simplified leading-term approximation, and the actual configurations used in experiments deviate from this simplified form because the paper adjusts for practical considerations (matching exact rather than leading-term FLOPs, or using different gating widths).

The paper's actual methodology for the 3-head variant is better understood directly from the configuration table: they set $d_{\text{mhmoe}} = 512$ , use top-3 gating, and arrive at $E_{\text{mhmoe}} = 96$ . The intermediate dimension $512$ is $2d/3$ (since $d=768$ ), which is substantially smaller than the standard MoE's $2048$ . With top-3 gating, each token activates 3 experts out of 96, compared to 1 out of 8 in the baseline. The total activated parameters are $3 \times 2 \times (768/3) \times 512 = 3 \times 2 \times 256 \times 512 = 786,432$ for MH-MoE versus $1 \times 2 \times 768 \times 2048 = 3,145,728$ for standard MoE—but remember, these numbers need to account for the head and merge layers in the MH-MoE total.

Key takeaway from this derivation exercise: the paper's Equation 8 and Equation 9 provide the conceptual framework (reduce intermediate dimension, increase expert count), but the actual experimental configurations are selected to achieve matched total parameters and matched leading-term FLOPs simultaneously, which requires jointly solving for $d_{\text{mhmoe}}$ , $E_{\text{mhmoe}}$ , and sometimes $k$ . The specific numbers (2-head: $d_{\text{mhmoe}}=768$ , $k=2$ , $E=40$ ; 3-head: $d_{\text{mhmoe}}=512$ , $k=3$ , $E=96$ ) are designed to land on integer configurations that approximately satisfy both constraints.

The MH-MoE Forward Pass: Token-Level Walkthrough

What happens to a single token. To understand the architecture concretely, let's trace one token through the MH-MoE layer in the 3-head variant ( $h=3$ , $d=768$ ):

Input: A token representation $x \in \mathbb{R}^{768}$ arrives at the MH-MoE layer from the preceding Transformer component (typically the attention output after residual connection and layer norm).
Head projection: The token is multiplied by $W_{\text{head}} \in \mathbb{R}^{768 \times 768}$ , producing $\hat{x} = xW_{\text{head}}$ , also in $\mathbb{R}^{768}$ . This linear transformation mixes all feature dimensions, preparing the representation for the subspace decomposition that follows. It is analogous to how the combined QKV projection in multi-head attention doesn't yet separate into heads—it creates a representation where head-wise splitting will be meaningful.
Splitting into sub-tokens: The 768-dimensional vector $\hat{x}$ is divided into $h=3$ contiguous segments, each of 256 dimensions: $\tilde{x}_1, \tilde{x}_2, \tilde{x}_3 \in \mathbb{R}^{256}$ . These are conceptually arranged as if they were three separate tokens in the batch (the batch dimension effectively gets multiplied by $h$ for the duration of the expert processing step).
Gating (per sub-token, per head): A gating function $G: \mathbb{R}^{256} \to \mathbb{R}^{96}$ computes a score for each of the 96 available experts. With top-3 gating, only the three highest-scoring experts are activated per sub-token. The gating scores are typically computed via a learned linear projection followed by a softmax over experts, with the top- $k$ operation setting all non-selected expert scores to zero (and renormalizing). Importantly, the gating function operates independently on each sub-token: $\tilde{x}_1$ might route to experts {4, 17, 83}, while $\tilde{x}_2$ routes to {17, 42, 91} and $\tilde{x}_3$ to {8, 17, 55}. The original token thus engages with up to 9 distinct expert computations across its three sub-tokens (fewer if the same expert is selected for multiple sub-tokens).
Expert computation: For each activated expert, the 256-dimensional sub-token is processed through a SwiGLU FFN. The SwiGLU activation (Shazeer, 2020) uses three weight matrices instead of two: $W_{\text{gate}} \in \mathbb{R}^{256 \times 512}$ , $W_{\text{up}} \in \mathbb{R}^{256 \times 512}$ , and $W_{\text{down}} \in \mathbb{R}^{512 \times 256}$ . The computation is: $\text{Expert}(\tilde{x}) = (\text{SiLU}(\tilde{x}W_{\text{gate}}) \odot (\tilde{x}W_{\text{up}})) W_{\text{down}}$ where SiLU (Sigmoid Linear Unit, also called Swish) is $x \cdot \sigma(x)$ and $\odot$ is element-wise multiplication. The gating mechanism inside the FFN (not to be confused with the expert-routing gate) uses $W_{\text{gate}}$ to compute a multiplicative filter that is applied to the up-projected representation—this is the "Gated Linear Unit" family that SwiGLU belongs to. The output is a 256-dimensional vector.
Weighted combination: For each sub-token, the outputs of its $k=3$ activated experts are combined via a weighted sum, using the normalized gating scores as coefficients. This yields three 256-dimensional output vectors: $\tilde{y}_1, \tilde{y}_2, \tilde{y}_3$ .
Concatenation: The three 256-dimensional outputs are concatenated in their original order to form a single 768-dimensional vector $\hat{y} = [\tilde{y}_1; \tilde{y}_2; \tilde{y}_3]$ .
Merge projection: This concatenated vector is multiplied by $W_{\text{merge}} \in \mathbb{R}^{768 \times 768}$ , producing the final output $y = \hat{y}W_{\text{merge}} \in \mathbb{R}^{768}$ . This linear transformation allows information from different sub-token paths to interact—crucially, the expert processing for different sub-tokens of the same original token occurred independently, and the merge layer is where cross-subspace integration happens.

Why splitting before routing matters. In standard MoE, the gating function sees the token in a single 768-dimensional representation. All routing decisions are based on whatever information is most salient in that single view. In MH-MoE, the head projection followed by splitting means that different subspaces of the token's representation can emphasize different features, and the gating function can make different routing decisions for each subspace. A token that contains both a mathematical expression and a natural language description might route one subspace to math-specialized experts and another to language-specialized experts—something impossible in standard MoE, where the token goes to at most $k$ experts total, and all routing is based on a single representation.

Why the merge projection matters. Without it, the concatenated expert outputs would remain segmented by subspace—each chunk of 256 dimensions would contain information processed independently of the others. The merge projection mixes across these chunks, enabling the layer to learn which combinations of subspace-specific computations are useful. This is again analogous to the output projection in multi-head attention, where the independently computed attention outputs from different heads are linearly combined.

The Shared Expert Variant (Residual MoE Configuration)

The paper also evaluates MH-MoE in a residual MoE configuration following the approach of DeepSeekMoE (Dai et al., 2024). This design pattern, which has been empirically shown to improve MoE training stability and performance, adds a shared expert that processes all tokens regardless of routing decisions.

How it works. In a residual MoE layer, alongside the $E$ routed experts (which only process tokens assigned to them by the gating function), there is one additional expert—the shared expert—that processes every token. The outputs of the routed experts and the shared expert are summed:

$y = \sum_{p \in \Phi} G(x) \cdot \text{Expert}_p(x) + \text{Expert}_{\text{shared}}(x)$

The shared expert has the same architecture as the other experts but is exempt from the gating mechanism: it is unconditionally activated for every token. In the paper's experiments, the shared expert's hidden dimension is set to 2048 (matching the standard MoE expert dimension) for all models, and its parameters are counted toward the total parameter budget.

Why this helps. The shared expert serves as a "default" computation that captures knowledge and patterns common across all tokens. The routed experts can then specialize in handling token-specific nuances—the shared expert handles the heavy lifting of general language modeling, while routed experts provide task- or context-specific refinements. This reduces the pressure on the gating function to make perfect routing decisions, because even if a token is poorly routed, it still receives competent processing from the shared expert. It also mitigates the representation collapse problem documented by Chi et al. (2022), where a small number of experts receive most tokens and others are underutilized.

MH-MoE with shared experts. When MH-MoE is combined with a shared expert, the shared expert must also be adapted to the multi-head framework. The paper states that "a shared expert with the same size (hidden dimension is set to 2048) is applied to all MoE models." This means the shared expert operates on the full-dimension tokens (not the sub-tokens), processing the input $x$ at dimension $d=768$ with intermediate dimension 2048. Its output is added after the merge projection of the MH-MoE path. The shared expert's parameters are accounted for in the total parameter budget, so the number of routed experts $E_{\text{mhmoe}}$ is adjusted accordingly to maintain parameter parity with the baseline (which also has a shared expert).

Experimental motivation. The inclusion of shared expert experiments in Tables 1 and 2 (the paper's primary results tables) serves as a robustness check: it demonstrates that MH-MoE's benefits are not specific to the non-residual configuration and that the multi-head mechanism composes with established MoE architectural improvements. The consistent perplexity improvements of MH-MoE over baselines in both settings strengthen the claim that the gains come from the multi-head routing mechanism itself, not from interactions with specific architectural choices.

Specific Experimental Configurations and Hyperparameters

The paper's experimental setup (Section 3) uses a consistent baseline architecture and training procedure across all comparisons to isolate the effect of the MH-MoE design.

Model architecture. All models are decoder-only Transformers with 12 layers and a model dimension $d = 768$ . This is a relatively small model by contemporary standards (comparable to GPT-2 small-medium scale), chosen to enable systematic experimentation within reasonable compute budgets. The MoE layers replace dense FFN layers every two layers—a common cadence in MoE Transformers that balances the overhead of expert routing against the benefits of increased capacity. Layers that are not replaced with MoE retain standard dense FFNs.

Expert FFN details. All experts use the SwiGLU activation function (Shazeer, 2020), which has become the de facto standard in modern Transformer architectures due to its superior performance over ReLU and GELU variants. The SwiGLU FFN has three internal weight matrices (gate, up, down) as described earlier. The intermediate dimension $d_{\text{moe}}$ for the baseline SMoE experts is set to 2048.

Baseline SMoE configuration. Top-1 gating with 8 experts. This means each token is routed to exactly one expert per MoE layer, and there are 8 possible experts to choose from. The number of activated experts per token is $k=1$ , and the total expert count is $E_{\text{moe}}=8$ .

Fine-grained SMoE configuration. The intermediate dimension is halved to 1024, and the number of experts is doubled to 16. This maintains the same total expert parameters ( $8 \times 2048 = 16 \times 1024 = 16,384$ parameter units per dimension of model width, approximating the total) while providing finer-grained expert specialization. Top-1 gating is preserved.

MH-MoE 2-head variant. With $h=2$ , the head and merge layers each contribute $768^2 = 589,824$ parameters. To maintain FLOPs parity, the expert intermediate dimension $d_{\text{mhmoe}}$ is set to 768 (which is $d$ , compared to $2048 \approx 2.67d$ for the baseline). The gating is switched from top-1 to top-2 ( $k=2$ ), meaning each sub-token activates 2 experts. The number of experts is increased to $E_{\text{mhmoe}} = 40$ . With $h=2$ , each sub-token has dimension $d/h = 768/2 = 384$ . Each expert has $2 \cdot 384 \cdot 768 = 589,824$ parameters (for the two weight matrices; the SwiGLU splits this differently but the total is equivalent). Total expert parameters: $40 \cdot 589,824 = 23,592,960$ . Adding head and merge: $23,592,960 + 2 \cdot 589,824 = 24,772,608$ . For comparison, the baseline SMoE with 8 experts of dimension $768 \to 2048 \to 768$ has $8 \cdot 2 \cdot 768 \cdot 2048 = 25,165,824$ expert parameters—approximately matched, with small differences due to the SwiGLU parameterization and rounding.

MH-MoE 3-head variant. With $h=3$ , the sub-token dimension is $768/3 = 256$ . The intermediate dimension $d_{\text{mhmoe}}$ is set to 512 ( $2d/3$ ). Top-3 gating ( $k=3$ ) is used, meaning each of the three sub-tokens activates 3 experts. The number of experts is increased to $E_{\text{mhmoe}} = 96$ . Per-expert parameters (approximate): $2 \cdot 256 \cdot 512 = 262,144$ . Total expert parameters: $96 \cdot 262,144 = 25,165,824$ . Head and merge layers add $2 \cdot 768^2 = 1,179,648$ . Total: $26,345,472$ —slightly above the baseline, but the paper describes all models as "matched in terms of parameters and computation," suggesting that the differences are small enough to be negligible or are compensated by small adjustments not fully enumerated in the text.

Training hyperparameters. All models are pre-trained for 100,000 steps on the RedPajama dataset (Together Computer, 2023), an open-source reproduction of the LLaMA training data consisting of approximately 1.2 trillion tokens from web crawls (CommonCrawl, C4), code repositories (GitHub), academic papers (ArXiv), books, Wikipedia, and StackExchange. Each training batch contains 0.5 million tokens (approximately 650 sequences of length 768, or 520 sequences of length 960, depending on the exact sequence length used). The paper states that "the same code base, training parameters, and pre-training tasks" are used across all experiments, ensuring that differences in perplexity are attributable to architecture rather than training recipe variations. However, specific optimizer settings (learning rate, schedule, weight decay, Adam betas) are not reported in the paper—a notable omission that limits full reproducibility.

Evaluation. Perplexity is reported on three held-out validation sets: the RedPajama validation split, English Wikipedia, and C4 (Colossal Clean Crawled Corpus). Perplexity is the exponentiated negative log-likelihood $\exp(-\frac{1}{T}\sum_t \log p(y_t|y_{<t}))$ , measuring how well the model predicts the next token—lower is better. Results are reported at both 50,000 steps (mid-training) and 100,000 steps (final), providing a view of relative convergence rates as well as final performance.

Integration with 1-Bit Quantization (BitNet)

Section 3.2 explores whether MH-MoE remains effective when the model is quantized to 1-bit precision using BitNet (Ma et al., 2024). This is a forward-looking experiment motivated by the growing interest in extremely compressed LLMs for edge deployment.

What BitNet changes. In standard Transformer training, weights and activations are stored in floating-point formats (typically FP16 or BF16 during training, with optional post-training quantization to INT8 or INT4 for deployment). BitNet replaces all linear layer weights with ternary values $\{-1, 0, +1\}$ (1.58 bits per parameter on average, using the $-1, 0, +1$ representation) and quantizes activations to 8 bits. During training, the model maintains full-precision "latent weights" that are quantized to ternary values on the forward pass, with the straight-through estimator used to propagate gradients through the non-differentiable quantization operation.

Why this test matters for MH-MoE. The head and merge layers in MH-MoE are $d \times d$ linear projections that add parameters and computation not present in standard MoE. Under aggressive quantization, one might worry that these additional components—which are essential to the multi-head mechanism's benefits—would suffer disproportionate degradation, negating MH-MoE's advantages. Alternatively, the multi-head decomposition might provide some robustness to quantization by distributing representational capacity across multiple narrower paths. The experiment tests which of these hypotheses holds.

Experimental setup for 1-bit experiments. The same model architectures from the main experiments are trained using the BitNet quantization procedure. All models are "matched in terms of parameters and computation" as before, with the only change being the quantization-aware training. The training dataset, batch size, step count, and evaluation benchmarks remain identical.

Results interpretation. Table 3 shows that MH-MoE continues to outperform both standard SMoE and fine-grained SMoE under BitNet quantization. For example, at 100K steps on RedPajama, MH-MoE (head=3) achieves 26.47 perplexity versus 26.78 for SMoE and 26.68 for fine-grained SMoE. The relative improvements are comparable in magnitude to the non-quantized setting, suggesting that the multi-head mechanism's benefits are robust to extreme weight compression.

Performance gap from quantization. The paper acknowledges that all models perform substantially worse under BitNet than in full precision: the dense model degrades from 12.13 to 30.04 perplexity on RedPajama at 100K steps, and the best MH-MoE variant degrades from 10.51 to 26.47. The authors attribute this to the small model scale (768-dimensional, 12 layers), noting that "when the model size is relatively small, BitNet tends to degrade performance, a finding that aligns with the conclusions reported in the original BitNet paper." This is an important caveat: the quantization results establish compatibility but do not claim that 1-bit MH-MoE at this scale is practically useful; rather, they suggest that at larger scales where BitNet's degradation is less severe, MH-MoE's benefits would compound with quantization efficiency gains.

Ablation Study Design: Isolating Head and Merge Layer Contributions

Section 3.3 presents ablation experiments that are essential for establishing that the multi-head mechanism—not merely the added parameters or computation from the head and merge projections—is responsible for MH-MoE's improvements.

The core experimental question. If you add $d \times d$ projection layers (head and merge) to a standard SMoE architecture, you increase FLOPs unless you compensate by shrinking experts. The ablation asks: does the benefit come from the multi-head splitting and independent routing (the architectural novelty), or does it come from simply having additional linear projections that mix representations in a way that could be achieved without the multi-head structure?

Ablation 1: Adding head and merge layers to SMoE without multi-head splitting (Table 4). The authors take the standard SMoE and fine-grained SMoE models and add head and merge layers ( $W_{\text{head}}$ and $W_{\text{merge}}$ , each $\mathbb{R}^{768 \times 768}$ ) that operate on the full-dimension tokens without splitting into sub-tokens. The forward pass becomes: $x \to xW_{\text{head}} \to \text{MoE layer}(xW_{\text{head}}) \to \text{output} \cdot W_{\text{merge}}$ . There is no splitting, no multi-head routing—the expert layer still sees full 768-dimensional tokens and routes them as before. This inevitably increases FLOPs (the head and merge layers add $4Bd^2$ to the total, uncompensated by expert shrinkage), so the comparison is not FLOPs-matched. The results show "only marginal gains in performance": SMoE with head/merge achieves 11.84 vs. 11.87 without them on RedPajama (50K steps), and fine-grained SMoE achieves 11.67 vs. 11.68. These are essentially noise-level differences.

Why this result is important. It establishes that the $d \times d$ projections themselves are not the source of MH-MoE's gains. If they were, SMoE + head/merge would substantially outperform SMoE alone. The fact that it doesn't suggests the benefit requires the combination of projections AND splitting/routing in subspaces.

Ablation 2: MH-MoE without head and merge layers (Table 4, bottom row). Removing head and merge from MH-MoE eliminates the multi-head mechanism entirely: the tokens are split into sub-tokens, processed through experts, and concatenated—but without the initial projection that creates meaningful subspaces or the final projection that reintegrates them. This variant achieves 11.71 perplexity on RedPajama, which is worse than the full MH-MoE (11.46) and only marginally better than fine-grained SMoE without head/merge (11.68). This confirms that both projections are necessary for the multi-head mechanism to work.

Ablation 3: Separate contributions of head and merge layers (Table 5). This ablation uses four configurations: neither layer (✗, ✗: 11.97 perplexity), head only (✓, ✗: 11.74), merge only (✗, ✓: 11.84), and both (✓, ✓: 11.60). The head layer provides a larger improvement (11.97 → 11.74, a drop of 0.23) than the merge layer (11.97 → 11.84, a drop of 0.13). This aligns with the intuition: the head projection is what creates the diverse subspaces that make multi-head routing meaningful, while the merge layer primarily integrates already-computed information. The head layer is the enabling component; the merge layer refines.

Design implications from ablations. These results validate the paper's core architectural thesis: multi-head routing through different representational subspaces is genuinely beneficial, but it requires both the subspace creation mechanism (head projection) and the reintegration mechanism (merge projection) to realize those benefits. Without the split, the projections are just extra matrix multiplies that don't change the routing behavior. Without the projections, the split is just an arbitrary partition of feature dimensions that doesn't correspond to meaningful representational subspaces. The combination—and only the combination—produces consistent improvements.

4. Key Insights and Innovations

Innovation 1: Reframing MH-MoE As a Constraint Satisfaction Problem Rather Than a New Architecture

The most distinctive intellectual move in this paper is not the proposal of a new mechanism—the multi-head routing concept is credited entirely to Wu et al. (2024)—but rather the reframing of MH-MoE implementation as a constraint satisfaction problem with a closed-form solution. This transforms MH-MoE from a promising but impractical idea into a drop-in replacement for standard MoE layers at identical computational cost.

What the field assumed before this paper. The original MH-MoE paper demonstrated that multi-head routing improves language modeling performance, but it did so at a 4× FLOPs premium over standard sparse MoE. The implicit message from that work, whether intended or not, was that the multi-head mechanism's benefits were partly attributable to increased computation—you get better results, but you pay for them with more FLOPs. This left the community with an unresolved question: does multi-head routing provide genuine architectural advantages, or is it simply a more expensive way to achieve what could be obtained by scaling up a standard MoE model? The field had no way to disentangle these hypotheses because no implementation existed at computational parity.

What this paper does differently. Rather than proposing yet another gating mechanism or expert allocation scheme, the authors treat the problem as one of resource accounting: given a fixed budget of FLOPs and parameters (inherited from a standard MoE configuration), how can the head and merge projections be "paid for" within that budget? The answer—reducing each expert's intermediate dimension while increasing the total number of experts according to the derived Equations 8 and 9—is conceptually simple but was not obvious before this analysis. The paper shows that the head and merge cost can be exactly offset by making experts narrower, and the resulting parameter deficit can be closed by adding more (smaller) experts. This is not a new architectural idea; it is a new configuration methodology that makes an existing architecture viable.

Why this reframing matters beyond performance numbers. By establishing FLOPs parity as a hard constraint rather than an aspirational goal, the paper changes the evaluation landscape for MoE architectural innovations. It implicitly argues that any proposed modification to the MoE layer should be assessed at matched FLOPs, not just matched parameters—a standard that was inconsistently applied before. This is analogous to the shift in neural architecture design where topping leaderboards with 10× more compute became less compelling than showing gains at equal cost. The paper's approach also provides a reusable template: for any future modification that adds overhead to the MoE layer (new projections, additional gating networks, auxiliary losses), the same FLOPs-matching methodology could be applied to determine the necessary compensating adjustments to expert dimensions and counts.

Tying to evidence. The ablation in Table 4 is the crucial piece of evidence that validates this reframing. When the head and merge layers are added to a standard SMoE model without the compensating expert adjustments (that is, without honoring the FLOPs parity constraint), the performance gains are marginal (11.87 → 11.84 on RedPajama). Only when the full MH-MoE configuration is used—with the head and merge layers AND the reduced expert dimensions AND the increased expert count—do meaningful improvements emerge (11.46 for 2-head MH-MoE). This demonstrates that the benefit is not from the projections alone, but from the joint configuration that the constraint-satisfaction methodology produces. The performance gains are thus attributable to the architectural pattern that emerges from satisfying the constraints, not to any single component.

Incremental or fundamental? This is a fundamental reframing of an existing architecture, not an incremental tweak. While the mechanism (multi-head routing) existed before, the paper's contribution is establishing the conditions under which it can be fairly evaluated and providing the analytical tools to achieve those conditions. This transforms MH-MoE from a lab curiosity into a practically deployable architecture.

Innovation 2: The Multi-Head Mechanism Improves Expert Utilization Without Changing the Gating Function

A subtle but important finding embedded in the experimental results is that MH-MoE improves performance without any modification to the gating mechanism itself. Standard MoE, fine-grained MoE, and MH-MoE all use the same fundamental gating architecture: a learned linear projection followed by softmax and top-k selection. The difference is what the gating function operates on—full-dimension tokens in standard MoE versus subspace projections in MH-MoE.

The dominant prior assumption. Work on improving MoE performance has largely focused on two strategies: making experts more specialized (fine-grained MoE, DeepSeekMoE's shared vs. routed expert decomposition) or improving the gating function itself (load-balancing losses, auxiliary losses to prevent collapse, learned routing strategies beyond simple top-k). The implicit assumption was that routing quality is primarily limited by the gating mechanism's architecture—make the gate smarter, and routing improves. This paper suggests an orthogonal axis: keep the gate architecture fixed and change the representation it operates on.

What's distinctive about this finding. The MH-MoE gating function receives sub-tokens of dimension d/h rather than d. In principle, this is less information per routing decision—each sub-token contains only a fraction of the original token's features. Yet routing quality improves, as evidenced by the consistent perplexity gains over standard MoE at matched FLOPs. This is counterintuitive: less information per routing decision leads to better overall routing. The resolution, as suggested by the multi-head attention analogy, is that different subspaces emphasize different features, and making h independent routing decisions on these specialized views is more expressive than making one routing decision on the full representation. A token that might be ambiguous in the full 768-dimensional space (is it more mathematical or more linguistic?) can be clearly mathematical in one subspace and clearly linguistic in another, enabling both aspects to be routed appropriately.

This reframes the routing problem from "make a single correct decision per token" to "make multiple correct decisions per token across different representational subspaces"—a fundamentally different objective that is not achievable through any improvement to the standard gating function alone, no matter how sophisticated.

Significance beyond raw performance. This finding opens a design space that was previously unexplored: representation-space engineering for routing. Rather than asking "how do we build a better gate?", the question becomes "how do we structure the input representation so that simple gating works better?" The head projection matrix W_head is learned, which means the model can discover which subspace decompositions are most useful for routing. This shifts the burden of routing quality from the gating architecture (which must be simple for computational efficiency) to the representation learning (where we have more flexibility and capacity).

Tying to evidence. Table 1 shows that the 3-head variant (more subspaces, finer-grained routing decisions) consistently outperforms the 2-head variant (10.51 vs. 10.70 on RedPajama at 100K steps). Since both use the same basic gating mechanism (just with different k values), the improvement must come from the increased number of representation subspaces available for routing. Additionally, the head-layer ablation in Table 5 shows that adding only the head projection (which creates the subspaces but doesn't reintegrate them) provides a larger gain (11.97 → 11.74) than adding only the merge layer (11.97 → 11.84). This supports the interpretation that subspace creation—what the gating function operates on—is the primary driver of improvement.

Incremental or fundamental? This is a conceptual shift in how to think about routing quality in MoE models. It is not a new mechanism but a new understanding of where the leverage lies: representation design, not gating sophistication.

Innovation 3: FLOPs Parity As a Methodological Standard for MoE Architecture Evaluation

The paper establishes, by example, a methodological standard that the original MH-MoE work failed to meet: architectural innovations in MoE should be evaluated at matched FLOPs, not just matched activated parameters. This might sound like an obvious engineering consideration, but it represents a significant tightening of evaluation norms in the MoE literature.

The status quo ante. The sparse MoE literature has traditionally emphasized activated parameters as the primary efficiency metric. The narrative is: "Our model has X total parameters but only activates Y per token, so it's more efficient than a dense model with X parameters." This framing correctly captures the memory vs. compute tradeoff—you need to store all experts but only compute a subset—but it has an important blind spot. When architectural modifications add computation outside the expert FFNs (as MH-MoE does with its head and merge projections), the activated parameter count remains unchanged while the actual FLOPs increase. A model could have the same activated parameters as a baseline while consuming substantially more compute per token, and the standard efficiency narrative would not detect this.

The original MH-MoE paper fell into exactly this trap: it matched activated parameters with baseline MoE but consumed 4× more FLOPs. Readers who focused only on the parameter counts would have missed the computational premium they were paying for the reported performance gains.

What this paper establishes. By making FLOPs parity the central constraint in its configuration derivation, the paper effectively argues that FLOPs, not activated parameters, should be the primary unit of computational cost in MoE architecture comparisons. This is a methodological contribution, not an architectural one, but it has significant implications for how the field evaluates new MoE variants. The paper's Equation 8 and Equation 9 serve as both configuration tools and evaluative standards: any MH-MoE configuration derived from them can be directly compared to its baseline MoE counterpart without caveats about computational cost.

Why this matters for future work. The paper implicitly critiques a broader pattern in the efficient architecture literature where innovations are evaluated with favorable FLOPs accounting. For example, adding auxiliary modules, extra normalization layers, or more sophisticated routing functions can improve performance while incrementally increasing FLOPs in ways that are easy to overlook if only parameter counts or FLOPs-leading-terms are reported. By providing an explicit, reproducible methodology for achieving exact (leading-term) FLOPs matching, the paper raises the bar for what constitutes a fair comparison.

Tying to evidence. The entire experimental section is structured around this standard. Tables 1 and 2 explicitly state "all models are matched in terms of parameters and computation." The ablation in Table 4 reinforces the point by showing what happens when this standard is violated—adding head/merge to SMoE without FLOPs compensation yields minimal gains, but those gains would be misleadingly attributed to the architectural change if FLOPs were not controlled for. The consistency of the evaluation framework across standard MoE, fine-grained MoE, and both MH-MoE variants, all at matched cost, is itself the evidence that this methodological standard is being applied rigorously.

Incremental or fundamental? This is an incremental refinement of evaluation methodology, but one with outsized practical significance. The standard of FLOPs-matched comparison is not new to the field as a whole (it is standard in efficient architecture design for vision models, for example), but applying it systematically to MoE architectural variants—and deriving the analytical tools to do so—is a genuine contribution that the original MH-MoE work lacked.

Innovation 4: The Multi-Head Mechanism Survives Aggressive Quantization, Suggesting Representational Robustness

The 1-bit quantization experiments (Section 3.2) produce a finding that is easy to overlook in a performance-table scan but has significant implications: MH-MoE's benefits persist under BitNet quantization, and the relative ordering of architectures (MH-MoE > fine-grained SMoE > SMoE > dense) is preserved despite a substantial absolute degradation in perplexity.

Why this is surprising. Quantization to 1-bit (ternary) precision is an extreme form of compression that fundamentally changes the nature of the computations being performed. The head and merge projections in MH-MoE are large d×d matrices that, under BitNet, are reduced to {-1, 0, +1} values. One might expect these projections to be particularly vulnerable to quantization because they are responsible for creating meaningful representational subspaces—a task that seems to require the expressive capacity of full-precision weights to learn useful feature transformations. The fact that they remain effective under ternary quantization suggests that the subspace decomposition they learn is structurally simple—capturable with low-precision weights—rather than requiring fine-grained weight patterns that quantization would destroy.

What this implies about the multi-head mechanism. The robustness to quantization hints at a deeper property: the multi-head decomposition may function as a form of implicit regularization that encourages the model to learn representations that are compressible. By forcing the routing to operate on multiple low-dimensional subspaces rather than one high-dimensional space, MH-MoE may discover feature decompositions that are more modular and less dependent on precise weight values. This is speculative—the paper does not investigate this mechanism—but the empirical result raises the question and provides initial evidence for it.

Significance for deployment scenarios. The combination of MH-MoE and BitNet is practically relevant because MoE models, despite their parameter efficiency, are challenging to deploy on memory-constrained devices due to the need to store all experts. 1-bit quantization dramatically reduces the memory footprint, but only if the architectural benefits survive the compression. The paper's results suggest that MH-MoE and BitNet are compatible in a way that compounds benefits: MH-MoE improves the quality-efficiency tradeoff at full precision, and BitNet compresses the result without destroying the multi-head advantage. This is not a given—many architectural modifications interact negatively with quantization, as the paper acknowledges by noting the overall degradation at small scales.

Tying to evidence. Table 3 shows that MH-MoE (head=3) achieves 26.47 perplexity at 100K steps under BitNet, versus 26.78 for standard SMoE and 26.68 for fine-grained SMoE. The ranking is identical to the full-precision setting (Table 1), and the gaps are proportionally similar. The paper explicitly notes the absolute degradation (10.51 → 26.47 for the same architecture) but frames this as a model-scale effect rather than an MH-MoE-specific vulnerability, citing consistency with the original BitNet paper's findings.

Incremental or fundamental? This is an incremental but practically significant finding. It does not change our understanding of how MH-MoE works, but it establishes an important compatibility property that expands the range of scenarios where MH-MoE can be deployed. The suggestion of representational robustness to quantization is intriguing but under-explored—it remains a hypothesis rather than an established insight.

Innovation 5: The Head Layer, Not the Multi-Head Routing, Is the Primary Driver of Gains

The ablation analysis in Section 3.3 yields a finding that refines our understanding of where the multi-head mechanism's value actually comes from. Table 5 shows that the head layer alone (without the merge layer) provides most of the perplexity improvement over the no-projection baseline (11.97 → 11.74 for head-only, versus 11.97 → 11.84 for merge-only). This is diagnostically important because it localizes the benefit to the subspace creation step rather than the subspace integration step.

Why this is a non-obvious finding. The multi-head attention analogy might lead one to expect that both the input projection (creating query/key/value subspaces) and output projection (merging head outputs) are equally important—after all, both are essential in multi-head attention. The fact that the merge layer contributes less than the head layer suggests that the MH-MoE mechanism is doing something different from multi-head attention at a functional level. In attention, each head computes independent attention patterns, and the merge projection is where the independently attended information is combined into a coherent output—both steps are crucial because the heads compute qualitatively different operations. In MH-MoE, the "operations" are expert FFN computations, which are more generic than attention weightings. The finding suggests that the primary benefit comes from enabling the gating function to make better routing decisions (via the head projection's subspace decomposition), while the merge projection's integration is beneficial but secondary.

This has design implications: if most of the gain comes from the head layer, future work might explore alternative merge strategies that are cheaper than a full d×d projection, potentially reducing the FLOPs overhead and enabling more aggressive multi-head configurations (more heads) within the same compute budget.

Tying to evidence. Table 5 is the primary evidence, with three key comparisons: (a) head-only vs. baseline: 11.74 vs. 11.97, a 0.23 perplexity drop; (b) merge-only vs. baseline: 11.84 vs. 11.97, a 0.13 perplexity drop; (c) both vs. head-only: 11.60 vs. 11.74, a 0.14 incremental gain from adding merge. The head layer accounts for approximately 60% of the total improvement (0.23 / 0.37), with the merge layer contributing the remaining 40% (and approximately half of that 40% being redundant with the head's contribution, since merge-only alone provides only 0.13 gain).

Incremental or fundamental? This is a diagnostic finding—it advances understanding of an existing mechanism rather than introducing a new one. It has practical value for future architecture design (prioritize head-layer optimization over merge-layer optimization) and theoretical value for understanding what multi-head routing actually accomplishes (better gating decisions, not better output integration).

5. Experimental Analysis

Evaluation Methodology

Dataset. All models are pre-trained on the RedPajama dataset (Together Computer, 2023), an open-source reproduction of the LLaMA training corpus comprising approximately 1.2 trillion tokens sourced from CommonCrawl, C4, GitHub, ArXiv, books, Wikipedia, and StackExchange. The paper reports validation perplexity on three held-out evaluation sets: the RedPajama validation split, English Wikipedia, and C4 (Colossal Clean Crawled Corpus). No test set distinct from these validation sets is used—the reported perplexity numbers serve as both the optimization target (implicitly, through training on the RedPajama training split) and the primary evaluation metric. The paper does not specify the exact size of each validation split or how they were constructed from the source datasets.
Base model(s). All experiments use a decoder-only Transformer architecture with 12 layers and a model dimension of d = 768. This is a relatively small model by contemporary standards—comparable in scale to GPT-2 small-medium—chosen, implicitly, to enable systematic architectural comparisons within a tractable compute budget for 100,000-step pre-training runs. The base architecture uses standard dense FFN layers with SwiGLU activation (Shazeer, 2020). For the MoE variants, these dense FFN layers are replaced with MoE layers every two Transformer layers, following the common cadence established in prior MoE Transformer work. The paper does not specify the number of attention heads, the sequence length used during training, or the vocabulary size—these architectural details, while standard, are not reported. The paper states that "the same code base, training parameters, and pre-training tasks" are used across all experiments, ensuring that differences in perplexity are attributable to architecture rather than training recipe variation. However, specific optimizer hyperparameters (learning rate, schedule, weight decay, Adam betas, dropout, gradient clipping) are entirely absent from the paper. This is a notable reproducibility gap: without these details, the reported perplexity numbers cannot be independently replicated even with access to the described architectures and dataset.
Metrics. The sole evaluation metric is perplexity, computed as the exponentiated negative log-likelihood: exp(−(1/T) ∑_t log p(y_t | y_{<t})), where T is the total number of tokens in the evaluation set. Lower perplexity indicates better next-token prediction performance. Perplexity is reported at both 50,000 training steps and 100,000 training steps, providing a view of relative convergence rates across architectures as well as final performance. The paper does not report any downstream task evaluations (e.g., few-shot accuracy on reasoning benchmarks, zero-shot performance on standard NLP tasks), which means the perplexity improvements cannot be mapped to practical task-level gains. This is a significant limitation for assessing whether the multi-head mechanism's benefits translate to user-facing applications.
Baselines. The paper compares against four baseline architectures:
- Dense: A standard 12-layer decoder-only Transformer with dense SwiGLU FFN layers throughout (no MoE layers at all). This serves as the lower-bound reference point for what MoE architectures are supposed to improve upon.
- SMoE (Sparse Mixture-of-Experts): The canonical MoE architecture described in Section 1, using top-1 gating with 8 experts and intermediate dimension d_moe = 2048. MoE layers replace dense FFN layers every two layers. This is the primary baseline that MH-MoE aims to outperform.
- Fine-grained SMoE: A variant of standard SMoE where the intermediate dimension is halved to 1024 and the number of experts is doubled to 16, maintaining approximately the same total expert parameters while providing finer-grained expert specialization. Top-1 gating is preserved. This baseline tests whether simply increasing the number of smaller experts (without multi-head routing) provides comparable benefits to MH-MoE.
- SMoE with shared expert (Table 2 only): All MoE variants (SMoE, fine-grained SMoE, MH-MoE) are also evaluated in a residual MoE configuration with a shared expert of hidden dimension 2048, following DeepSeekMoE (Dai et al., 2024). This tests whether MH-MoE's benefits persist when combined with established MoE architectural improvements.
Generation budget / compute accounting. The paper does not use "generation budget" in the sense of inference-time compute scaling—this is a pre-training study, not an inference-time study. Instead, compute parity is enforced architecturally through the scalar multiplication analysis in Section 2.2 and the configuration equations in Section 2.3. All models are described as "matched in terms of parameters and computation" (Tables 1, 2, 3 captions), meaning the total number of scalar multiplications per forward pass (leading term) and the total number of trainable parameters are designed to be equal across the MoE variants being compared. The paper provides the analytical framework for this matching (Equations 8 and 9) and specifies the resulting configurations (e.g., 40 experts for 2-head MH-MoE, 96 experts for 3-head MH-MoE), but does not empirically verify the FLOPs matching by measuring actual runtime or profiling hardware utilization. The matching is purely analytical and based on leading-term scalar multiplication counts, which ignores lower-order terms and implementation-level factors like kernel launch overhead, memory bandwidth effects, and communication costs in distributed training. For the dense baseline, FLOPs are not matched—the dense model serves as a lower-bound reference and is expected to have different computational characteristics.
Cross-validation / statistical protocol. There is no cross-validation, statistical significance testing, or uncertainty quantification reported in the paper. Perplexity is reported as a single number per model per evaluation set per training step, with no confidence intervals, standard deviations, or multiple training runs to assess variance. The 100,000-step training runs are substantial enough that training noise is unlikely to dominate the architectural comparisons for a 768-dimensional model, but the absence of any statistical protocol means the reader cannot assess whether a difference of 0.1-0.2 perplexity (e.g., MH-MoE head=3 at 10.51 vs. MH-MoE head=2 at 10.70 on RedPajama at 100K steps) is reliable or within the range of run-to-run variation. Given that some of the claimed improvements are small (e.g., fine-grained SMoE at 10.74 vs. MH-MoE head=2 at 10.70), the lack of variance estimates is a meaningful limitation.

Main Quantitative Results

The paper's experimental results are organized around two primary comparisons (MH-MoE vs. baselines without shared experts in Table 1; with shared experts in Table 2), one quantization experiment (Table 3), and a set of ablation studies (Tables 4 and 5). I walk through each in turn.

Language Modeling Without Shared Experts (Table 1)

Headline result. MH-MoE with 3 heads achieves the lowest perplexity across all three evaluation sets at both 50,000 and 100,000 training steps, outperforming both standard SMoE and fine-grained SMoE at matched FLOPs and parameters. The improvements are consistent across training duration and evaluation domain.

Detailed comparisons at 100,000 steps (final performance). On the RedPajama validation set, MH-MoE (head=3) achieves 10.51 perplexity, compared to 10.90 for standard SMoE and 10.74 for fine-grained SMoE. This represents a 0.39 perplexity improvement over standard SMoE (a ~3.6% relative reduction) and a 0.23 improvement over fine-grained SMoE (~2.1% relative reduction). The 2-head MH-MoE variant achieves 10.70, which is better than fine-grained SMoE (10.74) but by a smaller margin (0.04). The dense baseline achieves 12.13—substantially worse than all MoE variants, confirming the expected benefit of sparse expert architectures even at this modest scale.

On English Wikipedia, the pattern is similar but the gaps are smaller in absolute terms: MH-MoE (head=3) at 9.18 versus SMoE at 9.68 (0.50 improvement) and fine-grained SMoE at 9.38 (0.20 improvement). MH-MoE (head=2) at 9.26 is again between the two baselines.

On C4, MH-MoE (head=3) achieves 13.63 versus SMoE at 14.35 (0.72 improvement) and fine-grained SMoE at 13.97 (0.34 improvement). The relative improvements on C4 are larger than on RedPajama or Wikipedia, suggesting that the multi-head mechanism may be particularly beneficial for the web-crawl text distribution represented by C4.

Performance at 50,000 steps (mid-training convergence). The 50,000-step results in the upper half of Table 1 show that the relative ordering of architectures is already established by mid-training: MH-MoE (head=3) at 11.45 on RedPajama versus SMoE at 11.87 and fine-grained SMoE at 11.68. The gaps are proportionally similar to the 100,000-step results, suggesting that MH-MoE's advantage is not merely a convergence speed effect (where it reaches a given perplexity faster but plateaus at the same point)—rather, the architecture appears to be genuinely more parameter-efficient, producing better predictions at matched compute throughout training.

Head count scaling. Across all three datasets and both training durations, the 3-head variant consistently outperforms the 2-head variant. On RedPajama at 100K steps: 10.51 (3 heads) vs. 10.70 (2 heads), a 0.19 gap. The improvement from increasing head count from 2 to 3 is smaller than the improvement from introducing the multi-head mechanism at all (SMoE at 10.90 vs. MH-MoE head=2 at 10.70, a 0.20 gap), suggesting potential diminishing returns to further head increases—though the paper does not explore head counts beyond 3.

An important baseline question: how much of the SMoE → MH-MoE improvement is from multi-head routing vs. from the increased expert count? MH-MoE (head=3) uses 96 experts with top-3 gating, while standard SMoE uses 8 experts with top-1 gating. The MH-MoE configuration thus has both more experts total (96 vs. 8) and more activated experts per token (3 vs. 1). Fine-grained SMoE partially controls for the expert count increase (16 experts instead of 8) but still uses top-1 gating. The fact that MH-MoE (head=2) with 40 experts outperforms fine-grained SMoE with 16 experts suggests that the multi-head mechanism provides benefits beyond simply having more experts. However, the paper does not include a crucial control experiment: a variant with the same number of experts and same gating width as MH-MoE, but without the head/merge projections and multi-head splitting. For example, a non-multi-head SMoE with 96 experts and top-3 gating, dimension-adjusted to match FLOPs. Without this control, we cannot conclusively attribute the gains to the multi-head mechanism rather than to the increased expert count and gating width independently. The ablation in Table 4 partially addresses this concern by showing that adding head/merge layers to standard SMoE without adjusting expert count or dimension yields minimal gains—but this tests the wrong direction (adding projections without changing experts) rather than the right direction (changing expert count and gating to match MH-MoE's configuration, but without the multi-head structure).

Language Modeling With Shared Experts (Table 2)

Headline result. Adding a shared expert (residual MoE configuration, following DeepSeekMoE) improves perplexity for all MoE architectures, and MH-MoE maintains its advantage over baselines in this setting. The relative ordering of architectures (MH-MoE head=3 > MH-MoE head=2 > fine-grained SMoE > SMoE) is preserved.

Detailed comparisons at 100,000 steps. On RedPajama, MH-MoE (head=3) with shared expert achieves 10.28, versus SMoE at 10.66 and fine-grained SMoE at 10.41. The gap between MH-MoE (head=3) and SMoE is 0.38—nearly identical to the 0.39 gap in the non-shared-expert setting. This consistency suggests that the multi-head mechanism's benefits are additive with the shared expert's benefits: both contribute independently rather than redundantly. On Wikipedia, MH-MoE (head=3) achieves 8.72 (vs. 9.44 for SMoE, a 0.72 gap). On C4, it achieves 13.49 (vs. 14.30, a 0.81 gap).

Comparing shared vs. non-shared within the same architecture. The shared expert consistently improves performance. For SMoE: 10.66 (shared) vs. 10.90 (non-shared) on RedPajama, a 0.24 improvement. For MH-MoE (head=3): 10.28 (shared) vs. 10.51 (non-shared), a 0.23 improvement. The absolute benefit of the shared expert is approximately constant across architectures (~0.23-0.24 on RedPajama), reinforcing the additivity interpretation: the shared expert provides a uniform performance lift, and MH-MoE's multi-head mechanism provides an additional lift on top.

Mid-training results. At 50,000 steps with shared expert: MH-MoE (head=3) at 11.26 vs. SMoE at 11.76 on RedPajama, a 0.50 gap. This is larger than the 0.42 gap at 50K steps without shared expert (11.45 vs. 11.87), hinting that the combination of multi-head routing and shared experts may accelerate early-training convergence—though without multiple training runs, this could also be noise.

1-Bit Quantization Results (Table 3)

Headline result. Under BitNet 1-bit quantization, MH-MoE continues to outperform both standard SMoE and fine-grained SMoE, with the same relative ordering as in full precision. However, all models suffer substantial absolute perplexity degradation compared to their full-precision counterparts.

Detailed comparisons at 100,000 steps. On RedPajama, MH-MoE (head=3) achieves 26.47 under BitNet, versus SMoE at 26.78 (0.31 gap) and fine-grained SMoE at 26.68 (0.21 gap). The dense baseline degrades even more severely: 30.04, a gap of 3.57 from the best MH-MoE. On Wikipedia: MH-MoE (head=3) at 21.06 vs. SMoE at 21.54. On C4: 29.14 vs. 29.73.

Quantifying the quantization penalty. Across all architectures, BitNet training approximately doubles the perplexity compared to full precision. For MH-MoE (head=3): 26.47 (BitNet) vs. 10.51 (full precision) on RedPajama—a factor of ~2.5. For SMoE: 26.78 vs. 10.90—a factor of ~2.46. The proportional penalty is similar across architectures, meaning the relative benefits of MH-MoE (the gaps between architectures) are preserved in absolute perplexity terms. The paper attributes the overall degradation to the small model scale, citing consistency with the original BitNet paper's finding that "when the model size is relatively small, BitNet tends to degrade performance." This is plausible—768-dimensional models are far below the scale where 1-bit quantization typically shows its strongest results—but the paper provides no larger-scale BitNet experiments to verify that the architecture ranking holds at sizes where BitNet's penalty is smaller.

The gap between MH-MoE and baselines under quantization is proportionally similar to full precision. On RedPajama at 100K steps, the SMoE → MH-MoE (head=3) improvement is 0.31 perplexity (BitNet) vs. 0.39 (full precision). As a fraction of the absolute perplexity, this represents roughly 1.2% improvement in both settings, suggesting that the multi-head mechanism's benefits are not disproportionately affected by quantization—neither amplified (which would suggest quantization reveals the mechanism's robustness) nor destroyed (which would suggest incompatibility).

Ablation Results: Head and Merge Layer Contributions (Tables 4 and 5)

Headline result. Adding head and merge projections to standard SMoE or fine-grained SMoE without the compensating expert configuration changes (reduced expert dimension, increased expert count) yields only marginal, likely noise-level improvements. In MH-MoE, removing the head and merge layers substantially degrades performance (Table 4). Separately, the head layer contributes more to MH-MoE's performance than the merge layer (Table 5).

Ablation: Projections without multi-head architecture (Table 4, rows 2 and 4). When head and merge layers are added to standard SMoE, perplexity on RedPajama drops from 11.87 to 11.84—a 0.03 improvement. For fine-grained SMoE, the drop is from 11.68 to 11.67—a 0.01 improvement. These differences are well within the range of training noise for a single run, and the paper appropriately characterizes them as "only marginal gains in performance." On Wikipedia and C4, the patterns are similar: SMoE with head/merge achieves 10.48 (vs. 10.51 without) and 15.61 (vs. 15.63 without). These results demonstrate that the head and merge projections, in isolation and at added computational cost, do not provide meaningful benefits.

Ablation: MH-MoE without head and merge (Table 4, row 5 vs. row 6). Removing both the head and merge layers from MH-MoE (head=2) degrades perplexity from 11.46 to 11.71 on RedPajama—a 0.25 increase, which is substantial relative to the original MH-MoE advantage over SMoE (0.17, comparing 11.46 for MH-MoE to 11.63 mean of SMoE and fine-grained SMoE). On Wikipedia, the degradation is from 9.98 to 10.16 (0.18); on C4, from 14.89 to 15.23 (0.34). The degraded MH-MoE (row 5, 11.71) performs comparably to fine-grained SMoE (11.68) and worse than fine-grained SMoE with head/merge (11.67). This confirms that the head and merge layers are necessary for MH-MoE's performance advantage—without them, the architecture reduces to something close to fine-grained SMoE (many small experts with a particular token-splitting scheme, but no learned subspace projections).

Ablation: Separate contributions of head and merge (Table 5). This ablation uses the MH-MoE (head=2) configuration and varies the presence of each projection:

Head layer	Merge layer	RedPajama perplexity
✗	✗	11.97
✓	✗	11.74
✗	✓	11.84
✓	✓	11.60

The head layer alone provides a 0.23 improvement (11.97 → 11.74). The merge layer alone provides a 0.13 improvement (11.97 → 11.84). Having both provides a 0.37 improvement (11.97 → 11.60). The head layer is responsible for approximately 62% of the total gain (0.23 / 0.37); the merge layer accounts for approximately 35% (0.13 / 0.37). The incremental benefit of adding the merge layer when the head layer is already present is 0.14 (11.74 → 11.60), suggesting that the merge layer's contribution is partially but not entirely redundant with the head layer's.

Why this asymmetry is diagnostically informative. If the multi-head mechanism's benefit came primarily from increased representational capacity (more parameters devoted to transforming the token representation before and after expert processing), the head and merge layers would be expected to contribute roughly equally, since they have identical parameter counts (d × d each). The asymmetry—head matters more than merge—suggests instead that the primary benefit is from improving the gating function's input representation (what the head layer does: creating subspaces for better routing) rather than from improving the output integration (what the merge layer does: combining subspace outputs). This aligns with the paper's framing of MH-MoE as enabling "collective attention to information from various representation spaces within different experts"—the representational diversity is created by the head projection, and the merge projection mainly reassembles what was already effectively computed.

Ablation Studies and Robustness Checks

What the ablation studies cover. The paper presents two sets of ablation experiments in Section 3.3: (1) whether adding head and merge projections to standard MoE architectures (without the full MH-MoE configuration) provides benefits (Table 4), and (2) the separate contributions of the head and merge layers within the full MH-MoE configuration (Table 5). These ablations directly address the questions: "Would standard MoE benefit from just having extra projection layers?" (no) and "Which component of MH-MoE matters most?" (the head layer). The coverage is targeted but narrow—the paper does not ablate other architectural choices, training hyperparameters, or design decisions.
The configurational ablation that is missing. The paper does not ablate the specific expert-dimension and expert-count choices that emerge from the FLOPs-matching equations. For instance, with h=2 and the derived configuration of d_mhmoe = 768, E = 40, k = 2, the paper does not test variants that slightly deviate from this exact parity point: What happens if d_mhmoe is set to 1024 with proportionally fewer experts? What happens if top-1 gating is used with a different expert dimension? Without exploring the neighborhood around the FLOPs-parity configuration, the reader cannot assess whether the reported configuration is genuinely optimal or simply happens to be the one the equations produce. It is possible that a configuration with slightly different expert dimensions and counts—still FLOPs-matched but not exactly following Equation 8—would perform even better. The paper's analytical framework provides a single point in the configuration space but does not validate that this point is locally or globally optimal.
No ablation on head count. The paper tests two head counts (2 and 3) and finds that 3 outperforms 2, but does not test 1 head (which would reduce to standard MoE with a particular configuration), 4 heads (the configuration used in the original MH-MoE paper), or higher head counts. The extrapolation from 2→3 improvement to potential 3→4 improvement is entirely unknown. The original MH-MoE paper (Wu et al., 2024) used h=4, so the fact that this paper caps at h=3 raises the question of whether h=4 was attempted and found to underperform, or simply not tested. If 4 heads were tested and found to be worse, that would be an informative negative result about diminishing returns or negative returns to excessive subspace splitting. The absence of any h=4 result—despite the original paper using exactly that configuration—is a notable gap.
No ablation on the MoE layer placement cadence. The paper places MoE layers "every two layers" of the 12-layer Transformer. This is a reasonable default but is not ablatted. The interaction between multi-head routing and MoE layer frequency is unexplored: perhaps MH-MoE's benefits are larger or smaller depending on how many layers use expert routing. In a denser MoE configuration (every layer instead of every two), the head and merge projection cost would be incurred more frequently, potentially changing the optimal balance of expert dimension reduction vs. expert count increase. Conversely, in a sparser MoE configuration (every four layers), the overhead of the projections would be amortized over fewer expert computations, potentially making FLOPs matching easier or changing the optimal h.
No ablation on training hyperparameters. The paper states that "the same code base, training parameters, and pre-training tasks" are used across all experiments, but does not report what those training parameters are. This makes it impossible to assess whether the results are robust to learning rate, batch size, sequence length, optimizer choice, or training duration. For example, if MH-MoE's multi-head splitting effectively increases the number of routing decisions per token (by a factor of h), it might benefit from different learning rate scaling or different expert load-balancing loss coefficients compared to standard SMoE. None of these interactions are explored.
No robustness check across random seeds. All reported perplexity values are from single training runs. The paper provides no evidence that the perplexity differences between architectures—particularly the small gaps (e.g., 10.70 vs. 10.74 for MH-MoE head=2 vs. fine-grained SMoE on RedPajama at 100K steps)—are robust to random initialization and data ordering. For a 100,000-step training run on a 768-dimensional model, run-to-run variance in final perplexity is likely non-negligible, and the absence of multiple seeds means that some of the claimed improvements may be within the noise floor.
No validation of the FLOPs matching through empirical measurement. The paper's entire contribution hinges on the claim that MH-MoE and the baseline architectures are "matched in terms of parameters and computation." The parameters matching is straightforward to verify analytically. The FLOPs matching, however, is based on leading-term scalar multiplication counts—an analytical approximation that ignores lower-order terms, activation function costs, gating network costs, normalization layer costs, attention costs (which are shared across all variants and thus controlled for), and most importantly, implementation-level efficiency factors. On real hardware, two operations with the same scalar multiplication count can have substantially different wall-clock times due to memory access patterns, parallelism, and kernel fusion opportunities. The 3-head MH-MoE with 96 experts, each operating on 256-dimensional sub-tokens, has a fundamentally different computation and memory access pattern than standard SMoE with 8 experts operating on 768-dimensional tokens. Without profiling actual training throughput (tokens per second) or inference latency, the claim of computational parity remains an analytical convenience rather than an empirically validated fact. A simple table reporting training step time or tokens-per-second for each architecture would substantially strengthen the paper's central claim.
The BitNet results constitute a robustness check for the architecture, not a claim about practical deployment. The 1-bit experiments demonstrate that MH-MoE's benefits survive quantization, which serves as evidence that the multi-head mechanism does not rely on high-precision weight representations. However, the absolute perplexity values under BitNet (26-30 range) are too poor for the 768-dimensional models to be practically useful, as the paper acknowledges. The robustness check is thus conceptual—"the mechanism is robust to quantization"—rather than practical—"you can deploy quantized MH-MoE today." At larger model scales where BitNet's degradation is less severe, the conclusion might change, but no such experiments are provided.

Critical Assessment

This paper makes one central empirical claim and several supporting claims. Here, I evaluate each against the evidence presented.

Central Claim: MH-MoE, when implemented with the derived FLOPs-and-parameter-matching configuration, outperforms both standard SMoE and fine-grained SMoE on language modeling perplexity.

Assessment: Supported, but with important caveats about what constitutes a fair comparison. The evidence in Tables 1 and 2 consistently shows MH-MoE (both 2-head and 3-head variants) achieving lower perplexity than SMoE and fine-grained SMoE across three evaluation datasets, two training durations, and both non-shared and shared-expert configurations. The improvements are directionally consistent and appear at both 50K and 100K steps, ruling out the possibility that MH-MoE merely converges faster but plateaus at the same level.

However, the comparison is not fully controlled in a way that isolates the multi-head mechanism per se. MH-MoE (head=3) uses 96 experts with top-3 gating, while SMoE uses 8 experts with top-1 gating. The architectures differ along at least three axes simultaneously: (1) the presence of head and merge projections, (2) the number of experts, and (3) the gating width (k). The paper's FLOPs-matching equations tie these axes together—you cannot change one without changing the others while maintaining parity—but this means the "multi-head mechanism" is being evaluated as a package deal rather than an isolated variable. A reader could reasonably ask: what if I simply increased the number of experts and the gating width in a standard SMoE to match the 96-expert, top-3 configuration, without adding head/merge projections or multi-head splitting? Would that architecture also outperform the 8-expert baseline? The paper does not answer this question. The ablation in Table 4 shows that adding head/merge to 8-expert SMoE doesn't help—but this is a different question from whether the full 96-expert, top-3 routing (without multi-head) would help.

The paper's argument for why the multi-head mechanism specifically—rather than increased expert count or gating width—is responsible for the gains rests on two pieces of evidence: (1) the ablation in Table 4 showing that projections without multi-head splitting don't help standard SMoE, and (2) the comparison of MH-MoE (head=2, 40 experts) vs. fine-grained SMoE (16 experts), where MH-MoE outperforms despite the fine-grained variant having more experts than standard SMoE. This evidence is suggestive but not definitive. The fine-grained SMoE has different gating width (top-1) and different intermediate dimensions (1024 vs. 768) compared to MH-MoE (head=2). A truly isolated test would hold expert count, gating width, and intermediate dimensions constant across the multi-head vs. non-multi-head comparison—this is impossible under the FLOPs-matching constraint, which is precisely why the paper's analytical framework is both its contribution and its limitation: it ties the variables together in a way that precludes clean isolation.

Claim: MH-MoE is compatible with 1-bit quantization (BitNet) and maintains its advantage over baselines under compression.

Assessment: Supported, with strong caveats about practical relevance. Table 3 demonstrates that MH-MoE outperforms SMoE and fine-grained SMoE under BitNet quantization, and the relative gaps are proportionally similar to the full-precision setting. This establishes compatibility—the multi-head mechanism does not catastrophically fail under quantization—which is the claim being made. However, the absolute perplexity values are poor enough (26-30 range) that the result is primarily a proof-of-concept rather than a practical deployment recommendation. The paper's attribution of the degradation to small model scale is plausible but unverified at larger scales. A skeptical reader might note that if the multi-head mechanism's benefits are fundamentally tied to precise alignment of representational subspaces (fine-grained weight patterns in the head projection), then quantization might be expected to degrade MH-MoE more than standard SMoE—the fact that it doesn't is the interesting result, but the mechanism is unexplained and the evidence is limited to a single (small) scale.

Claim: The head layer is the primary driver of MH-MoE's performance gains (Section 3.3, Table 5).

Assessment: Supported. The ablation in Table 5 cleanly isolates the head and merge contributions within the MH-MoE (head=2) configuration, showing that the head layer provides approximately 62% of the total improvement. The experimental design here is appropriate: all four configurations (neither, head-only, merge-only, both) are tested within the same expert configuration, so the only variables are the presence or absence of the two projections. The asymmetry between head and merge contributions is a genuine empirical finding, not forced by the experimental design.

Missing experiments that would have strengthened the paper:

A non-multi-head MoE with matched expert count and gating width. The most important missing baseline is an SMoE variant that uses the same number of experts, same gating width, and same per-expert dimensions as an MH-MoE configuration, but without the head/merge projections and without splitting tokens into subspaces. This would require matching FLOPs by some other means (e.g., reducing the model dimension d, or accepting a slight FLOPs difference), but it would isolate whether the multi-head mechanism provides benefits beyond simply having more experts and broader gating.
Head count h=4 experiments. Given that the original MH-MoE paper used 4 heads, testing this configuration in the FLOPs-matched framework would provide direct continuity with prior work and reveal whether the benefits saturate, reverse, or continue to grow with additional heads.
Empirical wall-clock time or throughput measurements. Reporting tokens-per-second during training and inference for each architecture would validate the analytical FLOPs-matching claims and reveal any hidden efficiency costs (e.g., from the increased number of smaller expert kernels, which may have worse GPU utilization).
Downstream task evaluations. Perplexity improvements on validation sets do not always translate to proportional improvements on downstream tasks (few-shot reasoning, knowledge-intensive QA, code generation). Demonstrating that MH-MoE's perplexity gains translate to task-level accuracy improvements would substantially strengthen the practical case for the architecture.
Larger-scale experiments. The 768-dimensional, 12-layer models tested are roughly GPT-2 Small scale (117M parameters for the dense version; the MoE variants have more total parameters but comparable activated parameters). The original MH-MoE paper and the MoE literature more broadly are primarily concerned with scaling to much larger models (Mixtral 8×7B, DeepSeekMoE at 16B and 145B). Whether MH-MoE's benefits persist, grow, or diminish at scales where MoE is actually deployed is entirely unknown from this paper's experiments.
Multiple training runs with variance reporting. Confidence intervals or standard deviations across at least 3-5 random seeds would allow readers to assess whether differences of 0.1-0.2 perplexity are statistically reliable or within the noise floor. This is particularly important for claims involving small gaps (e.g., MH-MoE head=2 vs. fine-grained SMoE on RedPajama: 10.70 vs. 10.74).

Overall assessment of experimental rigor. The experiments are carefully designed to test the specific claim the paper makes—that a FLOPs-matched MH-MoE implementation outperforms standard MoE baselines on language modeling perplexity—and the results consistently support this claim within the tested regime. The ablation studies are targeted and informative, particularly the head vs. merge decomposition. However, the paper's experimental scope is narrow: one model scale, one training dataset (RedPajama), one evaluation metric (perplexity), single training runs, and a limited exploration of the configuration space around the analytically derived parity point. The most significant limitation is the absence of empirical FLOPs validation—the claim of computational parity is purely analytical and based on leading-term scalar multiplication counts, which may not reflect actual hardware efficiency. For a paper whose primary contribution is establishing FLOPs parity as a viable implementation strategy, the lack of even a single wall-clock time measurement is a notable omission. The results should be interpreted as demonstrating that MH-MoE can be configured to match the analytical FLOPs of standard MoE while improving perplexity, with the practical efficiency implications remaining to be validated through implementation-level benchmarking.

6. Limitations and Trade-offs

6.1 The FLOPs Parity Claim Is Purely Analytical—No Empirical Throughput Measurements Validate It

The assumption or constraint. The paper's entire contribution rests on the claim that MH-MoE can be configured to consume identical FLOPs to standard sparse MoE. This claim is defended entirely through leading-term scalar multiplication counting in Section 2.2 and the derivation of Equations 8 and 9 in Section 2.3. The paper explicitly states:

"we only consider the leading term of the equation" (Section 2.3)

and characterizes models as "matched in terms of parameters and computation" (Table 1, 2, 3 captions). At no point does the paper report actual wall-clock training time, tokens-per-second throughput, or inference latency for any architecture.

The consequence. Leading-term FLOPs matching is a useful analytical tool but does not guarantee practical computational parity on real hardware. The 3-head MH-MoE configuration uses 96 experts each operating on 256-dimensional sub-tokens with top-3 gating, while the baseline SMoE uses 8 experts on 768-dimensional tokens with top-1 gating. These configurations have fundamentally different memory access patterns, kernel launch overheads, and parallelism characteristics. The MH-MoE variant launches many more small matrix multiplications (each expert processes d/h = 256-dimensional inputs rather than d = 768), which may underutilize GPU tensor cores optimized for larger matrix operations. The gating network must compute scores over 96 experts rather than 8, and the top-3 selection operates on a much larger set. These implementation-level factors could make MH-MoE substantially slower in practice than the analytical FLOPs count suggests, which would undermine the economic argument for adopting it over standard SMoE. If MH-MoE is, say, 20-40% slower in wall-clock time despite matched scalar multiplication counts, the perplexity improvements must be weighed against real latency costs that the paper does not quantify.

What evidence exists in the paper. None. The paper provides no throughput or latency measurements whatsoever. This is the single most significant omission for a paper whose stated purpose is to demonstrate that MH-MoE can be made computationally competitive with standard SMoE.

Mitigation status. Not addressed. The paper does not acknowledge this as a limitation or suggest profiling experiments as future work. The entire evaluation framework treats analytical scalar multiplication parity as synonymous with practical computational parity, which is an assumption that the systems and architecture communities would not accept without empirical validation.

6.2 Single Model Scale (768-Dimensional, 12 Layers) Leaves Scaling Behavior Completely Unknown

The assumption or constraint. All experiments use a 12-layer decoder-only Transformer with model dimension d = 768. This is a small model by contemporary standards—comparable to GPT-2 Small (~117M parameters for the dense variant; the MoE variants have more total parameters but comparable activated parameters). The paper does not test any larger configurations. This scope is not explicitly acknowledged as a limitation; the paper simply presents results at this scale without commenting on whether they might generalize.

The consequence. The entire practical motivation for MH-MoE comes from large-scale MoE deployments (Mixtral 8×7B, DeepSeekMoE at 16B and 145B parameters). At 768 dimensions with 12 layers, the model is far below the regime where MoE architectures are typically deployed. Several aspects of MH-MoE's behavior may not scale:

The head and merge layer cost as a fraction of total FLOPs scales differently with d. The head and merge layers contribute 4Bd² to the leading term, while expert computation contributes 4Bd d_mhmoe k. As d grows, the head/merge overhead grows quadratically in d, while expert cost grows only linearly in d (multiplied by d_mhmoe). At larger model dimensions, the FLOPs ratio between projections and experts may shift, potentially requiring different d_mhmoe adjustments that could make FLOPs matching harder or alter the optimal head count.
The optimal number of heads h may depend on d. At d = 768, splitting into 3 sub-tokens yields 256-dimensional subspaces. At d = 4096 (a more typical large model dimension), 3 heads would give 1365-dimensional subspaces—a very different information bottleneck. The optimal h likely scales with d, but the paper provides no evidence on how.
The BitNet results are particularly scale-dependent. The paper acknowledges that BitNet's degradation at small scales is consistent with prior work, but the practical case for 1-bit MH-MoE deployment depends entirely on demonstrating that the benefits persist at scales where BitNet's penalty is small. At d = 768, all models under BitNet achieve perplexities of 26-30, which are too poor for any practical application. Whether the architecture ranking (MH-MoE > fine-grained SMoE > SMoE) holds at 7B+ scale with BitNet or other quantization methods is unestablished.

What evidence exists in the paper. Only the single-scale results in Tables 1-3 and the BitNet degradation noted in Section 3.2: "when the model size is relatively small, BitNet tends to degrade performance, a finding that aligns with the conclusions reported in the original BitNet paper." The paper acknowledges this for BitNet specifically but does not extend the scaling concern to the full-precision results.

Mitigation status. Not addressed for the full-precision experiments. The BitNet scale limitation is acknowledged but not resolved. The paper does not propose scaling experiments as future work or discuss what scale would be needed to make deployment-relevant claims.

6.3 The Multi-Head Mechanism Is Not Isolated From Confounding Architectural Variables

The assumption or constraint. The FLOPs-matching methodology ties together several architectural variables—head count, expert intermediate dimension, number of experts, and gating width—into a single configuration package. The paper compares MH-MoE against baselines (standard SMoE, fine-grained SMoE) that differ along multiple of these axes simultaneously. There is no experiment that holds expert count and gating width constant while varying only the presence of multi-head splitting and projections.

The consequence. The paper's central claim—that multi-head routing improves performance over standard MoE—is confounded with other architectural differences:

MH-MoE (head=3) uses 96 experts with top-3 gating, while the baseline SMoE uses 8 experts with top-1 gating. The MH-MoE variant activates 3 experts per sub-token × 3 sub-tokens = 9 expert computations per original token (though some experts may be selected for multiple sub-tokens), while SMoE activates exactly 1. This dramatically increases the combinatorial expressiveness of expert combinations per token, independent of the multi-head mechanism.
The increased expert count (96 vs. 8) means MH-MoE has access to a much larger "vocabulary" of specialized computations. Fine-grained SMoE partially controls for this (16 experts instead of 8), but the gap between 16 and 96 is vast.
The paper does not include a crucial control experiment: a non-multi-head MoE configured with the same number of experts, same gating width, and matched FLOPs through some other dimensional adjustment (e.g., reducing model dimension or expert intermediate dimension without splitting into subspaces). Without this control, the improvement from SMoE to MH-MoE could be explained by "more experts and wider gating" rather than "multi-head routing" per se.

The ablation in Table 4 partially addresses this concern by showing that adding head/merge projections to standard SMoE (without changing expert count or gating) doesn't help. However, this tests the wrong direction: it asks whether projections help without the expert configuration changes, rather than whether the expert configuration changes help without the projections. The relevant control would be: a non-multi-head MoE with 96 experts, top-3 gating, and dimensions adjusted to match FLOPs—does it match MH-MoE's performance?

What evidence exists in the paper. The ablation in Table 4 (adding projections to SMoE/fine-grained SMoE without expert reconfiguration) and the comparison between MH-MoE and fine-grained SMoE (which increases expert count from 8 to 16, but not to 40 or 96). These provide suggestive but not definitive evidence that the multi-head mechanism contributes independently to the gains.

Mitigation status. Not addressed. The paper does not acknowledge this confound or propose the missing control experiment. The framing treats the FLOPs-matched package as the atomic unit of comparison without interrogating which elements of the package are necessary versus sufficient.

6.4 No Downstream Task Evaluations—Perplexity Improvements May Not Translate to Practical Gains

The assumption or constraint. The sole evaluation metric is language modeling perplexity on validation sets from RedPajama, Wikipedia, and C4. The paper does not evaluate any downstream tasks: no few-shot reasoning (e.g., ARC, HellaSwag, MMLU), no code generation, no question answering, no dialogue or instruction-following. This is standard for pre-training architecture papers, but it leaves a critical gap for practitioners.

The consequence. Perplexity improvements on held-out validation data do not always translate to proportional improvements on downstream task performance, particularly for architectural modifications that change how the model represents and routes information. Several failure modes are possible:

The multi-head mechanism might overfit to surface-level statistical patterns in the pre-training data that improve next-token prediction but don't help with reasoning or knowledge-intensive tasks. If MH-MoE's gains come from better capturing local syntactic or stylistic regularities (which dominate perplexity measurements), the architecture might not improve—and could even hurt—performance on tasks requiring deeper semantic understanding.
The 96-expert, top-3 gating configuration might achieve better perplexity through a form of mixture sparsity that harms the model's ability to generalize to out-of-distribution prompts. Standard MoE models with aggressive expert specialization are known to sometimes underperform dense models on certain downstream tasks despite better perplexity.
The perplexity improvements are modest in absolute terms. On RedPajama at 100K steps, MH-MoE (head=3) achieves 10.51 vs. 10.90 for SMoE—a 3.6% relative improvement. Whether this translates to, say, a 1% or 5% improvement on MMLU few-shot accuracy is anyone's guess. Practitioners deploying MoE models care primarily about task performance, not pre-training perplexity, and the paper provides no evidence on this translation.

What evidence exists in the paper. None beyond perplexity. The paper does not even mention downstream evaluation as a limitation or future work direction.

Mitigation status. Not addressed. This is standard scope for a methods paper at venues where pre-training perplexity is accepted as the primary evaluation, but it represents a significant limitation for practitioners trying to assess whether to adopt MH-MoE in their systems.

6.5 The Difficulty Estimation Cost for Orchestrating Multi-Head Routing Is Not Accounted For

The assumption or constraint. This limitation is more structural than the paper's explicit scope: standard MoE uses a gating network that scores E experts per token (e.g., E=8 for the baseline). MH-MoE splits each token into h sub-tokens, each of which must be independently gated over E_mhmoe experts (e.g., E_mhmoe=96 for the 3-head variant). This means the gating computation scales with h × E_mhmoe rather than E_moe. For the 3-head configuration: 3 × 96 = 288 expert scores must be computed per original token, compared to 8 for the baseline SMoE—a 36× increase in gating computation.

The consequence. The paper's FLOPs analysis in Section 2.2 counts scalar multiplications for the head layer, expert FFNs, and merge layer, but does not include the gating network's cost. The gating function G: R^(d/h) → R^E_mhmoe is typically implemented as a linear projection from the sub-token dimension to the number of experts, followed by softmax and top-k selection. For the 3-head MH-MoE with d/h = 256 and E_mhmoe = 96, this is a 256 × 96 matrix multiplication per sub-token, performed h=3 times per original token. The total gating FLOPs are 3 × 2 × 256 × 96 = 147,456 scalar multiplications per token for MH-MoE, versus 2 × 768 × 8 = 12,288 for the baseline SMoE—a 12× increase in gating cost. While this is small relative to the expert FFN cost (which is in the millions of scalar multiplications), it represents an unaccounted computational overhead that grows with both h and E_mhmoe.

More importantly, the gating computation requires materializing a (B × h) × E_mhmoe logit matrix in memory, which can be large. For the 3-head configuration with E_mhmoe = 96 and large batch sizes typical in pre-training (the paper uses 0.5M tokens per batch), this matrix has 0.5M × 3 × 96 = 144M elements—substantial memory pressure that is not reflected in the parameter or FLOPs counts.

What evidence exists in the paper. The explicit FLOPs analysis in Equations 6 and 7 does not include gating network terms. The paper does not discuss gating cost as a factor in the FLOPs matching or acknowledge it as an omission. The analytical framework thus systematically underestimates MH-MoE's true computational cost relative to baselines.

Mitigation status. Not addressed. The paper neither accounts for gating cost in the FLOPs analysis nor discusses its potential significance. For the 3-head, 96-expert configuration, the gating overhead is measurable and should be included for a genuinely fair comparison.

6.6 Single Training Run Per Configuration—No Statistical Evidence That Improvements Are Reliable

The assumption or constraint. Every perplexity value in Tables 1-5 is reported as a single number per configuration per checkpoint, with no confidence intervals, standard deviations, or evidence of multiple training runs. The paper states that all models use "the same code base, training parameters, and pre-training tasks" across experiments, but does not report what those training parameters are (learning rate, schedule, optimizer, etc.), nor does it mention multiple random seeds.

The consequence. The paper claims improvements over baselines based on differences as small as 0.04 perplexity (MH-MoE head=2 at 10.70 vs. fine-grained SMoE at 10.74 on RedPajama at 100K steps). At this margin, run-to-run variance from random initialization and data ordering could easily exceed the claimed improvement. Without multiple seeds, the reader cannot assess whether these small gaps are reliable or within the noise floor.

The problem is compounded by the fact that the paper does not report training hyperparameters. If, for example, the learning rate was tuned on the baseline SMoE and then applied unchanged to MH-MoE, the MH-MoE configurations might be operating at a suboptimal learning rate due to their different architectural characteristics (more experts with smaller dimensions, different gradient flow through the head and merge projections). The paper's claim that "the same code base, training parameters, and pre-training tasks" are used across all experiments sounds like a control, but it actually introduces a confound: the hyperparameters were likely designed for standard Transformer training and may not be optimal for MH-MoE's different optimization landscape. Performance differences could reflect hyperparameter sensitivity rather than architectural superiority.

What evidence exists in the paper. Single-run perplexity values with no uncertainty quantification. No hyperparameter sweep results, no learning rate sensitivity analysis, no evidence of hyperparameter tuning for fairness across architectures.

Mitigation status. Not addressed. The paper does not report random seeds, variance estimates, or hyperparameter details. This is a meaningful gap for a paper whose claimed improvements include margins that are small relative to typical training variance in language model pre-training.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper makes a methodological contribution that shifts the evaluation standard for MoE architecture research rather than introducing a new mechanism or achieving a new state-of-the-art. Its primary effect is to establish that the multi-head routing mechanism proposed by Wu et al. (2024) can be implemented at FLOPs parity with standard sparse MoE—a constraint the original work violated by roughly 4×—thereby transforming MH-MoE from an intellectually promising but computationally impractical idea into a viable drop-in replacement for standard MoE layers.

This is a reframing, not a paradigm shift. The multi-head mechanism itself is unchanged from the original proposal. What changes is the analytical toolkit—the derived Equations 8 and 9 that specify exactly how to adjust expert intermediate dimension and expert count to offset the added cost of the head and merge projections. This reframing matters because it converts a qualitative architectural idea ("let tokens route through multiple representation subspaces") into a quantitative resource allocation problem with a closed-form solution. Before this paper, a practitioner wanting to use MH-MoE faced an uncomfortable choice: accept a 4× FLOPs penalty (and hope reviewers don't notice, or argue the performance gains justify it), or hand-tune expert dimensions and counts through expensive trial-and-error. After this paper, there is a recipe: plug your baseline MoE configuration into the equations, and the resulting MH-MoE variant is guaranteed (analytically, at leading-term) to match FLOPs and parameters. This lowers the barrier to adoption from "research project" to "configuration change."

The paper resolves a specific ambiguity left by the original MH-MoE work. Wu et al. (2024) demonstrated that MH-MoE outperforms standard SMoE, but the performance gains were potentially explainable by the 4× FLOPs increase—a confound the original paper did not control for. The implicit question hanging over that work was: "Does multi-head routing actually help, or did you just build a more expensive model and call it better?" This paper answers that question cleanly: the multi-head mechanism provides genuine performance improvements at matched FLOPs, with MH-MoE (head=3) achieving 10.51 perplexity vs. 10.90 for standard SMoE on RedPajama at 100K steps (Table 1). The ablation in Table 4 closes the loop by showing that adding head/merge projections to standard SMoE without the compensating expert adjustments yields only marginal gains (11.87 → 11.84), confirming that the projections alone don't help—the complete configuration matters.

The work redirects attention from gating-mechanism design to representation-space engineering for routing. The dominant prior approach to improving MoE routing has been to make the gating function more sophisticated: load-balancing losses, auxiliary objectives, learned routing strategies beyond simple top-k, stochastic routing, and so on. MH-MoE takes the opposite approach: keep the gating function simple (standard top-k on a linear projection) and instead engineer the input representation that the gate operates on. The head projection creates h distinct representational subspaces, each of which can route the token to different experts based on different features. This is a fundamentally different axis for improving expert utilization, and the finding that it works—combined with the FLOPs-parity implementation that makes it practical—opens a design space that was previously underexplored. Future MoE research may shift from asking "how do we build a smarter router?" toward asking "in what representational space should the router operate?"

The paper also establishes, by example, an implicit methodological standard: architectural modifications to MoE layers should be evaluated at matched FLOPs (not just matched activated parameters), and any added computational overhead should be accounted for through compensating adjustments to the expert configuration. While the concept of FLOPs-matched comparison is not new to the field broadly, applying it systematically to MoE architectural variants—and providing the analytical tools to do so—raises the bar for future MoE architecture papers. A new gating mechanism that improves performance by 2% but adds 10% to the per-layer FLOPs budget would now need to argue why that tradeoff is worthwhile rather than simply reporting the accuracy gain.

One critical thing this paper does not do is demonstrate that MH-MoE's benefits scale to model sizes where MoE architectures are actually deployed. The experiments use 12-layer, 768-dimensional models—roughly GPT-2 Small scale. Mixtral 8×7B, DeepSeekMoE, and other production MoE models operate at scales 50-100× larger in activated parameters. Whether the multi-head mechanism's perplexity improvements persist, grow, or diminish at those scales is entirely unknown. The paper therefore changes the conversation about MH-MoE (from "can it work?" to "how much does it help, and at what scale?") but does not itself provide deployment-relevant evidence. A skeptical reader could reasonably conclude that MH-MoE is now proven viable at research-prototype scale but remains unvalidated for production use.

Follow-Up Research This Work Enables

Scaling MH-MoE to 1B+ parameter models to determine whether the multi-head benefit persists, saturates, or grows with model size. The paper's experiments use 768-dimensional, 12-layer models where the absolute perplexity improvements are modest (0.39 on RedPajama for 3-head MH-MoE over standard SMoE at 100K steps). The original MH-MoE paper demonstrated gains at similar small scales. The critical open question is whether the mechanism provides proportional or greater benefits at the 7B-70B activated parameter scales where MoE is deployed in practice. A strong follow-up would train MH-MoE and baseline SMoE configurations at matched FLOPs on a 1B-7B scale model (e.g., using the LLaMA or Pythia architecture families) on a standard pre-training corpus (C4, The Pile, or RedPajama at larger scale) and evaluate both perplexity and downstream tasks. The key measurement would be whether the perplexity gap widens or narrows as model dimension d increases—theory suggests the head/merge layer overhead scales as d² while expert computation scales as d × d_mhmoe, which could change the FLOPs-matching tradeoff at larger d and potentially shift the optimal head count.

Isolating the multi-head mechanism from the increased expert count and gating width through a controlled baseline experiment. The paper's MH-MoE (head=3) uses 96 experts with top-3 gating, while the baseline SMoE uses 8 experts with top-1 gating. The architectures differ simultaneously on expert count, gating width, and multi-head splitting. A critical control experiment would construct a non-multi-head MoE with matched expert count and gating width: for example, a standard SMoE with 96 experts and top-3 gating, where FLOPs parity is achieved through a different dimensional adjustment (reducing model dimension d or expert intermediate dimension). If this configuration achieves perplexity close to MH-MoE, then the gains are attributable to increased expert count and gating width rather than multi-head routing per se. If it underperforms MH-MoE substantially, the multi-head mechanism's contribution is isolated. This experiment is conceptually straightforward but was not performed in the current paper, and it would decisively answer the question of whether multi-head routing provides benefits beyond what is achievable through simply having more experts and broader gating in a standard architecture.

Exploring the head count scaling curve—specifically, testing h=4, h=8, and h=16 at matched FLOPs to find the optimal subspace granularity. The paper tests h=2 and h=3, finding that 3 outperforms 2. The original MH-MoE paper used h=4 (though without FLOPs matching). The shape of the scaling curve beyond h=3 is unknown: do gains continue to increase, saturate, or reverse? Each additional head increases the overhead of the head/merge projections (fixed cost) and reduces the per-sub-token dimension d/h, which may eventually make the subspaces too narrow for meaningful expert computation. A systematic sweep of h ∈ {1, 2, 3, 4, 6, 8} at a fixed model dimension (e.g., d=768 or d=1024), with FLOPs-matched expert configurations derived from Equations 8 and 9 for each h, would map out the optimal subspace granularity. The hypothesis that gains saturate or reverse at high h is plausible—splitting a 768-dimensional token into 8 sub-tokens of 96 dimensions each might create subspaces too impoverished for the experts to learn useful specializations—but it needs empirical testing. If the optimal h turns out to be model-dimension-dependent (e.g., h_opt ≈ d / 256), that would be a valuable design principle.

Empirical wall-clock throughput benchmarking to validate the analytical FLOPs parity claim on actual GPU hardware. The paper's central contribution is analytical FLOPs matching through leading-term scalar multiplication counts (Equations 6-9). However, the 3-head MH-MoE with 96 experts operating on 256-dimensional sub-tokens has fundamentally different computational characteristics than the 8-expert baseline operating on 768-dimensional tokens—many more small matrix multiplications versus fewer large ones, different memory access patterns, 12× more gating computation, and different expert load-balancing dynamics. A thorough systems study would measure training tokens-per-second and inference latency on representative GPU hardware (e.g., A100, H100) for each architecture, controlling for implementation quality (e.g., using the same MoE kernel library across all variants). This would reveal whether the analytical FLOPs matching translates to actual throughput parity, or whether MH-MoE incurs hidden efficiency costs from kernel launch overhead, poor GPU utilization on narrow experts, or communication overhead in distributed expert parallelism. Such a study would also inform practical deployment decisions: even if MH-MoE achieves better perplexity at matched analytical FLOPs, a 15-20% wall-clock slowdown might make it less attractive than standard SMoE for latency-sensitive applications. A negative result (MH-MoE significantly slower in practice than analytical FLOPs suggest) would motivate research into more efficient implementations of narrow-expert, many-head MoE layers, or alternative merge strategies that reduce overhead.

Testing whether MH-MoE's benefits transfer to downstream tasks beyond language modeling perplexity. The paper evaluates only pre-training perplexity. The practical value of MH-MoE for practitioners depends on whether the perplexity improvements translate to better performance on tasks they actually care about: few-shot reasoning (ARC, HellaSwag, MMLU), code generation (HumanEval, MBPP), mathematical reasoning (GSM8K, MATH), and instruction following. A strong follow-up would take the best MH-MoE and baseline SMoE configurations from this paper's pre-training setup, continue training to a larger token budget if needed to reach reasonable downstream performance, and evaluate on a standard suite of zero-shot and few-shot tasks. The key question is whether the multi-head mechanism's improved routing—which helps with next-token prediction—also improves the model's ability to perform multi-step reasoning, follow instructions, or retrieve factual knowledge. It is possible that MH-MoE's gains are concentrated in capturing surface-level statistical patterns that dominate perplexity (local syntax, common phrases) but don't help with deeper semantic understanding. A dissociation between perplexity improvement and downstream task improvement would be an important negative result that would refine our understanding of what the multi-head mechanism actually contributes.

Combining MH-MoE with other established MoE improvements—particularly DeepSeekMoE's fine-grained expert segmentation and shared expert isolation—to test for additive or super-additive benefits. The paper already incorporates shared experts (Table 2) as one form of combination, showing additive benefits (MH-MoE + shared expert outperforms either alone). But DeepSeekMoE (Dai et al., 2024) introduced a more refined decomposition: besides shared experts, they split routed experts into finer-grained units and introduce a dedicated knowledge-sharing mechanism. A natural extension would be to apply the MH-MoE multi-head routing to the fine-grained expert decomposition—that is, split tokens into subspaces, route each subspace through a large pool of very small experts, and combine with shared experts. The hypothesis is that multi-head routing and fine-grained expert segmentation provide complementary benefits: fine-graining increases the combinatorial expressiveness of expert combinations, while multi-head routing enables those combinations to be made with awareness of multiple representational subspaces. The risk is that both mechanisms push toward smaller per-expert capacity (fine-graining reduces intermediate dimension to increase expert count; MH-MoE also reduces d_mhmoe to pay for projections), and combining them might produce experts too narrow to learn useful functions. A well-designed experiment would systematically vary expert granularity and multi-head splitting while maintaining FLOPs parity, mapping the frontier of this tradeoff.

Investigating whether the learned head and merge projections in MH-MoE transfer across tasks, domains, or even model architectures. The head projection W_head learns to create representational subspaces that are useful for expert routing. If these subspaces capture general properties of how language should be decomposed for specialized processing, they might transfer across domains (a head projection learned on RedPajama pre-training might improve routing quality on code or math fine-tuning data without retraining) or even across model scales (a head projection from a smaller model might initialize a larger model's head projection). A transfer experiment would pre-train an MH-MoE on general text, extract W_head and W_merge, and test whether using these as initialization (rather than random) for fine-tuning on a specialized domain accelerates convergence or improves final performance compared to random initialization of the projections. If transfer works, it suggests the projections learn something structurally general about optimal subspace decomposition for routing. If it doesn't, it suggests the projections are tightly coupled to the specific data distribution and expert configuration they were trained with, making MH-MoE less flexible than it might appear.

Stress-testing MH-MoE under distribution shift and adversarial inputs to assess whether multi-head routing introduces new failure modes. Standard MoE models are known to sometimes exhibit brittle routing behavior under distribution shift—tokens from out-of-distribution inputs may be routed to inappropriate experts, causing degraded performance. The multi-head mechanism, by routing each token through multiple subspaces simultaneously, might be more robust (because the token has multiple "chances" to be routed correctly across its subspaces) or more fragile (because each subspace receives only partial information about the token, making routing decisions based on less context). A stress-test would evaluate MH-MoE and baseline SMoE on datasets with systematic distribution shifts: different text domains (scientific papers vs. social media vs. code), different languages (low-resource language evaluation), and adversarially perturbed inputs (typos, word substitutions, prompt injections). The measurement would be the relative perplexity degradation under shift compared to in-distribution performance. If MH-MoE degrades less than SMoE, it would suggest the multi-head mechanism provides routing robustness as a side benefit—a practically significant finding that the paper does not address.

Practical Applications and Downstream Use Cases

Pre-training pipelines for organizations building MoE language models from scratch. For teams developing custom MoE LLMs on proprietary data, the MH-MoE implementation provides a drop-in architectural improvement that requires no changes to the training pipeline, data processing, or optimizer configuration. The paper shows a 0.39 perplexity improvement on RedPajama at 100K steps for the 3-head variant over standard SMoE at matched FLOPs and parameters. If this improvement holds at production scales (1B+ activated parameters), it translates to better language modeling quality for the same training and inference budget. The practical adoption path is straightforward: take an existing MoE training codebase, replace the MoE layer implementation with the MH-MoE variant described in Section 2, adjust d_mhmoe and E_mhmoe according to Equations 8 and 9, and proceed with training as before. The paper's finding that the shared expert configuration (Table 2) is compatible with MH-MoE means teams already using DeepSeekMoE-style residual MoE can adopt multi-head routing without abandoning their existing architectural choices.

Quantization-aware deployment of MoE models on edge and consumer devices. The BitNet compatibility results (Table 3) demonstrate that MH-MoE's benefits survive 1-bit quantization, with the 3-head variant achieving 26.47 perplexity vs. 26.78 for standard SMoE on RedPajama at 100K steps under BitNet. While the absolute perplexity values at 768 dimensions are too poor for practical deployment, the architecture ranking is preserved. If this finding scales to larger models where BitNet's degradation is smaller (the original BitNet paper shows much smaller gaps at 1B+ scale), then MH-MoE + BitNet could enable on-device MoE models that combine the parameter efficiency of MoE (many total parameters, few activated per token) with the memory efficiency of 1-bit quantization (dramatically reduced storage for expert weights) and the representational benefits of multi-head routing. The compounding benefit is substantial: MoE already reduces inference FLOPs per token relative to a dense model with the same total parameters; MH-MoE improves the quality-efficiency tradeoff of the MoE architecture itself; and BitNet compresses the result for deployment. Each piece addresses a different dimension of the quality-cost-memory frontier.

Expert specialization analysis and interpretability research. The multi-head mechanism creates a natural framework for understanding what different experts specialize in. In standard MoE, each token is routed to 1-2 experts based on its single-representation gating scores, and interpreting expert specialization requires analyzing which tokens activate which experts. In MH-MoE, a single token can activate different experts through different representational subspaces, and the head projection determines which features of the token are emphasized in each subspace. This creates a richer signal for interpretability research: by analyzing the head projection weights, one can identify which feature dimensions are grouped together into subspaces, and by analyzing which experts are activated through each subspace, one can characterize expert specialization in terms of representational perspectives rather than just token types. A token that routes to a "math expert" through one subspace and a "language expert" through another is revealing that the model has learned to decompose the token's processing along these functional lines. For organizations deploying MoE models and needing to understand their behavior (debugging routing failures, auditing for bias, ensuring coverage across domains), MH-MoE's multi-subspace routing traces provide more diagnostic information than standard MoE's single-routing traces.

Cost-efficient inference for applications with mixed-difficulty query distributions. This application is more speculative but follows from the paper's finding that MH-MoE improves the quality-efficiency tradeoff of MoE architectures. In production settings where LLMs serve queries of varying complexity—some requiring deep reasoning (where expert capacity matters), others requiring broad knowledge (where routing diversity matters)—MH-MoE's ability to route different subspaces of the same input to different experts means that a single query can simultaneously engage specialists in multiple domains. A medical query like "What are the drug interactions between Lisinopril and potassium supplements, and what should patients watch for?" might route one subspace to a pharmacology expert and another to a patient-communication expert within the same MH-MoE layer, combining specialized knowledge that in a standard MoE would require separate tokens or layers to integrate. While this paper provides no direct evidence for this application (there is no domain-specialization experiment), the architectural capability is implicit in the multi-head mechanism, and a deployment team with domain-labeled expert routing data could test whether MH-MoE improves answer quality for cross-domain queries without increasing inference cost (since FLOPs are matched to baseline).