Towards Interpretability Without Sacrifice: Faithful Dense Layer Decomposition with Mixture of Decoders

We are grateful to the Reviewer JDFb for praising the idea to move to layer-level sparsity as interesting, and highlighting the stability of MxDs to varying levels of sparsity. We address each of the raised concerns and questions below:

[W1] Advantages for interpretability

”Since the ultimate goal of this line of research is to enable better interpretation of MLPs, this approach does not particularly provide an advancement”

We respectfully disagree with the reviewer--throughout the abstract [L6] and introduction [L56:L65], we argue that the benefit of MxDs lies in overcoming the accuracy trade-off plaguing sparse layers. MxDs consequently do advance interpretability, by removing the large cost to model performance incurred when adopting sparse layers in place of original model components [1]. Two reviewers highlight this key contribution:

• Reviewer-zUPk praises that MxDs help(s) address a very core issue in interpretability work where sparse approximations greatly underperform, constituting a potentially viable replacement for the original component.
• Reviewer-9H1M: tackles the core trade-off between performance and interpretability and central challenge in AI safety [...] making the research highly relevant and important

Would any additional experiment or metric convince the reviewer about the value of the proposed method? We are open to additional experiments if the reviewer has a specific hypothesis to test.

[W2 & Q2] Shared experts

We respectfully clarify that "A key claim of the paper is that MxD improves reconstruction faithfulness under high sparsity via expert specialization" is not the claim we make in our submission. In [L76, L145] we attribute the reconstruction faithfulness specifically to the full-rankness of the weights that MxD provides (independent of specialization).

Please see Sect. C.2 and Fig. 21-22 in the appendix, where we show that a single shared expert can be avoided if desired by training with a randomized top-K selection, at very small cost to reconstruction loss. For the same GPT2 $K=64$ model, we tabulate below the Frobenius norms of differences between the original model’s decoder $\mathbf{D}^*$, our learnt base decoder $\mathbf{D}$, and the “shared expert” weight matrix \mathbf{W}_\text{shared}=\mathbf{D}\text{diag}\(\mathbf{c}\_{\text{shared}}) (with index $\text{shared}=19772$), all of which are matrices of shape $H\times O$. As can be seen, their difference is significant, showing that the shared expert does not simply learn the original decoder:

We are grateful to the reviewer for the chance to address this further, however--we indeed use no explicit loss discouraging degenerate solutions, finding it sufficient to rely on the random weight initialization alone. We will expand on our discussion of this on [L327] in the limitations as follows:

"Secondly, MoEs are prone to issues of expert imbalance [71] or collapse [88]. Through random weight initialiazation alone, we find MxDs in our experiments learn diverse experts without collapsing to the base decoder. However, we expect standard MoE load-balancing [2] or diversity losses [3] to be useful for MxDs should one need more explicit ways of encouraging expert diversity, or when training new interpretable MxD architectures end-to-end."

[W3] Fig. 1 presentation

We appreciate the constructive feedback about Figure 1--we will indeed use $H$ in place of $N$ to denote the number of specialized units in the centre subfigure in the revised manuscript. Would the reviewer have any more specific suggestions on how to improve readability?

[Q1] Can the authors clarify how MxD can be used as a better interpretation tool by future research?

By showing the feasibility of alternative layer designs that largely overcome the sparsity-accuracy trade-off, we argued that MxDs are an important step towards interpretable layers capable of replacing dense MLPs of base models, distilling specialized computations directly into the new model. One immediate application for future research that we are excited about is always-on feature monitoring. With the updated base model computation consisting of tens of thousands of specialized experts, safety-relevant features could be tracked, ablated, and manually tweaked at inference time without any additional post-hoc cost.

Secondly, MxDs provide more promising evidence that the dominating Mixture of Experts (MoE) paradigm [4] can yield interpretable subcomputations [5,6,7]. Given that a large number of new foundation models are MoE models [8,9], we hope the study in the MxD paper might advocate for the research community to focus its efforts on interpretable units of computation in existing foundation models, and/or design new interpretable-by-design architectures for training end-to-end.

---

• [1]: Elhage, et al., "Softmax Linear Units", Transformer Circuits Thread, 2022.
• [2]: Wang, Lean, et al. "Auxiliary-loss-free load balancing strategy for mixture-of-experts." arXiv preprint arXiv:2408.15664 (2024).
• [3]: Liu, Boan, et al. "Diversifying the mixture-of-experts representation for language models with orthogonal optimizer." arXiv preprint arXiv:2310.09762 (2023).
• [4]: Shazeer, Noam, et al. "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer." arXiv preprint arXiv:1701.06538 (2017).
• [5]: Park, Jungwoo, et al. "Monet: Mixture of monosemantic experts for transformers." ICLR 2025.
• [6]: Oldfield, James, et al. "Multilinear mixture of experts: Scalable expert specialization through factorization." NeurIPS 2024.
• [7]: Yang, Xingyi, et al. "Mixture of experts made intrinsically interpretable." ICML 2025.
• [8]: Meta AI. “The Llama 4 herd: The beginning of a new era of natively multimodal AI innovation.” 2025.
• [9]: Guo, Daya, et al. "Deepseek-r1: Incentivizing reasoning capability in LLMs via reinforcement learning." arXiv preprint arXiv:2501.12948 (2025).

---

We welcome Reviewer JDFb to ask any additional questions that might raise their support of the paper--we will gladly answer them.

We are grateful to Reviewer-8ypg for acknowledging the robustness of MxDs to varying levels of sparsity and the strengths of the theoretical results provided. We address each of the concerns below:

[W1 & L3] Predictive performance vs faithfulness

We inherit the term “faithfulness” from prior work [1], which uses the term synonymously with predictive ability, e.g.:

"which seek to faithfully approximate a densely activating MLP layer with a wider, sparsely-activating MLP layer"

Furthermore, the reconstruction loss in [1] is named a “faithfulness loss”. We respectfully disagree with the reviewer's claim in the limitations that meeting this requirement does not guarantee that new model will learn the same representation as that learned by the original: through the reconstruction loss, the layer mapping is indeed trained to reconstruct the same MLP output vector $\mathbb{y}=\text{MLP}(\mathbf{x})\in\mathbb{R}^O$ as closely as possible; we are not training on the final downstream cross-entropy loss.

We are open to specific suggestions that the reviewer might have for how to refer to this property instead if they strongly believe in this enough to depart from the terminology of prior work.

[W2 & Q4 & L5] Adopted evaluations & LLM as a judge

We acknowledge that no evaluation in the paper is strong evidence in isolation. However, each experiment is designed to complement the other--probing evaluates features for which we have labels (frequently used to evaluate prior work [2,3]), whilst evaluations of steering with judge LLMs allow for a scalable evaluation of features for which we do not have labels. Taken as a whole, we argue that the experiments do substantiate the claims we make in the paper; that "MxDs learn similarly specialized features" on [L15].

We highlight that using LLM-as-a-judge is the standard evaluation of steering in a recent benchmarking suite AxBench [6] and in past work [7]. We will cite these works in the feature steering section to better ground our evaluations in standard protocol, and we will acknowledge their potential limitations in the conclusion:

Revised limitations [L330]

With regards to feature evaluations, our steering experiments rely on LLMs as judges. Whilst this is commonplace in recent popular steering benchmarks [6], there is mixed evidence emerging about the reliability of all base models [4,5]. Our experiments attempt to circumvent issues by reporting scores across two capable SOTA models, but we caution inferring too much into the absolute values of reported scores (rather than the relative performance between sparse layers).

[W3 & L4] Appendix results

Many additional results are delegated to the appendix due to the page limit. However, we note that the results in Fig. 10 specifically referenced by the reviewer showcase the same experiments as those in Fig. 3 (with cross-entropy loss used in [1,8]) in the main paper, only with an alternative metric (reconstruction loss), and thus are supplementary.

With the additional page available in the camera-ready revision, we will include results from Sect. B.1 in the main paper.

[Q1] Training complexity

We highlight that the forward pass cost of the proposed MxD layer is not significantly more complex--when parameter-matched to sparse MLP layers as in our experiments, each MxD forward pass has only an additional $O/2$ FLOPs.

Parameter count & FLOPs We first tabulate the theoretical parameter counts and FLOPs for MxDs vs Sparse MLPs below. To be consistent with the popular PyTorch library fvcore [9] we count one fused multiply-add as one FLOP, and the Hadamard product between two $d$-dimensional vectors as requiring $d/2$ FLOPs:

$I,O$ denote the input and output dimensions respectively, $H^*$ the original models’ hidden layer, $N$ the MxD expert count, and $H$ the sparse MLPs' width.

Empirical benchmarks We next run benchmarks for a Sparse MLP layer vs MxD with a batch size of 512, and dimensions: $I=H^*=O=1024$, and number of experts/features as $N=8192, H=9216$.

[Q2] Scalability to larger models

We argue that the computational benchmarks above provide evidence that MxDs will be no more difficult to scale than sparse MLP layers, with very similar FLOPs and latency costs. For initial empirical evidence of further scalability, we perform three additional large-scale experiments on the 27B Gemma2 model below, where we see MxDs continue to outperform on the sparsity-accuracy frontier:

Gemma2-27B, Layer 20, K=32, 50k iterations. To fit in memory, the model is loaded in half precision, and we use 1/8th the batch size and # buffers stored, and 1/2 the multiplier on the # of latent features (with values of 4, 16, and 16, respectively).

We further note that the MoE architecture itself [4] has been scaled successfully to hundreds of billions of parameters in the pre-training regime [10,11].

Revised claims We will revise the claims in our discussions of scalability as follows (based on the computational benchmarks performed), to highlight that we expect MxDs to be no more difficult to scale than sparse MLPs:

[L196]:

[...] we expect MxDs to scale just as well *as sparse MLPs* to models with tens of billions of parameters or more.

[L321]:

[...] Whilst we fully anticipate this trend to *scale just as well as with sparse MLPs* in even larger models, our experiments only provide direct evidence for LLMs with up to 3B parameters, given our limited resources.

[Q3] Relevance of citations [17,19]

We agree that these studies are only high-level motivation: [17] shows humans prefer explanations with $\leq 32$ concepts, suggesting smaller K may aid interpretability, while [19] provides complementary quantitative metrics. To avoid over-reliance on these priors, we'll remove their appeal from Sect. 3, citing them briefly in the introduction alone.

[L1 & L2] Additional related work

We are grateful to the reviewer for pointing out interesting related work [12, 13, 14]; the latter two of which appearing on ArXiv a month before the NeurIPS submission deadline. We will include a discussion of how they relate to the proposed method in the revised manuscript in addition to acknowledging alternatives to interpretability-via-sparsity.

Concretely, a new section following the discussion of sparse decompositions will read:

Interpretable layers: Whilst MxDs follow the paradigm of specialization-via-sparsity, a number of promising interpretable mechanisms have been recently proposed that do not rely on sparsity constraints. For example, recent work explore bilinear layers [12], compositional structure [13], and/or tensor networks [14] as alternative ways of designing interpretable architectures. Notably, [12] show how the Hadamard product between two linear transformations of the same input vector allows direct interpretation (such as analyzing input-output feature interactions). Whilst MxDs’ yields a similar Hadamard product forward pass, it involves two different input vectors, and establishes a theoretical equivalence to mixture of experts and conditional computation when sparsity is present in one of the operands (i.e., the expert coefficients).

---

References

• [1]: Dunefsky et al. "Transcoders find interpretable LLM feature circuits." NeurIPS 2024*.*
• [2]: Gao et al. "Scaling and evaluating sparse autoencoders." ICLR 2025*.*
• [3]: Gurnee, et al. "Finding neurons in a haystack: Case studies with sparse probing." TMLR 2023.
• [4]: Zheng et al. "Judging llm-as-a-judge with mt-bench and chatbot arena." NeurIPS 2023.
• [5]: Bavaresco et al. "Llms instead of human judges? a large scale empirical study across 20 NLP evaluation tasks." ACL 2025.
• [6]: Wu et al. "Axbench: Steering LLMs? Even simple baselines outperform sparse autoencoders." ICML 2025.
• [7]: Cao et al. "Personalized steering of large language models: Versatile steering vectors through bi-directional preference optimization." NeurIPS 2024.
• [8]: Gonçalo et al. "Transcoders beat sparse autoencoders for interpretability." arXiv 2025.
• [9]: fvcore (FAIR), detectron2 library.
• [10]: Meta AI. “The Llama 4 herd: The beginning of a new era of natively multimodal AI innovation.” 2025.
• [11]: Guo et al. "Deepseek-r1: Incentivizing reasoning capability in LLMs via reinforcement lerning." arXiv 2025.
• [12]: Pearce et al. "Bilinear MLPs Enable Weight-based Mechanistic Interpretability" ICLR 2025.
• [13]: Dooms et al. "Compositionality Unlocks Deep Interpretable Models" AAAI 2025 ColoRAI Workshop
• [14]: Adler et al. "Towards Combinatorial Interpretability of Neural Computation" arXiv 2025.

---

Might Reviewer 8ypg have any additional questions? If all concerns were addressed, we would otherwise be grateful if the reviewer raised their score.

We are grateful to the Reviewer zUPk for highlighting the “clear experimental evidence for [the paper’s] claims” and that the paper “address[es] a very core issue in interpretability”. Furthermore, we are glad the reviewer shares our view that the Hadamard factorization might even be “extensible to other applications outside of this paper”. We address each of the concerns below:

[W1 & Q1] Equivalence between experts/features

We thank the reviewer for the chance to further clarify the relationship between sparse MLP features and sparse expert coefficients. In terms of functional form, there is no fundamental difference between the two as feature extractors: both Sparse MLPs and MxDs learn features by a matrix multiplication followed by a TopK gating function, with $\mathbf{z}=`TopK`(\mathbf{E}^\top\mathbf{x})$ and $\mathbf{a}=`TopK`(\mathbf{G}^\top\mathbf{x})$ (for the Sparse MLP hidden units and MxD expert coefficients respectively). For this reason, we find them to learn a similarly high level of interpretable features.

The key difference between Sparse MLPs and MxDs is in how the sparse features are constructed to contribute to the output. Each non-zero neuron in Sparse MLPs adds a static column vector from the overcomplete dictionary basis to the output, through $\mathbf{y}=\sum_h \mathbf{d}_h z_h$ as in in Eq. 2. In contrast, each non-zero coefficient in MxDs’ expert coefficients adds a linear transformation of the input to the output: $\mathbf{y}=\sum_n a_n\mathbf{W}_n^\top\mathbf{x}$, as per Eq. 3; through features-as-linear-transformations, MxDs have greater capacity to combine the sparse features to well-reconstruct the MLP mapping.

We hope this clarifies the similarities between the existing neuron- and layer-level sparsity approaches: the key difference lies on the decoder side, not in how the sparse features/experts are selected.

[W2 & Q5] Hadamard trick & experimental details

We kindly refer the reviewer to Appendix A.3, where an alternative presentation of the method is given in terms of tensor methods, which might add additional insight. Furthermore, we highlight the Jupyter notebook linked to from [L162] (submitted in the anonymous github repository at the time of submission), which walks through various equivalent model forms in pytorch.

To enhance understanding even further, we will condense some of the content in this tutorial-style notebook into the appendix of the revised manuscript, inline.

We agree with the reviewer that it would be useful to explain experimental details inherited from past work (e.g. the Top-K activation) inline. We will add a detailed description of the TopK activation function and the reconstruction-style training setup to the appendix.

[Q2] Efficiency of MxDs vs Sparse MLPs

We thank the reviewer for the chance to compare the computational overhead and efficiency of the layers. We first report theoretical layer FLOPs, and then report empirical benchmarks. We then use the results here to address the following question about scalability.

Parameter count & FLOPs We first show the theoretical parameter counts and inference-time FLOPs for MxDs vs Sparse MLPs below. To be consistent with the popular PyTorch library fvcore [1] we count one fused multiply-add as one FLOP, and the Hadamard product between two $d$-dimensional vectors as requiring $d/2$ FLOPs:

$I,O$ denote the input and output dimensions respectively, $H^*$ the original models’ hidden layer, $N$ the MxD expert count, and $H$ the width of the sparse MLPs.

For a chosen expert count $N$, we set the width of Sparse MLPs (e.g. Transcoders) to $H:=N+H^*$ to parameter-match the models.

Empirical benchmarks We next run benchmarks for a Sparse MLP layer vs MxD with a batch size of 512, and dimensions: $I=H^*=O=1024$, and number of experts/features as $N=8192, H=9216$.

Both layers have a ReLU, followed by a TopK activation (for the sparse hidden layer and expert gating). MxDs have an additional ReLU and final Hadamard product--however, we see from the table’s results that the additional time and compute costs are minimal.

We thank the reviewer again for the chance to present the computational comparisons, and will include all tables here in the revised manuscript.

[Q3] Justification for scalability to larger models

We thank the reviewer for raising this concern--we will revise our claims that we expect MxDs to scale just as well as sparse MLPs (please see the bottom of this response to Q3), based on the similar computational costs of layers established above.

For strong initial evidence of scalability to even larger models, we perform three additional large-scale experiments on the 27B Gemma2 model below (in half-precision with a smaller batch size), where we see MxDs continue to outperform on the sparsity-accuracy frontier:

Normalized MSE on Gemma2-27B, Layer 20, K=32, after 50k iterations. To fit in memory, the model is loaded in half precision, and we use 1/8th the default batch size and number of buffers stored in memory, and 1/2 the multiplier on the number of latent features (with values of 4, 16, and 16, respectively).

We further note that the Mixture of Experts architecture itself [2] has been scaled successfully to hundreds of billions of parameters in the pre-training regime (e.g. 400B [3] and 685B [4] parameters, respectively).

Revised claims We will revise the claims in our discussions of scalability as follows (based on the computational benchmarks performed, showing that MxDs have similar computational costs), to highlight that we expect MxDs to be no more difficult to scale than sparse MLPs:

[L196]:

Whilst we do not have the computational resources to show similarly thorough experiments on even larger models, we expect MxDs to scale just as well *as sparse MLPs* to models with tens of billions of parameters or more.

[L321]:

Our experiments show MxDs outperform on the sparsity-accuracy frontier on 4 diverse LLMs. Whilst we fully anticipate this trend to *scale just as well as with sparse MLPs* in even larger models, our experiments only provide direct evidence for LLMs with up to 3B parameters, given our limited resources.

[Q4] Training details of sparse layers

We indeed follow past work [5,6] by training just the newly initialized layers to reconstruct the mapping from the MLP’s input to its output (we do not train end-to-end on the next-token cross-entropy loss for example).

We will add this clarification around [L199]:

we train sparse layers to minimize the normalized reconstruction loss between its output and that of the original MLP layer when fed the OpenWebText dataset. We use objectives of the form [...] , where f(.) denotes the various learnable sparse MLP layers. This follows the protocol of past work [1,2], where the new sparse layers' parameters alone are trained on the direct output of the MLP.

[Q5] Confusing order of equations

Our intention with Sect. 2.3 discussing the rank properties appearing before Sect. 2.4 with the factorized forward pass is to immediately connect the high-level motivation for moving to the layer-wise approach to the technical low-level benefits (the rank preservation). Whilst Sect. 2.4 is a crucial implementation detail, we felt that the flow of the presentation worked best with the current ordering of subsections for this reason. However, if the reviewer feels strongly that the re-ordering of these two subsections significantly aids clarity, we are more than happy to make this change in the revised manuscript.

---

References:

• [1]: fvcore (FAIR), detectron2 library.
• [2]: Shazeer et al. "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer." ICLR 2017.
• [3]: Meta AI. “The Llama 4 herd: The beginning of a new era of natively multimodal AI innovation.” 2025.
• [4]: Guo et al. "Deepseek-r1: Incentivizing reasoning capability in LLMs via reinforcement learning." arXiv 2025.
• [5]: Dunefsky et al. "Transcoders find interpretable LLM feature circuits." NeurIPS 2024.
• [6]: Gonçalo et al. "Transcoders beat sparse autoencoders for interpretability." arXiv 2025.

---

Would Reviewer zUPk have any further questions? We’d be happy to address them.

Dear Reviewer-zUPk,

We are grateful for your suggestions and positive support of the paper! If you later think of any additional ways we might further improve the paper's clarity, please do not hesitate to let us know, and we will aim to incorporate them into the camera-ready accordingly.

Sincerely,

Authors of Submission6888

We are grateful to the Reviewer 9H1M for praising many aspects of the paper; addressing a “critical, high-impact problem”, the “clear conceptual innovation” and “strong theoretical and empirical backing”. We address the concerns below:

[W1] Computational comparisons

We thank the reviewer for raising this point. In summary, there is minimal differences between layer computational cost. We first report theoretical layer FLOPs, and then report empirical benchmarks below.

Parameter count & FLOPs We first tabulate the theoretical parameter counts and inference-time FLOPs for MxDs vs Sparse MLPs below. To be consistent with the popular PyTorch library fvcore [1] we count one fused multiply-add as one FLOP, and the Hadamard product between two $d$-dimensional vectors as requiring $d/2$ FLOPs:

$I,O$ denote the input and output dimensions respectively, $H^*$ the original models’ hidden layer, $N$ the MxD expert count, and $H$ the width of the sparse MLPs.

For a chosen expert count $N$, we set the width of Sparse MLPs (e.g. Transcoders) to $H:=N+H^*$ to parameter-match the models.

Empirical benchmarks We next run benchmarks for a Sparse MLP layer vs MxD with a batch size of 512, and dimensions: $I=H^*=O=1024$, and number of experts/features as $N=8192, H=9216$.

Both layers have a ReLU, followed by a TopK activation (for the sparse hidden layer and expert gating). MxDs have an additional ReLU and final Hadamard product--however, we see from the table’s results that the additional time and compute costs are minimal.

We thank the reviewer again for the chance to present the computational comparisons of MxDs vs sparse MLPs, and will include all tables here in the camera-ready version of the paper.

[W2] Scalability to frontier models

We are grateful to the reviewer for acknowledging our experiments on 3B parameter models as respectable. We highlight that the 3B parameter base model used is already larger than any of those used by the papers of the baseline methods [2,3].

As we acknowledge on [L323] in the limitations, however, our experiments only provide direct evidence for MxDs’ ability on base models up to the 3B parameters range due to limited resources. However, we perform three additional large-scale experiments on the 27B Gemma2 model below (in half-precision with a smaller batch size), where we see MxDs continue to outperform on the sparsity-accuracy frontier--providing strong initial evidence of scalability to even larger models:

Normalized MSE on Gemma2-27B, Layer 20, K=32, after 50k iterations. To fit in memory, the model is loaded in half precision, and we use 1/8th the default batch size and number of buffers stored in memory, and 1/2 the multiplier on the number of latent features (with values of 4, 16, and 16, respectively).

We further note that the Mixture of Experts architecture itself [4] has been scaled successfully to hundreds of billions of parameters in the pre-training regime (e.g. 400B [5] and 685B [6] parameters, respectively).

Revised claims We will revise the claims in our discussions of scalability as follows (based on the computational benchmarks performed), to highlight that we expect MxDs to be no more difficult to scale than sparse MLPs:

[L196]:

Whilst we do not have the computational resources to show similarly thorough experiments on even larger models, we expect MxDs to scale just as well *as sparse MLPs* to models with tens of billions of parameters or more.

[L321]:

Our experiments show MxDs outperform on the sparsity-accuracy frontier on 4 diverse LLMs. Whilst we fully anticipate this trend to *scale just as well as with sparse MLPs* in even larger models, our experiments only provide direct evidence for LLMs with up to 3B parameters, given our limited resources.

---

References

• [1]: fvcore (FAIR), detectron2 library.
• [2]: Dunefsky, Jacob, Philippe Chlenski, and Neel Nanda. "Transcoders find interpretable LLM feature circuits." NeurIPS 2024.
• [3]: Paulo et al. "Transcoders beat sparse autoencoders for interpretability." arXiv 2025.
• [4]: Shazeer, Noam, et al. "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer." ICLR 2017.
• [5]: Meta AI. “The Llama 4 herd: The beginning of a new era of natively multimodal AI innovation.” 2025.
• [6]: Guo, Daya, et al. "Deepseek-r1: Incentivizing reasoning capability in LLMs via reinforcement learning." arXiv 2025.

---

If Reviewer 9H1M has any additional questions, we would be glad to answer them.