Advancing Expert Specialization for Better MoE

TL;DR

This paper introduces a simple method to improve expert specialization in Mixture-of-Experts models by adding orthogonality and variance losses, leading to significant performance gains without architectural changes.

Contribution

It proposes two new objectives that enhance expert specialization in MoE models, addressing issues caused by auxiliary load balancing loss.

Findings

Up to 23.79% performance improvement on benchmarks

Enhanced expert specialization demonstrated across models

Maintains load balancing without extra architecture

Abstract

Mixture-of-Experts (MoE) models enable efficient scaling of large language models (LLMs) by activating only a subset of experts per input. However, we observe that the commonly used auxiliary load balancing loss often leads to expert overlap and overly uniform routing, which hinders expert specialization and degrades overall performance during post-training. To address this, we propose a simple yet effective solution that introduces two complementary objectives: (1) an orthogonality loss to encourage experts to process distinct types of tokens, and (2) a variance loss to encourage more discriminative routing decisions. Gradient-level analysis demonstrates that these objectives are compatible with the existing auxiliary loss and contribute to optimizing the training process. Experimental results over various model architectures and across multiple benchmarks show that our method significantly enhances expert specialization. Notably, our method improves classic MoE baselines with auxiliary loss by up to 23.79%, while also maintaining load balancing in downstream tasks, without any architectural modifications or additional components. We will release our code to contribute to the community.