Mixture of Experts Made Intrinsically Interpretable

We sincerely appreciate Reviewer tfxC's suggestions. We have carefully incorporated them into the revised manuscript.

>>> Q1Related Work on MoE

>>> A1We truly appreciate R-tfxC for bringing the ReMoE [1] paper on sparse MoE to our attention. We will definitely include a citation in the revised version.

While both work utilize the properties of ReLU in MoE, there are key differences:

1. Purpose of study: ReMoE aims to make the routing function fully differentiable, while our focus is on enhancing the interpretability of the MoE.
2. Position where ReLU used: In our work, ReLU is applied within each expert, whereas ReMoE uses it in the gating function.

>>> Q2Feature Streering Experiments

>>> A2We truly appreciate the suggestion. While we cannot perform a full qualitative comparison within the rebuttal period, we provide an initial analysis by applying feature steering to MoE-X and evaluating the results.

Love − Hate Steering

Building on the Activation Addition method [A], we demonstrate how feature steering can influence model completions. Using MoE-X-S, we first extract feature vectors for the contrasting concepts "Love" and "Hate".

These vectors are then used to compute a steering adjustment v_Steered = v_Unsteered + coefficient * (v_"Love" − v_"Hate"). We apply the steering at MoE hidden feature at layer 6 with a coefficient of 5.

The table below compares unmodified model completions (left) with steered ones (right). The bolded text represents the prompt.

This demonstrates how targeted feature modifications can also steer model behavior in our MoE-X.

[A] Turner, Alexander Matt, et al. "Activation addition: Steering language models without optimization." arXiv e-prints (2023): arXiv-2308.

>>> Q3Interpretability and Performance

>>> A3The review is absolutely right -- there is indeed a trade-off between interpretability and performance. If we enforce extreme sparsity (e.g., pushing the $l_0$​ norm to zero), the model would have no activations and thus fail to learn.

To address this, we use ReLU-based sparsity rather than hard constraints like $l_1$​ regularization (as in SAEs). This allows the model to adaptively learn sparse patterns while still maintaining strong performance. As a result, our approach preserves interpretability without sacrificing the model's ability to fit the data effectively.

We sincerely thank R-nXvb for their valuable questions!

>>> Q1Evaluation Metrics

>>> A1The reviewer noted that metrics for interpretability is perplexity only. We appreciate the feedback, but would like to clarify two critical points

1. Misconception About Perplexity: Perplexity is not an interpretability metric; it measures language modeling accuracy.
2. Interpretability Metrics: Our work evaluates interpretability using three distinct metrics
   • Coverage/Reconstruction Score (Table 1): Measures the alignment between the hidden neurons and human-defined chess states.
   • Auto-Interp Score (Figure 8): Assesses the accuracy of predicting activations for unseen samples using extracted explanations.
   • Case Studies (Figure 7): Provide qualitative validation of interpretability. These metrics are commonly used in previous research.

Additionally, both R-K7C2 and R-tfxC acknowledged that our evaluation is reasonable. Based on this, we believe our evaluation is sufficient.

>>> Q2Comparing with Proto-lm [A]

>>> A2Really thank the reviewer for bring this great work! We will cite Proto-lm in our revision. While both methods share similar motivation, key differences make direct comparison challenging:

1. Base Models: Proto-lm focuses on encoder-only models (e.g., BERT), while MoE-X is designed for decoder-only causal LLMs like GPT.
2. Tasks: Proto-lm targets text classification (e.g., SST2), whereas MoE-X is for general language generation. So metrics like Simulatability cannot be directly applied to our model.
3. Layers. Proto-lm explains only the final layer's embeddings, while MoE-X provides interpretability across all layers.
4. Interpretability Categorization. Proto-lm uses prototype-based explanations for local, per-sample explaination. On the other hand, MoE-X focuses on mechanical interpretability, offering a global explanation of how the network functions by identifying the meaning of individual neuron.

>>> Q3User Study

>>> A3Thanks for the suggestion! As requested, we conduct a blinded human evaluation to assess neuron monosemanticity in MoE-X compared to GPT-2 and GPT-2 + SAE.

Specifically, 5 raters evaluate 20 random features from each model. Based on the top 5 activating samples, raters classified each feature as: Yes (clearly monosemantic), Maybe (partially interpretable), No (not monosemantic).

Results showed MoE-X’s features were most interpretable, with 70% (14/20) labeled as Yes.

>>> Q4Faithfulness and Simulatability Score

>>> A4As noted in Q2, directly comparing Faithfulness and Simulatability is not feasible due to differences in model, task, and setup. However, as suggested, we retrained the model and defined new evaluation metrics to make this comparison possible.

Specifically, we fine-tune MoE-X-Medium on SST2 using the last token for classification. We report the interpretability on the final layer. For Proto-lm, scores are extracted from the raw figure in its paper, which may have minor discrepancies. This provides reasonable, though not exact, comparability for interpretability analysis.

Faithfulness

To evaluate faithfulness, we define two metrics inspired by [A].

• Comprehensiveness (Comp): Measures the decrease in model confidence when top $k$% activated neurons are removed.
• Sufficiency (Suff): Measures the change in confidence when only top $k$% activated neurons are preserved.

We do evaluations on the SST2, with with $k\in$ {1,5,10, 20, 50}.

Results: The table shows Comp and Suff scores. Notably, for $k\geq 5$%, MoE-X achieves perfect scores (Comp=100%, Suff=0%) due to its extreme sparsity: MoE-X activates <100 neurons out of 8,192 in a layer. This ensures faithful explanations, outperforming Proto-lm.

Simulatability

After training MoE-X on SST-2, we identify each neuron's concept and select the top 5 activated concepts for each sample. Following the evaluation setup in [A], 3 evaluators categorize 50 SST-2 questions based on these concepts.

Results: MoE-X performs slightly worse than Proto-LM but outperforms other methods in [A]. Note that differences in question selection and human judgment make results not directly comparable to [A].

Special Note: The detection score in paper is a Simulatability score. It predicts activation patterns (not classes) using an LLM, not humans.

[A] "Proto-lm: A prototypical network-based framework for built-in interpretability in large language models." EMNLP (2023).

We truly thank the R-K7C2 for the nice comments.

>>> Q1Additional Application Scenarios

>>> A1We sincerely appreciate the suggestion! Expanding to the medical domain is valuable, but a key challenge is finding a benchmark that evaluates both interpretability and performance. While many medical datasets focus on performance, few assess interpretability, highlighting the need for a dedicated benchmark.

As a first step, we tested the small MoE-X model (8×124M) on MMLU’s medical subset using 5-shot prompting, comparing it to a similarly sized dense model (GPT-Neo 125M).

While MoE-X-S performs good, we’d need more research to evaluate how interpretable it is in specialized domains.

>>> Q2Human Evaluation

>>> A2Thanks for the suggestion! We conducted a blinded human evaluation to assess neuron monosemanticity in MoE-X compared to baseline models (GPT-2 and GPT-2 with SAE).

Following GatedSAE’s approach, 5 raters evaluated 20 random features from each model, presented in random order. Each rater evaluated 12 features (20 × 3 / 5). For each feature, raters were shown the top 5 activating samples. Raters classified each feature as: Yes (clearly monosemantic), Maybe (partially interpretable), No (not monosemantic).

MoE-X’s features were rated as ​most interpretable, with 70% (14/20) labeled Yes. While this is a small-scale study, it provides a promising preliminary insight.

>>> Q3Validate the Gaussian assumption

>>> A3Very luckily, a recent arXiv paper [A] has conducted a formal study on investigates the distribution of weights in LLMs. The authors find that the weights of all major open-source LLMs closely follow a Gaussian distribution, including LLaMA, Qwen, and the Vicuna family. They provide both statistical evidence and theoretical explanations for this phenomenon. We will cite this work in the revised version!

[A] Unveiling the Mystery of Weight in Large Foundation Models: Gaussian Distribution Never Fades https://arxiv.org/abs/2501.10661

>>> Q4MoE Feature Steering

>>> A4We truly appreciate the suggestion. While we cannot perform a full qualitative comparison within the rebuttal period, we provide an initial analysis by applying feature steering to MoE-X and include the results below.

Love − Hate Steering

Building on the Activation Addition method [B], we demonstrate how feature steering can influence model completions. Using MoE-X-S, we first extract feature vectors for the contrasting concepts "Love" and "Hate".

These vectors are then used to compute a steering adjustment v_Steered = v_Unsteered + coefficient * (v_"Love" − v_"Hate"). We apply the steering at MoE hidden feature at layer 6 with a coefficient of 5.

The table below compares unmodified model completions (left) with steered ones (right). The bolded text represents the prompt.

This demonstrates how feature addition can also steer model behavior in our MoE-X.

[B] Turner, Alexander Matt, et al. "Activation addition: Steering language models without optimization." arXiv e-prints (2023): arXiv-2308.

>>> Q5Formatting issue

>>> A5We sincerely apologize for the formatting and line number issues! We will make sure to correct them in the revision.

>>> Q6Impact of number of experts

>>> A6Thanks for the insightful question! As suggested, we explore two ways to increase the number of experts, and both show promising improvements:

1. Activating More Experts (with a fixed total number): In fact, this has been shown in Fig 5 in the paper, which we varying the number of activated experts $k \in \{1,2,4\}$. It shows that increasing the number of activated experts leads to better interpretability.

2. Expanding Total Experts (keeping the number of activated experts constant): We increased the total experts from 8 to 16. This improved interpretability. However, the improvements were modest compared to increasing activated experts, likely because the per-inference cost remained the same.

We sincerely appreciate R-a7ZR's thoughtful comments and suggestions.

>>> Q1Evaluation Setup and Additional Experiments

>>> A1We truly appreciate the suggestion! As R-a7ZR mentioned, our primary focus is on interpretable model design. In line with this goal, we use small-scale datasets, guided by 3 key factors:

1. Comparison with Prior Work: We follow the same evaluation protocol used in prior studies (SAE on GPT-2 Small) for fair comparison.
2. Interpretability Evaluation: We need a benchmark that evaluates both interpretability and performance. Small, controlled dataset like chess is good for this.
3. Computational Constraints: Our resources are limited to an 8×48G GPU server, which supports prototyping but not large-scale experiments.

Large-scale Evaluations

While large-scale training isn't feasible within the rebuttal period, we follow the reviewer’s suggestion and evaluate our trained checkpoints on additional standard few-shot benchmarks. The results are shown in the table below.

Though preliminary (due to the model’s small size), our MoE-X model performs comparably to the Switch Transformer. We plan to scale up the model in future work.

>>> Q2Fair Comparison between GPT-2 and MoE-X

>>> A2Thank R-a7ZR for this thoughtful suggestion! We completely agree that a parameter-matched comparison is critical for fairness.

As suggested, we re-train MoE-X Small with 1 activated expert (124M active) to match GPT-2 Small’s size. To speed up training, we fine-tune the old top-2 MoE-X, but activating only 1 expert per token. The training took ~8 hours. This top-1 MoE-X (124M) has the same active parameters as GPT-2 (124M).

Apples-to-Apples Comparison: We report performance using perplexity. As shown below, MoE-X-S (1 expert active) outperforms GPT-2 Small across all datasets while using ​identical active parameters. But as expected not as good as the MoE-X-S (2 from 8×124M) with around 50% more parameters.

​Interpretability Gains Persist: We evaluate interpretability using the Detection Score defined in the paper. Still, with the same activated parameters, MoE-X-S (1 from 8×124M) demonstrates a clear improvement in identifying more accurate concepts compared to GPT2-Small.

These results highlight that MoE-X’s advantages stem not just from scale but from its structured sparse design. We deeply appreciate the suggestion and will include in the revised paper!

>>> Q3Switch Transformer in Figure 8

>>> A3Great suggestion! As suggested, we incorperate Switch Transformer into Figure 8. We re-run the auto-interp experiment for Switch-S and compare with MoE-X-S using the average Detection Score.

While both models perform well, MoE-X-S consistently outperforms the standard MoE in interpretability across all quantiles.

>>> Q4SAE applied to MoEs

>>> A4That’s really an interesting direction! SAE can indeed be applied to MoEs, including MoE-X.

Due to time constraints, we conduct a quick experiment by training SAE on MoE-X and Switch Transformer using the Chess dataset for post-training interpretability. The SAE is placed at post-res layer 6 with a hidden size of 4096.

Three key observations:

1. SAE improves interpretability for MoE but increases validation loss.
2. Switch + SAE slightly outperforms MoE-X alone. But MoE-X + SAE fights back.
3. Improvements are modest due to SAE's hidden size matching MoE's channel width.

>>> Q5Typos and formating

>>> A5We sincerely thank the reviewer for the careful proofreading! We will correct the typos in the references and ensure proper formatting for terminology in the revision.