MoE Models — An Introductory Overview

To provide some background, the Mixtures of Experts framework was introduced in the 1991 paper titled "Adaptive Mixtures of Local Experts," co-authored by the renowned Geoffrey Hinton. The fundamental concept of MoE involves predicting an output ( y ) based on an input ( x ) by aggregating multiple "experts" ( E ), with their contributions managed by a "gating network" ( G ).

The gating network ( G ) employs a straightforward linear model:

Here, ( W ) is a learnable matrix that allocates training examples to specific experts. The learning objectives when training MoE models are twofold:

1. Experts learn to process their assigned inputs into optimal outputs (predictions).
2. The gating network learns to appropriately route training examples to the correct experts, effectively mastering the routing matrix ( W ).

MoE is particularly effective when computation is executed over just the single expert with the highest gating value. Thus, we can approximate ( y ) as follows:

This technique, known as “hard routing” or “sparse gating,” has been pivotal in breakthroughs like the Switch Transformer, allowing models to scale with a computational complexity of ( O(1) ).

With this context, let's explore some specific applications and the learning outcomes for the experts.

MoE in Vowel Discrimination

To gain further insights into the specialization of experts, we can reference the original 1991 MoE paper, which provides valuable clues. In this work, the authors developed a 4-expert MoE model to tackle a vowel discrimination task, distinguishing between sounds like [A] and [a] as well as [I] and [i] in voice recordings.

The following figure illustrates their data (i and I in the top left, a and A in the bottom right) based on formant values (acoustic features that characterize the sound of the vowel):

Interpreting this figure: - The plotted “points” represent the data: i, I, a, and A (the readability is slightly affected due to the age of the paper). - The lines labeled “Net 0,” “Net 1,” and “Net 2” depict decision boundaries learned by three of the four experts. The fourth expert, however, did not learn a useful decision boundary. - The line “Gate 0:2” illustrates the gating mechanism's decision boundary between directing inputs to expert 0 (to the left) and expert 2 (to the right).

In this case, Expert 1 focused on distinguishing [i] from [I], while Experts 0 and 2 specialized in differentiating [a] from [A], likely due to the greater complexity in classifying that data compared to [i] and [I].

The key takeaway here is that the gate effectively clusters the data, while experts define decision boundaries within each cluster. More challenging data regions tend to attract a greater number of experts, though some may contribute minimally.

MoE in Translation

Next, we examine another example that illustrates expert specialization effectively. This instance originates from the 2017 paper “Outrageously Large Neural Networks,” also from Hinton’s lab at Google Brain.

In this research, MoE was applied to the natural language processing challenge of translating sentences from English to French. The authors incorporated an MoE layer featuring 2048 experts between two LSTM modules, hypothesizing that different experts would specialize in varying types of inputs.

The following table presents the leading input tokens, organized by gate value, for three of the 2048 experts:

Once again, we observe a clustering pattern: Expert 381 concentrates on a word cluster encompassing terms like “research,” “innovation,” and “science,” while Expert 2004 focuses on words such as “rapid,” “static,” and “fast.”

Interestingly, there is also an expert, Expert 752, that does not seem to add significant value, as it solely examines the token “a.”

This behavior is intriguing: the model was not explicitly instructed to associate these words or to form clusters, nor were experts pre-assigned to specific words. This emerging behavior arose solely from the predetermined number of experts and the defined learning objective.

MoE in Synthetic Data

Finally, let's discuss a recent paper that greatly enhanced our understanding of the inner workings of the MoE layer, titled “Towards Understanding the Mixture-of-Experts Layer in Deep Learning,” by researchers Zixiang Chen et al. from UCLA.

In this study, a straightforward 4-expert MoE model was utilized on a synthetic dataset with four clusters of data points belonging to two classes. The objective was to differentiate these classes across all clusters. The experts in this model comprised 2-layer CNNs equipped with either linear or non-linear (cubic) activation functions.

The following visualization illustrates the training process, comparing the MoE model with non-linear activation at the top and linear activation at the bottom. This graph displays the data points representing the two classes (crosses and circles), the routing of each data point by the gate (yellow, blue, green, red), and the learned decision boundaries of the model (jagged lines).

Key insights include: 1. Specialization takes time: Initially, there is no specialization; all experts are dispersed. As training progresses, certain clusters become associated with specific experts. 2. Random expert allocation: There is no systematic assignment of clusters to experts; the process is random. For instance, the top-right cluster may have more data points routed to the “blue” expert, potentially leading to the entire cluster being identified as blue. 3. Non-linearity outperforms linearity: Linear experts underperform compared to non-linear ones, as evidenced by the comparison between the non-linear and linear plots. The decision boundaries from linear experts are less effective, and the clusters are not as distinctly separated, highlighting that expert non-linearity is crucial for MoE success.

Moreover, it is insightful to monitor the “dispatch entropy” of the gate, which peaks when each expert receives training examples from all clusters and reaches a low point (0) when each expert gets examples from a single cluster. As training progresses (from left to right in the plot below), dispatch entropy decreases until it stabilizes, indicating a near-1:1 correspondence between clusters and experts:

This again emphasizes the narrative: the gate learns to categorize the data into clusters, with experts specializing in their randomly assigned clusters.

Three studies, spanning three decades, converge on a single narrative.

Takeaways

With this comprehensive context, let's revisit and address the earlier posed questions: - Q: When is MoE effective?

A: MoE excels when data naturally clusters, as observed in vowel discrimination, translation tasks, and synthetic datasets.

• Q: Why doesn’t the gate send all training examples to the same expert?

  A: Such an approach would yield poor performance; a single expert cannot effectively learn the decision boundaries for all clusters, resulting in significant gradients that push the model in the opposite direction.

• Q: Why doesn’t the model collapse into a state where all experts are identical?

  A: Again, this would lead to suboptimal performance; varied expert specialization across different data regions enhances overall performance.

• Q: How do the experts specialize, and in what areas?

  A: Experts specialize in various regions of the data, with specialization being random, influenced by the gate's (random) initialization.

• Q: What does the gate learn?

  A: The gate learns to categorize the data into clusters and assigns each cluster to one or more experts.

MoE continues to be one of the most intriguing and practically valuable modeling approaches in machine learning, and we are just beginning to realize its implications for modern applications. Comprehending the internal workings is a vital step toward enhancing their efficacy.

Curious to deepen your knowledge about the latest advancements in ML? Subscribe to my Newsletter.