Part-Of-Speech Sensitivity of Routers in Mixture of Experts Models

We measure expert specialization at each layer $l$ by first ordering experts based on how many tokens they handle for each specific POS. We then calculate the proportion of tokens handled by the top-$k$ experts of this ranking:

where, for a given POS, $T_{\text{POS},\text{top-}k}$ is the number of tokens processed by the top-$k$ experts that handle the most tokens and $T_{\text{POS},\text{all}}$ represents the overall number of tokens for this POS. The top-$k$ experts are selected per layer, highlighting layer-specific routing dynamics.

The average specialization score for each POS across all layers is given by:

where $L$ is the number of layers.

The model’s global specialization score is:

where $13$ is the number of POS categories, excluding SYM and X which account for less than 1% of the total, reflecting overall specialization consistency across all layers and POS categories¹¹1See Section 4.1 for the tagset being used..

We compare these values to the expected token routing percentages, denoted as $(\mathcal{U})$, which represent the proportion of tokens that top-$k$ experts would handle under a uniform distribution. The difference between $Spec$ and $\mathcal{U}$, denoted as $\Delta\;\mathcal{U}$, quantifies the deviation from this expectation.

For each token we extract the $k$ routed experts at each layer to create a "path" of the token among the different experts of the model.

From this, we train a Multi-Layer Perceptron (MLP)²²2We used the MLPClassifier of the Scikit-Learn library with the default parameters except for ”max_iter” which has been increased to ensure convergence. to predict the POS of each token based on the expert routing information from each layer.

The input to the MLP is the flattened vector $\mathbf{r}\in\mathbb{R}^{S}$, ($S$ being 64 for Mixtral) which represents the list of experts chosen for the token at each layer. The target is a one-hot encoded POS vector $\mathbf{y}$ with $15$ possible classes. The MLP outputs a predicted vector $\hat{\mathbf{y}}$ that represents the probabilities for each POS tag. We evaluate the routers’ effectiveness by assessing the MLP’s performance in predicting the correct POS on the test data, based on the expert routing paths.

For the POS data, we use 5000 random English sentences from OntoNotes 5.0 Weischedel et al. (2013), corresponding to 116,379 tokens according to the Mixtral tokenizer. This corpus has the advantage of including other types of annotations aligned with the POS, which can be used to expand the analysis we conduct here. We cast³³3We follow the conversion table given on the UD website the POS tags from Penn Treebank POS tags to Universal Dependencies (UD) tags, grouping them into broader categories such as all nouns, verbs, etc.

Tokens are linked to their word’s POS by assigning each token the POS tag of the word it belongs to. For instance, in the case of "humanities," both tokens "human_" and "_ities" are tagged as NOUN.

The distribution of the proportion of POS in our subset is shown in Appendix A.

We split the data into a training set and a test set for the MLP training experiments described in Section 3.2 using a two-thirds/one-third split at the token level, allowing different tokens of the same word to appear in different sets.

We compared 6 models in our experiments: dbrx-base (132B parameters, 36B active, 40 layers, 4 among 16 experts per layer https://huggingface.co/databricks/dbrx-base), Mixtral-8x7B-v0.1 (46.7B parameters, 13B active, 32 layers, 2 among 8 experts per layer Jiang et al., 2024), Phi-3.5-MoE-instruct (41.9B parameters, 6.6B active, 32 layers, 2 among 16 experts per layer Abdin et al., 2024), deepseek-moe-16b-base (16.4B parameters, 2.8B active, 28 layers, 6 among 64 experts per layer, plus 2 shared experts Dai et al., 2024), Qwen1.5-MoE-A2.7B (14.3B parameters, 2.7B active, 24 layers, 4 among 60 experts per layer Team, 2024), and OLMoE-1B-7B (7B parameters, 1B active, 16 layers, 8 among 64 experts per layer Muennighoff et al., 2024).

Among the 6 tested models, only OLMoE-1B-7B can be considered fully reproducible, as it is the only model with a documented and accessible training dataset alongside open-source code for both training and inference.

All models were loaded in fp8, and when a base version (i.e. without preference tuning) of the model existed, this was used. Experiments on Mixtral confirmed that this did not alter prediction results. Tests with full precision, half precision, and with/without the instruct version yielded nearly identical expert routing.

If we look only at expert specialization, i.e. the percentage of tokens the top-$k$ experts retrieve from a given POS layer by layer, we observe a specialization. This approach ensures that we account for not only the top-1 logits but also consider the routing to multiple experts, reflecting the overall expert distribution. Table 1 presents the overall scores for each model. All models exhibit a higher degree of specialization than expected with a uniform distribution, with $Spec_{POS}$ ranging from +25.21% for Mixtral to +36.32% for OlMoE, compared to the expected percentages shown in parentheses.

These are average scores, with higher specialization observed at specific layers. Table 2 highlights the maximum specialization $\max_{l}(Spec_{\text{POS},l})$ for 4 POS on the Phi-3.5-MoE-instruct model. Results for all POS and models are provided in Appendix C.

Table 3 summarizes the KL divergence statistics across models. dbrx-base shows the lowest $\mu$ Mean (0.21), while Qwen shows the highest $\mu$ Max (1.92), suggesting stronger expert specialization in certain layers. Notably, OLMoE maintains a fairly balanced profile with low $\mu$ Min (0.04) but higher $\mu$ Max (1.52), reflecting both shared and distinct expert behaviors (see Appendix D for detailed KL at each expert and layer). However, KL divergence is less interpretable than the other metrics due to its unbounded scale $(0,\infty)$ and the difficulty in defining what represents a small or large divergence.

The results from MLPs trained with paths from different models are summarized in Table 1. All models show an accuracy between 0.79 and 0.88, except OlMoE, regardless of whether full routing information or just the top-expert information from the gating network is used. A simple baseline, predicting the most common POS for a word form, achieves an accuracy of 0.91. However, predicting POS from token-level routing paths is more challenging, as it relies solely on how the token was routed, without any explicit information about the token’s form. Despite this, MLPs still perform well, suggesting that the router effectively captures syntactic information.

Examining the confusion matrices (see Appendix B for all matrices) on the test set reveals which POS types are better predicted, indicating stronger clustering in the paths and experts used. This analysis highlights connections between categories; for example, symbols (SYM) are often confused with punctuation (PUNCT) or numbers (NUM), and adverbs (ADV) and adjectives (ADJ) with nouns (NOUN) or verbs (VERB). Notably, the matrix for OLMoE is more distinct, showing significantly greater confusion across classes such as SYM, PUNCT, and ADJ and an overall lower average accuracy compared to the other models.