Towards Interpretability Without Sacrifice: Faithful Dense Layer Decomposition with Mixture of Decoders
Authors: James Oldfield
Shawn Im
Sharon Li
Mihalis Nicolaou
Ioannis Patras
Grigorios Chrysos
Venue: NeurIPS 2025 poster ()
Keywords: interpretability
mixture of experts
Primary area: social_and_economic_aspects_of_machine_learning

Abstract: Multilayer perceptrons (MLPs) are an integral part of large language models, yet their dense representations render them difficult to understand, edit, and steer. Recent methods learn interpretable approximations via neuron-level sparsity, yet fail to faithfully reconstruct the original mapping--significantly increasing model's next-token cross-entropy loss. In this paper, we advocate for moving to layer-level sparsity to overcome the accuracy trade-off in sparse layer approximation. Under this paradigm, we introduce Mixture of Decoders (MxDs). MxDs generalize MLPs and Gated Linear Units, expanding pre-trained dense layers into tens of thousands of specialized sublayers. Through a flexible form of tensor factorization, each sparsely activating MxD sublayer implements a linear transformation with full-rank weights--preserving the original decoders' expressive capacity even under heavy sparsity. Experimentally, we show that MxDs significantly outperform state-of-the-art methods (e.g., Transcoders) on the sparsity-accuracy frontier in language models with up to 3B parameters. Further evaluations on sparse probing and feature steering demonstrate that MxDs learn similarly specialized features of natural language--opening up a promising new avenue for designing interpretable yet faithful decompositions. Our code is included at: https://github.com/james-oldfield/MxD.

Towards Interpretability Without Sacrifice: Faithful Dense Layer Decomposition with Mixture of Decoders
James Oldfieldm..0
Shawn Im
Sharon Li
Mihalis A. Nicolaou
Toannis Patras
Grigorios G Chrysos
University of Wisconsin-MMdiion Queen Mary University of London The Cyprus Institute
*University of Cyprus
Abstract
Multilayer perceptrons (MLPs) integral part of large language models, yet their dense representations render them difficult to understand, edit, and steer. Recent methods learn interpretable approximations via neuron-level sparsity, yet fail to faithfully reconstruct the original mapping- signiiiggggn increasing model's next-token cross-entropy loss. In this paper, we advocate for moving to layer-level sparsity to overcome the accuracy trade-off in sparse layer approximation. Under this paradigm, we introduce Mixture of Decoders (MxDs). MxDs generalize MLPs and Gated Linear Units, expanding pre-trained dense layers into tens of thousands of specialized sublayers. Through flexible form of tensor factorization, each sparsely activating Mxl) sublayer implements linear transformation with full- rank weights- the original decoders' expressive capacity even under heavy sparsity. Experimentally, we show that MxDs significantly outperform state-of-the-arr methods (c.g., Transcoders) on the sparsity-accuret frontier language models with up to 3B parameters. Further evaluations sparse probing and feature steering demonstrate that MxDs learn similarly specialized features natural
up promising new avenue for designing interpretable yet faithful decompositions. Our code is included at: https:/ / github. com/ ames oldf i eetfdelld///////
Introduction
One strategy for addressing concerns about large language models' (LLMs) 2, 3] behavior is via bottom-up approach to understanding and controlling the network internals-developing models of how and where human-interpretable features are represented in LLMs and how they affect the output [4, 5, 6]. Such a mechanistic understanding has proved helpful for a number of issues relating to safety and transparency, from controlling refusal of harmful requests [7] to detecting generation of unsafe code [6] and latent model knowledge [8].
However, developing models of LLMs' internals faces challenges due to the dense nature of their representations [9, 10]. Indeed, many studies have found that individual neurons MLP layers encode multiple distinct concepts. Rather than numan-interpretable features being neatly aligned with individual neurons, they are often distributed across many [11, 12]. As result, it is not straightforward to cleanly isolate specific concepts of interest in the models' latent token representations.
Traditionally, imposing constraints on model form has offered way to instill more predictable properties or structure. Indeed, there is a rich history of success with constraints in machine learningg from basts- representations through non negaatiivv [13, 14], structure through low-rankness or assumptions on geometry [15, 16). With the particular issues posed by dense representations
Cooresponding author: jimenalexaneerodldieeelllgg..... com. Work done whilst UW-Madison.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

MLPs
Neuran-leval sparsity
Layer-level sparsity
4.
de
C -R
DApA
Unspsciviied ookarra
Special zed columns Weighesd by:
Specialted linear transfomations of dense hidden units
Waihttd by:
Janse liddin undo
sparce hidden units
sparse expert units
Figure I: Units of specialization for sparse layer variants: Neuron-level sparsity of existing sparse MLPs [27, 26] (center) vs layer-level sparsity (right), which the proposed Mixture of Decoders (MxD) layer enables at scale. For GPT2, the dimensions are: O 768, H1 O H 32.
LLMs, specialization through sparsity has re-emerged as a dominating strategy for learning more interpretable representations. With prior work showing that sparser models both aid human explanation [17] and achieve higher scores on L.LM-based auto-interpretability metrics [18, 19], sparsity is often used proxy for interpretability [20, 21J. To this end, many recent works- such as sparse autoencoders [22, 23, 6]-take inspiration from traditional sparse dictionary learning methodologies [24, 25], re-writing pre-trained LLMs' activations as sparse, non-negative linear combinations of atoms in a learned overeomplete basis. However, as argued [26], such approaches do not learn the functional mechanisms of LLMs' layers, and their inherent post-hoe nature demands additional parameters and computation on top of the base models.
One alternative approach is to directly replace layers with more interpretable equivalents [28], such as with wide MLPs with sparsity constraints. Transcoders [27, 29, 30, 26] (TCs) are recent example of this, training new MLPs to mimic the functional behavior of MLPs with sparse hidden units, which have recently been shown to also learn more interpretable features [26]. Thus, instead of relying on external post-hoc analysis, sparse MLP layers offer way distill specialized features directly into the model's forward pass itself.
Both of the above methods for learning specialized features fall into the same category of what one may call 'neuron-level sparsity', Dictionary learning methods restrict the number of non-zero elements used from a learned dictionary, whilst sparse MLPs [27] limit the number of active rows used from learned 'decoder' matrix. heavily trading off accuracy for sparsity, poorly reconstructing the original
core, whilst this constraint is useful for interpretability,
is too restrictive-off model components [31, 28]. We argue that preserving the base models' performance is a crucial component of sparse MLP layer approximations for the following two key reasons:
1. Model faithfulness: sparse layers that poorly approximate the original layers risk missing critical intricacies of the base models' behavior or latent features [32]. Conversely, an accurate reconstruction (yielding similar downstream next-token loss) is some evidence that the combination of newly learned subcomputations faithfully emulates the base model.
2. Practical adoption: sparse layers that closely preserve base models' performance are capable of replacing the existing MLPs, directly integrating specialized computation into the native forward pass. Otherwise, downstream use of the sparse layers' features must run on top the base models' computation. This introduces additional inference-time cost to every forward pass, and restricts any analysis to post-hoc settings.
In this paper, we advocate for moving from neuron-level to layer-level sparsity (as illustrated in Figure 1) to address this. We propose the Mixture of Decoders (MxD) layer to overcome sparsity-accuracy trade-off through scalable, resource-efficient conditional computation. Rather than individual vectors, MxDs learn interpretable sublayers as atomic units of specialization, This faithfully mirrors the functional form of dense layer we wish to approximate, and allows MxDs to readily generalize to modern MLP variants (i.c., the Gated Linear Unit [33]).
At a technical level, MxDs are constructed via flexible tensor factorization [34] with the Hadamard product [35]. Through their parameter efficiency, MxDs scale the number of specialized layers far beyond what is feasible with classic sparse mixture of experts (MoEs) [36], and recover prior adapter- based MoEs [37, 38] as special case. Crucially, we prove that the proposed tensor factorization in MxDs leads to each 'expert' sublayer implementing a linear transformation with full-rank weights- allowing faithful reconstruction even under heavy sparsity. Empirically, we demonstrate that MxDs significantly outperform alternative sparse MLP layers such Transcoders [27] and Skip Transcoders

---

[26] on the accuracy frontier. In addition to their faithfulness, MxDs remain competitive with the SOTA on interpretability metries. Our contributions can be summarized as follows:
We propose Mixture of Decoders, an instance of a flexible class of parameter-effiient MoE through Hadamard product-factorized weight tensors.
We prove that each specialized MxD expert's 's weights inherit up to the same rank as the original MLP's decoder, providing faithful approximation even in very sparse models.
Across 108 sparse layers in LLMs (with up to 3B parameters) MxDs (i) pareto-dominat existing techniques the sparsity-accurar frontier yet (ii) remain competitive on 34 sparse probing and steering tasks, validating the interpretability of the learned experts.
Methodology
We first recall the technical details of language models' MLP layers and existing approaches sparse approximations in Section 2.1. I. We then introduce the proposed MxD in Section 2.2, outlining the attractive rank properties it inherits in Section 2.3 and factorized implementation in Section 2.4. We conclude with extensions to modern MLP layers in Section 2.5.
2.1 Preliminaries
Let x E IR' be the pre-MLP latent representation of a specific token at a given layer. Omitting bias terms throughout for brevity, the GPT2-styke MLP layer produces the output vector y RR as:
MLP(x) = T ERO with z = o
E
(1)
where E E II D E R xo are the learnable 'encoder' and 'decoder' parameters respectively, and <b-) is an activation function, often a GELU [39]. We use denote the weights/dimensions of the pre-trained base LLM.
Sparse approximations One approach to learning interpretable features in MLPs is to train new, wider MLPs with sparse hidden units to reconstruct the original layer's outputs [27, 26, 30, 29], reminiscent of dictionary learning techniques [25]. In general, sparse MLPs share the model form:
SMLP(x) D
Σ
24 d, E RR, with S(E ER.
where S(.) is pparsity-inducing function (such the top- [23] activation used in this paper). Here, the dimensionality sparse MLPs' learnable weights E E R 1x E IR #xO are set as > H* such that the hidden layer is significantly larger than that of the original MLP. The original post-MLP output vectors are approximated as a K-sparse, non-negative linear combination of the rows d, of newly learned decoder matrix. Whilst this model form has been shown to learn interpretable, specialized features 2n, in language models [27, 26], their poor reconstruction is of questionable faithfulness and limits their use as layer replacement in practice.
2.2 Mixture of Decoders
We now detail the proposed Mixture of Decoders (MxD) layer, which overcomes the sparsity- accuracy trade-off by treating sparsely activating linear layers the atomic unit of specialization. We approximate the original MLP with conditional combination N linear transformations:
MxD(x) = Σ an (W E R
where a = S R are sparse 'expert coefficients' from learnable gating matrix E R and z = (E x) E RR is the dense output from an encoder. Here, W E N is a third-order tensor of parameters collating all N experts' decoder weights W(n, W, E # In MxDs, we use large N scale the feature specialization, and set H = H to match the original MI.P's smaller hidden dimension.
3

---

With the gate routing each token to just its top-K experts, each W, n E R receives a gradient signal from only a specific set of semantically similar tokens. This implicit clustering naturally leads experts to specialize in seacure- subcomputations, while collectively covering the layer's full functionality. MxDs in Equation (3) also directly inherit the MI.P layers' original functional form, avoiding the need to impose sparsity and non-negativity constraints on the hidden units z E R# However, MxD decoders naively require prohibitive N HO aarameters-preventing N from scaling to tens of thousands of specialized components. To achieve parameter-efficiennyy whilst retaining layer capacity for faithful layer approximation, we parameterize MxDs' third-order weight tensor W Nx H'O specifically to yield full-rank expert weights, defined elementwise as:
Wnn, hn:) Cn d, RO, VN .I,...,NJ HE H..
(4)
where is the Hadamard product [34, 35], and Cn. d E Bº are the rows of learnable weights C E RN- D E RHx Intuitively, D implements a base transformation modulated by the N specialized units in C. Additional technical motivation for this parameterization with tensor methods can be found in Appendix A.3. This brings MxDs' parameter count down significantly to O-(N + H) from HHO in Equation (3) with N full decoders. One can then vary N parameter-matco sparse MLP layers. We next detail how this design (i) retains expressivity in each unit for faithful layer approximation under sparsity in Section 2.3 and (ii) yields a simple forward pass in Section 2.4.
2.3 MxDs are rank-preserving
In the original LLM, the linear transformation from the hidden units to the output constrained by the rank of the original MLP's decoder matrix D' R x0 Under only mild technical conditions, every expert's weight matrix in MxDs inherits the rank of D E IR: H x thus allowing it to match that of the original MLP's decoder, despite its parameter-effcienng
Lemma 1 (Decoder rank preservation). We can materialize linear expert n's 's weight matrix as W(n,;;) W, D diag (c,) RH Assuming diag (c,,) ROx0 is diagonal matrix with no zeros along its diagonal (and thus invertible), we then have
rank( W n) rank (D diag(c.)) rank(D).
The proof is found in Appendix 1, which first derives the matrix-valued expression for each expert from Equation (4) and then applies standard rank equality. At sparsity level of K, each MxD output vector is a weighted sum of K-many linear transformations (each with potentially full-rank weights) of the dense hidden units z. As a result, MxDs retain layer capacity even under high sparsity. Sparse MLPs' hidden units have only K non-zero elements in contrast-each output in Equation (2) therefore confined
-ddimensional subspace of 12 potentially limiting the capacity of sparse MLPs to faithfully approximate the original mapping in the small K regime desirable for interpretability (mirroring speculations by [26]). Further, whilst alternative soft linear MoEs achieve scalability through low-rankness [40], Lemma states that no such rank constraints are present MxDs. For approximating existing MLP layers where low-rank assumptions may not hold, MxDs are consequently a more suitable class of conditional layer.
2.4 Factorized forward pass
Mixture of Decoders (MxDs)
MxDs compute linear combination of N linear transfor- mations of the dense vector. With the proposed Hadamard- factorized weights, this yields simple implementation.
Confroent wegh
Lemma (Hadamard-factorrieed MoE. forward pass). Let E R and E R.N denote the MLP hidden units and expert coefficients respectively. Further, denote the decoder matrices as C E N D
MrD conditioonal branch: Madiliiee the aupput
Vanile MLP branch: Flecily interis the lawe
prodlct
Figure 2: Mixture of Decoders extends the base MIP/GLU layers with conditional 'expert' branch, modulating the MLP's outputs.
IR xx parameterizing W E NxHxO MxD's 's forward pass can be re-written as:
N
MxD(x)= = -an (W = (CTa) a) (D z).
(5)
The proof is found in Appendix A.2. We include a notebook at https: :/s: htt///gttttutt thub. com/ forer/fommmeeeee cettn--quivaleence ipynb showing the equivalence in Py- Torch. Further, please Appendix A.5 for discussion of how the Hadamard factorization relates

---

Table 1: Model formulations of related work: x E R', E R are the pre- and post-MLP representations respectively, the hidden units, and a the vector of the 'expert coefficients' for MxD. Model-specific encoders/decoders E, map between the hidden units and output.
MLPS
SAEs
Transcoders
Skip Transcoders
MxDs (Our)
[6]]
Model form Sparse component
y
z
y D
y=D'z y=D'z EE
y=
Og
==8 e
to prior parameter-efficien MoEs with element-wise scaling [37], and Appendix B.6 for perfor- mancefcomputational cost comparisons.
2.5 Extending MxDs to GLUs
In contrast to methods imposing neuron-level sparsity [22, 27, 26], MxDs do not make assumptions about the base layer's encoder architecture activation function. As result, MxDs readily generalize to alternative architectures such as the Gated Linear Units (GLUs) [33] used in recent LLMs [1. 2].
Recall that GLUs* hidden units computed ZGLU 7 GL [Ex)) E R. with additional GLU parameters EGLU RIxH H and GLU activation function & (e.g.. Swish [1]. By substituting in
the GLU hidden representations, MxDs straightforwardly extend the GLU model form too:
N MxDGI.uXx) = Σ an WT (EELU)) * E x)
(C a) DT((EEELUS (E (Exx))
n=1
GLU hidden units
where a S R are the expert units, W., D diag(c,) E R as before. For technical discussion of GLUs and their relationship to MxDs, refer readers to Appendix A.4- through the theoretical results developed in this paper, we show that GLU encoders themselves can be viewed mixture of rank-1 linear experts (in contrast to the rank- annk-rreeriving MxDs).
3 Experiments
The experimental section in the main paper is split into two parts. Section 3.1 first demonstrates how MxDs perform significantly better on the accuracy-sparsity frontier as sparse MLP layer ap- proximations on 4 LLMs. We then demonstrate in Section 3.2 that MxD's features retain the same levels of specialization through sparse probing and steering evaluations. Thorough ablation studies, experiments with matrix rank, and comparisons to low rank MoEs presented in Appendix B.
3.1 Sparse approximations of MLPs in LLMs
In this section, we perform experiments approximating LLMs' existing forw-forward layers with sparse MLPs, establishing that MxDs better navigate the sparsity- accuraay frontier, more faithfully approximating the base models' MLPs than the SOTA baseline methods.
Implementation details We train on 4 base models: GPT2-124M [3], Pythi Om, Pythia-1.44 [41], and Llama-3. 2-23B [1] with up to 80k experts/features. We train all sparse layers on total of 480M tokens of Open WebText [42], with learning rate 1e and a context length of 128, initializing the output bias as the empirical mean of the training tokens, and D in MxDs the zero-matrix (following [26]). We vary N in MxD layers to parameter-match Transcoders in all experiments, with parameter counts and dimensions shown in Table 2. For Llama3. 2-3B, we use the Swish-GLU variant of MxD and GELU-MLP MxDs for the other three models, matching the architectures of their base encoders. Through ablation studies in Appendix B.8 we show that MxDs using the GELU/GLU variants are much more accurate layer approximators than the ReLU variants. Full experimental details are included in Appendix D. 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 more.
5

---

Table 2: Sparse layer parameterss/imensionn:: H denotes the size of the layers' hidden units and N is the expert count. MxDs perform almost as many linear transformations as the baselines have features.
GPT2-124K
Pythia-410M
Pythi a-1.4
Llarra-3 2-38
Model
Params
Params
11
Params
11
Parums
N
Transcoders [27]
37.7M
24.376
67.1M
32,768
208.50
65.536
601M
98,301
Skip Transcoders [26)
38.4M 4M
24.376
86.2M
32,768
272. 7M
65,536
611M
86,304
4096
86,015
Pythia-410M
(Layer 8)
(Layer 15)
Mean cross- entropy loss
3.630
TC (3BM
Original
3.345
TC (67M params)
Original
STC (38M params]
STC (68M params)
3.625
MxDs (38M params)
3.340
MxDs (67M params)
3.335
3.615
3.610
3.605
3.600
16
128
256
16
32
128
256
Pythia-1.48
Liama3
(Layer 12)
3,0400
(Layer 12)
Mean Pross-entropy loss
3.120
Original LLM
Original
TC (269M params)
3.0375
TC (604M params)
STC (273M params)
STC (614M params)
MxDs (268M params)
3.0350
MxDs (604M params)
3.0325
3.110
3,0300
3.0275
3.105
3,0250
3.0225
3.100
3.0200
16
32
64
128
256
16
32
64
128
256
Sparsity level K
Sparsity level K
Figure 3: Model cross-entropy loss preserved when replacing MLPs with Transcoders [27], Skip Transcoders [26], and MxDs, as function of the number active units K (hidden neurons/expertss) We highlight that MxDs have consistently lower loss at all levels of sparsity.
Objective fuucctoo Given the frozen weights of the MLP. we train sparse layers minimize the normalized reconstruction loss between its output and that of the original MLP layer with objectives
of the form = Ex
[(x]] where f() denotes the various learnalle sparse MLP layers. [MTP((x]]2
This follows the protocol of past work [27, 26), where the sparse layers' parameters alone trained directly on the output of the MLP. To compare with recent work [26], we adopt the TopK activation function [23] for sparsity-inducho
function S(.), removing the need for additional
sparsity penalty. Please see Appendix A.6 for details on the TopK activation function.
3.1.1 Results: sparsity vs faithfulness
We train exhaustive set of 60 sparse MLP approximations across 4 diverse LLMs with up 3B parameters. We show Figure the resulting downstream base model cross-entropy loss when using the trained sparse layers in place of the original MLPs. As can be not only do the proposed

---

MxD layers outperform Transcoders [27] notably, but model performance is similarly preserved at all sparsity levels in MxD layers. Please also see Figure for results with normalized MSE. where MxDs' reconstruction emors are to order of magnitude smaller. Full results on additional layers are included in Appendix B.3 for 48 more trained sparse layers.
The recent "Skip Transcoders' (STCs) [26], introduce an additional 10 parameters with a skip connection I xx mapping the imput directly the output with y D + S'x STC layers thus have considerably more parameters (e.g., STCs on 1lama3. 2-3B have 10M more parameters than MxDs). Despite the smaller parameter counts, we find MxDs consistently outperform STCs on the apcrsity frontier, attesting to the benefits of MxDs' model form.
3.1.2 Faithfulness in output space
Next, we perform experiments comparing the faithfulness sparse layers as their computation propagates to the model output space. We sample 16 future tokens with the base model and then measure how similar the same generations are when target MLP layer is replaced with the sparse layers. We use 512 text snippets from OpenWebText, and take the first 4 words of each as the initial prompts, generating 16 future tokens after each prompt. We plot in Figures 4a and the percentage of the samples' continuations that are identical in the original LLM and hooked LLMs up to n future tokens ahead. We note that this is rather punishing task-any small deviations quickly compound as grows. Despite this, baselines, exhibiting more faithfulness in model output space (as well as in latent space).
that the MxDs match the future token generations far better than the
Please see qualitative examples of the first prompts and the subsequent 'diffs" (using Python 3's difflib) the generated tokens in the Appendix Figures 7 and
1.0
MxD
0.8
0.68
0.9
0.8
STC
TC
0.99
0.83
0.7
0.61
0.57
0.49
0.45
0.39
0.33
0.31
0.28
0.25
0.22
0.2
0.18
0.3
SAE
0.95
0.56
0.31
0.18
D.11
0.057
0.031
0.02
0.016
0.012
0.0039
0.0039
0.002
0.1
0.2
0.0
2
5
10
11
12
13
14
15
n-gram length
(a) Pythia-410m.
1.0
STC MXD
0.94
0.88
0.85
0.81
0.78
0.73
0.7
0.66
0.62
0.58
0.51
0.5
0.46
0.9
0.8
0.7
0.98
0.85
0.73
0.65
0.57
0.49
0.43
0.37
0.32
0.27
0.22
0.19
0.17
0.15
0.13
0.6
0.5
0.99
0.83
0.7
0.6
0.52
0.46
0.38
0.33
0.29
0.23
0.2
0.17
0.13
0.12
0.088
0.4
2
5
n-gram length
(b) GPT2-124m.
Figure 4: Proportion of 512 generated samples that contain predicted future words identical the original model's output when replacing the base L.M''s MLP layer with the sparse layers.
3.2 Feature evaluations
The accurate reconstruction of MxD models in Section 3.1 provides some evidence that MxDs are faithfully emulating the original MLP layers' functional mapping. However, for interpretability. we care equally about the extent which the learned features correspond to specialized, human- interpretable concepts. We confirm that MxD's features compete with the basclines quantitatively in

---

ag_ news
ag_
1.9
1.0
Test F1 score
Text
SP12
pythle 420m
McD
02-
SkipTCC Teck
2.0 World
Business
00
Word
Sports
Business
type
News
Figure 5: Highest Fl score probing for 'news category' [48] on individual features/experts. As expected, the MxDs remain competitive with the Transcoder baselines, outperforming TopK-SAEs.
two ways: through probing for known concepts in Section 3.2.1 and by steering the model using the learned features Section 3.2.2. For all experiments in this section, use the K 32 models.
Shared experts and specialization Interestingly, find MxDs naturally learn 'shared' expert performing a common base transformation- remaining K - active experts are thus free to dedicate their capacity modelling features unique to individual tokens. This emergent shared/private processing complements recent trends to use shared experts by design in MoEs [43, 44, 45, 46, 47] with [43] arguing this facilitates greater specialization. Furthermore, one may view the skip connection in STCs [26] as performing an analogous role to the shared expert. With MxDs, however, all units have the same high capacity accurately learn separate subcomputation regardless of the frequency or rarity of features.
We also observe that trained MxDs exhibit very few 'dead' experts, shown in Appendix C.I, with many experts contributing actively. Furthermore, initial ablations Appendix C.2 show that one train MxDs without shared experts if desired, small performance cost. Please qualitative results of activated tokens for particular experts in Appendix
3.2.1 Sparse probing with individual fcatures/experts
One challenge is that the sparse layers learn features in unsupervised manner. As pointed out in [23], we therefore do not know which high-level features we ought to expect the model to learn (or whether they exist in the Open WebText training data). Nonetheless,
reasonably expect useful unsupervised model to learn at least a handful of commonly occurring concepts and linguistic themes. We accordingly focus our evaluation the relative abilities of the sparse models to learn features well-predicting a variety of binary features used in the literature.
Concretely, to quantify the extent to which sparse layer features reliably fire in response to common high-level, interpretable concepts of natural language, we adopt the experimental settings [49, 23, 191. training binary probes the individual units of specialization (sparse hidden units 2n for TCs/SAEs and expert units a, for MxDs-all pre-activation). For probing of sample-level concepts, mean-pool activations across all non-padding tokens [19]. We train separate probes on 100 features with the largest mean difference between positive and negative activations, as per [49].
We perform experiments on all 24 binary probing tasks in the SAEBench suite [19]. Four of which are shown Figure 5, plotting the best Fl score (on a held-out set) for news topic classification in a 1-vs-all setting [48]. As can be seen, there exist individual MxD expert units that are predictive of various categories news articles, competitive with the baselines. We refer readers to Appendix B.7 for additional experiments 20 more sample-level probing tasks, 10 token-level probing tasks, and experimental details.
3.2.2 Feature steering
Specific features might reliably fire in response to interpretable patterns of the input, yet not contribute to the generation process. Here, we aim test this functional role features by steering the LLMs. We note that these experiments do not aim to establish TCs/MxDs as competitive with the SOTA for controllable L.LM generation. Rather, we aim to validate that the learned features contribute mechanistically to the LLM's forward pass in predictable way.

---

1.0
1.0
t
f
Textual coherence
0.8
Textual coherence
0.8
0.6
0.6
GPT2-124m
Pythia-410m MxD
0.4
MxD
0.4
TC
TC
0.2
Skip-TC
0.2
Skip-TC
TopK-SAE
TopK-SAE
0.0
0.0
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Concept steerability +
Concept steerability
Figure 6: Mean score along dimensions of "textual coherence' and 'steerability' of text generated by steering with the first 100 features of the sparse layers. Each sample is scored by 2 LLM judges.
Mechanisms for steering Let À E R be a hyperparameter controlling the desired 'strength' of the model edit. For TCs, we hook the forward pass at the relevant layer to increase the presence of target feature with Intuitively, increasing the weight an expert's contribution in the forward pass modulates the token representation in the direction of the learned specialization.
y AdA In contrast, MxDs be steered with y À (W, z).
We perform steering with the first 100 neurons/experts individually, using À == 100 for all experiments. We generate collection of 10 synthetic outputs for each neuron, each string consisting of 32 generated tokens to the prompt "Let's talk about " We then ask two LLMS to rate the collection of text along dimensions separately: (1) the extent which shared concept, theme, or linguistic pattern is present throughout the generated collection of text, and (2) the grammatical fluency of the text (please see Appendix D.1 for the full prompt). As can be seen from the mean scores over the 100 neurons shown in Figure 6, MxDs are competitive with the baselines, exhibiting similar trade-off between textual coherence and presence of concept we expect.
Results
4 Related work
Sparse decompositions
Learning sparse [50, 25], non-negative [51] features of a data signal has found many applications in computer vision [15, 52, 53, 54] and natural language processing [55, 56, 57]. motivated by the pursuit of interpretable, parts-based representations [13, 14]. In transformer-based language models [3], similar variants have been proposed for post-hoc analysis; sparse autoencoders (SAEs) combinations of atoms in a learned overcomplete dictionary. imposing either soft sparsity penalties [6, 22, 31] or thresholding activations directly [23, 58, 59]. Recent work aims to sparsify the existing layers of pretrained LLMs, either by learning new MIPs with sparse hidden units [29] (for circuit analysis [27] more interpretable yet faithful computation [26, 60],, by decomposing model parameters directly using attribution [61] and/or masking [62]. Despite the surge of interest in SAEs, many works emerging drawing attention to their immitaions-unddeperffome underperforming baselines for probing [63], unlearning [64], and steering [65], in addition to other pathologies [66, 32, 67, 68].
popular method that rewrites latent features non-negative
Conditional computation One natural alternative static fully connected layers is conditional computation [69, 70]. Tracing back to the early work of [71, 72], single dense layers are replaced with specialized Aobunits-conditiona on the input-as form of layer-level sparsity. The Mixture of Experts (MoE) architecture [36, 73, 74] is prominent example of conditional computation, breaking the link between parameter count and FLOPs. Consequently, MoEs have seen rapid adoption in SOTA models in recent wears-scaling to very large parameter counts [75, 76, 77, 78, 79]. For
We use gemini-2.. 1ammi-2..-ffaae 0-flash and www--ooutt----------e- 1et---ruu-ntttt two independent LLM judges.

---

parameter-efficien instruction tuning [37] introduces conditional (IA)* adapters [38], modulating the MLP hidden dimension with the Hadamard product. Our proposed formulation with factorized weight tensors yields 'Mo [37] as a less scalable special case (see Appendix A.5). In contrast, MxDs model the decoder output space directly for reconstruction, and also provide significantly more specialized units than [37], making MxDs more suitable for our goal of interpretability.
Whilst the primary focus of MoEs has been on their impressive capabilities, the literature has observed that individual experts often specialize in particular semantic patterns of the input data, despite not being trained do so [80, 81, 43, 82, 83]. For example, many works find that data that are in some sense similar are routed to the same apperts--pecialiiing to object shapes [84], texture [85], image category [86]. or semantic patterns in natural language [36]. In the context of large language models, this emergent property of specialization in MoEs has been primary focus of recent work: from encouraging monosemantic experts [87] or sparsity amongst experts' weights [88] to efficiently scaling the expert count for fine-grained specialization [40]. In contrast to these works exploring pre-training, we explore an efficient design of MoE to replace existing LLMs' dense layers.
Interpretability by design Whilst MxDs follow the paradigm of specialization- siaaiiionvvi--paasitttyy number of promising interpretable mechanisms have been recently proposed that do not rely on sparsity constraints. For example, recent work explores bilinear layers [89], combinatorial structure [90], and/or tensor networks [91] as alternative ways of designing interpretable architectures. Notably, [89] 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 yield similar Hadamard product forward pass, it involves two different input vectors, and establishes theoretical equivalence to mixture of experts and conditional computation when sparsity is present in one of the operands (i.e., the expert coefficients).
5 Conclusion
In this paper, we showed the benefits of decomposing dense layers' computations as a mixture of interpretable sublayers. We proposed the Mixture of Decoders (MxD) layer to achieve this at scale, proving that MxDs' linear experts preserve the matrix rank properties of the original decoders. Experimentally, we showed MxDs significantly outperform on the aparsiyyyacurray frontier when trained to replace dense MLP layers. Quantitative results on sparse probing and feature steering demonstrated MxDs nonetheless learn specialized latent features similarly to existing interpretability techniques. Crucially, MxDs reexamine the dominating neuron-level sparsity paradigm of popular techniques, providing evidence that specialization doesn't have to come with such a high cost to model performance. We believe MxDs (and specialization at the layer-level more generally) are an important step towards sparsity without sacrifice, and hope future work continues to build interpretable mechanisms that better preserve model capabilities. We are excited about future work exploring MxDs (and mixtures of linear transformations more generally) in alternative settings, such as for cross-layer features or transformations [92].
Limitations Our experiments show MxDs outperform on the sparsity-accurace 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. Furthermore, whilst the Topк activation can greatly reduce the decoders' FLOPs, the large encoders in sparse MLPs and the gating function in MxDs remain an additional inference- time cost. Future work could explore hierarchical structures [87, 36] and/or efficient retrieval [93] for further reductions in FLOPP. Secondly, MoEs are prone to issues of expert imbalance [73], or collapse [94]. Through random weight initialization alone, we find MxIDs in our experiments learn diverse experts without collapsing to the base decoder. However, we expect standard MoE load-balaneing [95] or diversity losses [96] be useful for MxDs should one need more explicit ways of encouraging expert diversity, or when training new interpretable MxD architectures end-to-end. With regards to feature evaluations, our steering experiments rely on LLMs as judges. Whilst this is commonplace in recent popular steering benchmarks [65], there is mixed evidence emerging about the reliability of all base models [97, 98]. 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).
10

---

Acknowledgments
James Oldfield is grateful Demian Till for reviewing the draft and providing valuable feedback and suggestions, and would also like to thank Markos Georgopoulos, Benjamin Hayum, and Wisconsin Al Safety Initiative's Safety Scholars for insightful discussions throughout the project. We are also grateful to the open-source Zulip platform for facilitating research discussion. Sharon Li is supported in part by the AFOSR Young Investigator Program under award number FA9550-23-1-01888 National Science Foundation under awards IIS-2237037 and (IS-23316) 1669, Office of Naval Research under grant number N00014-23, 1-2--4-- Schmidt Sciences Foundation, Open Philanthropy, Alfred P. Sloan Fellowship, and gifts from Google and Amazon. Shawn Im is also supported by the National Science Foundation Graduate Research Fellowship Program under Grant No. 2137424. Any opinions, findings, and conclusions recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation. Support was also provided by the Graduate School and the Office of the Vice Chancellor for Research at the University of Wisconsin-Maddionn with funding from the Wisconsin Alumni Research Foundation. Mihalis implemented
Nicolaou supported in part by the TensorICE project TEXCELLENCE/0520000 under the social cohesion programme "THALIA 2021-2027", co-funded by the European Union through the Research and Innovation Foundation.
References
[1] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Llama herd of models. arXiv preprint arXiv:2407.21783, 2024.
[2] Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhupatiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, et al. Gemma 2: Improving open language models at practical size. arXiv preprint wrXiv:2408.00111 2024.
[3] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenA/ blog, 188)9, 2019.
Zhengxuan Wu, Atticus Geiger, Thomas Icard. Christopher Potts, and Noah Goodman. In- terpretability at scale: Identifying causal mechanisms in Alpaca. In A. Oh, T. Naumann, Globerson, K. Saenko, M. Hardt, and Levine, editors, Adv. Neural Inform. Process. Syst. (NeurlPS), volume 36, pages 78205-78222 Curran Associates, Inc., 2023.
[5] Atticus Geiger, Zhengxuan Wu, Christopher Potts, Thomas Icard, and Noah Goodman. Finding alignments between interpretable causal variables and distributed neural representations. In Francesco Locatello and Vanessa Didelez, editors, Proceedings of the Third Conference on Causal Learning and Reasoning, volume 236 of Proceedings of Machine Learning Research, pages 160-187. PMLR, 01-03 Apr 2024.
[6] Adly Templeton. Scaling monosemanticity: Extracting interpretable features from Claude 3 Sonnet. Anthropic, 2024.
[7] Andy Arditi, Oscar Obeso, Aaquib Syed, Daniel Paleka, Nina Panickssery, Wes Gurnee, and Neel Nanda. Refusal in language models mediated by a single direction. arXiv preprint arXiv:2406.1 11717, 2024.
[8] Collin Burns, Haotian Ye, Dan Klein, and Jacob Steinhardt. Discovering latent knowledge in language models without supervision. Int. Conf. Learn. Represent (ICLR), 2023.
[9] David E. Rumelhart and James L. McClelland. General Framework for Parallel Distributed Processing, pages 45-76. 1987.
[10] Geoffrey E Hinton. Distributed representations. 1984.
[11] Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. Zoom In: An introduction to circuits. Distill, 2020. doi:
[12] Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac HatfieldDDodds, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. Toy models of superposition, 2022.

---

[13] Daniel D Lee and H Sebastian Seung. Learning the parts of objects by non-negative matrix factorization. Nature, 401(6755))788-75) 1999.
[14] Bruno A Olshausen and David Field. Emergence of simple-cell receptive field properties by learning a sparse code for natural images. Nature, (6133))6) 1996.
[15] Emmanuel J Candès, Xiaodong Li, Yi Ma, and John Wright. Robust principal component analysis? Journal of the ACM (JACM), 58(3):1-37, 2011.
[16] Jiankang Deng, Jia Guo, Niannan Xue, and Stefanos Zafeiriou. Arcface: Additive angular margin loss for deep face recognition. In IEEE Conf. Comput. Vis. Pattern Recog. (CVPR), pages 4690-4699, 2019.
[17] Vikram V. Ramaswamy, Sunnie S. Y. Kim, Ruth C. Fong, and Olga Russakovsky. Over- looked factors in concept-based explanations: Dataset choice, concept learnabllity, and human capability. JEEE Conf. Comput. Vis. Pattern Recog. (CVPR), pages 10932-10941, 2022.
[18] Caden Juang, Gonçalo Paulo, Jacob Drori, and Nora Belrose. Open source automated inter- pretability for sparse autoencoder features. https: //bleg///llo.......ee eleuther. atttintrrnt July 2024. EleutherAl Blog.
[19] Adam Karvonen, Can Rager, Johnny Lin, Curt Tigges, Joseph Bloom, David Chanin, Cal- lum McDougall, Yeu-Tong Lau, Eoin Farrell, Arthur Conmy, Kola Ayonrinde, Demian Till, Matthew Wearden, Samuel Marks, and Neel Nanda. SAEBench: A comprehensive benchmark for sparse autoencoders in language model interpretability.. In Int. Conf. Mach. Learn. (ICML), 2025.
[20] Zachary Chase Lipton. The mythos of model interpretability. Communications of the ACM, 61:36 - 43, 2016.
[21] Forough Poursabzi-Sangea Sangdeh, Daniel G. Goldstein, Jake M. Hofman, Jennifer Wortman Vaughan, and Hanna M. Wallach. Manipulating and measuring model interpretability. Pro- ceedings of the 2021 CHI Conference on Human Factors in Computing Systems, 2018.
[22] Robert Huben, Hoagy Cunningham, Logan Riggs Smith, Aidan Ewart, and Lee Sharkey. Sparse autoencoders find highly interpretable features in language models. In Int. Conf. Learn. Represent. (ICLR), 2023.
[23] Leo Gao, Tom Dupre la Tour, Henk Tillman, Gabriel Goh, Rajan Troll, Alec Radford, Ilya Sutskever, Jan Leike, and Jeffrey Wu. Scaling and evaluating sparse autoencoders. In Int. Conf. Learn. Represent. (ICLR), 2025.
[24] Bruno A Olshausen and David J Field. Sparse coding with an overcomplete basis set: A strategy employed by VI? Vision research, 37(23):3311-3325 1997.
[25] M. Aharon, M. Elad, and A. Bruckstein. K-SVD: An algorithm for designing overcomplete dictionaries for sparse representation. JEEE Transactions on Signal Processing, 54(11): 4311-4322, 2006. doi: 10.1109/TSP.2200688811 1999.
[26] Gonçalo Paulo, Stepan Shabalin, and Nora Belrose. Transcoders beat sparse autoencoders for interpretability. arXiv preprini arXiv:2501. 18823, 2025.
[27] Jacob Dunefsky, Philippe Chlenski, and Neel Nanda. Transcoders find interpretable LLM feature circuits. In Adv. Neural Inform. Process. Syst. (NeurIPS), 2024.
[28] Lee Sharkey, Bilal Chughtai, Joshua Batson, Jack Lindsey, Jeff Wu, Lucius Bushnaq, Nicholas Goldowsky-Dill, Stefan Heimersheim, Alejandro Ortega, Joseph Bloom, et al. Open problems in mechanistic interpretability. arXiv preprint arXiv:2501. 16496, 2025.
[29] Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Con- erly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, Zac Hatfield-Dodds, Alex Tamkin, Karina Nguyen, Brayden McLean, Josiah E Burke, Tristan Hume, Shan Carter, Tom Henighan, and Christopher Olah. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023. https:/ttrassfomme- html.
[30] Samuel Marks, Adam Karvonen, and Aaron Mueller. dection https://githir com/
lewrooonnrryyeennnng.g..a 2024.
12

---

[31]
Senthooran Rajamanoharan, Tom Lieberum, Nicolas Sonnerat, Arthur Commy, Vikrant Varma, János Kramár, and Neel Nanda. Jumping ahead: Improving reconstruction fidelity with jumprelu sparse autoencoders. arXiv preprint arXiv:2407. 14435, 2024.
[32] Joshua Engels, Logan Riggs, and Max Tegmark. Decomposing the dark matter of sparse autoencoders. arXiv preprini arXiv:2410. 14670, 2024.
[33] Noam Shazeer. GLU variants improve transformer. arXiv preprint arXiv:2002.05202, 2020.
[34] Tamara G Kolda and Brett W Bader. Tensor decompositions and applications. SIAM review, 51(3):455-500 2009.
[35] Grigorios G Chrysos, Yongtao Wu, Razvan Pascanu, Philip Torr, and Volkan Cevher. Hadamard product deep learning: Introduction, advances and challenges. IEEE Transactions on Pattern Analysis and Machine Intelligence, pages 1-20, 2025. doi:
[36] Noam Shazeer, * AZalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of- experts layer. In Int. Conf. Learn. Reprexent. (ICLR), 2017.
[37] Ted Zadouri, Ahmet Üstün, Arash Ahmadian, Beyza Ermis, Acyr Locatelli, and Sara Hooker. Pushing mixture of experts to the limit: Extremely parameter efficient moe for instruction tuning. Inf. Conf. Learn, Represent. (ICLR), 2024.
[38] Haokun Liu, Derek Tam, Muqeeth Mohammed, Jay Mohta, Tenghao Huang, Mohit Bansal, and Colin Raffel. Few-shot parameter-efficient fine-tuning is better and cheaper than in-context learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Adv. Neural Inform. Process. Syst. (NeurIPS), 2022.
[39] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (GELUs). arXiv preprint arXiv:1606. 001116666004415. 2016.
[40] James Oldfield, Markos Georgopoulos, Grigorios Chrysos. Christos Tzelepis, Yannis Pana- gakis, Mihalis Nicolaou, Jiankang Deng, and Ioannis Patras. Multilinear mixture of experts: Scalable expert specialization through factorization. In Adv. Neural Inform. Process. Syst. (NeurIPS), 2024.
[41] Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, al. Pythia: A suite for analyzing large language models across training and scaling. In Int. Conf. Mach. Learn. (ICML), pages 2397-2430. PMLR, 2023.
[42] Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. OpenWehText corpus. http://Skklion007. gi thub........ wttp:i/eeeeebeeeeoooooo 2019 2019.
[43] Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. DeepSeekMoE: Towards ultimate expert specialization in ofxepereoo-eeperts language models, 2024.
[44] Meta AI. The Llama 4 herd: The beginning of new era of natively multimodal AI innova- tion, 2025. URL https //ai meta. hom/blog//// blog/1lamaa 4- mom-mmmlmmomaala iitt--iiieneeeee...... Accessed: 2025-04-06.
[45] Qwen Team. Qwenl.5-MoB: Matching 7b model performance with 1/3 activated parameters", February 2024. URL https: / http://www/mmmmmmm tit.hb..u... wtt.t/////o/oo....oooooooooooeoo wttpo///llgggoo//ooooooeee qwen-moe.
[46] An Yang. Baosong Yang. Binyuan Hui, Bo Zheng. Bowen Yu, Chang Zhou Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang, Jian Yang. Jianhong Tu, Jianwei Zhang, Jianxin Ma, Jianxin Yang, Jin Xu, Jingren Zhou, Jinze Bai, Jinzheng He, Junyang Lin, Kai Dang, Keming Lu, Keqin Chen, Kexin Yang, Mei Li, Mingfeng Xue, Na Ni, Pei Zhang, Peng Wang, Ru Peng. Rui Men, Ruize Gao, Runji Lin, Shijie Wang, Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Xuejing Liu, Yang Fan, Yang Yao, Yichang Zhang, Yu Wan, Yunfei Chu, Yuqiong Liu, Zeyu Cui, Zhenru Zhang, Zhifang Guo, and Zhihao Fan. Qwen2 technical report, 2024.
[47] An Yang. Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoyan Huang, Jiandong Jiang, Jianhong Tu, Jianwei Zhang, Jingren Zhou, Junyang Lin, Kai Dang, Kexin Yang, Le Yu, Mei Li, Minmin Sun, Qin Zhu, Rui Men, Tao He, Weijia Xu, Wenbiao Yin, Wenyuan Yu, Xiafei
13

---

Qiu, Xingzhang Ren, Xinlong Yang, Yong Li, Zhiying Xu, and Zipeng Zhang. Qwen2.5-1n technical report, 2025.
[48]
Antonio Gulli. Ag corpus of news articles. http://groups di wttp.dd........gggll///////////e
corpus_ of news aarricleeee.... htmm,, 2005.
[49] Wes Gurnee, Neel Nanda, Matthew Pauly, Katherine Harvey, Dmitrii Troitskii, and Dimitris Bertsimas. Finding neurons haystack: Case studies with sparse probing. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
[50] Rodolphe Jenatton, Guillaume Obazinski, and Francis Bach. Structured sparse principal component analysis. In Yee Whye Teh and Mike Titterington, editors, Proceedings of the Thirteenth International Conference Artificial Intelligence and Statistics, volume 9 of Proceedings of Machine Learning Research, pages 366-373, Chia Laguna Resort, Sardinia, Italy, 13-15 May 2010. PMLR.
[51] Patrik O Hoyer. Non-negative matrix factorization with sparseness constraints. Journal of machine learning research, 5(Nov): 2004.
[52] Edo Collins, Radhakrishna Achanta, and Sabine Süsstrunk. Deep Feature Factorization for Concept Discovery, page 352-368. Springer International Publishing. 2018. ISBN 9783030012649, doi:
[53] James Oldfield, Christos Tzelepis, Yannis Panagakis, Mihalis Nicolaou, and Ioannis Patras. PandA: Unsupervised learning of parts and appearances in the feature maps of GANs. In Int. Conf. Learn. Represent. (ICLR), 2023.
[54] Yue Song, Thomas Anderson Keller, Yisong Yue, Pietro Perona, and Max Welling. Un- supervised representation learning from sparse transformation analysis. arXiv preprint arXiv:2410.05564 2024.
[55] Wei Xu, Xin Liu, and Yihong Gong. Document clustering based on non-negative matrix factorization. In Proceedings of the 26th Annual International ACM SIGIR Conference on Research and Development in Information Retrieval, SIGIR '03, page 267-273, New York, NY, USA, 2003. Association for Computing Machinery. ISBN 1581136461 doi: 10. 1...145/66445.6666 860485
[56] Da Kuang, Jaegul Choo, and Haesun Park. Nonnegative matrix factorization for interactive topic modeling and document clustering. Partitional clustering algorithms, pages 215-243, 2015.
[57] Sanjeev Arora, Yuanzhi Li, Yingyu Liang, Tengyu Ma, and Andrej Risteski. Linear algebraic structure of word senses. with applications polysemy. Transactions of the Association for Computational Linguistics, 6:483-495, 2018. doi: 10. 1621/a// a 00034.
[58] Alireza Makhzani and Brendan Frey. K-sparse autoencoders. arXiv preprint arXiv:1312.566 2013.
[59] Bart Bussmann, Patrick Leask, and Neel Nanda. BatchTopK sparse autoencoders. In NeurIPS
2024 Workshop on Scientific Methods for Understanding Deep Learning, 2024.
[60] Lucy Farnik, Tim Lawson, Conor Houghton, and Laurence Aitchison. Jacobian sparse autoencoders: Sparsify computations, not just activations, 2025.
[61] Dan Braun, Lucius Bushnaq. Stefan Heimersheim, Jake Mendel, and Lee Sharkey. Inter- pretability in parameter space: Minimizing mechanistic description length with attribution- based parameter decomposition. arXiv preprint S.ti:25501.19220 2025.
[62] Lucius Bushnaq, Dan Braun, and Lee Sharkey. Stochastic parameter decomposition. arXiv preprini arXiv:2506.2000
2025.
[63] Subhash Kantamneni, Joshua Engels. Senthooran Rajamanoharan, Max Tegmark, and Neel Nanda. Are sparse autoencoders useful? A case study sparse probing, 2025.
[64] Eoin Farrell, Yeu-Tong Lau, and Arthur Conmy. Applying sparse autoencoders to unlearn knowledge in language models, 2024.
[65] Zhengxuan Wu, Aryaman Arora, Atticus Geiger, Zheng Wang, Jing Huang, Dan Jurafsky, Christopher D. Manning, and Christopher Potts. AxBench: Steering L.MSs? even simple baselines outperform sparse autoencoders, In Int. Conf. Mach. Learn. (ICML), 2025.
14

---

[66] David Chanin, James Filken-Smith Tomáš Dulka, Hardik Bhatnagar, and Joseph Bloom. A is for absorption: Studying feature splitting and absorption sparse autoencoders, 2024.
[67] Patrick Leask, Bart Bussmann, Michael Pearce, Joseph Bloom, Curt Tigges, Noura Al Moubayed, Lee Sharkey, and Neel Nanda. Sparse autoencoders do not find canonical units of analysis, 2025.
[68] Gonçalo Paulo and Nora Belrose. Sparse autoencoders trained on the same data learn different features, 2025.
[69] Yizeng Han, Gao Huang, Shiji Song, Le Yang, Honghui Wang, and Yulin Wang. Dynamic neural networks: A survey, IEEE Trans. Pattern Anal. Mach. Intell. (TPAMI), 44(11):7436- 7456, 2021.
[70] Emmanuel Bengio, Pierre-Lue Bacon, Joelle Pineau, and Doina Precup. Conditional com- putation in neural networks for faster models. In Int. Conf. Mach. Learn. Worksh. (ICMLW), 2015.
[71] Robert A Jacobs, Michael Jordan, and Andrew G Barto. Task decomposition through compe- tition in a modular connectionist architecture: The what and where vision tasks. Cognitive science, 15(2)2219-250, 1991.
[72] Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures local experts. Neural computation, 3(1):79-87, 1991.
[73] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(1))) 2022.
[74] Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. ST-MoE: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202.08900 2022.
[75] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with conditional computation and automatic sharding. In Int. Conf. Learn. Represent. (ICLR). 2021.
[76] Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts. In Int. Conf. Mach. Learn. (ICML), pages 5547-5569. PMLR, 2022.
[77] Albert Q. Jiang Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie-Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mixtral of experts, 2024.
[78] Dee,Seek-AI, Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang. Han Bao, Hanwei Xu, Haocheng Wang, Haowei Zhang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Li, Hui Qu, J. L. Cai, Jian Liang, Jianzhong Guo, Jiaqi Ni, Jiashi Li, Jiawei Wang, Jin Chen, Jingchang Chen, Jingyang Yuan, Junjie Qiu, Junlong Li, Junxiao Song, Kai Dong, Kai Hu, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lean Wang, Lecong Zhang, Lei Xu, Leyi Xia, Liang Zhao, Litong Wang, Liyue Zhang, Meng Li. Miaojun Wang, Mingehuan Zhang. Minghua Zhang, Minghui Tang. Mingming Li, Ning Tian, Panpan Huang, Peiyi Wang, Peng Zhang, Qiancheng Wang, Qihao Zhu, Qinyu Chen. Qiushi Du, R. J. Chen, R. L. Jin, Ruiqi Ge. Ruisong Zhang, Ruizhe Pan, Runji Wang, Runxin Xu, Ruoyu Zhang, Ruyi Chen, S. S. Li, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shaoqing Wu, Shengfeng Ye, Shengfeng Ye, Shirong Ma, Shiyu Wang, Shuang Zhou, Shuiping Yu, Shunfeng Zhou, Shuting Pan, T. Wang, Tao Yun, Tian Pei, Tianyu Sun, W. L. Xiao, Wangding Zeng. Wanjia Zhao, Wei An, Wen Liu, Wenfeng Liang, Wenjun Gao, Wenqin Yu, Wentao Zhang, X. Q. Li, Xiangyue Jin, Xianzu Wang, Xiao Bi, Xiaodong Liu, Xiaohan Wang, Xiaojin Shen, Xiaokang Chen, Xiaokang Zhang, Xiaosha Chen, Xiaotao Nie, Xiaowen Sun, Xiaoxiang Wang, Xin Cheng, Xin Liu, Xin Xie, Xingchao Liu, Xingkai Yu, Xinnan Song,
15

---

Xinxia Shan. Xinyi Zhou, Xinyu Yang, Xinyuan Li, Xuecheng Su, Xuheng Lin, Y. K. Li, Y. Q. Wang, Y. X. Wei, Y. X. Zhu, Yang Zhang. Yanhong Xu, Yanhong Xu. Yanping Huang, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Li, Yaohui Wang, Yi Yu, Yi Zheng, Yichao Zhang, Yifan Shi, Yiliang Xiong. Ying He, Ying Tang, Yishi Piao, Yisong Wang, Yixuan Tan, Yiyang Ma, Yiyuan Liu, Yonggiang Guo, Yu Wu, Yuan Ou, Yuchen Zhu, Yuduan Wang, Yue Gong, Yuheng Zou, Yujia He, Yukun Zha, Yunfan Xiong, Yunxian Ma. Yuting Yan, Yuxiang Luo, Yuxiang You, Yuxuan Liu, Yuyang Zhou, Z. F. Wu, Z. Z. Ren, Zehui Ren, Zhangli Sha, Zhe Fu, Zhean Xu, Zhen Huang, Zhen Zhang, Zhenda Xie, Zhengyan Zhang, Zhewen Hao, Zhibin Gou, Zhicheng Ma, Zhigang Yan, Zhihong Shao, Zhipeng Xu, Zhiyu Wu, Zhongyu Zhang, Zhuoshu Li, Zihui Gu, Zijia Zhu, Zijun Liu, Zilin Li, Ziwei Xie, Ziyang Song, Ziyi Gao, and Zizheng Pan. Deepseek-v v3 technical report, 2025.
[79] Joan Puigcerver, Carlos Riquelme, Basil Mustafa, and Neil Houlsby. From sparse to soft mixtures of experts. In Int. Conf. Learn. Represent. (ICLR), 2024.
[80] Aya Abdelsalam Ismail, Sercan O Arik, Jinsung Yoon, Ankur Taly, Soheil Feizi, and Tomas Pfister. Interpretable mixture of experts. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
[81] Marmik Chaudhari, Idhant Gulati, Nishkal Hundia, Pranav Karra, and Shivam Raval. MoE lens an expert is all you need. In Sparsity in LLMs (SLLM): Deep Dive into Mixture of Experts. Quantization, Hardware, and Inference, 2025.
[82] Huy Nguyen, Xing Han, Carl Harris, Suchi Saria, and Nhat Ho. On expert estimation in hierarchical mixture of experts: Beyond softmax gating functions, 2025.
[83] Stefan Nielsen, Rachel Teo, Laziz Abdullaev, and Tan Minh Nguyen. Tight clusters make specialized experts. In Int. Conf. Learn. Represent. (ICLR), 2025.
[84] Brandon Yang, Gabriel Bender, Quoc V Le, and Jiquan Ngiam. CondConv: Conditionally pa- rameterized convolutions for efficient inference. Adv. Neural Inform. Process. Syst. (NeurIPS), 32, 2019.
[85] Basil Mustafa, Carlos Riquelme Ruiz, Joan Puigcerver, Rodolphe Jenatton, and Neil Houlsby. Multimodal contrastive learning with LIMoB: the language-image mixture of experts. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Adv. Neural Inform. Process. Syst. (NeurIPS), 2022.
[86] Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Ady. Neural Inform. Process. Syst. (NeurIPS), 34:8583-8599 2021.
[87] Jungwoo Park, Ahn Young Jin, Kee-Eung Kim, and Jaewoo Kang. Monet: Mixture of monosemantic experts for transformers. In Int. Conf. Learn. Represent. (ICLR), 2025.
[88] Xingyi Yang. Constantin Venhoff, Ashkan Khakzar, Christian Schroeder de Witt, Puneet K. Dokania, Adel Bibi, and Philip Torr. Mixture of experts made intrinsically interpretable, 2025.
[89] Michael T Pearce, Thomas Dooms, Alice Rigg, Jose Oramas, and Lee Sharkey. Bilinear MLPs enable beie mechanistic interpretability. In Int. Conf. Learn. Represent. (ICLR), 2025.
URL https: openr ewwer..eewww... www. f Id-IIOKKKUKKS
[90] Micah Adler, Dan Alistarh, and Nir Shavit. Towards combinatorial interpretability of neural computation. arXiv preprint arXiv:2504.0884/ 2025.
[91] Thomas Dooms, Ward Gauderis, Geraint A. Wiggins, and Jose Oramas. Compositionaliiy unlocks deep interpretable models. In AAAI'25 workshop CoLoRAI- - Comnecting Low-Rank Representations in AI, 2025. URL https:t///// hptp//vppeeeeeeeeeeeeeeeeeeeeeee rett/foorum?dd===Xtt552223
[92] Jack Lindsey, Brian Chen, Adam Pearce, Sasha Hydrie, and Thomas Conerly. Sparse mix- tures of linear transforms. https: / /t/t//ranfffrrrreeeeeeeeeee circuits. mab/2025/bull /ulpbuuuuuu///// index. html, 2025. Preliminary research update, Transformer Circuits.
[93] Xu Owen He. Mixture million experts, 2024.
[94] Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, Xian-Ling Mao, Heyan Huang, and Furu Wei. On the representation collapse of sparse mixture of experts. Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Adv. Neural Inform. Process. Syst. (NeurIPS), 2022.
16

---

[95] Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun. and Damai Dai. Auxiliaryloss--ree load balancing strategy for mixture-of-experts, 2024. URL 15664.
[96] Boan Liu, Liang Ding, Li Shen, Keqin Peng, Yu Cao, Dazhao Cheng, and Dacheng Tao. Diver- sifying the mixture-of-exppetss representation for language models with orthogonal optimizer, 2024. URL
/211:///w..rraaa///0.... 09762.
[97] Lianmin Zheng, Wei-Lin Chiang, Ying Sheng. Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing. Hao Zhang, Joseph E Gonzalez, and Ion Stoica. Judging LLM--a--a-jude with MT-Bench and chatbot arena. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Adv. Neural Inform. Process. Syst. (NeurIPS), volume 36, pages 46395-46663 2023.
[98] Anna Bavaresco, Raffaella Bernardi, Leonardo Bertolazzi, Desmond Elliott, Raquel Fernández, Albert Gatt, Esam Ghaleb, Mario Giulianelli, Michael Hanna, Alexander Koller, Andre Martins, Philipp Mondorf, Vera Neplenbroek, Sandro Pezzelle, Barbara Plank, David Schlangen, Alessandro Suglia, Aditya K Surikuchi, Ece Takmaz, and Alberto Testoni. LLMs instead of human judges? a large scale empirical study across 20 NLP evaluation tasks. In Wanxiang Che, Joyce Nabende, Ekaterina Shutova, and Mohammad Taher Pilehvar, editors, Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 238 255, Vienna, Austria, July 2025. Association for Computational Linguistics. ISBN 979-8-89176-252-1 doi: 10. 00.18653/v1/225. acl-short.20.
[99] James E. Gentle. Matrix Algebra: Theory, Computations, and Applications in Statistics. Springer, New York, 2nd edition, 2007.
[100] Nicholas D Sidiropoulos and Rasmus Bro. On the uniqueness of multilinear decomposition of n-way arrays. Journal of Chemometrics: Journal of the Chemometrics Society, 14(3):
229 209, 2000.
[101] Donghyun Lee, Jaeyong Lee, Genghan Zhang, Mo Tiwari, and Azalia Mirhoseini. CATS: Context-aware thresholding for sparsity in large language models. In First Conference Language Modeling, 2024.
[102] Frank Lauren Hitchcock. The expression of a tensor or a polyadic as a sum of products. Journal of Mathematics and Physics, 6: 6:14 1927.
[103] J. Douglas Carroll and Jih Jie Chang. Analysis of individual differences in multidimensional scaling via an n-way generalization of "Eckart- "Eckart-Young" decomposition. Psychometrika, 35: 283-319, 1970.
[104] fvcore: Flop counter for PyTorch models. httter//gttiab.cm/acccbbtreeeere fvcore.
[105] CodeParrot. Github code dataset. https///hnniinggeeecoo//tttttttt github-code, 2022.
[106] Yupeng Hou, Jiacheng Li, Zhankui He, An Yan, Xiusi Chen, and Julian McAuley. Bridging language and items for retrieval and recommendation. arXiv preprint arXiv:2403.03952. 2024.
[107] Philipp Koehn. Europarl: A parallel corpus for statistical machine translation. In Proceedings of Machine Translation Summit X: Papers, pages 79-86, Phuket, Thailand, September 13-15 2005.
[108] Maria De-Arteaga, Alexey Romanov, Hanna Wallach, Jennifer Chayes, Christian Borgs, Alexandra Chouldechova, Sahin Geyik, Krishnaram Kenthapadi, and Adam Tauman Kalai. Bias in bios: A case study of semantic representation bias in a high-stakes setting. In Proceed- ings of the Conference on Fairness, Accountability, and Transparency, FAT '19, page 120 128, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 97814 778145036122 255. URL ht / /doii. org/10. 10/1/114//222 3217/608776.0 .2877772
doi: 10. (0.1145/58875003227778
[109] Anthony Duong Joseph Bloom, Curt Tigges and David Chanin. Saclens. https:/ // /////ttttn /github. com/jbloomAus//55 /SAELess, 2024.
17

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The quantitative results in Section 3.1 support the claims that the MxDs better navigate the sparsity-accurace frontier in models up to 3B parameters. Section 3.2 quantitatively supports the claims that the features similarly specialized.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitations are stated explicitly in Section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specfiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
18

---

Answer: [Yes]
Justification: The proofs of Lemmas I and 2 are provided in the appendix in Appendix A.2, and referenced from the main paper. It made explicit in Lemma under exactly which technical conditions the rank property holds. A small summary of the proof of Lemma is given in the main paper, too.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof' provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Yes, details, configurations, and resources used train the models are included in Appendix D. Implementations of the models are provided in notebooks explicitly linked to from the main paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibilitty In the case of closed-source models, it may be that access to the model is limited in
19

---

reproducible, they some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: A link the code train the models is included in the submission.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
conoooocyllyllooolomooommooimlll for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nips. wwwppcpppiiccccoccccccc caablicgguides... /oope//Cooeumiii wembbbmsssiii sssisssssooooso ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Details about the main training runs are provided in Table 10, and the details about the sparse probing experiments are provided in Appendix B.7.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: The largest model runs take the order of multiple days per value of K-thus, training multiple models not feasible on our resource budget.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
20

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative eтror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes, this is provided in Table 10, including the specific GPU types and wall-clock runtime.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: //tprr//////...eeeeeee cc cctppblcc//cccccccccccccccccccct cettciie/eieeeeeeeeeeeeeeeeeeeee SetiiiceGel
Answer: [Yes]
Justification: No violations of the code of ethics are made. In particular, the paper focuses designing more transparent and interpretable mechanisms, and thus we don't consider the paper to cause any direct negative societal issues. However, we acknowledge that there is often an inherent dual-use nature much interpretability work, including our own. Furthermore, whilst our largest training runs take multiple days, each uses a single GPU, therefore the environmental impact is limited.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
21

---

Justification: We discuss how our design of increasingly performant yet interpretable models may provide better understanding, debugging, and editing of pre-trained models. We do not explicitly discuss the downsides of the work.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
The proposed work only explores existing pre- mr-ttrineed models that are openly accessible.
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging. and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We include links to all used models' huggingface pages in Table 10.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
22

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi wwwrsswithhoeeee thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's ' creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new assets are released.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the datasct/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: No human subjects are used.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human subjects used.
Guidelines:
23

---

The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [Yes]
Justification: We use LLMs via public APIs for our steering evaluation. The full prompt is provided in Appendix D.I, and the exact model endpoint names are provided for repro- ducibility.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: : /nnuriiss.....// cc/Conf e /LLM) for what should or should be described.
24

---

Appendix
Table of Contents
A Proofs and additional technical results
25
A.1 Proof of rank equality
25
A.2 Proof of MxD forward pass equivalence
26
A.3 Intuition weight parameterization through the lens tensor methods
26
A.4 GLU encoders mixture of rank-1 linear experts
27
A.5 fadamard-factorizee factorized tensors generalize MoVs
28
A.6 TopK activation function...
28
B Additional quantitative results and ablations
29
B.1
Faithfulness in output space
29
B.2 Additional reconstruction metrics
30
B.3 Results on additional layers
30
B.4
Initial results Gemma2-27B
31
B.5
Expert rank ................................
32
B.6
Computational cost of sparse layers
34
B.7
Sparse probing
34
B.8
Ablations
35
Feature balance and shared experts
40
C.I. Expert/feature balance
40
C.2 Shared experts
42
Detailed experimental setup
Feature steering details
E Additional qualitative results
44
Proofs and additional technical results
A.1 Proof of rank equality
Proof of Lemma 1. We first derive the expression for expert n's weight matrix W, D diag(c,, e Ifxx and then show the rank equality that follows. First, recall that we have the third-order weight tensor defined as
W(n, h,:) = c, d₁ E R
for matrices C NxO D E Hx0, We can express each element of the tensor W E N in terms of elements of the two matrices as
W(n, = Cno dno = (D) ho Cno-
(6)
Equation (6) shows that for fixed expert n, the nh row c,, E Bº essentially scales the columns of matrix D R #x0 This is equivalent to right-mut multiplying matrix D by a diagonal matrix formed from C, E _ Indeed, the (h, o) entry of such matrix iroduct
D diag(cn))no =
Σ (D) diag(c.)io i=1
(7)
hD) )an
(8)
d Cnor
(9)
25

---

since all off off-diagonal terms (i.e., i * o) in Equation (7) vanish and diag(Cn)ou Cno by con- struction. Comparing Equation (6) and Equation (9) shows that, for every {1,2.... H and E (1,2....,0) O} we have
(n,, h, o) [D diag(c.))no
Hence, indexing into the first mode of the tensor alone gives us the matrix-valued expression for expert n as claimed:
W(n W, -ddia((..) R HxO
Finally, a standard result in linear algebra [99] has that rank (AB) (An for any A E # and invertible matrix B _xxo Since matrix diag (c.) E R invertible by assumption in Lemma 1, setting A = D and B = diag(c.) yields the rank equality.
A.2 Proof of MxD forward pass equivalence
Recall we have input vector z E R expert coefficients a E RR and layer weights W E RR~xHxo, The weights are defined in Equation (4) element-wise through the Hadamard product
Wnn, h, = Cn * dj E RO, A' (I....,, NJ, HE H},
for learnable parameters E Nx D R#xo Lemma states that MxD's forward pass can be equivalently expressed
Σ an (W = (C z).
Proof of Lemma 2. The LHS can first be re-written as an explicit sum over the hidden dimension
N
N
H
W
2
(Wnn:Rh) E RO.
(10)
Plugging in the definition of Wich: E R.O from Equation (4) then yields
an (wnh:zh)
(11)
N
H
an (C,, * dii 21h)
(12)
n= -=
N
CnCn
Σ
(13)
n=l
h=1
(C a) (D (D))),
(14)
which exactly the RHS Equation (5), showing the MxD forward pass is equivalent to the
Hadamard product of C and D z.
A.3 Intuition for weight parameterization through the lens of tensor methods
A second complementary way of viewing the MxD layer's parameterization (and its full-rank propertics) is through the lens of tensor methods [34]. A tensor-based motivation for MxD's weight tensor parameterization and forward pass is presented in Appendix A.3.1 and Appendix A.3.2, respectively.
Notation and definitions A brief primer first included below, based on [34] (and can be safely skipped for those already familiar):
26

---

The mode-n fibers of an N order tensor X' R/xx1xxxxx are the dn-dimensional column vectors obtained by fixing every index except that of the n mode (e.g. X,i R are the mode-1 fibers of third-order tensor X E (1))). Stacking all mode-n fibers column-wise yields the so-called mode-n unfolding (n) E ER4xk with number of columns given by the product of remaining dimensions I, -1144
The Khatri-Rao product (denoted by )) between two matrices A E RRI and RJ*K is the column- wise Kronecker product (denoted by @):
A @ B == a @ b, ... a @ b. E R(IJ)XK
The mode-n (vector) product of tensor X E R with a vector E R& is denoted X' x, and has entries (X' X, Siatttii
la Di==1 Bigig..dy Wi..
A.3.1 MxD weight tensors through the Khatri-Rao product
MxDs construct the collective weight tensor through the Khatri-Rao product @ [34] of the two factor matrices C E RN D HRx Concretely, the mode-3 unfolding of the third-order weight tensor W E RN in MxDs from Equation (4) is alternatively given by:
W ()) == (C o D) O
(15)
Given that the factor matrices are learned end-to-end without constraints, they are likely of full column-rank, i.e rank(D) rank(C) O (as N > O, H 4.0 O in practice given the MLP layers' larger bottleneck). Consequently, their Khatri -ao product parameterizing the collective N experts' weights will be of maximum rank O too, through Lemma 1 of [100]. As result, parameterized this way, the O-dimensional fibers likely span the full output space.
A.3.2 Tensorized MxD forward pass
Furthermore, the layer's forward pass can then be viewed as performing two tensor contractions between the third-order weight tensor W NxxOOO (collecting all N experts' H O-dimensional matrices) and expert coefficients a E RR and hidden activations z E R.. This can be expressed in terms of the so-called mode-n product (denoted by x,) [34] follows:
(W.,z)
On
N
H
N
H
W 2nh
a 2 Wnh
n=
h=1
A=11
W X
(16)
A.4 GLU encoders are a mixture of rank-1 linear experts
Both the proposed MxDs and Gated Linear Units (GLUs) [33] share similar functional form, using the element-wise product. However, there are crucially important differences between GLUs and MxDs that make both their interpretation and model capacity different.
In short, the technical results here our paper show that GLUs" encoder can be viewed as linear mixture of expert layer with rank-1 experts. Furthermore, GLUs can be modified and extended to MxDs with two additions to their model form as detailed at the end of this subsection. First, recall that the GLU encoder [33] computes:
Youu (E (E (ETx)= H
x)
(17)
for input vector x E RR learnable weights EGLU, E E R and activation function v.).. To transform Equation (17) into the same model form MxDs, first pre-multiply the LHS by the identity matrix to match the MxD model form of Equation (5), yielding:
YOLU = (Iaa) * E x),
(18)
1 'which is simply a reshaping of a higher-order order tensor into matrix, arranging all N expert matrices' column vectors along the columns of new matrix.
27

---

where a ( E R"' and I E R is the H- H-dimensional identity matrix. Next, we can write this explicitly in terms of linear MoE with expert weights W,, E R as follows:
YGLU
(I
a)
E (Exx
(19)
H
Σ
On [W,x)) x)
(20)
H
Σ
On (Ediag ((I),)) x),
(21)
where (T), E R is the nh row of the H-dimensional identity matrix (i.e. one-hot vector with its only non-zero element at index n). We draw particular attention to how the n expert's 's matrix W,, n E diag (TT),) R |xH essentially picks out the nº column of E, leaving all remaining H columns as zero vectors. Therefore, GLU encoders compute a MoE with linear expert weights of (at most) rank 1. This relationship between GLUs and conditional computation is consistent with prior work interpreting individual GLU column vectors as experts [101]. Whilst GLUs' encoders' model form does not put any inherent restrictions on the total number of rank-1 terms that can contribute to the output, the sparsity necessary for specialization does.
We conclude this section by summarizing the two technical changes needed to transform GLUs into full-rank linear MoEs based on the Hadamard product:
1. Replace Il Equation (18) with learnable, non-diagonal weight matrices for full-rankness.
2. Choose (-) produce non-negative, sparse coefficients to encourage specialization through sparsity among the experts (for example, a softmax function, or a ReLU activation followed by TopK).
The first of the steps above provides full-rankness, whilst the second brings the sparsity and non- negativity needed for specialization. We include a notebook showing this connection in PyTorch at: https://gt thub. com/ http://ammmlllleddd/dddoddo
wwwwoeled/dd//bbeee/eaeaaae http///loommmmmam.//////mm googggllssoooooo to-eees...emmmmmm-oo-oemooooeee. ipynb.
A.5
Hadamard-factorta tensors generalize MoVs
Prior work [37] proposes to linearly combine N many (IA)² adapters [38] for parameter-efffcient MoEs for instruction fine-tuning. The implementation results in very similar functional form to the factorized forward-pass in MxDs. Interestingly, the Hadamard product parameterization of the third-order weight tensor in Equation (4) provides more general framework through which one can also derive MoVs' model form. shedding light on the relationship to the proposed MxDs and their benefits. Concretely, factorizing the weight tensor instead along the second mode as W(n,; = Cn * d, E RR# in our framework immediately recovers MoV [37] as a special case. In particular, in contrast to the MxD in Appendix A.3 whose weight tensor can be parametrized equivalently through its mode-3 unfolding [34], MoV's implicit weight tensor can be given in terms of its mode-2 unfolding in terms of similar Khatri- -ao product of two factor matrices.
Instead, MoVs in analogy would vield expert weights by pre-multiplying D as: W,, = diag (c.) D E xO for much larger C E RR &RN Due to H >> O, our proposed MxD formulation yields around 4x the number of specialized units as MoVs with the same parameter budget (yet MoVs' experts are of no higher rank than MxDs'), making MxDs much more suitable and efficient class of layer for our goal of scalable specialization. We therefore see that the proposed lens of tensor methods for unification provides valuable insights about how design more interpretable layers with the minimum trade-off to capabilities.
A.6 TopK activation function
Formally, for activations RH (possibly with the ReLU activation applied), the TopK activation function can be formulated elementwise as the following:
2h. if 24 2 (x)z)
(2)
10
otherwise
28

---

where (z( E I returns the value of the K largest element of vector z E R". Intuitively, the TopK activation retains only the largest K f H elements of the H neurons, setting the rest zero.
We note here that the definition above does not handle ties (although this is very unlikely in practice, given the activations are continuous).
B Additional quantitative results and ablations
B.1 Faithfulness in output space
Experiments in Section 3.1.2 of the main paper report the number of future predicted words that are identical from the original model and from the LLM with the sparse layer replacements. We also show qualitative examples of the first 8 prompts and the subsequent 'diffs' (using Python 3's difflib) of the generated tokens in Figures 7 and 8-we see MxDs' superior ability to preserve model functionality as it propagates through the output space of future tokens.
Generation 1, Pythia- - -41m
GT:
your daily routine and just you need to take break from your daily routine and just There are times when you need to be ablo to do something that is not casy.
There times you nead take break
There are times
STC:
There are times when you need to take a break from your dally routine and just SAE: There times when need to
TC:
the
in high-speed
Generation 2, Pythia-410m
GT:
Hunanitarian chief warns capacity of US to be tested By Staff reports Published: Friday, May 38,
MyD:
Huranitarian chief warns capacity of us
tested By Starf reports Published: Friday, May 36,
STC:
Huranitarian chief
capacity
tested By Associated Press Publiished: Tuesday, March
Huranitarian chief wams capacity
tested By Associated Press January 24, 2013
SAE:
Huranitarian chief warns capacity of the U.S. military to be a threst The U.S. military 15
Generation 3, Pythia-410m
GT:
buring Trurp addinistrationn, Departuent of Homeland Security bean tasked with oversesing ianigration enforcement.
During the Trune adrinistration, Department of Honelant Security (185) been tasked protecting Avericans
STC: During the Trune adrinistration, the Departeent of Justice has been accused of using "bribery" to Luence
TC: During the Trunp administration, the Department of looneland Security has been tasked with everseeing imigratioo enforcement.
SAE: During the trurp adhinistration. the president-elect" compaign was "crisis" that could solved by
Generation 4, Pythia-410m
GT: Racandan neuspaper ro. cites report from the Hinlstry of Interier and Security (MIS) that shows
MxD: Romanian newspoper 2F. cstes report from the Hinistry or Interior and Secursty (MIS) that shows
STC:
Ronanian
that Ronania Ministry of Foreign Affairs Trade (MFA) Romanian newsppper 2t cites report that the country's police officers are not allowed to use
cites report the European Connission
TC:
Romanian nevepaper
SAE:
Generation 5, Pythia-410m
GT:
Denocratic Virginia 6ov.
McAuliffe
Tuesday said
not seek re-election
Donocratic Virginta Gov. Terry NcAulstte (0) Tuesday sasd will not seek re-election 2629,
STC:
Denocratic Virginta Gov. Terry Nctuliffe (0) 15 running for president in the 2028 election, but he's
TC:
Derocratic Virginia Gev. Terry Nctuliffe (D) is facing challenge from group of Denocratác state
SAE:
Desocratic Virginie Gov. Terry Not. Gingis Desocrat, but he's 's not the rost popular governor
Generation 6, Pythia-410m
GT: federal court has ruled that the Trump adainistration's travel ban on people fros seven Mustia-majority
MxD: rederal court has ruled that the Truns administrationns travel bon pcople fros seven Nus onjority
STC: tederal court has ruted that Trump administration's travel
fros Mus injority people fros seven Mastime-agoritty
federal court ruled that Trump administration's travel
SAE:
federal court has ordered fomeer U.S. anbessador to the United Nations, Willias H. Taylor,
Generation 7, Pythia-410m
GT:
British cotumbia takes in British Columbia takes in 551.5 billion the sale of its ail and gas reserves The British Colunbia takes in X51.5 billion in tax breaks The British Columbia government has announced that
billion the sale of its eil and gas
STC:
TC:
British Columia takes in the world The British Columbia geve mesntt has announced that it will spend a plan to expand its coal-wining operations, but
SAE:
British Columbia takes in the third
Generation 8, Pythia-410m
Eartier this nonth, I wrote about the upcoming release of the first episade of The Walking MXD: Earlier this nonth, I wrote about the upcoming release of the first episode of The Walking
GT:
STC:
tarlier this nonth. wrote about the upconing relesse of the "Sonic Boom" soundtrack. The albux farlier
TC:
wrote about upconing release of Farlier this nonth, was invited to conference in city Torento. The event
Mania" gane, gane
SAE:
Figure 7: Pythia-410m: The first few generated tokens from the base model ('GT') and the corresponding tokens from the model when the sparse layers replace the target MLP. Red denotes tokens that are removed, orange denotes newly inserted tokens, and green denotes matching tokens.
29

---

Generation 1, GPT2-124m
GT:
There are tires when you need to be little more careful with your food. You
MxD:
There are tixes when you need to be on the lookout for new job. But,
STC:
There are times when need to get out of bed and go to the bathrous. been told There are times when you need to get tittle bit more than your usual. You're
TC:
There tixes
feel Like I'n being judged.
grades
SAE:
Generation 2, GPT2-124m
GT:
Humaniterian chief warns capacity to handle refugees is "usprecedented" The uN refugee
has warned that
MxD:
risk The rafugie agency has Husanitarian chief warns capacity to handle mnargenciss is at risk The govermment has warned that
Husinitarian chief
capacity
refugaas
STC:
TC: Husanitarian chief wamss capaeity to handle evergencies is at risk The UK govemmeent has warned that
SAE:
Hunanitarian chief warns capacity to hold The BBC political correspondent, Peter Robinson. says the governrent
Generation 3, GPT2-124m
GT:
buring the Trump adsinistration the White House been working make sure that its innigration buring the Trunp advinistration, the white House has been working to make sure that its innigration
MxD:
STC:
burzing the Trunp adkinistration, the white House has been working make sure that it i5
burting Trump adrinistration, white SAE: burang the Trusp adainistration, white House been working on plan to
TC:
been working
that
Generation 4, GPT2-124m
GT:
Rosanian newspaper 2F.re cites report the European connission that the country's 's government is considoring report the European Connission that the country's governeent considering country's government is considering Romenion nevspaper 2F ro cites report by the Curopean Connission that the EJ is considering inposing
MxD:
Romanian
2F.
STC:
Rosanian newspaper
cites report by the European Conniazion
TC:
SAE:
Romanian nevspaper 2F ro cites the "revelation" of the "death of the country's econony" a key
Generation 5, GPT2-124m
GT:
Desocratic Virginia Gev. Terry Nolaliffe (0) said he would not seek re-election in 2018, but
MxD:
Denccratic Virginta 6ov Terry Rctuliffe (0) said he vould not seek revelection in 2018. but he
STC:
Denocratic Virginia Gev. Terry Rctulitfee (D) has said he will not support bill that would
TC:
Desocratic Virginia Gov. Terry Nefuliffe (R) has seid he will not seek rexelcction in 2015, but
SAE:
besocratic Virginia Terry Nouliffe (R) signed
would allow state
Generation 6, GPT2-124m
GT: A federal. court has ordered the Department of Justiie to pay 1$1.5 miltionn to group
MxD: federal
has ordered Departaent Justice to pay X$1.5 million varrant court ordered Department Justice to stop using
STC:
TC: federal court has ordered Department Justice to stop using search varrant
SAE:
federal court has ordered the Department of Justice to pay 1$1.5 rillion to forer
Generation 7, GPT2-124m
billion annually from federal goversment, only about MxD: British Colurbis takes in X51.5 billion foreign aid ennually, sccording to the govermentts latest report
GT:
British Calurbia
STC:
British Colurbia takes in V$1.5 billion in foreign aid anmullly, according to the Canadian Council of
TC:
aritish colurbia takes in 141.5 billion annually fron the federal governeent, but it's not a big fron province's tourisn industry. according
SAE:
British Coturbia takes 151.5 million
Generation 8, GPT2-124m
tarlier this mosth, I wrote about the state of our industry. We been in bit MxD: farlier this month. wrote about the state of our industry. We've been in state Fartier morth.
GT:
STC:
state
in building solid foundation
TC:
Eartier this month, I wrote about the state the for creating a custos 30 tarlier this morth, I was asked to write a piece about the *anti: Somitic" rovement in Germany.
SAE:
Figure 8: GPT2-124m: The first few generated tokens from the base model ('GT') and the corre- sponding tokens from the model when the sparse layers replace the target MLP. Red denotes tokens that are removed, orange denotes newly inserted tokens, and green denotes matching tokens.
B.2 Additional reconstruction metrics
To highlight the scale of difference in the reconstructions between MxDs and the baselines, we also plot in Figure 9 the normalized MSE at the end of training for all models and LLMs. At the smallest values of K (which we care about most for interpretability). MxDs' normalized MSE is up to an order of magnitude smaller than Transcoders'.
B.3 Results on additional layers
We also fully train all models and baselines (with 4 different values of K) on different target layers for each model. The results are shown in Figure 10 for 48 additional trained layers for the same setup in the original paper, using different colours to highlight that these are new results. As can be seen, the same trend holds; MxDs significantly outperform the baselines at small K in all LLMs.
30

---

0.017
0.0018
Pythia-410M
(Layer 8)
(Layer 15)
0.015
TC (3BM params)
0.0016
(67M params)
0.014
STC (38M params)
STC (68M params)
Nomalized MSE
0.013
MxDs (38M params)
0.0014
MxDs (67M params)
0.0013
0.012
0.0012
0.011
0.0011
0.010
0.009
0.0009
0.008
0.0008
0.007
0.0007
0.006
0.0006
0.0005
0.005
0.0004
0.004
0.0003
0.003
0.0002
0.002
0.0001
0.001
0.0000
16
32
64
128
256
32
64
128
256
Pythla-1.48
0.00075
Llama3.2-3
0.0024
(Layer
(Layer
0.0022
TC (269M params)
TC (604M params)
Normalized aleed MSE
0.0020
STC (273M params)
MxDs (268M params)
0.00060
STC (614M params)
MxDs (604M params)
0.0018
0.00055
0.0016
0.00050
0.00045
0.0014
0.00040
0.0012
0.00035
0.0010
0.0008
0.0006
0.00020
0.0004
0.00015
0.00010
0.0002
0.00005
0.0000
0.00000
64
128
256
64
128
256
Sparsity level
Sparsity level
Figure 9: Normalized MSE at the end of training Sparse MLP layers, as a function the number of active units (i.e., hidden neurons experts); with differences as large order of magnitude in еггоr.
B.4 Initial results on Gemma2-27B
In the main paper. we perform experiments on base models with up to 3B parameters. For initial evidence of scalability even larger models, perform here additional large-scale experiments the 27B Gemma2 model (in half-precision, and with a smaller batch size to fit into memory). We MxDs continue to outperform the aparsity-accuracy frontier, shown by the (normalized) reconstruction loss in Table 3.
Table 3: Normalized MSE (L) on Gemma2-27B. Layer 20, K 32, partially trained for 50k iterations. To fit in memory, the model is loaded in half precision. We use 1/8 the batch size and buffers stored, and 1/2 the multiplier on the number of latent features (with values of 16, and 16. respectively).
Iterations
10k
20k
30k
40k
50k
Transcoder
0.147
0.132
0.128
0.123
0.119
Skip Transcoder
0.128
0.107
0.102
0.100
0.093
MxD
0.098
0.096
0.093
0.086
0.069
31

---

3.645
Pythia-410M (Layer 12)
(Layer 10)
3.3425
Mean cross-Entropp loss
3.640
Original
Original
TC (3BM params)
3.3400
TC (67M paramss)
3.635
STC (38M params)
STC (68M params)
MxDs (38M params)
MxDs (67M params)
3.630
3.3350
3.3325
3.3300
3.3250
3.3225
32
64
128
256
32
64
128
256
3.038
Pythia-1.48 (Layer 10)
LJama3.2-38
(Layer 10)
3.036
Original LLM
Original
Mean cross- entropy loss
3.1150
(269M params)
3.034
TC (604M params)
STC (273M params)
STC (614M params)
3.1125
MxDs (268M params)
3.032
MxDs (604M params)
3,030
3.1075
3.026
3.1050
3.024
3.1025
3.022
3.1000
3.020
32
64
128
256
32
64
128
256
Sparsity level K
Sparsity level K
Figure 10: Additional layer results: model cross entropy loss preserved when replacing MLPs with Transcoders [27]. Skip Transcoders [26]. and MxDs, as a function of the number of active units (hidden neurons/experts). These results complement those in the main paper, but here train new set of additional models on different layers.
B.5 Expert rank
This section concerns the matrix rank of the linear experts in parameter-efficient MoEs. We first compare to low-rank MoEs in Appendix B.5.1 to demonstrate the benefits of full-rankness, and then follow up Appendix B.5.2 by confirming that the learned MxD expert ranks are close maximum in the trained models.
B.5.1 Comparisons to low-rank MoEs
In this section, study the impact of expert rank on the ability of efficient MoE layers to reconstruct pre-trained MLP layers' mappings. One compelling alternative to MxDs for efficient conditional computation is the μМоΕ layer [40], which imposes low-rankness
expert weights to achieve Whilst μлоEs are found to perform competitively in the pre-training setting, the impact of low-ratd oow-rakkness approximations of existing layers will determine their suitability in the sparse layer approximation setting studied in this work.
We therefore compare to juMEE layers, which we use to compute linear MoE in place of the MLP's decoder. CPμMo N experts' weight matrices are jointly parameterized through low-rank tensor structure with the CP decomposition [102, 103] for chosen rank R E N+. With the same learnable encoder and expert gating matrices producing the expert coefficients E R&N and hidden units IR generated the same way in the main paper, train μMoE layers to approximate the
32

---

original MLP layer's output with:
H ΣΣΣ 10
N
R
xM)E(x)
On D(r, h) C(r,n) W(ir)) E RD;
(22)
where C E RRXN D E R W E ROXR are the learnable low-rank terms of the implicit third-order tensor parameterizing all N collective experts' weights.
We match the MxD experimental configuration as closely as possible for a fair comparison. For the encoders, we mirror MxDs and use the GELU activation function, which we find through ablations in Appendix B.8 to perform the best. We initialize the parameters the same as MxDs and Skip Transcoders: use the standard PyTorch linear layer initialization for D, C (and the encoder layers), and initialize W as the zero matrix.
3.611
Rank ablation GPT2-small
Rank ablation
3.334
pythia-410m
MxD uMoo
MrD MMDE
cross-entropy
Mean cross coss-entrogy
3.332
3,609
3.608
3.330
3.607
3.328
3.606
3.605
3.326
3.604
0.5
0.6
0.7
0.8
0.9
1.0
0.5
0.6
0.7
0.8
0.9
1.0
Normalized expert rank (mean)
Normalized expert rank (mean)
Figure 11: Comparisons to μMoEs for various choices of (normalized) rank: high rank weights best-preserve the models' downstream cross-entropy loss.
We vary the (MoE. layer rank R. training fully 3 sparse approximation layers for K = 32 active experts, varying the total number of experts N to keep the parameter count the same-isolating the impact of the choice of rank. As with the main experiments, we record the downstream model loss when we splice in the trained layer to replace the MLP layers, shown in Figure 11.
As can be seen, the μMоE layers perform well when they are close to full-rank (i.e. when the normalized rank A - 1). Crucially, however, performance drops off notably when the rank reduced. Whilst μMоEs still perform far better than neuron-level sparsity methods (i.e. the corresponding CE loss results in Figure 3), we observe that full-rankness is necessary for the most faithful layer approximations
the proposed MxDs provide by design.
As motivating example, for why Sparse and SoftMoFS are not practical: Sparse [36] and SoftMoFs [72] require 2.16 trillion parameters for a single layer, for the same 86k experts we use for Llama-3. 2-BB.. This is orders of magnitude more parameters than the entire base network itself, making it prohibitively costly for SparseMoEs to scale to sufficiently high expert counts.
B.5.2 MxD empirical expert rank
Next, show experimentally that the learned experts' matrices W. D diag(c,)) "xO very nearly full-rank in practice, corroborating the propertics of expert matrices shown theoretically in Lemma 1. We compute the mean 'normalized rank', which take for MxDS to be the empirical matrix rank of the learned expert's weights, over the maximum possible rank given the dimensions:
N rank (W,) min H. O}
(23)
N
We show in Table 4 the normalized rank across all 4 base models: MxD's learned experts exhibit no rank deficiencies, providing further evidence of the large potential capacity of MxD layers despite their sparsity constraints on the expert-level.
33

---

Table Mean normalized expert matrix rank of Equation (23) across models for the first 2k experts in K = 32 trained MxDs the learned expert matrices are very close to full column rank.
GPT2-124M
Pythia-410M
Pythia-
Llama-3.2-3B
0.99 + 0.005 0.99 + 0.007 0.99 + 0.005 0.99 + 0.002
B.6 Computational cost of sparse layers
We show here that there is minimal difference between the computational cost of sparse layer variants. We first report theoretical layer FLOPs, and then report empirical benchmarks.
Parameter count & FLOPs We first tabulate in Table 5 the theoretical parameter counts and inference-time FLOPs for MxDs vs Sparse MLPs. To be consistent with the popular PyTorch library [104] we count one fused multiply-ado d-dimensional vectors as requiring d/2 FLOPs. Note that, For a chosen expert count N, we set the width of Sparse MLPs (e.g. Transcoders) H == N H* to parameter-match the models.
one FLOP, and count the Hadamard product between two
Table 5: Here I, O denotes the input and output dimensions, H the original model's hidden width, N the MxD expert count, and H the width of the sparse MLPs.
Parameter count
FLOPSs
Sparse MLP MxD
H (1++) N N ( H (I+O)
H (1++))
+ H H
Empirical benchmarks We next run benchmarks for Sparse MLP layer vs MxD with batch size of 512, and dimensions: == H 1024. and number of experts/features as N 8192, 9216. The results are tabulated in Table 6, where we see performance (and cost) is similar for the two layers.
Table Here I, O denotes the input and output dimensions, H" the original model's hidden width, N the MxD expert count, and H the width of the sparse MLPs.
Peak memory usage (MiB) Latency (ms)
Parameter count Reported FLOPSs ([104)
Sparse MLP MxD
386.50
1.394
18.874.368 18,874,368
18,874.368 18,874,880
389.50
1.457
B.7 Sparse probing
Sample-level probing Here, we follow the SAEBench [19] evaluation protocol. In this 'sample- level' setting, each text string is labeled with a binary concept at a global level (e.g., the language of the snippet, its sentiment). This in contrast to what we refer to as 'token-level probing', where each token within the text samples is labeled individually (e.g., whether a word is certain part of speech). We perform experiments on a total of 24 sample-level sparse probing tasks with the same 'maximum mean difference" feature filtering applied in [19]. The details of the datasets used are summarized in Table 7.
Token-level probing We also explore sparse probing for 10 features defined at the token-level. For this, we follow [49], and include experiments training probes on the mean feature activations under tokens spanning the surnames of the individuals. We note that this is a significantly harder task, and makes even stronger assumptions about the features the dataset includes, but is nonetheless some additional weak evidence about the relative feature- reature-learning abilities of the sparse models. Through various surnames, we probe for occupations of individuals, whether or not individuals are alive, and individuals' labeled gender. We also experimented with probing for compound words as in [49], but found no predictive features in our trained models. Details of the surname token-level probing datasets (and the total training examples the tokenizers could parse) are included in Table 8.
34

---

Table 7: Details of sample-level sparse probing datasets used.
Dutaset
Training examples
Test examples
Classification task description
Number of classes
fancyzho/ag news [45]
16,000
article topic
ondeparooggthhb-oodd [105]
20,000
5,000
Progrumming language
amaaznn, reviewseevwwwee.ee mecsssy land5 sentimmntt [106]
8,000
2,000
Posiive/ncaative review sentiment
Hebsink--NL PlPleropal [107]
European language
LabHChias jn_ bios [108]
32,000
8,000
Profession from bio
Table 8: Details of token-level sparse probing datasets used.
Dataset
# Training examples
Test examples
Classification task description
Number of classes
Occupation [49]
4784
1195
Occupation of individual
6
Is alive? [49]
4800
1 199
Are they alive
2
Gender [49]
4800
1200
Labeled gender
2
Experimental setup For sample-level probing, we truncate the input strings to the first 128 tokens for all datasets but for the Gi Cithub dataset, where we take the last 128 tokens to avoid license headers [19, 49]. For token-level probing, we instead take only the last 128 tokens, where the final token contains the surname of the individual in question in the datasets of [49].
Binary probes trained 80% of the training data (randomly shuffled) with the sklearn library's Logisti cRegressionn module with parameters:
class ses_-weeitttttatnnee balanced
penalty=112'
solver= wower==nwwtoo-cooleskky
max iter=200
A random seed of 42 is used throughout the code to ensure reproducibility.
B.7.1 Sparse probing results
We show in Figure 12 results on 20 additional (sample-level) sparse probing tasks, where MxDs remain competitive with the baselines. We also plot the expert activation (of the single expert with the highest F1 test set score) for the positive/negative classes for all tasks split across Figures 13 and 14. One observe certain degree of separability between the two semantic clusters of data given by the expert coefficient, thus confirming that individual experts are learning to specialize to particular high-level features.
We also include results on 10 token-level probing tasks in Figure 15, with the corresponding activation densities displayed in Figure 16. Whilst MxDs appear to perform slightly less well here on average, they remain competitive as expected.
B.8 Ablations
We turn next to ablation studies to explore the value of the various model components below:
B.8.1 Choice of sparsity constraint
We first train a variety of MxDs on GPT2 models with the TopK activation function [23] and instead train models with ReLU followed by an explicit A/|-|11 sparsity penalty on the specialized components addition to the reconstruction loss [22]. We show the results in Figure 17, where, similarly to [26], we find the TopK activation to dominate on the sparsity-accuracy frontier-we thus use the TopK activation for all experiments.
B.8.2 Choice of MxD encoder
Secondly, we show in Figure 18 the benefits of MxDs' flexibility in inheriting the original MLP layer's encoder form/activation function. All models here are trained from scratch for the same number of tokens and with the same experimental setup in Section 3.1, with K 32. In the
35

---

europari
0.8
Test F1 score
0.6
0.5
2p52 MaD
pythi--1100
0.2
6:2-
European language
European language
(a) Europarl dataset, on GPT2
(b) Europarl dataset, on Pythia-441 410mm
github code
github code
1.0
3.0
0.0-
2.8
Fest Fattlscore score
F1
aythi 430m
Mio
VsD
0.2
6.2
TO
ShipTC
HTMLL
Jova
PHP
Python
HTML
Java
PHiP
Python
Pragramming language
Programming language
(c) Github code dataset, on GPT2
(d) Github code dataset, on Pythia-410m
wwtpn//reiissscuuuee....oom sentiment
amazon_ aacioo_reviews__meeeee has___maalll__yyaee land5 sentiment
1.0
2.9
0.8
C.8
Test F1 score
0.4
0.4
Ma3
VuD
0.2
TC
Skip
TockS4z
Negative
Pasitive
Negative
Powitive
Amazon review sentiment
Amazon reviow sentiment
(c) Amazon review sentiment dataset, on GPT2
(f) Amazon review sentiment dataset, on Pythia-4 10m
10m
bias_ in__inn__i bios
bla_ii__ioo
Test Fl score
Fl
aythie 410m
MEO
VSD
attomey
tentrt
jumalist
photographer
physrsin
profeseor
piychologitt
beachee
annomme
deriist
iurmaiet
photoprapher
atysccian
proteaor
prycholoyist
Inacher
Profession
Profossion
g Bias Bios dataset, on GPT2
(h) Bias in Bios dataset, on Pythia-410m
Figure 12: Sample-leve sparse probing results individual experts/features; the best Fl score held out set is presented.
36

---

Sports
Business
1.0
1.0
1.0
0.5
0.5
0.5
0.5
0.0
-3
2.5
0.0
2.5
5.0
-2
33
Pre-activttion Value
Pre-actication value
Pre-ectiontion value
Pre-activation value
(a) AG news dataset, GPT2
World
Sports
Business
Tech
Dendty
0.5
0.00
0.25
value
(b) AG news dataset, on Pythia-410m
de
en
85
fr
nl
Density
1.0
0.2
1.0
Deusity
0.5
0.5
0.1
0.5
0.0
0.0
0.0
0.0
10
1.0
-2
artaactivaionn
ne-activadonvaa
Pre-activtion value
Pro Pro-activaion
(c) Europarl dataset, on GPT2
Decstt
1.5
-10
-0.
-1.0
-055
0.0
0.5
0.0
0.5
1.0
activationnnoooooo value
Pro octivotton value
Pro activation value
Pro activation value
Pro vec-actvation value
(d) Europarl dataset, on Pythia-410m
HTML
Java
PHP
Python
Dendty
1.0
1.0
Dundky
0.5
0.5
0
0.0
0.0
Presac value
Peeaactiation
Pri-activation value
Pre-activation
Pre-activation
(c) Github code dataset, on GPT2
HTML
Java
PHP
Python
Density
Deauity
-0.5
0.0
0.5
Pro activation
activation
activattoon
Pro-activaiionaa value
Pretcctvation
(f) Github code dataset, on Pythia-410m
Negative
Positive
Donsty
2
Pre-actination value
Presactivaiion Vallue
Amazon review sentiment dataset, on GPT2
Negative
Posttive
Denshy
5.0
2.5
2.5
0.0
0
-0.6
-0.2
Pre-actrattion valuo
Pre-acitvalion
(h) Amazon review sentiment dataset, on Pythia-410m
Figure 13: [1/2] Sample-level sparse probing results on individual experts for MxDs; here we plot the values of the expert pre-activation for positive/other classes (in the 1-vs-all setting).
37

---

attomey
dentist
joummaist
photographer
1.5
1.0
0.5
6,5
6.0
30
0
Pro coivotion value
Pro activation value
Pro activation value
Pro ree-sciivationn value
physician
professor
psychologist
teacher
Deesky
1.0
0.0
0.0
Pe-activation value
Phevactivation
Pre-activation value
Pre-activation
(a) Profession from biography, on GPT2
attomny
dentist
journalist
photographer
4
Dordtt
1.2
1.0
0.8
0.5
0.4
0.2
0.50
0.25
0.00
0.25
0.50
0.4
0.2
0.0
0.2
0.4
1.0
0.8
0.6
0.4
0.2
Phesasciration value
Prevaclivation
Pre-activadion value
Pre-atkivatione Value
physician
professor
psychologist
Daraty
-0.8
-0.6
-044
-0.2
0.0
1.00
-0.75
0.00 -..50
-0.25
-1.25
1.00
-0.75
-0.50
-0.25
-1.2
-1.00
-0.75
-0.50
-0.5
Pro wo--actvation volue
Pro activattiiinno volue
Pro act.aaddooonoooo value
Pre-activation value
(b) Profession from biography, on Pythia-410m
Figure 14: [2/2] Sample-level sparse probing results on individual experts for MxDs; here we plot the values of the expert pre-activation for positive/other classes (in the 1-vs-all setting).
occupation
occupation
1.0
1.0
0.8
0.8
Fe score
0.6
Fl
0.4
0.4
gpt2
pythia-410m
MXD
MyD
0.2
TC
0.2
TC
TC
Skip-TC
TopKSAE
Tap
0.0
Singer Researcher
Actor
Athiltt
Politician Journalstt
Singer Researcher
Actor
Athiete
Politician Journalist
Occcpation
Occupation
(a) Occupation surname probing, GPT2
(b) Occupation surname probing, Pythia-410m
2014
(c) "Is alive?" surname "Is alive?" surname probing, on GPT2
(e) Gender surname prob-
(f) Gender surnam prob- ing, on Pythia-410m
probing, on Pythia-410m ing, on GPT2
Figure 15: Token-level sparse probing results on individual experts/features; the best F1 score held out set is presented.
38

---

Singer
Researcher
Actor
Athete
Politician
Journalist
0.4
0.50
0.50
0.2-
0.25
0.25
100
-5
0
550
-5
0.0
Pre-activation value
Pre- Pre-activation value
Pre acticattion valuee
Pre-ectivation value
Pre cet--ctiiationnooo value
Pre-activation value
(a) Occupation surname probing, GPT2
Singer
Researcher
Actor
Athlete
Politician
Joumalistt
1.0
Deratt
1.0
0.5
0.5
0.0
0.0
Presectiv
Pre-ectivation valuu
Presectival
Presictivalln Prevactivatiion
Prescciiaaion value
(b) Occupation surname probing, on Pythia-410m
Yes
No
Doreny
0.2
0.25
Pre-actiration value
(c) Aliveldead surname probing, GPT2
Denatty
11
Pre-actiattion value
Presectvation value
(d) Alive/dead surname probing, on Pythia 410m
Dondty
6.25
0.00
Pro setivottion value
Pro act.actttioonooooo value
(e) Gender surname probing, GPT2
M
Doratty
(f) Gender surname probing. on Pythia-410m
Figure 16: Token-level sparse probing results on individual experts for MxDs; here plot the values of the expert pre-activation for positive/other classes (in the 1-vs-all setting).
39

---

GPT-2 ablation: penalty Topk activation
1.5e
Transcoders (L1)
3,640
Transcoders (TopK)
Ours (L1)
3.635
Mean cross-entropy loss
(TopK)
Original model
3.630
0m
3.625
3.620
7.0e5
4.0e4
3.615
3.0e-
3.610
1-1 4.15e-xx
3.605
3,600
2
2
Mean to
Figure 17: ReLU+TopK activation function [23] vs ReLU w/ LI sparsity penalty [22]: both MxDs and Transcoders better recover the cross entropy loss with the TopK activation.
Pythio-410N
GPT2-124M
300010
Pythio-1.48
Liama-5.2-5
0.0025
000025
3,20070
Nommllze ASE
0.0055
000015
5.00010
2.20010
0.020
000010
5,00005
000005
3. 0.20000
0.0000
0.00000
5,00000
Rotu IMLS)
d-GELD JNLPI
(~RicU (MLF)
(-CEU (ML))
(=RoLU (MU1
Figure 18: Encoder architecture ablation: MSE loss when using ReLU activation vs the GELU used by the base models; and MLPs vs GLUs for Llama (rightmost subfigure).
first 3 left-most subfigures, see the Normalized MSE is low half when using GELU vs the non-native ReLU activation.
We next ablate the impact of inheriting the same encoder as the Llama-3 3-2-BB base model. In the rightmost subfigure of Figure 18, we train MxDs with ReLU-MLP. GELU-MLP, and Swish-Gi GLU encoders. As can be seen, using GLU with a Swish activation model (matching the base model architecture) yields Normalized MSE almost an order of magnitude smaller than MLPs with GELU/ReL.U.
C Feature balance and shared experts
C.1 Expert/feature balance
Following the code of [27, 109], we log how often each unit of specialism/feature used, over a fixed window of ~ 1M tokens. We show in Figure 19 the feature frequency at the end of training, where we observe that MxDs see a similar healthy balance experts to the frequency of usage of features in the baselines.
Interestingly, we observe small peak of experts that fire more frequently MxDs (e.g., around -2 on the -errhas specializing in common patterns and primitives in natural language.
40

---

Feature usage: GPT2-124M MxD
Feature usage: Pythia-410M MxD
3000
Number of features
2500
Number of features
3000
2000
2000
1500
1000
1000
500
-10
-8
-6
-4
-2
-10
-8
-6
-4
-2
0
Logl0 feature firing frequency
Log10 0 feature firing frequency
Feature usage: GPT2-124M
Feature usage: Pythia-410M TC
4000
Number of features
Number of features
3000
2000
2000
1000
0
-10
-2
-6
Log10 feature firing frequency
Log10 feature firing frequency
Feature usage: GPT2-124M STC
Feature usage: Pythla-410M STC
5000
4000
Number of features
Number of features
4000
3000
3000
2000
2000
1000
1000
-10
-8
-6
-2
-10
-8
-6
-2
Log10 feature firing frequency
Log10 feature firing frequency
Figure 19: log feature sparsity (following [27, 109]); MxDs' experts are well-balanced, similar to the baselines" features.
41

---

C.2 Shared experts
We find that, by default, our MxD models naturally learn to use shared expert, with the remaining experts exhibiting strong specialization in a wide range of themes and linguistic patterns. The use of a shared expert is becoming an increasingly popular design choice, including in the latest Llama models [44]-we therefore allow this pattern to emerge naturally in our base models, further justified through the evidence in [43] that shared experts can enhance specialization among the remaining experts [43]. We highlight, however, that a simple trick of sampling K ~ Unif[ K K/a, K+ + K/a) for the Top-K K activation at train-time (for e.g. a == 2) is sufficient to remove the dominating shared-expers at minimal hit to reconstruction performance, if desired.
We train two sets of models with a base = 32 on GPT2-small and pythia-410m, using a == We first show in Figure 20 the indices of the top-activating experts for the 2 model variants on template prompt, after training has finished. On the left-hand side of Figure 20, the models route all tokens through the same shared expert at position 1. However, we on the right-hand side that training with the 'random-K" strategy breaks the dependence on a shared expert in position I. Furthermore, we include in Figure 21 the corresponding train-time MSE loss for the 4 models here as ablations -olatioos-observing observing that the random-K K strategy also brings comparable performance. Based on these experiments, recommend this simple training strategy if one desires MxD models without shared experts.
Expert diversity/collapse To what extent is the shared expert simply learning the base decoder itself? In initial attempt to study this, take the same MxD trained with K 64 on GPT2, and tabulate the Frobenius norm of differences between the original model's decoder D', our learned base decoder D, and the "shared expert" weight matrix W shared D diag (shared) (with index shared = 19772), all of which are matrices of shape (H x O). As can be seen in Table 9, their difference is significant, showing that the shared expert does not simply learn the original decoder.
As discussed in the limitations section, we highlight that MxDs as formulated contain no explicit loss terms to discourage expert collapse, or encourage expert balance diversity. Whilst we find no need for these in our experiments, they may prove useful (or even necessary) in other settings or future work.
Table 9: Distance between learned shared expert. base decoder, and original decoder
shared D* D...s
W shared D'
|
-
(9 elements of the difference matrix)
70.09
0.03
0.067
208.08
221.6934
0.21
0.12
0.07
00.17
0.15
0.11
D Detailed experimental setup
We list in Table 10 the resources used for each experiment: the GPU and the indicative run-time for a single model. The expansion_fee input dimension to generate the MLP width in the sparse layers (i.e. H == I 1--mmppeee iospansion_fctttin
column refers to the expansion factor applied to the
Table 10: Total training time and resources used to produce the k- 32 experiments (the required compute being roughly the same across models trained with different k).
Model
GPU use
VRAM
Traiiin
fuctoe
Asset link
GPT2 124m
x GeForce RIX
2015B
Nb 34m
32
kdtpp: //taggggggggggeeeeeeee
Pythia
xi GeForce RIX 3050
24GB
8h 35m 176
buspa /taagg/aaaaeeeeo wo//Beutbeall 420ss
Pythis1
A100
80G8
antpe www.g////gggggnge ptw./aaateee///eeaaa-....
Llams-12-38
80G8
50s
atspe://nooginttte ww/pste- www.a///////a//.....
42

---

Model trained with fixed K
Model trained with random K
2nd highhest
4th highest
highest index
highest
3rd highhest
exspert
expprt
Token
[10160
10962
19772
9610
526
18499
7257
8244
Token 21
197722
15461
2630
8228
Token 2
116992
3344
17100
7388
Token
19772
18694
7385
3494
Token 19829
10864
7720
5507
Taken
[19772
19466
10619
970]
Token [20001
15277
1905
11387
Prompt: "Who is the president of the USA?"
GPT2-small
Model trained with fixed
Model trained with random K
1st: high experrt index
2nd highest expert index
bighhst expert indlex
4th highest expert
highest expert
highest expert index
3rd hiih expert Index
highest expert indes
Taken
128104
1694
18149
2013
Token
7412
13294
3097
19430
Taken 2
28104
1163
5124
11890
Token
13439
24209
13723
18099
Taken 3
28114
5124
27687
3657
Token 3
9587
3857
10715
6198
Token
22810
4126
12814
23628]
Token
2809
3378
25799
9435]
Prompt:
"Who is the president the USA?"
Pythia-410m
Figure 20: Top-activating experts template prompt with and without using a randomized value of K at train-time for TopK expert selection: randomization largely prevents a shared expert. Shown are the leading tokens and expert indices.
GPT2-small
Pythla-410m
TC
TC
0.035
STC
0.006
MxD
MxD
MxD (random
MxD (random K)
Normalized MSE
0.025
Normalized MSE
0.004
0.020
0.003
0.015
0.002
0.001
0.005
20K
40K
40K
80K
100K 120K
Training steps
Training steps
Figure 21: MxD performance with random sampling: Normalized MSE loss a func-
tion of training steps using fixed Top K 32 expert selection and when sampling IS Unif { K
K K + K
2,
43

---

D.1 Feature steering details
For the steering experiments, we use two LLM judges to grade generations on two axes. The full template prompt we feed to gemini-2 -2. 0 O-flash and 1lama-4-scooe 4-cout----t----- --s-cuut-1b-1660 1 iebo-inntructt is as follows (note that line breaks and emphases included here only to aid visualization):
Prompt given to LLM judges
You are an expert evaluator of synthetic text.
TASK: Rate a collection of [num samples) samples along independent axes. AXIS 1 CONCEPT COHERENCE:
0.00 shared concepts/themes/sylle.
0.25 faint overlap.
0.50 some overlap or similar structure.
0.75 mostly the same concepts or structure; a few partial drifts.
1.00 all snippets clearly share the same concepts, themes, style, or structure.
AXIS 2 GRAMMATICAL FLUENCY:
0.00 incomprehensible.
0.25 dense errors; meaning often obscured.
0.50 frequent errors; meaning still mostly recoverable.
0.75 minor errors that rarely hinder comprehension.
1.00 completely grammatical and natural.
(Do not penalise fluency if snippet starts ends abruptly.). SCORING: Choose any real value in [O, 1] for each axis.
OUTPUT FORMAT: Respond with exactly two numbers formatted '0.00, 0.00' in the order [coherence, fluency] and no other text or symbols. TEXT TO EVALUATE: (samples)
E
Additional qualitative results
We show in Figures 22 and 23 tokens activating the first experts as they appear numerically. We sample bins of expert coefficient value to show both tokens that highly activate the experts and those that do so only mildly. As can be seen, both high- and low-level specializations emerge in both GPT and Pythia models.
Whilst we observe specializations to a range of concepts (such as punctuation, MMO games, words in specific contexts), we do not notice any systemic differences between the types of expert special- izations that emerge between the two models in MxD layers.
44

---

P10hi Tokens routed to specific experts
Feature 0
Feature 1
Feature2 2
Artivating 3.06 1.63
3,61 9.57
3.61
conslanity
Artivation betomes 2.00 and 1.66
Activation betweea 3,65 and 3.81
Activation betweee 2.41 and 2,01
Wesian Jalowess
Dh
201 soat rea peris atl Test
2.01
Activttaon betwore 2.52 and 3,00
wells kor corea
free all the
Activation and 2:32
Activation of aeniee
Activation
Purk sentt
deeado afte
M
rinder mat
and
Activatann briwers 1.77 and 1.54
As Activatinn letwere 6. 1.30
Activation between 40 1:20
fer eiter for desl aestert ailiipilen
Laggter My 11860)
acttried
of
between 0,40
Aierie vertes aitict supeinn
wis unile eateritte em
sendcheri -i perioe.
Feature
Feature
Feature
Attivatina betweees 5.20 and 3.54
As Mitation letwern 3.00 and 5.71
Artisation letwern 3.51 and 4.21
accorl pcoectar 1TTEREIE adiing
siediteccio sillinee vot
Activatin
2.47
Activation
2:31
gearally
raain to 3C receena
or sateuer
Activation
Activation bethoes 1:31 and 3.97
Hutting caed
ME 3/ Mare retter ucadyy
Fraamsiinn lestac racaa Zie
srupptts
05 sibe sne Welore aaterty
Artivarion betooes 8.66 and 1.31
Six
coa enry
aroad TV
Activation between 0.00 and 0.52
bete rol: 0.70
Day
Activelin belevea 8.00 and 8.
ritic
sciatin: Milonn Hares
2iardcteen: aerrer hellereyy.
delicgg is Irirg beatts book opurall Soomso
Feature
Feature? 7
Feature 8
Suppor
2.19
Everecttyy ctiomia
Activation Weterts 5.00
Activation between 243
Srstie 1888
acstactdd maint
voess 70 Galleern 3.22
Mie Bos fuelee let live medd Best Satia
Acthvation
Activation Genett Centat
Activation between 1,22 1.2
nost
lesitinn 1 Je0.lee Mage:
Site suli (atse Stadio the
Activation betwren 6.50 and 1.41
Artisation
2n1. a0d. coltrabiie and ca
dredaci
wai Gote
Activadios betanen 3.60 and 2.14
1au7e 24
intaraaiin
R Rouster
reaaectionsa sicamnn
and uttcinn inttind lam
2e protinnd
Activation between 0.41 und 0.51
mic anectie do caraint an qucertertt thet the
Siant. Pret nechia
100 faar
iecroniele cowe ware encvll rivrs with
Activation
wee mavirals sefocciaon watered tye
Figure 22: Tokens activating the first 9 numerical experts on MxDs with K = 32 trained on Pythia-410m Om; we sample 6 bands of activations show both tokens that highly activate experts and those that activate them only mildly. Magnitude of activation is denoted by the orange highlight. Moderate specialism emerges, e.g., to MMO games, abbreviations, and words in specific contexts.
45

---

G2-2-124m: Tokens routed to specific experts
Feature 0
Feature 1
Feature 2
Artivarion bettreen and 6.25
Activation bettowen 5.70 and 6.95
Activation between 5.35 and 6,12
Acthodion bebwwen 6.17 and 5.21
Activation between 4.03 and 5.79
Acthvaton between 4.20 and 5.35
More
1artar 11/ A09
42g ressaarts
Lecerieta
returikurg haate often Clumeritt ror the
Clcadl Nerli or Mers orltn cmenrc the decrtatrnl
sienk/tects kanwers geotrege belly eeturon the
Artication between 3,21 and 0.20
dectcety 11r9
Activasion between 2.13
2017
Activation mebe
2:32
si purteer wn werth
Acthaaion betheen and
sall my custsase 68
opted
caryarry Now lea
Activatius belnnes 1.18 and 2.32
1077 MA
Activation betwere
caaisinn sons clack 9o
nol sereria 11.
(nd ded crecede
Activation bettenen and
and feree
Activation
rec-el juct may, Ligpect kntleat stenoltaate free sesie.
slate <iadeftucty
Feature 3
Feature 4
Feature 5
Activation between 1.26 2.13
Activation butweee 4.31 and 1.85
Activation befreeen 8.50 and 7.85
Activtton between 1.2 and 1.78
Activation belineus 5.57 and 4.21
Activation befwern and 6.35
teme
intrectionl:
sigeriioe Hoset Cor andeec
Arthvation scally cor
Actiaatinn between 3.60 and 5.34
the
cevesing car
Activation bottonen 3.52 and 3.30
Activation
Arthvetion between 0.73 and 1
720ats iedtabeeed seary cite
The Wooi ding emtrin reted Eeselas
parr leattt
byse accian TM thme citis me
betteenn 1x dtanne recetiis rasoctas
Activetion betwees and
ceratats
Beitaddd is 11st or nettoone wetter a4
Activation
sepory
W
parteer Due 11 colectivv free
Juser remeddag
their
2151l lox
the- the-0
Activatton butwoes
<laerfattt Foll 20100.
verraar, soultry Latooiae the 21 2014
wiing
Feature 6
Feature 7
Feature 8
clatrs cal lor (iiiees or USMS oatery
Acthvation
and
Activation betwous B.04 6.03
Activation 5.07
a05 Tay secar (1att fitical vee
180 -00 Benlery 1o/e served
dack S0tT
Acti-ation belween 3.80 and 3.07
Acthotion betwers 1.51 and 271
and
Actisation borwoee and LB1
-100-414
tintte sinem:
only Mastonptens
5 rekittenleit live daartee for.
Acthvaiin between 0.00 0.99
mith not Were ato daalery Brom Tay
T0we: befure why Talkyss vaite is Gex enciciand -390
Artivation
vole de Dout
Pr14s001 Sickue
Activation 6, and 1,27
3a1 Capaciia
Activativn beiwwrn 111 and
st1ff 1ar0e Ineesee 1ffa:
offfce
1OA 15 sin no 1001 <jenstiast Fin
Figure 23: Tokens activating the first 9 numerical experts on MxDs with K 32 trained on GPT2- 114m; we sample 6 bands of activations to show both tokens that highly activate experts and those that activate them only mildly. Magnitude of activation is denoted by the orange highlight. Moderate specialism emerges, e.g., to punctuation, names, and months.
46


Reviews:
[{"review_id": "KyxxPXVlIH", "reviewer": "Reviewer_9H1M", "summary": "This paper introduces Mixture of Decoders (MxDs), a novel method that overcomes the accuracy-interpretability trade-off in large language models by using layer-level sparsity to expand dense layers into thousands of specialized sublayers, faithfully preserving the model's performance while learning interpretable features.", "questions": "See Weaknesses.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\n\n1.  **Addresses a Critical, High-Impact Problem:** The paper tackles the core trade-off between performance and interpretability in large language models. This is a central challenge in AI safety, transparency, and alignment, making the research highly relevant and important.\n\n2.  **Clear Conceptual Innovation:** The proposed shift from \"neuron-level\" to \"layer-level\" sparsity is a powerful and intuitive conceptual leap. By arguing that individual neurons are not the right unit of analysis and proposing sublayers instead, the paper offers a compelling new way to think about model decomposition.\n\n3.  **Focus on \"Faithfulness\" and Practicality:** The authors rightly emphasize that an interpretable model is only useful if it faithfully reproduces the behavior of the original. This practical focus on maintaining low next-token loss ensures that the resulting model is not just a simplified, inaccurate curiosity but a potentially viable replacement for the original component.\n\n4.  **Strong Theoretical and Empirical Backing (as claimed):** The paper claims to provide a formal proof that its method preserves the expressive capacity (full-rank weights) of the original layers, lending it strong theoretical credibility. This is supported by claims of extensive and rigorous experiments (\"108 sparse layers in 4 LLMs,\" \"348 sparse probing and steering tasks\") that show it \"pareto-dominates\" existing techniques.\n\n### Potential Weaknesses\n\n1.  **Unaddressed Computational Overhead:** The paper expands a single layer into \"tens of thousands\" of sublayers. While it mentions parameter efficiency, it says nothing about the potential computational costs. This raises key questions about inference latency (how fast is the conditional routing?) and training complexity (how long and how much memory does it take to train these MxD layers?), which are critical for practical adoption.\n\n2.  **Scalability to Frontier Models:** The experiments are conducted on models \"up to 3B parameters.\" While respectable, this is significantly smaller than current state-of-the-art open models (e.g., 70B+ parameters). It's an open question whether the training and memory dynamics of this method will scale effectively to these much larger and more complex systems.", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n## [W1] Computational comparisons\n\nWe 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.\n\n**Parameter count & FLOPs** \nWe 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:\n\n|                | Parameter count                                                       | FLOPs                                                                                            |\n| -------------- | --------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------ |\n| **Sparse MLP** | $H (I+O)$ *(encoder+decoder)*                                    | $H (I+O)$ *(encoder+decoder)*                                                               |\n| **MxD**        | $N (I+O) + H^* (I+O)$ *(conditional branch) + (encoder+decoder)* | $N (I+O) + H^* (I+O) + O/2$ *(conditional branch) + (encoder+decoder) + (hadamard product)* |\n\n> $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.\n\nFor a chosen expert count $N$, we set the width of Sparse MLPs (e.g. Transcoders) to $H:=N+H^*$ to parameter-match the models.\n\n**Empirical benchmarks**\nWe 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$.\n\n|                | Peak memory usage (MiB) | Latency (ms) | Parameter count | Reported FLOPs per sample (fvcore [1]) |\n| -------------- | ----------------------- | ------------ | --------------- | --------------------------- |\n| **Sparse MLP** | 386.50                  | 1.394        | 18,874,368         | 18,874,368                     |\n| **MxD**        | 389.50                  | 1.457        | 18,874,368         | 18,874,880                     |\n\nBoth 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.\n\nWe 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.\n\n## [W2] Scalability to frontier models\n\nWe 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].\n\nAs 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:\n\n| **Step**                                  | **9999**  | **19999** | **29999** | **39999** | **49999** |\n| ----------------------------------------- | --------- | --------- | --------- | --------- | --------- |\n| **Transcoder** Normalized MSE (↓)      | 0.147     | 0.132     | 0.128     | 0.123     | 0.119     |\n| **Skip Transcoder** Normalized MSE (↓) | 0.128     | 0.107     | 0.102     | 0.100     | 0.093     |\n| **MxD** Normalized MSE (↓)             | **0.098** | **0.096** | **0.093** | **0.086** | **0.069** |\n\n> **Normalized MSE on Gemma2-27B, Layer 20, K=32, after 50k iterations.**\n> 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).\n\nWe 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).\n\n**Revised claims**\nWe 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*:\n\n[L196]: \n\n    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.\n\n[L321]:\n\n    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.\n\n\n----------\n\n**References**\n\n- [1]: fvcore (FAIR), detectron2 library.\n- [2]: Dunefsky, Jacob, Philippe Chlenski, and Neel Nanda. \"Transcoders find interpretable LLM feature circuits.\" NeurIPS 2024.\n- [3]: Paulo et al. \"Transcoders beat sparse autoencoders for interpretability.\" arXiv 2025.\n- [4]: Shazeer, Noam, et al. \"Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.\" ICLR 2017.\n- [5]: Meta AI. “The Llama 4 herd: The beginning of a new era of natively multimodal AI innovation.” 2025.  \n- [6]: Guo, Daya, et al. \"Deepseek-r1: Incentivizing reasoning capability in LLMs via reinforcement learning.\" arXiv 2025.\n\n---\n\nIf **Reviewer 9H1M** has any additional questions, we would be glad to answer them."}, {"review_id": "WZlOU3UpZP", "reviewer": "Reviewer_zUPk", "summary": "This paper proposes to learn an interpretable approximation of LLMs via sparse learning, but at the layer level rather than previously studied neuron level. The method replaces the computation in Transformer MLPs with a k-sparse sum of linear transformations on the hidden intermediate representation, and uses a Hadamard-product based parameterization to reduce the footprint of this sparse replacement. Through a special hadamard product trick, the mixture of linear transformations becomes much more computationally feasible, and experimental results show much closer performance to baseline models, while preserving interpretability capabilities.", "questions": "1. Can you provide some high level justification for how we can think of the equivalence between a single expert in MxD and a single neuron in SAE as the core unit of decomposition? \n2. Despite efficiency of the sparse method not being the core importance in interpretability work, and despite the Hadamard trick improving feasibility, can you comment on the efficiency of this method versus SAE and transcoders?\n3. Lines 197-198 suggest that this work would scale well to larger models, but what is the justification for this claim based on the MxD construction? \n4. How are sparse layers trained? Is it via layer replacement and then training just the layers? Some training details seem glossed over perhaps as a result of being in line with prior work, but these details are important for reproducibility. \n5. The factorized forward pass in equation 5 is critical to understand how equations 3 and 4 relate to each other, but the rank preservation text appears between equations 4 and 5. I might suggest reordering sections 2.4 and 2.3 for reader understanding, but also additional clarity while introducing the parameterization could help.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths\n1. This paper presents clear experimental evidence for its claims, and presents numerous additional experiments for completeness in the appendix, like removing the shared expert or finding dead experts. \n2. The utility of the induced sparsity is well tested with probing experiments as well as steering experiments. Experimental results seem consistent with popular prior work on interpretability. \n3. The faithfulness results as measured by loss are far better than prior work in preserving original model performance. This helps address a very core issue in interpretability work where sparse approximations greatly underperform the corresponding LLM they intend to approximate. \n4. The Hadamard product trick is interesting and may be extensible to other applications outside of this paper, and seems to be a key insight to make this work feasible given the complexity of a naive implementation. \n\nWeaknesses\n1. The equivalence between sparse neurons from prior work and a sparse expert from this work is not clear to me as the same level of decomposition. While the steering experiments show that there is decomposition that can be steered, and probing experiments that show that different $a_n$ coefficients activate with different topics, the granularity seems quite different beyond these evaluations. In fact, even addressing potential limitations of this larger unit of sparsity would be helpful to situate this work in the context of prior work looking at the sparse neuron level. \n2. Some key details in the paper could be clarified for easier understanding; for example, more explanation of the Hadamard trick would be helpful as the explanation is quite dense and it takes some convincing of oneself as the reader that this works, some training details may be missing for how MxDs actually substituted and trained using OpenWebText, and some details about top-k activation training might also be necessary (despite it coming from prior work, it is important enough to mention its role in this work).", "was_revised": true, "final_justification": "The current work is clearly written, uses extensive experimental evidence to support their claims, and makes a measured contribution within the AI safety/interpretability subfield that I believe others within this field will find useful.", "rebuttal": "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:\n\n\n## [W1 & Q1] Equivalence between experts/features\n\nWe 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}=\\texttt{TopK}(\\mathbf{E}^\\top\\mathbf{x})$ and $\\mathbf{a}=\\texttt{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.\n\nThe 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.\n\nWe 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.\n\n\n## [W2 & Q5] Hadamard trick & experimental details\n\nWe 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.\n\nTo enhance understanding even further, we will condense some of the content in this tutorial-style notebook into the appendix of the revised manuscript, inline.\n\nWe 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. \n\n## [Q2] Efficiency of MxDs vs Sparse MLPs\n\nWe 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.\n\n**Parameter count & FLOPs** \nWe 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:\n\n|                | Parameter count         | FLOPs                         |\n| -------------- | ----------------------- | ----------------------------- |\n| **Sparse MLP** | $H (I+O)$             | $H (I+O)$                   |\n| **MxD**        | $N (I+O) + H^* (I+O)$ | $N (I+O) + H^* (I+O) + O/2$ |\n\n> $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.\n\nFor a chosen expert count $N$, we set the width of Sparse MLPs (e.g. Transcoders) to $H:=N+H^*$ to parameter-match the models.\n\n**Empirical benchmarks**\nWe 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$.\n\n|                | Peak memory usage (MiB) | Latency (ms) | Parameter count | Reported FLOPs per sample (fvcore [1]) |\n| -------------- | ----------------------- | ------------ | --------------- | --------------------------- |\n| **Sparse MLP** | 386.50                  | 1.394        | 18,874,368         | 18,874,368                     |\n| **MxD**        | 389.50                  | 1.457        | 18,874,368         | 18,874,880                     |\n\nBoth 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.\n\nWe thank the reviewer again for the chance to present the computational comparisons, and will include all tables here in the revised manuscript.\n\n\n## [Q3] Justification for scalability to larger models\n\nWe 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. \n\nFor 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:\n\n| **Step**                                  | **9999**  | **19999** | **29999** | **39999** | **49999** |\n| ----------------------------------------- | --------- | --------- | --------- | --------- | --------- |\n| **Transcoder** Normalized MSE (↓)      | 0.147     | 0.132     | 0.128     | 0.123     | 0.119     |\n| **Skip Transcoder** Normalized MSE (↓) | 0.128     | 0.107     | 0.102     | 0.100     | 0.093     |\n| **MxD** Normalized MSE (↓)             | **0.098** | **0.096** | **0.093** | **0.086** | **0.069** |\n\n> **Normalized MSE on Gemma2-27B, Layer 20, K=32, after 50k iterations.**\n> 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).\n\nWe 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).\n\n**Revised claims**\nWe 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*:\n\n[L196]: \n\n    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.\n\n[L321]:\n\n    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.\n\n\n## [Q4] Training details of sparse layers\n\nWe 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).\n\n**We will add this clarification around [L199]:**\n\n    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.\n\n\n## [Q5] Confusing order of equations\n\nOur 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.   \n\n\n----------\n\n**References:**\n\n- [1]: fvcore (FAIR), detectron2 library.\n- [2]: Shazeer et al. \"Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.\" *ICLR 2017.*\n- [3]: Meta AI. “The Llama 4 herd: The beginning of a new era of natively multimodal AI innovation.” 2025.\n- [4]: Guo et al. \"Deepseek-r1: Incentivizing reasoning capability in LLMs via reinforcement learning.\" arXiv 2025.\n- [5]: Dunefsky et al. \"Transcoders find interpretable LLM feature circuits.\" NeurIPS 2024.\n- [6]: Gonçalo et al. \"Transcoders beat sparse autoencoders for interpretability.\" arXiv 2025.\n\n---\n\nWould **Reviewer zUPk** have any further questions? We’d be happy to address them."}, {"review_id": "VSJsvY2wx2", "reviewer": "Reviewer_8ypg", "summary": "The paper proposes an inherently interpretable method that complements existing neuro-sparsity based methods, e.g. sparse overcomplete autoencoders (SAEs) and Transcoders, by imposing sparsity constrains at the layer level.  \n\nMore specifically, proposed Mixtture of Decoders (MxD) method, replaces MLP components present in standard LLMs by an autoencoder where the decoder part is defined as a sparse mixture of decoders who match in number of parameters/capacity the original MLP it aims to replace. In doing so the proposed method is capable of learning specialized expert (decoder) branches that are interpretable.\n\nExperiments on 4 LLMs (~3B parameters), i.e. GPT2-124M [3], Pythia-410m, 187 Pythia-1.4b [41], and Llama-3.2-3B [1], show that the proposed method are comparable with the state of the art methods (SAEs and Transcoders) while having a performance that is significantly robust to varying levels of sparsity.", "questions": "- As indicated in the manuscript the proposed methods introduces the capability of learning a representation among a mixture of sparse decoders. In this regard the training process of the proposed method is significantly more complex that of the considered competitors (SAEs and TCs). May you provide some details on how the training processes of these methods compare w.r.t. each other?\n    \n- In L. 196-198, it i stated that there is the expectation that the proposed MxDs are capable of scaling to significantly larger models. May you provide some theoretical insights or perhaps an intuition of why this would be the case? The manuscript would benefit from having this pointer further\n    \n    elaborated.\n    \n- In l.209-210, it is stated that [17,19] find sparse solutions to be more interpretable. Are there any guarantees, or perhaps clear indicators, that the findings from those articles are transferable to the proposed setting?\n    \n- In Sec. 3.2.2, two LLMs are used as gold standard for measuring properties of the output produced by feature steering. Since these are models that are not perfect, is there is pointer to existing work showing the validity of this practice/protocol?", "limitations": "In addition to the points listed as weaknesses I would stress the following aspects:\n\n- In the related work section (Sec. 4), the aspect related to Sparse Decomposition, i.e. the core direction of this paper, limits itself to listing some existing efforts. This section would benefit with an explicit positioning on how the proposed method compares w.r.t. existing efforts.\n    \n    In addition, The proposed method seems to bear resemblance to the work from [Pearce et al, 2025], [Dooms et al., 2025] [Adler et al., 2025] (please see below for complete references). These works do not seem to be discussed in the manuscript.\n    \n- One of the metrics used for the evaluation of the proposed method is “faithfulness” to the base model composed by standard MLPs. In this regard, the evaluation seems to be equating comparable predictive performance with faithfulness. While there should be the requirement to meet (or improve) the performance already achieved by the original model, meeting this requirement does not guarantee that new model will learn the same representation as that learned by the original.\n    \n- Empirical evidence, e.g reconstruction result (Fig 10), output space similarity (Sec. B.1), etc., used to support statements made in the paper are actually  delegated to the Appendix. This gives the impression of the paper not being self-contained. A better balance needs to be achieved here as to ensure statements made are paired with proper evidence (or pointers to the literature).\n    \n- In Sec. 3.2.2, two LLMs are used as gold standard for measuring properties of the output produced by feature steering. Since these are models that are not perfect, this casts doubts on the observations made from this results.\n    \n\n**References**\n\n- Pearce et al. \"Bilinear MLPs Enable Weight-based Mechanistic Interpretability\" ICLR 2025.\n    \n- Dooms et al. \"Compositionality Unlocks Deep Interpretable Models\" AAAI 2025 ColoRAI Workshop\n    \n- Adler et al, \"Towards Combinatorial Interpretability of Neural Computation\" arXiv:2504.08842", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n- To a good extent the manuscript has a very clear structure and the way content is presented, to a good extent, provides a very good and clear flow in the presentation of content.\n    \n- The proposed method seems to provide comparable if not superior performance to that of state of the art methods.\n    \n- The theoretical proofs put in place show, the proposed method has the capability to provide the desired sparsity properties without sacrificing information in the process, e.g. as is the case with SAEs which lose some information due to imperfect reconstruction processes.\n    \n- Several resources are released in order to ease the testing of the proposed method.\n    \n\n**Weaknesses**\n\n- The proposed evaluation seems to equate equal predictive performance to “faithfulness” to the original (standard MLP) model.\n    \n- Some of the adopted evaluation procedures (model faithfulness and feature steering) do not seem to be adequate.\n    \n- Some relevant parts of the manuscript have been delegated to the appendix.", "was_revised": true, "final_justification": "The clarification on the faithfulness loss effectively clarified my initial misunderstanding w.r.t. [W1 & L3]. \n\nThe suggested Revised Limitations text and other actions indicated to be taken to address [W2 & Q4 & L5] seem adequate. Likewise for my question (Q3) regarding the relevance of [17,19], i.e. the suggested course of action seems transparent enough and adequate.\n\nRegarding [W3 & L4], thanks for the indicated modification (of moving fig. 10 into the manuscript), I believe this would be  good complement to the results reported on predictive performance (fig.3). I would encourage the authors to include the pointers on the training complexity (Q1), present in the rebuttal, in the main body of the paper. Beyond that, I  understand it can be challenging to present complex and/or extensively-validated ideas in the given page limits. However, that is a constraint we need to deal with and for which a proper trade-off must be found.  Sometimes, aspects like this make the difference between two papers with two good ideas.\n\nRegarding [L1 & L2], positioning with respect to related work on sparse decomposition methods, I acknowledge two of the references are very recent. While I do not consider the  missing positioning with respect to those two as grounds for rejection, I would still encourage the authors to include them in their related work section for completeness. The work from Pearce et al., is a bit more older, I would be more strict on the inclusion of that one.", "rebuttal": "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:\n\n## [W1 & L3] Predictive performance vs faithfulness\n\nWe inherit the term “faithfulness” from prior work [1], which uses the term synonymously with predictive ability, e.g.:\n\n    \"which seek to faithfully approximate a densely activating MLP layer with a wider, sparsely-activating MLP layer\"\n\nFurthermore, 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.\n\nWe 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. \n  \n## [W2 & Q4 & L5] Adopted evaluations & LLM as a judge\n\nWe 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].\n\nWe 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:\n\n**Revised limitations**\n[L330]\n\n    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).\n\n\n## [W3 & L4] Appendix results\n\nMany 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.\n\nWith the additional page available in the camera-ready revision, we will include results from Sect. B.1 in the main paper.\n\n## [Q1] Training complexity\n\nWe 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.\n\n**Parameter count & FLOPs** \nWe 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:\n\n|                | Parameter count         | FLOPs                         |\n| -------------- | ----------------------- | ----------------------------- |\n| **Sparse MLP** | $H (I+O)$             | $H (I+O)$                   |\n| **MxD**        | $N (I+O) + H^* (I+O)$ | $N (I+O) + H^* (I+O) + O/2$ |\n\n> $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.\n\n**Empirical benchmarks**\nWe 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$.\n\n|                | Peak memory usage (MiB) | Latency (ms) | Parameter count | Reported FLOPs per sample (fvcore [1]) |\n| ---| --- | --- | --- | --- |\n| **Sparse MLP** | 386.50                  | 1.394        | 18,874,368         | 18,874,368                     |\n| **MxD**        | 389.50                  | 1.457        | 18,874,368         | 18,874,880                     |\n\n## [Q2] Scalability to larger models\n\nWe 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:\n\n| **Step**                                  | **10k**   | **20k**   | **30k**   | **40k**   | **50k**   |\n| --- | --- | --- | --- | --- | --- |\n| **Transcoder** Normalized MSE (↓)      | 0.147     | 0.132     | 0.128     | 0.123     | 0.119     |\n| **Skip Transcoder** Normalized MSE (↓) | 0.128     | 0.107     | 0.102     | 0.100     | 0.093     |\n| **MxD** Normalized MSE (↓)             | **0.098** | **0.096** | **0.093** | **0.086** | **0.069** |\n\n> **Gemma2-27B, Layer 20, K=32, 50k iterations.**\n> 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).\n\nWe 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].\n\n**Revised claims**\nWe 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*:\n\n[L196]: \n\n    [...] we expect MxDs to scale just as well *as sparse MLPs* to models with tens of billions of parameters or more.\n\n[L321]:\n\n    [...] 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.\n\n\n## [Q3] Relevance of citations [17,19] \n\nWe 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.\n\n\n## [L1 & L2] Additional related work\n\nWe 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.\n\n**Concretely, a new section following the discussion of sparse decompositions will read:**\n\n    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).\n\n---\n\n**References**\n\n- [1]: Dunefsky et al. \"Transcoders find interpretable LLM feature circuits.\" NeurIPS 2024*.*\n- [2]: Gao et al. \"Scaling and evaluating sparse autoencoders.\" ICLR 2025*.*\n- [3]: Gurnee, et al. \"Finding neurons in a haystack: Case studies with sparse probing.\" *TMLR 2023.*\n- [4]: Zheng et al. \"Judging llm-as-a-judge with mt-bench and chatbot arena.\" NeurIPS 2023.\n- [5]: Bavaresco et al. \"Llms instead of human judges? a large scale empirical study across 20 NLP evaluation tasks.\" ACL 2025.\n- [6]: Wu et al. \"Axbench: Steering LLMs? Even simple baselines outperform sparse autoencoders.\" *ICML 2025.*\n- [7]: Cao et al. \"Personalized steering of large language models: Versatile steering vectors through bi-directional preference optimization.\" NeurIPS 2024.\n- [8]: Gonçalo et al. \"Transcoders beat sparse autoencoders for interpretability.\" arXiv 2025.\n- [9]: fvcore (FAIR), detectron2 library.\n- [10]: Meta AI. “The Llama 4 herd: The beginning of a new era of natively multimodal AI innovation.” 2025.\n- [11]: Guo et al. \"Deepseek-r1: Incentivizing reasoning capability in LLMs via reinforcement lerning.\" arXiv 2025.\n- [12]: Pearce et al. \"Bilinear MLPs Enable Weight-based Mechanistic Interpretability\" ICLR 2025.\n- [13]: Dooms et al. \"Compositionality Unlocks Deep Interpretable Models\" AAAI 2025 ColoRAI Workshop\n- [14]: Adler et al. \"Towards Combinatorial Interpretability of Neural Computation\" arXiv 2025.\n\n---\n\nMight **Reviewer 8ypg** have any additional questions? If all concerns were addressed, we would otherwise be grateful if the reviewer raised their score."}, {"review_id": "Lj3lkHvwx0", "reviewer": "Reviewer_JDFb", "summary": "The paper proposes the Mixture of Decoders (MxD) for post-hoc interpretation of MLPs in LLMs. The architecture achieves sparsity by approximating MLP with $N$ sparse transformations (or experts), which is unlike the existing sparse neuron approximations. The paper also proposes the Hadamard weight factorization, which allows for simplified implementation.", "questions": "1. Can the authors clarify how MxD can be used as a better interpretation tool by future research?\n2. Can the authors clarify the role of the shared expert?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The idea of sparse experts at the layer level is interesting and distinct from prior neuron-level sparsity formulations. It is nice that the paper also includes discussions about MxD's implementation and extension.\n2. MxD achieves a better sparsity-faithfulness trade-off across four models with sizes ranging from 124M to 3B. It also performs stably across different sparsity levels.\n3. The paper includes both faithfulness and interpretability evaluation of its approach, which is a fair experimental setup.\n\nWeaknesses:\n1. MxD does not show advantages over baselines on interpretability, based on the experimental results of probing and steering. Since the ultimate goal of this line of research is to enable better interpretation of MLPs, this approach does not particularly provide an advancement. On the other hand, the observation that MxD achieves higher faithfulness than existing approaches but only comparable interpretability also points out its technical weakness, i.e., its loss function does not particularly advocate better interpretability.\n2. The role of the \"shared expert\" in enabling the faithfulness gain needs more clarification. A key claim of the paper is that MxD improves reconstruction faithfulness under high sparsity via expert specialization. However, the presence of an emergent shared expert makes it unclear whether the gain stems from true specialization or this general-purpose expert. In an extreme case, it could happen that this shared expert replicates most of the original MLP and thus presents high faithfulness. However, relevant to Weakness 1, there's no particular architecture or loss function design that prevents this situation. \n3. Presentation: while the overall paper writing is clear, I find Fig 1 particularly hard to understand (or not intuitive enough without reading the technical details in the main text). For the neuron-level sparsity figure, the notation of $N$ corresponds to $H$ in Eq (2). It may be better to use the notations more consistently across the paper.", "was_revised": true, "final_justification": "During the rebuttal, my major concerns (no advantage in terms of interpretability evaluation and the potential collapse to shared experts) are addressed.", "rebuttal": "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:\n\n\n## [W1] Advantages for interpretability\n\n\n> ”Since the ultimate goal of this line of research is to enable better interpretation of MLPs, this approach does not particularly provide an advancement”\n\nWe 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**:\n\n\n- **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.`\n- **Reviewer-9H1M:** `tackles the core trade-off between performance and interpretability` and `central challenge in AI safety [...] making the research highly relevant and important`\n\nWould 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.\n \n\n## [W2 & Q2] Shared experts\n\n \nWe 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).\n\nPlease 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: \n\n| $\\vert\\vert \\mathbf{W}_{\\text{shared}} - \\mathbf{D}^* \\vert\\vert_F$ | $\\vert\\vert \\mathbf{D} - \\mathbf{D}^* \\vert\\vert_F$ | $\\left(\\vert \\mathbf{W}_{\\text{shared}} - \\mathbf{D}^* \\vert\\right)\\_{:3,:3}$ (9 elements of the difference matrix) |\n| --------------------------------------------------------------------- | ----------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------- |\n| 208.08                                                                | 221.6934                                              | [[0.09, 0.03, 0.06],  [0.21, 0.12, 0.07],  [0.17, 0.15, 0.11]]                                                    |\n\nWe 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:\n\n\n    \"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.\"\n\n\n## [W3] Fig. 1 presentation\n\nWe 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?\n\n## [Q1] Can the authors clarify how MxD can be used as a better interpretation tool by future research?\n\nBy 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.\n\nSecondly, 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.\n\n\n----------\n- [1]: Elhage, et al., \"Softmax Linear Units\", Transformer Circuits Thread, 2022.\n- [2]: Wang, Lean, et al. \"Auxiliary-loss-free load balancing strategy for mixture-of-experts.\" *arXiv preprint arXiv:2408.15664* (2024).\n- [3]: Liu, Boan, et al. \"Diversifying the mixture-of-experts representation for language models with orthogonal optimizer.\" *arXiv preprint arXiv:2310.09762* (2023).\n- [4]: Shazeer, Noam, et al. \"Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.\" *arXiv preprint arXiv:1701.06538* (2017).\n- [5]: Park, Jungwoo, et al. \"Monet: Mixture of monosemantic experts for transformers.\" *ICLR 2025.*\n- [6]: Oldfield, James, et al. \"Multilinear mixture of experts: Scalable expert specialization through factorization.\" *NeurIPS 2024.*\n- [7]: Yang, Xingyi, et al. \"Mixture of experts made intrinsically interpretable.\" *ICML 2025.*\n- [8]: Meta AI. “The Llama 4 herd: The beginning of a new era of natively multimodal AI innovation.” 2025.\n- [9]: Guo, Daya, et al. \"Deepseek-r1: Incentivizing reasoning capability in LLMs via reinforcement learning.\" arXiv preprint arXiv:2501.12948 (2025).\n\n---\n\nWe welcome **Reviewer JDFb** to ask any additional questions that might raise their support of the paper--we will gladly answer them."}]

Author rebuttal:

Meta-review: ## Summary
This paper introduces Mixture of Decoders (MxDs), a method for interpretable large language models that applies layer-level sparsity instead of neuron-level sparsity. Using a Hadamard-product factorization, MxDs expand dense MLPs into sparse sub-decoders while preserving full-rank capacity. The approach is validated with extensive experiments across multiple LLMs, showing strong faithfulness and scalability potential.

## Strengths
- Moves from neuron-level to layer-level sparsity, a significant conceptual shift.  
- Demonstrates high faithfulness while avoiding the performance drop common in sparse interpretability methods.  
- Provides theoretical guarantees and strong empirical validation.  
- Rebuttal added further evidence (computational benchmarks, experiments on larger models, clarification of shared experts).  

## Limitations and Weaknesses
- Computational and scalability implications for frontier-scale models remain partially open.  
- Interpretability gains over baselines (probing/steering) are not always stronger, though faithfulness is significantly better.  
- Some clarifications (training details, evaluation protocols) were only added after rebuttal and should be made explicit in the final version.  

## Rebuttal and Discussion
The rebuttal was thorough and addressed reviewers’ main concerns. Additional benchmarks demonstrated minimal computational overhead, while experiments on larger (27B) models supported scalability claims. Clarifications on interpretability evaluation and the role of shared experts resolved outstanding doubts, and reviewers updated or confirmed high scores after discussion.

## Recommendation
The reviewers converged on a positive view after rebuttal, with multiple scores at the top end and one reviewer explicitly raising their rating. The paper offers a novel and well-validated approach to interpretability that balances theory and practice, with strong evidence that it advances the field. Despite some remaining limitations, the strengths and overall contribution are clear.

**Final Decision: Accept.**
Decision: Accept (poster)