Multilinear Mixture of Experts

• Multilinear Mixture of Experts (MMoE) is a neural module that distributes computations across tensor-decomposed experts with sparse entmax gating.
• It employs various tensor factorization methods (CP, Tucker, Tensor Train) to achieve scalability and efficient inference in high-dimensional applications.
• The framework enables expert specialization and model editing, improving performance in tasks like vision benchmarks and bias correction.

A Multilinear Mixture of Experts (MMoE) is a neural module in which the computation of unified outputs is distributed across a parameterized set of “experts,” with the aggregation and selection of these expert contributions performed via multilinear mappings and differentiated, sparse gating mechanisms. MMoE architectures generalize classical Mixture of Experts (MoE) models by arranging the weight structure as a high-order tensor and employing tensor factorization for scalability, especially suitable for vision tasks and other high-dimensional applications. The MMoE framework enables training and deployment of models with a large—potentially exponential—number of implicit experts, while keeping inference- and memory-cost tractable by never instantiating the full expert ensemble directly (Oldfield et al., 19 Feb 2024).

1. Definition and Structure of MMoE Layers

In the standard dense Mixture of Experts, an input $z\in\mathbb R^I$ is processed by $N$ expert weight matrices $W_{:,:,n}\in\mathbb R^{O\times I}$, each weighted by a gating coefficient $a_n$:

$$y = \sum_{n=1}^N a_n(W_{:,:,n}z), \quad a = \mathrm{softmax}(G^\top z).$$

The MMoE generalizes this by collecting all $N$ expert matrices into a single $(E+2)$-way tensor:

$$\mathcal W \in \mathbb R^{O\times I\times N_1\times\cdots\times N_E},$$

where $E$ is the number of expert hierarchy levels. The output is computed via a sequence of mode-$n$ tensor-matrix products:

$$y = \mathcal W \times_2 z \times_3 a_1 \times_4 \cdots \times_{E+2} a_E \in \mathbb R^O,$$

with each gating vector $a_e=\phi(G_e^\top z)\in\mathbb R^{N_e}$ and $\phi$ typically chosen as $\text{entmax}_{1.5}$ for sparse, convex combinations. This nested structure enables hierarchical expert aggregation.

For $E=1$, this structure recovers the usual MoE forward pass but allows for factorizations that avoid materializing the full tensor $\mathcal W$ for large $N$.

2. Multilinear Tensor Factorizations for Scalability

The principal innovation of MMoE is to represent $\mathcal W$ in a compressed format using tensor decompositions, drastically reducing parameter count and computation:

• CP-decomposition (CPMMoE):

$$\mathcal W = \sum_{r=1}^R g_r^{(1)}\circ g_r^{(2)}\circ g_r^{(3)},\quad g_r^{(1)}\in \mathbb R^O,\,g_r^{(2)}\in\mathbb R^I,\,g_r^{(3)}\in\mathbb R^N$$

$$y = \sum_{r=1}^R g_r^{(1)} \cdot ( G^{(2)\top}z )_r \cdot ( G^{(3)\top}a )_r$$

with $O\bigl(R(I+N+O)\bigr)$ parameters/FLOPs.

• Tucker Decomposition (TuckerMMoE):

$$\mathcal W = \mathcal Z \times_1 G^{(1)} \times_2 G^{(2)} \times_3 G^{(3)}$$

$$y = \bigl( \mathcal Z \times_2 (G^{(2)\top}z) \times_3 (G^{(3)\top}a)\bigr) \times_1 G^{(1)}$$

with total cost $O(R_II + R_NN + R_OO + R_OR_IR_N)$.

• Tensor Train / Tensor Ring (TTMMoE/TRMMoE):

$$\mathcal W(o,i,n) = \mathrm{tr}[ \mathcal G_1(:,o,:) \mathcal G_2(:,i,:) \mathcal G_3(:,n,:) ]$$

with similar scaling, supporting very large $N$.

All formats generalize to multiple expert (hierarchy) modes. The factorization enables practical deployment of models with tens of thousands of experts, achieving monosemanticity (expert specialization), and eliminating the need to instantiate all expert matrices explicitly (Oldfield et al., 19 Feb 2024).

3. Expert Routing and Gating Mechanisms

MMoE replaces conventional top-$k$ or discrete softmax gating with fully differentiable, sparse “soft” gating based on entmax:

$$a_e = \phi(G_e^\top z),\quad \phi = \mathrm{entmax}_{1.5}$$

The entmax operator yields sparse convex mixtures, permitting direct backpropagation through both gating and factorized parameter tensors. Batch normalization is applied to $G_e^\top z$ prior to entmax, encouraging balanced expert utilization and reducing “dead” experts. This mechanism obviates additional load-balancing losses and retains efficient gradient-based optimization (Oldfield et al., 19 Feb 2024).

4. Inference and Training Complexity

Parameter and FLOP requirements are sharply reduced relative to dense or sparse MoE. For a CPMMoE with $I=O=768$, $N=512$, $R=512$:

• Naive dense MoE: $\approx 393\times 10^6$ FLOPs/sample
• CPMMoE: $\approx 1.1\times 10^6$ FLOPs/sample

Hierarchical CPMMoE (e.g., $E=4$, $N=8,192$) requires $\sim$1.08M parameters versus $\sim$6.3B for the equivalent dense model, with minimal accuracy loss (Oldfield et al., 19 Feb 2024).

Training employs AdamW, learning-rate warmup, cosine decay, and $\ell_2$-normalized factor initializations. Factorized MMoE layers fit as plug-in replacements for linear heads in vision backbones (CLIP ViT-B/32, DINO ViT-S/16) with immediate scalability and minimal change in accuracy.

5. Empirical Results and Visualizations

Qualitative and quantitative evidence supports substantial expert specialization as the number of experts increases:

• With $N=32$, experts cover polysemantic clusters (e.g., “gators” + “limos” + “quilt”).
• At $N=256$ or $2048$, most experts attend to one or two closely related object classes (“balloon” images alone), assessed via causal counterfactual intervention on expert weight slices: $$p^{(n)} = \|\,d^{(n)} - e_{\arg\max_c d^{(n)}_c} \|_2$$ Mean class-level polysemanticity $p^{(n)}$ decreases monotonically as $N$ increases.

On seven vision benchmarks, MMoE layers (CP, Tucker, TR) typically match or outperform standard and high-rank linear heads, with performance gains of $0.5$–$1.0\%$ in accuracy. An MMoE-augmented MLP-Mixer (all MLP blocks replaced by CPMMoE, $e=128$, $\text{CP-rank}\approx3.25\times$ channel-dim) increases ImageNet validation accuracy from $50.75\%$ to $59.85\%$ for similar parameter counts (Oldfield et al., 19 Feb 2024).

6. Applications: Conditional Bias Correction and Model Editing

The fine-grained control of expert activation enables direct intervention on model behavior:

• To mitigate subpopulation bias in attribute classification (CelebA), collect the average gate vector $\bar a$ on a specific subpopulation. The logit for head $o$ is then modified: $\tilde y_o = y_o + \lambda\,\bar a^\top a$ This reweights the output in favor of experts most engaged on the subpopulation. Expert-conditional bias correction via this formula improves fairness metrics (equality of opportunity, standard-deviation bias, max–min fairness, Model Rewriting Score) over unconditional thresholding, oversampling, and adversarial debiasing, with negligible loss in general accuracy (Oldfield et al., 19 Feb 2024).

7. Comparison With Related Models

MMoE recovers the standard linear MoE as a special case when no factorization is used. By employing tensor factorization and differentiable soft gating, MMoE avoids the excessive FLOP/memory cost of “soft” MoE (dense routing), and sidesteps the nondifferentiable selection and training instabilities associated with “sparse” MoE (top-$k$ routing). MMoE also enables a continuum between full soft and purely sparse gating regimes, with empirical ablations indicating entmax gating yields sparser and more specialized experts relative to softmax (Oldfield et al., 19 Feb 2024).

These architectural properties make MMoE particularly well-suited as a final layer or as a global replacement for standard MLPs in transformer-style architectures. The capacity to scale to hierarchically organized expert ensembles, while supporting efficient optimization and inference, marks the MMoE framework as a significant advance in modular, interpretable, and specialized neural computation.