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

TL;DR

EMoE introduces an eigenbasis-guided routing mechanism for Mixture-of-Experts models, promoting balanced expert utilization and diversity without auxiliary loss, thereby addressing key limitations of existing approaches.

Contribution

The paper proposes EMoE, a novel architecture using a learned eigenbasis for routing, which inherently balances load and encourages expert specialization without auxiliary losses.

Findings

Promotes balanced expert utilization

Encourages diverse, specialized experts

Eliminates need for auxiliary load-balancing loss

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.