Table of Contents
Fetching ...

EMoE: Eigenbasis-Guided Routing for Mixture-of-Experts

Anzhe Cheng, Shukai Duan, Shixuan Li, Chenzhong Yin, Mingxi Cheng, Shahin Nazarian, Paul Thompson, Paul Bogdan

TL;DR

MoE models in vision face load imbalance and expert homogenization, which throttle efficiency and specialization. EMoE introduces an eigenbasis-guided router that projects patch features onto the top-$r$ principal directions via a learnable $U in mathbb{R}^{D\times r}$ (with $r \ll D$) and gates to eight expert MLPs without a load-balancing loss, using an orthonormal regularizer to stabilize $U$; the class token bypasses routing. This geometrically principled routing eliminates conflicting auxiliary objectives while promoting balanced utilization and diverse specialization. Empirically, EMoE achieves competitive ImageNet accuracy with improved load balance, and demonstrates strong few-shot representations as well as beneficial brain-age predictions on 3D MRI data, indicating robust applicability to high-dimensional, heterogeneous domains.

Abstract

The relentless scaling of deep learning models has led to unsustainable computational demands, positioning Mixture-of-Experts (MoE) architectures as a promising path towards greater efficiency. However, MoE models are plagued by two fundamental challenges: 1) a load imbalance problem known as the``rich get richer" phenomenon, where a few experts are over-utilized, and 2) an expert homogeneity problem, where experts learn redundant representations, negating their purpose. Current solutions typically employ an auxiliary load-balancing loss that, while mitigating imbalance, often exacerbates homogeneity by enforcing uniform routing at the expense of specialization. To resolve this, we introduce the Eigen-Mixture-of-Experts (EMoE), a novel architecture that leverages a routing mechanism based on a learned orthonormal eigenbasis. EMoE projects input tokens onto this shared eigenbasis and routes them based on their alignment with the principal components of the feature space. This principled, geometric partitioning of data intrinsically promotes both balanced expert utilization and the development of diverse, specialized experts, all without the need for a conflicting auxiliary loss function. Our code is publicly available at https://github.com/Belis0811/EMoE.

EMoE: Eigenbasis-Guided Routing for Mixture-of-Experts

TL;DR

MoE models in vision face load imbalance and expert homogenization, which throttle efficiency and specialization. EMoE introduces an eigenbasis-guided router that projects patch features onto the top- principal directions via a learnable (with ) and gates to eight expert MLPs without a load-balancing loss, using an orthonormal regularizer to stabilize ; the class token bypasses routing. This geometrically principled routing eliminates conflicting auxiliary objectives while promoting balanced utilization and diverse specialization. Empirically, EMoE achieves competitive ImageNet accuracy with improved load balance, and demonstrates strong few-shot representations as well as beneficial brain-age predictions on 3D MRI data, indicating robust applicability to high-dimensional, heterogeneous domains.

Abstract

The relentless scaling of deep learning models has led to unsustainable computational demands, positioning Mixture-of-Experts (MoE) architectures as a promising path towards greater efficiency. However, MoE models are plagued by two fundamental challenges: 1) a load imbalance problem known as the``rich get richer" phenomenon, where a few experts are over-utilized, and 2) an expert homogeneity problem, where experts learn redundant representations, negating their purpose. Current solutions typically employ an auxiliary load-balancing loss that, while mitigating imbalance, often exacerbates homogeneity by enforcing uniform routing at the expense of specialization. To resolve this, we introduce the Eigen-Mixture-of-Experts (EMoE), a novel architecture that leverages a routing mechanism based on a learned orthonormal eigenbasis. EMoE projects input tokens onto this shared eigenbasis and routes them based on their alignment with the principal components of the feature space. This principled, geometric partitioning of data intrinsically promotes both balanced expert utilization and the development of diverse, specialized experts, all without the need for a conflicting auxiliary loss function. Our code is publicly available at https://github.com/Belis0811/EMoE.
Paper Structure (4 sections, 3 equations, 3 figures, 2 tables)

This paper contains 4 sections, 3 equations, 3 figures, 2 tables.

Figures (3)

  • Figure 1: Overview of EMoE. (a) Training schematic: the eigenbasis is recomputed from routers with frozen ViT weight (b) EMoE-augmented ViT block: input tokens convert into eigenbasis tokens, then pass through LayerNorm and multi-head self-attention, followed by a MoE feed-forward layer.
  • Figure 2: one-shot classification for the proposed model on CIFAR10. We reported $\pm$2% std.
  • Figure 3: Expert–class routing heatmaps across (a)CIFAR and (b)ImageNet for EMoE. Each panel shows average routed tokens per class (x) and expert (y).