Wasserstein Distances, Neuronal Entanglement, and Sparsity
Authors: Shashata Sawmya
Linghao Kong
Ilia Markov
Dan Alistarh
Nir N Shavit
Venue: ICLR 2025 Spotlight (2025)
Keywords: Polysemanticity
Disentanglement
Wasserstein Distance
Sparsity
Large Language Models
Primary area: interpretability and explainable AI
TL;DR: We show that the Wasserstein distance of a neuron's output distribution to a Gaussian is a pertinent indicator for its degree of entanglement, and propose a new framework for better investigating disentaglement under sparsity.

Abstract: Disentangling polysemantic neurons is at the core of many current approaches to interpretability of large language models. Here we attempt to study how disentanglement can be used to understand performance, particularly under weight sparsity, a leading post-training optimization technique. We suggest a novel measure for estimating neuronal entanglement: the Wasserstein distance of a neuron's output distribution to a Gaussian. Moreover, we show the existence of a small number of highly entangled "Wasserstein Neurons" in each linear layer of an LLM, characterized by their highly non-Gaussian output distributions, their role in mapping similar inputs to dissimilar outputs, and their significant impact on model accuracy. To study these phenomena, we propose a new experimental framework for disentangling polysemantic neurons. Our framework separates each layer's inputs to create a mixture of experts where each neuron's output is computed by a mixture of neurons of lower Wasserstein distance, each better at maintaining accuracy when sparsified without retraining. We provide strong evidence that this is because the mixture of sparse experts is effectively disentangling the input-output relationship of individual neurons, in particular the difficult Wasserstein neurons.

Published as a conference paper at ICLR 2025
WASSERSTEIN DISTANCES, NEURONAL ENTANGLEMENT, AND SPARSITY
Shashata Sawmya"", Linghao Kong"", Ilia Markov Dan 1istarh233,4 Nir Shavit MIT FIT Austria Neural Magic *Red Hat (shashata, Linghao, shani @mit. (ilia markov, aam.allstarreeiss....
ABSTRACT
Disentangling polysemantic neurons is at the core of many current approaches interpretability of large language models. Here we attempt to study how disentan- glement be used to understand performance, particularly under weight spar- sity, a leading post-training optimization technique. We suggest a novel measure for estimating neuronal entanglement: the Wasserstein distance of neuron's 's out- put distribution to a Gaussian. Moreover, we show the existence of a small number of highly entangled "Wasserstein Neurons" in each linear layer of an LLM, char- acterized by their highly non-Gaussian output distributions, their role in mapping similar inputs to dissimilar outputs, and their significant impact on model accu- racy. To study these phenomena, we propose new experimental framework for disentangling polysemantic neurons. Our framework separates each layer's inputs to create a mixture of experts where each neuron's 's output is computed by mix- ture of neurons of lower Wasserstein distance, each better maintaining accuracy when sparsified without retraining. We provide strong evidence that this is be- cause the mixture of sparse experts is effectively disentangling the input-output relationship of individual neurons, in particular the difficult Wasserstein neurons.
INTRODUCTION
Disentangling polysemantic neurons into their component, human Humannuunerrstnddabll features has been a longstanding goal of machine learning interpretability research (Olah et al., 2020:
Jermyn et al.)
2022 Elhage et al. 2022; Gurnee et al. 2023] [Tmmllttoo] 2024 Gurnee et al. 2024). While neurons are the basic building blocks of neural network architectures, they do not map one-to-one with specific features. Instead, neurons frequently engage in polysemantic representations, where they are activated by multiple, unrelated concepts and detect diverse features (Arora et al., 2018] Mu & Andreas 2020). It is suspected that every neuron is polysemantic some degree (Lecomte et al. 2023). and so we will refer to all neurons as polysemantic in this work.
Due to the importance of highly polysemantic neurons in network's computation (Bricken al., 2023]. the question of whether these neurons require more parameters naturally arises. However, the effects of polysemanticity on network performance under weight sparsity has not been well explored. Weight sparsification (Hoefier et al. (2221) aims to reduce the number of executed param- eters large language models (LLMs) by setting certain weight values to zero to improve efficiency. Various sparsification algorithms have been developed for this process (Han al.) 2015 Sun al. 2023 Frantar & Alistarh 2023). This paper investigates the relationship between an individual neu- ron's degree of entanglement (which we will formally define in a later section) and its ability to be sparsified in real-world models. To the best of our knowledge, this is the first work to explore this crucial perspective of entanglementt-epeedde
model sparsification.
To better understand the impact entanglement on sparsification, we introduce a novel metric that quantifies neuron's 's degree of entanglement. This metric is the Wasserstein distance between
"Equal contribution. Author order determined by coin toss.
Code available attttps: //gi thub. com/ Swam/Shaiiiai wttpsshaattttaaaaaaaaaaaaeea exxaasiion

---

Published as a conference paper at ICLR 2025
neuron's output distribution and Gaussian (Equation D. high Wasserstein distance (Figure Td. A are crucial for the performance of a network and very sensitive to pruning. We provide evidence that neuron's 's Wasserstein distance is related to its ability to distinguish similar inputs to different outputs through its dot product, and we refer to these neurons as especially entangled (Equation 2p. Akin previous works investigating special types of neurons Gurnee et al., [2023; Stolfo et al. 2024, Gurnee et al. 2024), this work explores the role of crucial neurons with implications for interpretability, specifically in the context of network sparsity.
We find that neurons with a particularly
30:
OxtpalNAMe
Figure I: The output distributions of neurons in Llama-2-7B computed densely and at 90% sparsity on Wikitext-2. WD refers the Wasserstein distance of the output distribution Gaussian. RI refers to the relative improvement of Sparse Expansion over SparseGPT. (a) The dense output dis- tribution of random neuron with a WD of 0.050 is well captured by SparseGPT, and (b) expanding this neuron via Sparse Expansion imparts only small (18%) increase in performance. (c) The clus- ter outputs are all concentrated in close proximity to each other. (d) SparseGPT struggles to capture the dense distribution of an entangled neuron with a WD of 0.524. (e) Following expansion, the sparse output the entangled neuron much better captured, leading to more improvement (77%). (f) Each expert specializes over different portion of the distribution.
To analyze the phenomenon of neuronal superposition under sparsity in greater detail, we create an experimental framework, which we dub Sparse Expansion. It expands a model into a mixture of sparse experts by clustering input embeddings layer-wise. Based on this clustering, Sparse Expan- sion utilizes the input-aware nature of the SparseGPT (Frantar & Alistarh (023) pruning algorithm to specialize different sparse experts to different sets of inputs, starting from the same base weights. Through Sparse Expansion, we are able to analyze the entangled neurons in much more detail, since now different subgroups of the inputs are being computed with different edges (Figure IF. A8 We find that as neuron lose edges, its output distribution tends to shift toward Gaussian distribution (Figure A9). However. through Sparse Expansion, the original output distribution can be better pre- served under sparse computation (Figure 3 A8e). We relate our findings to recent theoretical work on the bounds of neural computation under superposition (Hänni et al. 2024 Adler Shavit 2024).
Our main technical contribution is a detailed study of how model accuracy under sparsity is related to its degree of neuronal entanglement. In every LLM, there exist neurons that have striking, irregular output distributions (Figure 2p. ATD. These neurons have an outsized effect model performance and seem to be responsible for differentiating similar input vectors (Figure 2). We believe that the existence of these neurons is manifestation of polysemanticity in real-world language models. We find that the Wasserstein distance to a Gaussian is a strong indicator of such neurons.
In the next section explain such "Wasserstein neurons", neuronal entanglement, and the im- plication of ablating Wasserstein neurons in LLMs in detail. We then formulate our experimental framework Sparse Expansion and show how effectively disentangle the input-output relationship of neurons through Sparse Expansion, as well as some empirical computational bounds. Finally, we present some results showing its performance relative to other -ftte-of-the-art one-shot compression techniques in the hopes of inspiring future sparsification algorithms.
2
WASSERSTEIN NEURONS
2.1
CHARACTERIZING NON-GAUSSIAN NEURONAL OUTPUT DISTRIBUTIONS
We investigate the output distributions of individual neurons in all linear layers of transformer feed- forward networks (FFNs) during inference. Specifically, consider a linear operation Y = W + b, where Y E R is the output matrix, W E RR is the weight matrix, b E RD is the bias vector,
2

---

Published as a conference paper at ICLR 2025
broadcasted across all neurons, and X R= is the imput matrix, where each column represents an input vector. Each neuron is an individual row of W. and we collect individual scalar elements from the corresponding row in Y as the output distribution for that neuron.
We focus our analysis in Pythia- 1.4B (Biderman et al. 2023). Llama-2-7B (Tovroon et al., 2023),
and Llama-3-8B (Dubey et al. 2024). Most neurons exhibit reasonably Gaussian output distribu- tion after their dot product with the input vector (Figure Dh. 2b. However, we find the existence of a small group neurons with highly non-Gaussian outputs (Figure Ii. 2+) in all FFNs (Figure[A])
To characterize the degree of difference in terms of the shape of these distributions- the non- Gaussian output distributions of certain neurons with the Gaussian-like output distribution of most neurons- wwe considered several metrics, such as entropy. However, the Wasserstein distance (WD) (Kantorovich] 2006] Villani et al. 2009) proved be the most effective metric for quantifying this difference. In optimal transport theory, the WD measures the minimal transportation cost between two distributions, taking their geometry real space into account.
To find the WD of every neuron to the Gaussian N, we crucially first normalize the output distribu- tions of each neuron n to have zero mean and unit variance, and compare this normalized distribution n' to (0, 1). This normalization is performed because the range of neuron output distributions is quite variable, and we wanted to prioritize the differences in the shape of the distributions, rather than other properties. We use the 1-Wasserstein distance in one dimension, as shown in Equation I
Wi(n', --) (z)ldd.
| (2)
F-1 and 4p are the inverse cumulative distribution function of n' and (0, 1), respectively, which can be approximated with empirical data. To compute the WD of every neuron efficiently. we use the SciPy implementation (Virtanen al., 2020). When computing the difference metric in this way, we find that our originally observed neurons (Figure [.. (A8)) have been designated correctly with high WD to N. We thus term these neurons "Wasserstein neurons." " We also observe little overlap between neurons with high mean weight magnitudes and Wasserstein neurons (Figure A4a).
We additionally analyze Pythia-1.4B across its training, from network initialization to the final step. We find that Wasserstein neurons do not seem to receive more weight updates than other neurons (Figure A2)). Interestingly, we also find that Wasserstein neurons arise relatively early on in training, within 10-20 billion tokens (Figure A2 p). This phenomenon is likely related to and manifestation of other observations that fundamental model training dynamics rapidly stabilize, such as the rank of the gradient or the largest eigenvalue of the loss hessian (Gur-Ari al. 2018 Zhao et al. 2024 Noci et al. 2024). We leave further investigations into this crucial training period to future work.
2.2
WASSERSTEIN NEURONS AND ENTANGLEMENT
Here, we define and study the notion of entanglement of these Wasserstein neurons in greater detail by positing a new avenue to investigate entanglement. According to superposition theory, as the number of features increases relative to the number of neurons, features are forced become non- orthogonal in order to represent more of them, thus increasing entanglement (Elhage et al.) 2022). Consider neurons that must attend to multiple of these features. As the number of features increases, and different features are forced to become more similar in direction, such neurons must still manage to distinguish between them. Therefore, in this context, neurons that are highly entangled have the task of differentiating between similar input vectors, and mapping them to different output values.
To mathematically explore this concept, we study the input-output (IO) relationship of individual neurons. We introduce the metric "mapping difficulty" (MD), which measures how often a neuron must generate dissimilar outputs from similar inputs through its dot product computation. The MD for a particular neuron, given its weights and set of inputs, is calculated as follows (Equation 2):
|1 {
vill
[x
x;
MD(w, X) x)= = mean
15i<j5n
N,
Nu
(2)
Bi, By
E X,
y: = w Bi.
Ti = X
N₂ = max 15i<J5n
[||2; - 2y ||},
No = median (lle- - 1/1
15i<s5n
3

---

Published as a conference paper at ICLR 2025
2; and z; represent two distinct input vectors from the set of inputs X. Y4 and yy represent the two output scalars as a result of the dot product of an individual neuron's 's weights w with the inputs. For every pair of imputs, zero and one using the maximum norm N₂ We then compute the I,2 norm of the difference their corresponding outputs, and normalize them with the median norm Ny- More details on the rationale behind the normalizing factors can be found in Appendix [A.8] The MD of neuron can thus be calculated the average of the ratio between the normalized difference in outputs the normalized difference in inputs. Intuitively, greater MD means that neuron generally increases the separation of similar inputs into more dissimilar outputs.
compute the L norm of their difference, then scale the norms between
(e) Random Neuron Quipuut Distribution
(b) Random Nouron
(c) Entangled Neuroa Outgut Distribution
(4) Eatangled Neuron
(c) Entangleont Degres
Ic Relitionship
10 Relationship
Magging Diliealty
Dendry
94
03
02
04
Deiput
Input
Outgut
Ieput
Figure 2: A measure of neuronal entanglement. (a) The output distribution of a random neuron. (b) The normalized 2 plot of random neuron's 's pairs of inputs and outputs. (c) The output distribution of a Wasserstein neuron. (d) The normalized L plot of a Wasserstein neuron's 's pairs of inputs and outputs. This neuron must map fairly similar inputs dot product operation. The neurons are from the up projection matrix of the second FFN block in Pythia-1.45 (e) The MD of neuron is highly correlated with its WD. The selected random and Wasserstein neurons are highlighted in their respective colors.
outputs that very far apart through its
For the two neurons we have selected before, we plot the normalized L for pairs of inputs ([|x; - 2,|l/Na) and outputs (ly yj ll/N,), as defined in Equation 2 These imputs and out- puts were collected over the course of running the Wikitext-2 dataset (Merity et al.) (016)) through Pythia-1.4B. For the random neuron, as the difference between inputs decreases, so too does the dif- ference between outputs FFigur@pp). However, for the Wasserstein neuron, this is not the relatively similar inputs are mapped to outputs almost as far apart as the entire range of the neuron (Figure 21). A clear trend between the MD of neuron and its WD emerges (Figure 2c), and the two measures are highly correlated. Thus, we propose the WD of a neuron's output distribution to a Gaussian novel metric of entanglement, with Wasserstein neurons being particularly entangled.
2.3 EFFECT OF HIGH WASSERSTEIN NEURONS ON SPARSIFICATION
In the previous section, have related Wasserstein neurons novel formulation of entanglement. Now, we show that such neurons also have a substantially outsized effect on model performance un- der sparsity. In Llama-3-8B, if just 3% of all neurons aeurons-those with the highest WD)-an sparsified via SparseGPT in every FFN, model performance significantly degrades. This degradation is far more severe than when 3% of random neurons are sparsified, and remains true when compared to sparsifying the same number of other important neurons, such as those with the greatest mean and variance in their output distributions and even those with the greatest mean weight magnitude. As compression increases, this effect becomes more obvious (Figure 3a). Therefore, Wasserstein neu- rons are crucial for maintaining accuracy and are severely limited in their ability to be compressed.
To better understand which specific capabilities are impacted by neuron entanglement, we evalu- ate the Llama-3-8B model with its Wasserstein neurons sparsified across several language model evaluation benchmarks. We select five tasks spanning four broad categories, similar to the original Llama-3 work (Dubey et al. 2024). For reading comprehension, we use the 1-shot variant of the SQuAD 2.0 dataset (Rajpurkar et al. 2018). To assess knowledge reasoning and mathematical capa- bilities, we evaluate the model on the 5-shot TriviaQA-Win Wiki (Joshi et al.) 2017) and 5-shot GSM8K (Cobbe et al. 2021) datasets, respectively. Finally, to evaluate general reasoning, on two benchmarks: an easy task. 5-shot MMLU (Hendrycks et al. 2020), and a more challenging task, 3-shot Chain-of-Though (CoT) Big Bench Hard (BBH) (Suzgun et al. 2022
test the model
4

---

Published as a conference paper at ICLR 2025
(a) Sparsifying Wasserstein neurons hurts the most
(o) Ablatoo spersity: 90%
(c) Ablatooe sparsity: 96%
(d) Abletion spansity 9/%
XMLI
NKKU
MMLL
ligh
Trinus
TrisiQAA
Hig Neas Welga:
Randox
Perpicatt
Teein Vorie
BOEL
BBH
2011
SOWAD
SQUAD
SQUAD
70%
75%
00%
25%
309
55%
GSMEE
GSMBK
OSMBK
Spardiy
Figure 3: Entangled neurons are much more sensitive to compression. In Alama-3-8B,, 3% neu- rons from every FFN linear layer sparsified via SparseGPT in unstructured manner with subset of the Wikitext-2 train dataset as calibration data. (a) Sparsifying Wasserstein neurons (blue) impairs the model more than sparsifying neurons with the highest output distribution means (orange) and variances (green), those with the highest average mean weight magnitude (purple), and consid- erably more than random neurons (red). Perplexity measured on the Wikitext-2 test dataset. (b-d) Sparsifying the Wasserstein neurons (blue) affects general and mathematical reasoning much more than random neurons (red), shown the capability charts. At higher levels of neuron sparsity (2 95%), ablating Wasserstein neurons leads to collapse in model performance, which does not occur with random neurons.
Our findings reveal that when just small fraction of neurons (the top 3% Wasserstein neurons) are sparsified, the model's performance on complex tasks involving general reasoning and mathematical understanding is significantly impacted. However, when the same level of sparsification is applied to random neurons, the model is able to preserve most of its capabilities effectively. Additionally, as a neuron increasingly sparsified, the output distribution becomes more Gaussian (Figure A9 A10.. This in turn places even more stress upon the neuron- not only is it contending with de- creasing mean and variance of the output distribution (Figure (AT)). but also with the less expressive analyzing, the
distribution shape. Thus, seems that, especially the higher sparsities that irregular shape the entangled neurons is much more challenging to model with fewer weights than Gaussian-like distribution. Furthermore, partially due to their slightly lower mean weight magnitudes (Figure A44), Wasserstein neurons are actually sparsified more by SparseGPT during unstructured sparsity, compounding this issue (Figure A4p). However, keeping Wasserstein neurons dense at the cost of sparsifying all other neurons even more also does not seem to be the solution (Appendix A.7). To investigate the difficulty of sparsifying entangled neurons and the relationship between superposition and performance, we introduce Sparse Expansion.
AN EXPERIMENTAL FRAMEWORK TO STUDY DISENTANGLEMENT
To better study Wasserstein neurons and the phenomena between entanglement, sparsity, and per- formance that we observe, we create the experimental framework Sparse Expansion. It is inspired by recent work on the theoretical limits of computation within superposition (Hänni et al. 2024 Adler & Shavit 2024). Sparse Expansion was designed to achieve two goals in real-world models. First, it must originate from a trained dense model and not be retrained. This way, the dynamics of a single neuron, in particular Wasserstein neurons, can still be studied in depth after the model has been expanded. Second, from a theoretical perspective, it must test how varying the number of effective features in the input affects the number of required weights. Therefore, the relationship between superposition and sparsity the can be further understood.
3.1 SPARSE EXPANSION IN DETAIL
Sparse Expansion clusters the inputs to each layer into separate groups via an optional PCA di- mensionality reduction and K-means clustering. Each expert is then sparsified via the SparseGPT algorithm (Algorithm AT). Briefly, SparseGPT approximates the optimal sparse matrix of a layer with the Hessian of the error relative to the parameters the layer Y X + b. Doing so yields
H X where H is the Hessian matrix.
5

---

Published as a conference paper at ICLR 2025
Sparse Expansion Expert Creation
Sparse Expansion FFN Inference
Hossian H,
Sparse
SparseGPT Pruning
Sparse Expert
Reduction
Dimensionality
Clustering
H,
Gate Proj
Down Proj Router
Sparee Down Proj. Expert
Cluster
Sparss Esport W,
Sparse Proj. Expert
SparseGPT Pruning
Figure 4: The Sparse Expansion process. One-shot expert creation process of Sparse Expansion (left). Inference process in a FFN of an expanded model (right).
During inference, each input will be passed through the PCA and K-means model to decide its expert, then routed to the corresponding expert for the matrix multiply (Algorithm A2). As the routing is done via K-means multiply operation, both are inexpensive to add on to normal LLM inference. Furthermore, routing in this manner prevents the need to train and run a more expensive router.
a lower dimension, and the PCA is very low dimension matrix
Our design explicitly achieves the goals we set out. First, by starting from a dense model, we are able to study how the separation of inputs affects individual neurons, which we would not be able to do for the same neuron index across experts in a MoE model such as Mixtral (Jiang et al. 2024 and DeepSeek (Guo et (025). Second, by utilizing SparseGPT, each expert has its weights sparsified and tailored to a subset of inputs, testing the theoretical limits of how many weights are necessary to model given number of features.
3.2 SPARSE EXPANSTON DISENTANGLES NEURONS
We revisit the output distributions of neurons to determine the effect that clustering has in a sparse setting. First, we repeat the sparsification experiment conducted in Fn Wikitext-2 in Llama- 3-8B. Now, for just the neurons we pruned, we expand them into 16 experts and measure the recovery performance. Sparse Expansion able to recover significant performance following Wasserstein neuron sparsification, much more than it does during random neuron expansion. However, the recovery in performance for random neurons is not noticeable, because these neurons were not under significant entanglement initially (Figure 5p.. Furthermore, both the weighted cluster WD and weighted cluster MD of the majority of neurons decreases as a result of Sparse Expansion. The weighted cluster WD and MD are calculated as the average WD and MD within each cluster, weighted by cluster size. This is especially true for Wasserstein neurons, where 98% of neurons have decrease in weighted WD by median of 42% per neuron (Figure 5p), and where 96% of neurons have a decrease weighted MD by a median of 9% per neuron (Figure
For Llama-2-7B (Figure and Pythia- 1.4B (Figure A8). both models and both neuron types-random and maaseerteei-iimrre through Sparse Expansion, with the entangled neuron showing greater improvement. Furthermore, for the random neuron and especially for the entan- gled neuron, the geometry of the sparse output distribution in Sparse Expansion much more closely matches that of the dense distribution.
We also provide a visualization for the specialization of each cluster. Figure ] and Figure A8 each show the sparse output distributions of each individual cluster, with different color per expert. For the randomly selected neurons, there is still an improvement, although each expert is for the most part responsible for approximately the same range and shape. For the entangled neurons, there is significant specialization for different parts of the distribution further away from the mode.
However, Sparse Expansion is not limited to improving just Wasserstein neurons. Across different sparsities and across different models, all but tiny fraction of neurons improve through Sparse Expansion (Figure A12). Thus, like other work regarding polysemanticity (Lecomte et al. 2023 Bricken et al.) 2023), we believe that, in fact, every neuron is to some extent in entanglement. However, Wasserstein neurons are the most obviously entangled ones, and they benefit more from sparse disentanglementt especially at higher sparsities. Finally, note that other metries of the dense neuronal output distribution, such as their means and variances, fail to act as a predictor of neuronal improvement the degree that the Wasserstein distance Gaussian does (Figure 7

---

Published as a conference paper at ICLR 2025
5e) Recovery in Sparns Dxprsssion
231 Sgame Taperokin redeces WD
00 Spams Depazason rodaces
Kendow Tanden:
04
Figure 5: Sparse Expansion recovers performance of Wasserstein neurons. (a) Although Wasserstein neurons are penalized more under sparsity, they also recover better in Sparse Expansion compared to random neurons. We quantify this recovery using normalized perplexity relative the dense model. Data from Llama-3-8B. (b) As a result of Sparse Expansion, the median decrease in WD per neuron 19%. Although a few neurons with an initially low dense WD exhibit a higher average weighted WD, the majority (68%) of all neurons show a decrease in weighted WD. This is especially true in the top 10% of neurons with an originally high WD the Wasserstein neurons. (c) Sparse Expansion also decreases the weighted MD by a median of 2% per neuron. 70% of all neurons and 96% of Wasserstein neurons show a decrease weighted MD, the latter with a median decrease of 9% per neuron. (b, c) Data collected from of the up projection matrix in the second FFN of Pythia-1.4DD
A13). Thus, we believe that the WD to normal for a neuron's 's output distribution a very suitable and intuitive metric of entanglement within neuron.
3.3 MORE SPARSE EXPERTS BETTER FIT THE OUTPUT DISTRIBUTION
The complex dense output distribution of highly entangled neurons is difficult model with single sparse expert, as the case of SParseGPT. In Figure we show the output distribution of Wasser- stein neuron in both dense and sparse computation. As the number of sparse experts increases, the output distribution of the sparse computation more closely matches that the dense computation, as measured in the WD between the two distributions. Furthermore, the relative improvement (RI) of Sparse Expansion over SparseGPT increases. In this paper, RI is measured as the ratio of the RMSE between the SparseGPT sparse computation and the dense computation, to the RMSE between the Sparse Expansion sparse computation and the dense computation.
Clostee
Clisters
Clisters
Clusters
Closters
WD:
0.16
WD1.C.12
Rt: 1.00
Rt: 1.24
Ri: 1.31
1.60
Bl: L.87
Erequeiny
Chtppt Velun
Figure 6: Modeling recovery with more experts. The sparse computation output distribution (red) better matches the dense one (blue) with more clusters. Sparsity set to 90% for each expert. Here, WD refers to the Wasserstein distance between the Sparse Expansion sparse and dense output distributions, rather than to a Gaussian. RI represents relative improvement of Sparse Expansion (n clusters) over SparseGPT (n cluster). This is the same neuron from Figure 2
3.4 WASSERSTEIN DISTANCE BEST EXPLAINS IMPROVEMENT
So far, have claimed that Wasserstein distance is not only pertinent indicator of neuronal entan- glement, but also a predictor of its improvement in Sparse Expansion over SparseGPT. To test this idea, compare how well the RI is modeled by neuron's output WD, mean output magnitude, and output variance. Of these metrics, a neuron's Wasserstein distance is most correlated with its improvement in sparse computation from disentanglement (Figure A13).

---

Published as a conference paper at ICLR 2025
Wasserstein distance best explains improvement Sparsa Expansion
2.0
R*2 0.00
R=
200
20-
18-
1.0
1.
1.6
1.2
1.2
1.2
0.4
Optmmal
Figure 7: Wasserstein distance best explains improvement among tested metrics. The RI of each neuron in Pythia-1.44 .4B was calculated as before and compared against the optimal number of Gaus- sians needed to model its output distribution (gray), the average magnitude of its output distribution (orange), the variance of its output distribution (green), and the Wasserstein distance of its output distribution to normal (blue). For each metric, the line of best fit is calculated, and the coefficient of determination R² is found. For each optimal number of Gaussians, the mean improvement is marked. Of these metrics, the Wasserstein distance best correlates with relative improvement. Data collected from of the second up projection layer in Pythia-I.4B.
In addition, we test whether the estimated number of components in a Gaussian mixture model (GMM) is enough explain the improvement as a result of disentanglement. Specifically, given a neuronal output distribution, we applied Gaussian mixture modeling to determine the optimal number of Gaussians required to model the distribution, using the Bayesian Information Criterion (BIC) for evaluation. BIC is a metric that penalizes model complexity and tries to identify the minimum number of Gaussians which can optimally model the distribution. However, when testing the optimal number of Gaussians between one and sixteen models, our findings indicated almost no correlation (R² < 0.001) between the optimal number of Gaussians and the relative improvement in the Sparse Expansion setup, as seen in Figure 7 Thus, in our experiments, we find that the Wasserstein distance is a better indicator than others that we have tested.
3.5 THEORETICAL IMPLICATIONS OF SPARSE EXPANSION
Recent theoretical work (Hänni et al. 2024 Adler & Shavit 2024) investigates the algorithmic up- per and lower bounds of polysemantic neuronal computation in toy examples. To explore empirical evidence along this body of work for real-world models, we investigate the improvements made by Sparse Expansion in 1ythia-1.4BB in 80% unstructured sparsity. We estimate the approximate num- ber of effective features set of inputs has by applying PCA to the set and finding the minimum number of components required to reach 90% explained variance. As expected, the average mini- mum required components for the inputs to the experts, weighted by the number of inputs in each group, decreases after clustering for every FFN weight matrix (Figure (h).
To provide empirical evidence on the bounds of computation under entanglement, explore mod- eling ability compared to the number of input features. To identify a bound for minimum error under sparse computation, we consider the RMSE of each clustered sparse output the dense output, nor- malized to the overall RMSE for that layer as a proxy for computational ability. We compare this the number of required PCA components for said cluster before. Across all clusters in all layers of the network, there is a linear front that emerges in log-log scale: as the number of required components increases, so too does the minimum error (Figure (8)). Next, we consider the bound on maximum improvement in sparse computation under entanglement. When a cluster has fewer effective features, since each expert has the same number of parameters, Sparse Expansion allocates relatively more parameters to model these features than SparseGPT does, as the latter must ac- count for all inputs. However, when a cluster has many required components, Sparse Expansion and SparseGPT allocate a similar amount of parameters, leading to relatively lower improvement. There- fore, performance improvements increase with fewer effective features. However, beyond a certain point, adding more features no longer yields further performance gains. This trend is also visible in the linear frontier of the log-log plot (Figure Kp). Thus, we provide some empirical demonstrations of the existence of bounds of both loss and improvement of sparse computation under entanglement.
8

---

Published as a conference paper at ICLR 2025
Ce) Clusterin decrnoons required cotpuaenis
thi Empertea bornd for ucatauun creve
to connectiall boind for bostieuu zugeweeaen)
"ip peodecion Jave
Up prt Tunn (ee)
Bowi pnd
180 380 2:30
Regared tocgeansto
Figure 8: Empirical demonstrations of performance bounds. (a) As result of clustering, the weighed average minimum number components to capture 90% of the explained variance de- creases for every layer. (b) As the number of required components for particular cluster increases, so too must the error. (c) As the number of required components for a particular cluster decreases, Sparse Expansion improves more over SparseGPT, but up to bound. Data collected in Pythia- 1.4B.
3.6 SPARSE EXPANSION PERFORMANCE
We evaluate how well Sparse Expansion performs against other competitive one-shot pruning tech- niques, including in terms of inference speed (Table[A3). Despite its leading evaluation performance (Figure 9 Table A1 A5 this method is likely not practically implementable without further opti- mizations counteract the increase in memory footprint, including tuning the number of clusters neuron (Figure A3). Nevertheless, we hope that Sparse Expansion serves as an inspiration for future sparsification techniques that address entanglement for better performance.
3.6.1 MODELS, DATASETS, AND SETUP
We use the Pythia series of pre-trained LLMs to evaluate how Sparse Expansion performs across model sizes, from Pythia- 70M to Pythia-12B. We further evaluate Sparse Expansion across the entire Llama-2 family. We use subset the Wikitext-2 train dataset as calibration data for input-aware pruning and evaluate using the corresponding test through the perplexity metric. Furthermore, to evaluate the performance of Sparse Expansion in out-of-distribution (OOD) data, we evaluate the sparse model in 5 zero-shot standard benchmark tasks in both Llama and Pythia. For our perfor- mance benchmarks, we use 16 clusters at each level of routing in Sparse Expansion. We rely upon the RAPIDS library (Raschka et al., 2020) to accelerate the PCA and K-means models by orders of magnitude. We utilize and build upon the SparseGPT GitHub repository.
3.6.2 PERFORMANCE ACROSS SCALES
We evaluate the performance of Sparse Expansion against other one-shot pruning techniques across a range of model sizes in Pythia and sparsities in Llama-2-7B (Figure 9). Across all model sizes of Pythia, Sparse Expansion outperforms all other pruning techniques at 50% unstructured sparsity, approaching dense performance as model size increases. Moreover, for Llama-2-7B, across all levels of sparsity, Sparse Expansion outperforms all other techniques. At higher levels of sparsity, the gap in performance between the techniques grows. We run further experiments on the entire Llama 2 family as well, and Sparse Expansion similarly outperforms other methods (Table AS). Finally, our experiments show Sparse Expansion outperforming contemporary pruning algorithms in OOD settings as well (Table A1 A2
4 RELATED WORK
Polysemanticity There is plethora of ongoing research contributing to the understanding of pol- ysemanticity in neural networks from a mechanistic interpretability perspective (Bricken et al. [202]] Huben et al. 2023 Lecomte al. 2023 Templeton 2024). These efforts primarily rely upon sparse autoencoders to disentangle output activations into human- -uman-interpretable features, losing information specific to individual neurons in the process. As we focus neurons due to their direct role in network pruming, we derive our own formulation of entanglement as an extension of prior notions. There are also other works that investigate individual neuronal responses directly utilizing tech-

---

Published as a conference paper at ICLR 2025
(e) Performaace Across Sizes
Performance Acroes Sparsities
4.1
Hagaltude
Hugtizate
Wardi
Worth
Sprnctrt
Spuns Expaniie
Spane Expanteee
Cross Looe
2D
10"1
1t
101
Parameters
Ettective ML2 Spunsity
Figure 9: Sparse Expansion across model sizes and sparsities. (a) Performance comparisons on Wikitext-2 perplexity between Magnitude Pruning (MP). Wanda, SparseGPT, and Sparse Expansion on Pythia models from sizes of 70M parameters to 12B parameters. Every FFN in each model was sparsified to 50% sparsity. Each star represents particular model size the dense curve, and the corresponding sparsified model is the marker directly to its left on the sparse curves. (h) Performance for Llama-2-7B at different levels of sparsity for MP, Wanda, SparseGPT, and Sparse Expansion. The x-axis points in both graphs take into account the cost of routing.
niques such as sparse probing (Gurnee et al., 2023), well as those that identify special neuron types LLMs, such as Universal neurons (Gurnee et al. 2024) and Confidence Regulatory neu- rons (Stolfo et al. 2024). However, there is no recent literature tying polysemanticity and neuronal entanglement to sparse network performance.
Compression A multitude of advanced weight pruning algorithms, such as Wanda (Sun et al. 2023) and SparseGPT (Frantar & Alistarh 2022), and quantization algorithms (Kim et al. 2023
Dettmers al., 2022, Ashkboos et al. 22224, Egiazarian al.] 2024, Dettmers et al., 2023 /hao et al. 2023 Lin et al. 2024) exist. Most advanced algorithms are input-aware so as specialize
the weights to the most important input features. Other pruning approaches, such as SWAP YYou &
Cheng, 2024) and WD-based channel pruning (Duan & Li 2020), have also used WD, though for
the gradient of the loss or for channel similarity. rather than for analyzing neurons. While outliers in the features and weights are known to be the among the most challenging factors to address when quantizing to extremely low bits, no equivalent understanding has been made for high sparsities.
5 CONCLUSION AND DISCUSSION
In this work, for the first time demonstrate the impact of neuronal entanglement network performance under weight sparsity, a previously unexplored avenue. From our work and others, we suspect that every neuron is some extent entangled, but that this entanglement of features is easier for some neurons to resolve than it for others. We explore this notion of entanglement through our metric of mapping difficulty, and find that Wasserstein distance is novel, highly pertinent indicator of entangled neurons that must differentiate similar inputs into different outputs. Furthermore, as Wasserstein neurons in particular are incredibly sensitive to sparsification, we posit that the robust- ness of neuron to sparsity is directly dependent on its degree of entanglement. Finally, we have shown that our experimental framework Sparse Expansion is an effective way to disentangle the complex entangled state of sparse neuron, and use it to explore computational bounds in empir- ical weal-world models. The disentanglement provided by Sparse Expansion benefits Wasserstein neurons the most, providing further support that such neurons are the most entangled.
In future work, we plan to study Wasserstein neurons the framework mechanistic interpretabil- ity to understand what circuits they form. From our insight that more entangled neurons are harder to sparsify, we will investigate creating efficient, entanglement aware sparsification algorithms to pre- serve performance at higher sparsities. Looking forward, perhaps just as outlier features and weights are well understood to be one the most significant challenges when quantizing to fewer bits, so too neuronal entanglement be understood as the challenge of pruning to higher sparsities.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The authors would like to extend their gratitude to Lori Leu for her insightful comments on the application of the Wasserstein distance metrie. We also wish to thank Elias Frantar for his help in working with the SparseGPT implementation and his advice for the project. Additionally, we would like to thank Tony Tong Wang and Thomas Athey for their valuable feedback and constructive discussions.
This work was supported by an NIH Brains CONNECTS UOI grant and AMD's Al & HPC Fund.

---

Published as a conference paper at ICLR 2025
REFERENCES
Micah Adler and Nir Shavit. On the complexity of neural computation in superposition. arXiv preprint arXiv:2409.15316 2024.
Sanjeev Arora, Yuanzhi Li, Yingyu Liang, Tengyu Ma, and Andrej Risteski. Linear algebraie struc- ture of word senses, with applications to polysemy. Transactions of the Association for Compu- tational Linguistics, 6:483-495, 2018.
Saleh Ashkboos, Amirkeivan Mohtashami, Maximilian Croci, Bo Li, Martin Jaggi, Dan Alistarh, Torsten Hoefler, and James Hensman. Quarot: Outlier-free 4-bit inference in rotated llms. arXiv preprint arXiv:2404. 00456, 2024.
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 International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.
Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, et al. Towards monosemanticity: Decom- posing language models with dictionary learning. Transformer Circuits Thread, 2, 2023.
Roberto Lopez Castro and Dan Alistarh. Sparse marlin: a fast sparse plus 4-bit kernel for generative inference. https:// com. com DASLab/ Maparse-Marlin 2024.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10.1668,, 2021.
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Gpt3. int8 O: 8-bit matrix multiplication for transformers at scale. Advances in Neural Information Processing Systems, 35: 30318-30332, 2022.
Tim Dettmers, Ruslan Svirschevski, Vage Egiazarian, Denis Kuznedelev, Elias Frantar, Saleh Ashk- boos, Alexander Borzunov, Torsten Hoefler, and Dan Alistarh. Spqr: A sparse-quantized repre- sentation for near-lossless Ilm weight compression. arXiv preprint arXiv:2306. 03078, 2023.
Haoran Duan and Hui Li. Channel pruning for accelerating convolutional neural networks via wasserstein metric. In Proceedings of the Asian Conference on Computer Vision, 2020.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, al. The llama 3 herd of models. arXiv preprint arXiv:2407.2 1783, 2024.
Vage Egiazarian, Andrei Panferov, Denis Kuznedelev, Elias Frantar, Artem Babenko, and Dan Al- istarh. Extreme compression large language models via additive quantization. arXiv preprint arXiv:2401.06118, 2024.
Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodds, Robert Lasenby, Dawn Drain, Carol Chen, al. Toy models of superposi- tion. arXiv preprint arXhv:2209.1063 2022.
Elias Frantar and Dan Alistarh. Optimal brain compression: A framework for accurate post-training quantization and pruning. Advances in Neural Information Processing Systems, 35:4475-4488, 2022.
Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot. In International Conference on Machine Learning, pp. 10323-10333 10337. PMLR, 2023.
Elias Frantar and Dan Alistarh. Marlin: a fast 4-bit inference kernel for medium batchsizes. https:
//git /ggithub.com com IST DASLab
Btt.//IISTDAALLLAAAA mar rin 2024.
Elias Frantar, Salch Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210 10. /7323, 2022.
12

---

Published as a conference paper at ICLR 2025
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501. 2025.
Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace. arXiv preprint arXiv: 1812.04754, 2018.
Wes Gurnee, Neel Nanda, Matthew Pauly, Katherine Harvey, Dmitrii Troitskii, and Dimitris Bert- simas. Finding neurons in haystack: Case studies with sparse probing. arXiv preprint arXiv:2305.00 016000 2023.
Wes Gurnee, Theo Horsley, Zifan Carl Guo, Tara Rezaei Kheirkhah, Qinyi Sun, Will Hathaway, Neel Nanda, and Dimitris Bertsimas. Universal neurons in gpt2 language models. arXiv preprini 1rXiv:2401.1218/ 2024.
Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. Advances in neural information processing systems, 28, 2015.
Kaarel Hänni, Jake Mendel, Dmitry Vaintrob, and Lawrence Chan. Mathematical models of com- putation in superposition. arXiv preprint arXiv:2408.00 08.4055,,,,,, 2024.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint 2.Xiv:2009.03300 2020.
Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks. Journal of Machine Learning Research, 22(241):1 1-24, 2021.
Robert Huben, Hoagy Cunningham, Logan Riggs Smith, Aidan Ewart, and Lee Sharkey. Sparse autoencoders find highly interpretable features in language models. In The Twelfth International Conference on Learning Representations, 2023.
Adam S Jermyn, Nicholas Schiefer, and Evan Hubinger. Engineering monosemanticity in toy mod- els. arXiv preprint arXiv:221 2022.
Albert Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bam- ford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.04088. 2024.
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint 1xX.v:1705.0335 2017.
Leonid V Kantorovich. On the translocation of masses. Journal of mathematical sciences, 133(4): 1381- 2006.
Sehoon Kim, Coleman Hooper, Amir Gholami, Zhen Dong, Xiuyu Li, Sheng Shen, Michael W Mahoney, and Kurt Keutzer. Squeezellm: Dense-and-sparse quantization. arXiv preprint arXiv:2306.07629. 2023.
Victor Lecomte, Kushal Thaman, Trevor Chow, Rylan Schaeffer, and Sanmi Koyejo. Incidental polysemanticity. arXiv preprint arXiv.2312.03096, 2023.
Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Wei-Ming Chen, Wei-Chen Wang, Guangxuan Xiao, Xingyu Dang, Chuang Gan, and Song Han. Awq: Activation-aware weight quantization for on-device IIm compression and acceleration. Proceedings of Machine Learning and Systems, 6: 87-100, 2024.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint 17x1v:1609.07843, 2016.
Jesse Mu and Jacob Andreas. Compositional explanations of neurons. Advances in Neural Informa- fion Processing Systems, 33:17153 2020.
13

---

Published as a conference paper at ICLR 2025
Lorenzo Noci, Alexandru Meterez, Thomas Hofmaann, and Antonio Orvieto. Why do learning rates transfer? reconciling optimization and scaling limits for deep learning. arXiv preprint arXiv:2402.17457, 2024.
Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. Zoom in: An introduction to circuits. Distill, 5(3):c000240001, 2020.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoe-Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Compulational Linguistics (Volume I: Long Papers). pp. 1525-1534, 2016.
Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't 't know: Unanswerable questions for squad. arXiv preprint arXiv: /806. 1806.0882,, 2018.
Sebastian Raschka, Joshua Patterson, and Corey Nolet. Machine learning in python: Main develop- ments and technology trends in data science, machine learning, and artificial intelligence. arXiv preprint arXiv:2002.0488 04803, 2020.
Alessandro Stolfo, Ben Wu, Wes Gurnee, Yonatan Belinkov, Xingyi Song, Mrinmaya Sachan, and Neel Nanda. Confidence regulation neurons in language models. arXiv preprint arXiv:2406. 16254, 2024.
Mingjie Sun, Zhuang Liu, Anna Bair, and Zico Kolter. simple and effective pruning approach for large language models. arXiv preprint arXiv:2306. /695, 2023.
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V Le, Ed H Chi, Denny Zhou, et al. Challenging big-bench tasks and whether chain-of-thought can solve them. arXiv preprint 1rXiv:2210.09261, 2022.
Adly Templeton. Scaling monosemanticity: Extracting interpretable features from claude sonnet. Anthropic, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 2023.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
Pauli Virtanen, Ralf Gommers, Travis Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, et al. Scipy 1.0: funda- mental algorithms for scientific computing in python. Nature methods, 17(3):261-272, 2020.
Johannes Welbl, Nelson F Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. W-NUT 2017, pp. 94, 2017.
Lei You and Hei Victor Cheng. SWAP: Sparse entropic wasserstein regression for robust network pruning. In The Twelfth International Conference on Learning Representations, 2024. URL
htttps://ppenreee ew net fott::/oorumm2idd1 zuuBuy
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient Ilm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507, 2024.
Yilong Zhao, Chien-Yu Lin, Kan Zhu, Zihao Ye, Lequn Chen, Size Zheng, Luis Ceze, Arvind Krishnamurthy, Tianqi Chen, and Baris Kasikei. Atom: Low-bit quantization for efficient and accurate Ilm serving. arXiv preprint arXiv:2110.1910, 2023.
14

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 PSEUDO-CODE FOR SPARSE EXPANSION
Algorithm AT describes the sparsification process of Sparse Expansion. The sparse experts are created in a layer-wise sequential fashion for each linear layer of every FFN transformer block to create the sparse model. Algorithm A2 refers the inference procedure of Sparse Expansion once the model is pruned following the methods described in Algorithm Al and Section 3.1
Algorithm A1 Sparse Expansion model generation. The following layerwise procedure can be repeated for each linear layer in the transformer.
1: procedure LAYERWISE SPARSE EXPANSION SPARSIFICATION PROCESS
2: {a:) < 244 E RR
//set of calibration inputs to layer
3: < m x n Mayer weights
4:
//number of clusters
5:
//factor to reduce dimensionality by
6:
R < PCA(")) //new PCA object with components
7:
R.fi(((x)) //fit R to inputs
8:
K + Kmeans(c) //new K-means object with initial centroids K.fit((R(z))) //fit K to dimensionality reduced inputs
9:
for j 2.2.3...c do
11:
X; - [r]((R(z)) j} //rropp {z} into its component clusters
W, W //make a copy of the original weight matrix
13:
Sj + SparseGPT //make SparseGPT object
14:
W + (W,, //sparsify W, using X,
Algorithm A2 Sparse Expansion inference. The following layerwise procedure is repeated infer- ence time for each clustered layer.
1: procedure LAYERWISE SPARSE EXPANSION INFERENCE PROCESS
2:
{x} < T4 R." //set of inputs layer
3: [W] //set of experts
4: R. //PCA model
5:
K //K-means model
6:
for = 1,2, 3... do
7:
j K (((()))
//find the cluster assignment of x //run inference with the correct expert
8:
Yj < (xx
15

---

Published as a conference paper at ICLR 2025
A.2 DISTRIBUTION OF WASSERSTEIN DISTANCES ACROSS ALL LLAMA-2-7B FFN LAYERS
After collecting the Wasserstein distance to the normal distribution for every neuron, we find that all up and gate projection matrices in each Llama-2-7B FFN block have high WD neurons. We also find that certain down projection matrices also have high WD neurons, though most do not.
(a) Wasserstein Distances for Neurons in the Up Projection Matrix in Llama 2-7B
Wassersteln Distance
MPN Block
(b) Wasserstein Distances for Neurons in the Gate Projection Matrix in Llama 2-7B
Wiscerstetn Discance
FEN Block
(c)
Wasserstein Distances for Neurons the Down Projection Matrix in Llama 2-7B
Visscrrsttln Distance
16
PEN Block
Figure A1: High Wasserstein distance neurons each layer. Many neurons with a high WD to the Gaussian distribution exist in every FFN block, and in every up (a) and gate projection (b) specifically. Certain down projection layers also have high WD neurons (c). The box plots show the
range of non-outliers, well as the first quartile, the median, and the third quartile of neuronal WD. The outliers are defined as 1.5 times the interquartile range less than the first or more than the third quartile and represented by the points.
16

---

Published as a conference paper at ICLR 2025
A.3 SPARSE EXPANSION PERFORMANCE OUT-OF-DISTRIBUTIO DATA
We evaluate the performance of Llama-3.2-11B (Table A1 and Pythia-1.4B (Table A2) on range of natural language modeling tasks, including ARC-e (Easy) and ARC-c (Challenge) for arithmetic reasoning, Lambada (Paperno et al. (0166) for contextual word prediction, SciQ (Welbl et al. (017) for scientific question answering. and MMLU for multitask general knowledge assessment.
dense Pythia-I.4B does not score better than random chance on MMLU, we do not benchmark on this task. We compare various pruning algorithms 50% sparsity, including Magnitude Pruning (MP), Wanda, SparseGPT, and Sparse Expansion with 16 clusters. to the dense baseline. Sparse Expansion consistently excels across both models, achieving the highest scores on tasks among sparsification algorithms.
Table Al: Performance of Llama-3.2-IB under different pruning algorithms.
Algorithm
Sparsity
ARC-e
ARC-c
Lambada
SciQ
MMLU
Dense
0%
65.488
31.314
53.969
91.4
37.701
Magnitude
50%
45.244
22.354
4.677
67.1
23.493
Wanda
50%
50.800
23.635
31.457
85.2
25.428
SparseGPT
50%
55.640
24.403
31.613
86.8
25.046
Sparse Expansion 50%
57.713
26.962
35.807
87.5
28.729
Table A2: Performance of Pythia- under different pruning algorithms.
Algorithm
Sparsity
ARC-e
ARC-e
Lambada
SciQ
Dense
0%
61.742
27.389
48.981
86.9
Magnitude 50%
42.003
19.198
1.533
69.0
Wanda
50%
54.630
23.976
45.041
85.7
SparseGPT
50%
56.608
24.061
44.615
85.8
Sparse Expansion
50%
58.449
25.720
46.424
86.3
17

---

Published as a conference paper at ICLR 2025
A.4
NEURONAL ENTANGLEMENT TRAJECTORY ACROSS TRAINING IN LLMS
Enlanglement over the course of LLM training in Pythia-1.44
(A) Progression of Wasserstein Distance Across Steps
Wussersteln Distance
Too 3% ND Nourons
Botton Nourens
Aucroge
8000
10000
15000
200000
30000
40000
50000
amn0
700000
80a00
00800
1oce00
110000
120000
130000
140000
Traning Steps
(B) Change of Neuron Weighls Across Sleps
Norallzzd Welght Direeeeos
1.50
1.00
676
025
Bottors 3% ID Neurens
Average
0000
Ia00p
15000
21000
30000
40000
50000
08000
700000
90000
90040
10s000
110000
120000
130000
140000
Training Strpe
Figure A2: Analyzing neuronal entanglement during training. The intermediate checkpoints of Pythia-1.4P available over the course of its training, from initialization completion. Thus, collect data from 17 different checkpoints over the course of its training, first intervals of 5,000 steps, then at intervals of 10,000 steps after step 20,000. (a) We calculate each neuron's output distribution WD to Gaussian as before in Equation T We do so for each training step. From the WD of neurons in the last training step, separate out the top 3% of neurons with the highest WD and the bottom 3% of neurons with the lowest WD. We also find the average WD across all neurons. The progression of neuronal WD across training reveals that all neurons initially exhibit Gaussian- expected, but some neurons rapidly differentiate into entangled neurons with
like distribution, very high WD and within just 5,000 steps (corresponding to approximately 10 billion tokens). The WD of such neurons then levels off afterward. (b) Using the same groups as in (a), we visualize the change neuronal weights. We calculate the I,2 norm between each neuron's weights each training step and its weights at model initialization (step 0), and normalize this value by the L norm of the neuron's s weights at initialization. Notably, neurons with high WD not demonstrate more changes in their weights over the course of training than the average neuron, neurons with low WD. Error bars represent one standard error of the mean. Neurons are from the up projection matrix in the second FFN block of Pythia-1.4B.
18

---

Published as a conference paper at ICLR 2025
A.5 OPTIMIZATIONS FOR PRACTICAL IMPLEMENTATION
To evaluate the inference latency of Sparse Expansion we implemented a Sparse Expansion layer based PyTorch and optimized sparse-quantized GPU kernels called Sparse Marlin (Frantar & Alistarh 2024] Castro & Alistarh] 2024), which supports the INT4 + 2:4 sparsity format. To better utilize the compression kernel, we use both sparsification and quantization demonstrate speedup. We use a linear layer of appropriate size as an upper bound approximation for our router cost, which followed by 4 bit, 2:4 sparse matrix multiplication. We have run the layer-wise benchmarks for the typical layers sizes from Llama models on a single RTX3090 GPU. We can see in Table A that Sparse Expansion allows us to get up to a 4.8x speedup over the dense baseline. The speedup comes from the cegpeycoompressed linear layer representation. Although there is overhead compared to regular compressed matrix due to the presence of the router, such overhead decreases as layer size increases.
Table A3: Sparse Expansion inference speedup. Layer-wise single batch inference latency (in ji)).
The layer sizes are chosen specifically to match the layers of Llama-2-7B and Llama 2 70B.
Layer Size
4k x 12k
4k x 22k
11k x 4k
8k x 10k
8k x 57k
28k x 8k
Dense
132
227
114
220
1168
556
Sparse Expansion
76
76
75
76
241
138
Speedup
1.7x
3.0x
1.5x
2.9x
4.8x
4.0x
Sparse
26.8
44.7
24.4
42.3
216
109
Overhead
2,9x
1.7x
3.1x
1.8x
1.1x
1.3x
Additionally, we investigate how many experts different neurons need to improve performance. We find the relative improvement of each neuron, as defined Section 3.3 across a different number of total experts. Specifically,
choose 2, 4, 8, and 16 experts for Sparse Expansion, compared to SparseGPT with its single expert. In this setting, we analyze the top 3% of neurons with the high- est WD well as the bottom 3% of neurons with the lowest WD, defined before (Equation I We observe that Wasserstein neurons benefit far from Sparse Expansion than average for increasing clusters (left). Additionally, we split neurons into decile groups based on their relative improvement at 16 clusters. We find that, indeed, certain groups of neurons benefit very little from further addi- tional experts past eight experts (right). Thus, further optimizations can be made reach a balance between performance and memory increase.
Variation neuroncl Improvement across the number
1.30
1.30
Birton MD
1.25
1.25
Relelive imprevement
1.20
1.20
1.15
1.15
Ican E1)
103.
[71 Bt)
[86, (7)
(87. Ith
1.00
1.00
2
16
2
16
Number of clusters
Number of clusters
Figure A3: Improvement across clusters for different groups of neurons. (a) Wasserstein neurons benefit much more from Sparse Expansion than average with increasing clusters. (b) Different deciles of neurons have varying degrees of improvement from Sparse Expansion. Dn indicate the deciles from D1 D9. The decile groups decided by their relative improvement 16 clusters. For example, the first decile group consists of relative improvements between the minimum and DJ 16 clusters, the second decile group consists of relative improvements between DI and D2 at 16 clusters, and so on. Error bars represent one standard error of the mean. Neurons from the up projection matrix of the second FFN block of Pythia-
19

---

Published as a conference paper at ICLR 2025
A.6 WASSERSTEIN NEURONS DO NOT HAVE PARTICULARLY HIGH WEIGHTS
To understand whether Wasserstein neurons arise from having substantially higher weights than av- erage, We measure the mean weight magnitude for each neuron. We find that Wasserstein neurons do not have weight magnitudes that are particularly above average; anything, there seems to be slight negative correlation between neuron's WD and its mean weight magnitude (Figure (441). To investigate how this difference affects sparsification, we sparsify this layer to 80% unstructured sparsity via SparseGPT, calibrated with Wikitext-2. This takes into account both weight magnitude as well as the influence of the inputs via the Hessian matrix. Wasserstein neurons are especially sensitive to sparsity and have an outsized impact on model performance, even more so than neurons with high average weight magnitude (Figure 3). However, these neurons sparsified slightly more than average by this current advanced sparsificatio approach. This suggests that future sparsifica- tion schemes should take into account neuron's WD and degree of entanglement, rather than just its weights and the Hessian.
(b) Wasserstein neurons sparsified slightly more than average by current approaches
(a) Wasserstein neurons do not have particularly high weights
R-2 = 0.23
Meat weight magniuuee
80%
Sparsity
20%
60%
50%
40%
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3 0.4
0.5
Wasserstein distamon
Figure A4: Wasserstein neurons do not have particularly large weights, in terms of their average magnitude, and so are sparsified slightly more. (a) Neurons with high WD do not have large average weight magnitudes. Of the top 3% of neurons with the highest WD, just one is also within the top 3% of neurons with the largest average weight magnitudes. (b) Partially as a result of their lower than average weights, Wasserstein neurons tend to be sparsified slightly more than average in an unstructured setting. The top 3% Wasserestein neurons are sparsified 6% more than average. The neurons from the up projection of the second FFN in Pythia- 1.4B.
20

---

less Published as a conference paper at ICLR 2025
A.7 DIMINISHING RETURNS FOR KEEPING ENTANGLED NEURONS FULLY DENSE
As shown in Figure entangled neurons are particularly sensitive to pruning. We design an exper- iment to understand the opposite effect model performance, namely of keeping the Wasserstein neurons dense. In our setup, we selectively keep the top x% of Wasserstein neurons dense, while pruning each of the remaining neurons sparsity of 100-z % maintain an overall target spar- sity of s%. This approach is compared against a baseline where all neurons are pruned to the same sparsity percentage 5%, abbreviated as same sparsity per neuron (SSPN) in the table and equivalent to x 0%.
We conduct this experiment on Llama-2-7B. We use SparseGPT to sparsify the neurons, use part of the Wikitext-2 train set as calibration data, and evaluate on the Wikitext-2 test set. As illustrated in Table A4 keeping Wasserstein neurons dense at the cost of sparsifying every other neuron more to achieve a target sparsity does not enhance model performance. Additionally, model performance worsens progressively
the proportion of neurons kept dense (x%) increases, since tangled neurons are also being kept dense. This behavior is likely due to the fact that the benefit of allowing Wasserstein neurons to retain all of their weights is outweighed by the cost of sparsifying every other already sparse neuron even more.
Table A4: Perplexity of Llama-2-7B on Wikitext-2 while sparsifying to 8% overall and preserving x% of Wasserstein neurons.
$
50%
8 = 60%
$
70%
s 80%
8 = 90%
SSPN 0%)
6.219
7.420
12.73
33.26
366.0
3%
6.259
8.023
14.70
40.40
395.7
I 5%
6.345
8.131
16.03
46.67
629.5
7%
6.366
8.547
17.37
61.95
978.3
10%
6.522
9.232
19.48
79.53
8066
21

---

Published as a conference paper at ICLR 2025
A.8 DERIVING MAPPING DIFFICULTY AS METRIC OF ENTANGLEMENT
We show more reasoning behind the choice of the normalizing factors N₂ and Ng in Equation 2 First, we choose Nº maxisi<jsn i2; mソ|ر} to be the maximum L norm between a pair of inputs to simply scale all I,2 norms to be between and 1. Next, we choose Ny median iin (llb myll} to be the median based on the following observations. Specifically, we would like to preserve and highlight the fact that there are many IO pairs that have a relatively low difference in their inputs, but are mapped to very different outputs, one group of which is circled in purple in Figure 1.
First, we considered using the maximum L2, as did for No. However, note that there is a very small number samples that drives the maximum to be much further away from meaningful data points in both the random and Wasserstein neurons. Next, presence of the outlier data points of interest that have a much greater difference in output than expected, the mean is also driven much higher for the Wasserstein neuron. Indeed, observe that for the Wasserstein neuron, the mean is much greater than the mode, while they are much more similar for the random neuron. We therefore use the median normalize for inter-neuron differences in the expected range of their output differences, but to also be robust to outliers that would obfuscate the 10 pairs of interest through an inflated mean.
considered the mean. Due to the
Zondom Nearen Outpett Dtatribrtton
Randecs Nexan Relatrenhlp
(o Estangled Nouree Octpul Denttbttinn
00 Estangled Naurs Relatiensktp
Outpat tale
Batradion Sigar
Figure A5: Deriving a normalization constant for the difference in outputs. (a) The output distribu- tion of random neuron. (b) The non-normalized relationship between the I norm between pairs of inputs and the L norm between their corresponding outputs for the random neuron. (c) The out- put distribution of Wasserstein neuron. (d). The non-normalized relationship between the L norm between pairs of inputs and the L norm between their corresponding outputs for the Wasserstein neuron. Note how the mean is much higher than the median. One group that has a much higher output L norm than expected for its relatively low input L norm is circled in purple. These are the same neurons from Figure 2
22

---

Published as a conference paper at ICLR 2025
A.9 ADDITIONAL CLUSTERS IMPART MORE PERFORMANCE
To understand how Sparse Expansion scales with the number of experts per linear layer, we test its performance from to 32 experts. Interestingly, with 2 experts, very little performance benefits realized. However, with each doubling of experts following experts, we realize a nearly constant linear improvement perplexity.
Sparse Expansion performance
improves with more clusters
7.00
6.75
6.50
Perplexity
6.25
SparseCP
Sparse Expansion
8.00
Dense
3.75
3.50
21
25
Number of chusters
Figure A6: Increasing the number of clusters improves Sparse Expansion performance in Llama-2- 7B.
23

---

Published as a conference paper at ICLR 2025
A.10 FURTHER EXAMPLES OF DISENTANGLEMENT
We present additional evidence of neuronal disentanglement in Pythia- and Llama-2-7B. Figure A7 shows the recovery of the dense output distribution with increasing experts in Llama-2-7B. This Is analogous to Figure A9 where we see a similar trend in Pythia models. Clustering gradually decreases the WD of the sparse output to that of the dense, thus improving upon SparseGPT, equiv- alent to the single cluster case. Moreover, this results in direct improvement of model performance as depicted in figure A6
Chusder
Clusters
B Clasterr
Clasterx
RR. 1.00
BI 1.00
23. 1.43
M.. 1.38
R1-1.44 1.43
Modeling recovery with more experts in Llama-2-7B. Use of more experts can recover the dense output distribution even at very high sparsity, which is set to 90% for each expert. This is same neuron from Figure T..
Figure A7:
Analogous to Figure D we observe the effect of clustering inputs on a random neuron and an en- tangled neuron the gate projection of the second FFN of Pythia-1.49 SparseGPT fails to capture the output distribution of the high WD neuron as it does for a random neuron. With clustering via Sparse Expansion, both neurons improve, but the entangled neuron improves more. The granular analysis of the component clusters within both neurons reveals the specialization to vastly different parts of the output distribution the entangled neuron as compared to the normal neuron (Figure
A8
3:100
3.00
1.00
Rh
Figure A8: Sparse Expansion disentangles neurons in Pythia-I.4B. The dense output distribution of random neuron, along with its sparse via SparseGPT (a) and via Sparse Expansion (b) sparse output distributions. The dense output distribution
random neuron, along with its sparse via SparseGPT (d) and via Sparse Expansion (e) sparse output distributions. For both the random and entangled neuron, component clusters are shown a distinct color to visualize their range (c, e). WD represents the Wasserstein distance between the Sparse Expansion sparse output distribution and the dense distribution. RI represents relative improvement. These are the same neurons from Figure
24

---

Published as a conference paper at ICLR 2025
A.11 PERFORMANCE ACROSS THE LLAMA FAMILY
We analyze the performance of Sparse Expansion against other sparsification algorithms across all members of the Llama-2 family--Llama-2-15 Llama-2-13B, and Llama-2-70B; both under sparsity and joint sparsity-quamizzation compression (Touvron et al. 2023
Table A5: Sparse Expansion across the Llama family,
Sparsity
Bits
Llama-2-7B
Llama-2-133
Llama-2-70B
Dense
0%
16-bit
3.1168
4.5736
3.3192
MP
50%
16-bit
16.029
6.8270
4.9846
Wanda
50%
16-bit
6.7757
5.8527
4.0219
SparseGPT
50%
16-bit
5.7082
5.0521
3,9013
Sparse Expansion
50%
16-bit
5.5839
4.9728
3,8791
SparseGPT
2:4
16-bit
6.9767
5.9934
4,8002
Sparse Expansion
2:4
16-bit
6.4456
5.6255
4.6671
SparseGPT
2:4
4-bit
7.2759
6.1101
4.9036
Sparse Expansion
2:4
4-bit
6.5745
5.7151
4.7586
SparseGPT
2:4
3-bit
13.076
6.5055
5.2552
Sparse Expansion
2:4
3-bit
7.0757
5,9872
5.0588
Sparse Expansion outperforms all other pruning techniques for both 50% unstructured sparsity as well 2:4 sparsity in all Llama models (Figure A5). In addition to non-quantized sparsity, we consider how Sparse Expansion performs in the context of compression with 2:4 structured sparsity and quantization via GPTQ (Frantar al. 2022). We first sparsify each linear layer in each FFN block to 2:4 sparsity, then quantized to 3 and 4 bits. Our method outperforms SparseGPT across all models and across both conditions and in all models (Figure A5).
Across multiple model sizes, sparsity and compression levels, and advanced models, Sparse Expan- sion attains state-of-the-art performance for post-training one-shot sparsification when compared to other highly competitive pruning techniques. We do so by leveraging the powerful pruning algo- rithm of SparseGPT and combining with input specialization utilize the insights we gain from how entangled neurons behave under sparsity.
Because GPTQ (Frantar et 2022). a leading post-training quanization scheme, also relies upon the Hessian matrix for its algonith,, we combine it with SparseGPT for combined one-shot com- pression. Sparse Expansion also outperforms native SparseGPT and GPTQ across all compression settings.
25

---

Published as a conference paper at ICLR 2025
A.12 EFFECT OF SPARSITY ON NEURONAL OUTPUT DISTRIBUTIONS
With increasing sparsity, the sparse output distributions of the high WD neurons and random neurons converge toward the normal distribution (Figure ATO). A specific example of neuron is shown in Figure A9
Sparsity 0%
Sparsity - 50%
Sparsity - 90%
Sparsity- 95%
Sparsity 99%
WD: 0.44
WD: 0.44
WD: 0.34
WD: 0.28
WD: 0,08
Frequency
Öutput Value
Figure A9: Increasing sparsity induces normality. A highly entangled neuron's dense distribution (blue) and sparse distribution (red). As sparsity increases, the output distribution of the sparse neuron becomes progressively more Gaussian. WD represents the Wasserstein distance. This is the same neuron from Figure 2
Prigradeer Woseriia Dalanee Accoo Sjuratios
Ra 2% Neanse
0.25
310
1.51
Sp-ni
Figure A10: Output distributions become more normal under sparsity. The Wasserstein distance between neuron's normalized sparse output distribution and the Gaussian distribution is shown as sparsity increases for top 3% entangled Wasserstein neurons, the same number of bottom 3% WD neurons, and random sample of 3% of the neurons. For highly entangled neurons, the WD decreases significantly at higher sparsities whereas it remains more or less constant for the bottom 3% of neurons and for the random neurons. Range indicates maximum and minimum WD for group. Data collected from the second up projection matrix in Pythia- 1,4..
26

---

Published as a conference paper at ICLR 2025
Furthermore, with increasing sparsity, the magnitudes of the means and variances of each neurons' sparse output distribution both shift toward zero. This is reasonable, as with fewer nonzero weights to combine together features, both the mean and variance should decrease in magnitude.
Normaltood rsaca and vecsnce EAuEses
specitiea
1.
4.00
30 0.2
014 2.6 0.8 1.5 0.
Figure A11: Mean and variance shift toward zero under sparsity. Across all neurons, with increasing sparsity. the magniudee of the mean of output distribution (left) and the variance of the output distri- bution (right) both tend toward zero. Both mean and variance have been normalized to their dense values. Error bars represent one standard error. Data collected from of the second up projection matrix in Pythia-I.4B.
27

---

Published as a conference paper at ICLR 2025
A.13 ALL NEURONS IMPROVE, BUT ENTANGLED NEURONS IMPROVE MORE HIGHER SPARSITIES
Measuring the relative improvement of each neuron through Sparse Expansion, find that all neurons improve as a result of Sparse Expansion across both Pythia- and Llama-2-7B. Thus, believe that every neuron has some level of innate entanglement, and so all neurons can be and are improved. Interestingly, we note that, with increasing sparsity, highly entangled Wasserstein neurons tend to improve more.
Improvement 80% sparsity
Improvement 90% sparsity
22
2.0
2.0
1.6
1.6
Relative
1.4
1.0
10
0.0
0.1
0.2
0.5
0-1
0.5
0.6
0.0
0.1
0.2
0.1
0.5
0.6
Wasserstein
Wasserstein distance
Improvement at 80% sparsity
Improvement at 90% sparsity
8*2 = 0.07
R^2 0.14
2.0
1.0
1.6
Relative
1.2
1.0
0.2 Wesserstiei distance
Wassersstein distance
Figure A12: Entangled neurons improve more at higher sparsities. Relative improvement of each neuron the second up projection matrix in Pythia-1.44 1.4B (top row) and in the second gate projection matrix in Llama-2-7B (bottom row) with respect to their WD from the Gaussian. Two sparsity levels, 80% and 90%, are shown. Sparse Expansion improves expressibility of every neuron, thus improving performance. However, the entangled neurons improve more with higher sparsities, as visible the right column.
28

---

Published as a conference paper at ICLR 2025
A.14 THE WASSERSTEIN DISTANCE BEST CAPTURES WHICH NEURONS IMPROVE
We also consider whether the magnitude of the mean of the output distribution or the variance of the distribution would be good predictors of the degree of neuronal improvement through Sparse Expan- sion. However, across both Pythia-1.4B3 (Figure 7 and Llama-2-7B (Figure A13 the Wasserstein distance from the normal is better predictor of relative improvement, as defined previously. Though there is some correlation of the magnitude of the mean and variance of the output distribution with the relative improvement Pythia- 1.4B, that not the case Llama-2-7hh Furthermore, using the WD to predict neuronal improvement yields the highest coefficient of determination, R², across both models.
Wasseresteln distance best explains improvement from Sparse Expanston
h
R-2
Relative inprovvenenss
1.8
18-
1.6
1.-
16
1.4
1.4
10-
1.0
100
0.0
0.5
1.0
1.5
2.0
0.0
0.1
0.2
0.4
0.5
0.6
Magnctdde
Figure A13: Wasserstein distance best captures improvement. Relative improvement of each neuron in the second gate projection matrix in Llama-2-7B with respect to the magnitude of the mean, variance, and Wasserstein distance from normal of the dense output distribution. Neurons pruned to 90% sparsity.
29

---

Published as a conference paper at ICLR 2025
A.15 OUTPUT DISTRIBUTIONS OF ENTANGLED NEURONS IN PYTHIA AND LLAMA
Figures A414 and AIS show the non-trivial, non-Gaussian output distribution of a subset of neurons from the Pythia- 1. ..4B and Llama-2-7B models, illustrating examples of entangled neurons. We observe such neurons in every FFN block of the LLMs we investigated and believe that the existence of these neurons a global phenomenon in transformers.
1.00
1700
B00
600
13000
OM
BEB FATS
600
000
750
300
Seazce 1127 IVO 0.51
Strucce 2150 ID 0453
Neese 81200
Necres IVD 64D
Hotre 2012 1Wth 8:59
1,000
000
302
200-
Stacee E25L 640
Neurc 300 8:00
2003 ND.
Heurs E.D)
Hiurm 1145 203
800
7ah
1200
Figure A14: Dense output distributions of top 30 high WD neurons in 1ythia-1.4B. The distributions
are shown for the neurons of the up projection matrix in the second FFN block.
30

---

Published as a conference paper at ICLR 2025
4:40
400
300
200
Neara VIT 20th AnC
Misse VIAA [WS
2200
3000
400
Learn LOKE (AA) 638)
MMACc 966 (Wb
Mewrs E3S Ni5 8.37
Meuie 9V1 p63:
14 2031 (W0x 3.94
1100C
2000
MM2E1
FIVE
340
2003
Norry (NC
1W2:1330
900
800
GEC
Figure A15: Dense output distributions top 30 high WD neurons Llama-2-7B. The distributions are shown for the neurons of the up projection matrix in the sixteenth FFN block.
31


Reviews:
[{"review_id": "QIrCo2uwfD", "reviewer": "Reviewer_kR1B", "summary": "The paper proposes using the Wasserstein distance to select entangled neurons. The authors use the mapping difficulty to measure the entanglement and argue that the Wasserstein distance is strongly correlated with it. They further introduce a sparse expansion method that disentangle the Wasserstein neurons better.", "questions": "1. Directly normalizing the distribution to calculate the Wasserstein distance seems like a strong simplification. Is there a better approach that accounts for mean and variance?\n2. In Equation (2), why is the mean used for mapping difficulty, the maximum for $N_x$ , and the median for $N_y$?\n3. In Figure 4b, what does the \"normalized input/output L2\" represent?\n4. On line 282, there is a missing word after “the.”\n5. How is the weighted cluster distance calculated in Figure 6b?\n6. Are the K-NN and evaluation conducted on the same training set of wikitext-2?\n7. Could the authors explain why there is no correlation between the optimal number of Gaussians and relative improvement? Is this due to sparse expansion being applied directly to the weight matrix while the GMM is estimated for individual neurons? Is the relative improvement possibly influenced by the neurons with the largest GMM value? Could differences in measurement — one on the entire weight matrix, the other on individual neurons — lead to inconsistencies?\n8. How do outlier weights (or neurons) overlap with neurons that have large Wasserstein distances?\n9. How do neurons in the MoE model perform in terms of Wasserstein distance?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Understanding the relationship between sparsity and entanglement is an important research area.\n- The paper conducts extensive experiments to explore the connection between Wasserstein distance, entanglement, and their impact on sparsification.", "weaknesses": "- The concept of using the Wasserstein distance has been addressed in prior work [1]. The authors could benefit from a broader literature review. It would be helpful to discuss how other works evaluate entanglement.\n- The sparse expansion approach relies on prior knowledge of the distribution to support K-means clustering. How robust is this method for out-of-distribution data?\n- Practical aspects of the sparse expansion, such as runtime and inference speed (e.g., with 16 experts), should be discussed more thoroughly.\n- The writing could be improved; The reviewer finds that many phrases and sentences are difficult to parse.\n\n[1] You, Lei, and Hei Victor Cheng. \"Swap: Sparse Entropic Wasserstein Regression for Robust Network Pruning.\" The Twelfth International Conference on Learning Representations. 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "K5KQQVuTT1", "reviewer": "Reviewer_KmG9", "summary": "This paper examines the performance of models under weight sparsity by considering the entanglement of the neurons in the dense model. The authors first empirically show that neurons with high Wasserstein distance (WD) to a Gaussian distribution will more negatively impact the model's performance when sparsified. They then demonstrate that the WD metric has positive correlation with the Mapping Difficulty (MD) of each neuron, which serves as a heuristic for measuring neuron entanglement.\n\nMotivated by this observation, the authors propose an algorithm (Sparse Expansion - SE) that adjusts the SparseGPT (Frantar & Alistarh, 2023) algorithm for making the model sparse: For each layer, they learn $k$ sparse weights (experts), corresponding to $k$ clusters of input data, by the SparseGPT algorithm. The clusters at each layer are identified by $k$-means clustering on the low-dimensional representations derived from PCA on the original inputs. Then, during inference, each input is routed through the closest clusters. The paper claims that each cluster will help disentangle the neurons, leading to a decrease in WD in the sparse model. Finally, they show that the sparse models learned by SE have improved performance, as measured by perplexity.", "questions": "1.  On the Sparse Expansion (SE) algorithm:\n    - How is the number of clusters $k$ chosen? From Figure A2, it seems that increasing $k$ consistently improves performance (in terms of perplexity). Besides increased memory usage, are there other performance trade-offs when increasing $k$?\n    - If the clustering approach mainly addresses the issue with the entangled neurons, do all neurons need the same number of clusters? \n    - How much optimization/run-time overhead does the PCA+$k$-means step add to the sparseGPT algorithm when sparsifying the weights? \n2.  Since the link between entanglement and Wasserstein neurons is made through the MD metric, I am wondering if we see MD also decreases by the Sparse Expansion algorithm.\n3. Shouldn't the clusters learned through the SE algorithm depend on the structure of the data/task used for this purpose? For example, in Figure 10, is SE both performed and validated on Wikitext2? If so, would the sparse model still have improved performance on other downstream tasks/datasets (like the ones evaluated in Figure 2)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I am not fully familiar with the literature on neural disentanglement and the interpretability results. However, I find the paper’s approach to improving the performance of sparsified models by analyzing neuron interpretations interesting.\n\nThe authors take a step-by-step approach to justify their claims/arguments and to make connections between the different concepts discussed in the paper.", "weaknesses": "I find that many of the arguments in the paper are either based on intuition or established by showing correlations between metrics. In some cases, it is not very clear to me whether there is some sort of causation also involved. For example,  I don't understand why a neuron with a smaller WD to Gaussian should be less entangled (and thus more interpretable?). In the paper, this conclusion relies on the correlation between WD and MD, which is itself an intuitive metric for measuring entanglement. Nevertheless, the arguments are still motivating for further research in this direction.\n\nI think the flow of the paper can be slightly improved: For example, in section 2, Wasserstein neurons and entangled neurons are used interchangeably before properly establishing their connection in section 2.3. This was a bit confusing on the first read, requiring some back and forth through the sections to clarify their connections.\n\n(Minor) The paper's generally well-written, but it would benefit from a revision as there are several typos throughout the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vR5wiS3VKV", "reviewer": "Reviewer_uAy8", "summary": "The authors propose to compute the distance between the output distribution of a neuron in a neural network and a gaussian using Wasserstein distance as a measure of entanglement, i.e. a neuron could be highly polysemantic and thus entangled. It’s shown that some of the neurons in LLMs are highly entangled under the Wasserstein distance metric and sparsifying them leads to significant performance degradation compared to sparsifying random neurons. The authors propose sparse expansion by separating the input vectors and creating a mixture of neurons with low Wasserstein distance. The proposed sparse expansion methods outperform all other baselines under sparsity / pruning.", "questions": "See weaknesses", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "The proposed sparse expansion seems to outperform SparseGPT and all other baselines, though I’m not in this subfield and I don’t know what are the strong and weak baselines here.", "weaknesses": "- It’s very unclear what do you mean by the gaussian output distribution of a single neuron. For example, the introduction section says that a column of a weight matrix is a neuron. Here it’s unclear what’s the matmul operation between input and weights. Is it right or left multiply? Are you saying that the output scalar $y=\\mathbf{w}^\\top \\mathbf{x}$ is a sample from the gaussian distribution? As a reader I have to read between the lines to figure out what’s the proper definition. Please write down the precise mathematical quantity.\n- There is a confounding factor here regarding the idea of highly entangled neurons. If I understand properly, the output distribution of a neuron at initialization is basically a gaussian. This is because the weights are gaussian initialized unless otherwise specified. Then perhaps highly entangled neurons are basically neurons with proper feature learning, i.e. enough gradient updates, and neurons with low wasserstein distance to a gaussian could just be a neuron that didn’t receive enough gradient update. This is especially true when you consider LLMs under width scaling, as some of the parameters might not be learned properly. See https://arxiv.org/abs/2203.03466. Then it might not make sense to prune out or sparse highly entangled neurons since they could actually be the useful neurons or neurons with proper feature learning. Would like to hear what the authors think about this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WVA8hrdWMI", "reviewer": "Reviewer_2v2T", "summary": "In this work, the relevance of individual neurons in LLMs is estimated by comparing their normalized output distribution before the activation function to a Gaussian using the Wasserstein distance measure. \nThey find that the output of a small group of neurons deviates significantly from the gaussian distribution, and that these neurons contribute significantly to the model's performance. They define that a neuron is particularly entangled when it has to produce different outputs for similar inputs. They find that neurons that deviate from a gaussian distribution are particularly entangled.\nThey provide a framework that disentangles the neurons to some extent, that is, they reduce the Wasserstein distance measure. Using this framework, they obtain a lower loss on Wikitext-2 compared to other techniques when creating a sparsified model.", "questions": "1) The way I understand your sparsification algorithm in Figure 2 is this: For 3% of the neurons, you set 95% of the weights going into them to zero. Is this correct? And if so, how do you choose these weights? \n\n2) Figure 10: Which model is used? The caption speaks of many, but the figure shows only one, as far as I understand it. \n\nLine 422 typo: \"the the\"\nLine 428 typo: \"as\" -> has", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "1) The relevance of individual neurons is studied in detail, which is interesting from a technical point of view when reducing the model size, but also points to areas that explainable AI researchers should look for.\n\n2) A potential application in terms of model reduction is provided.", "weaknesses": "1) The exact sparsification algorithm used to generate Figure 2 is not given in the manuscript.\n\n2) There are some problems with the Figures: Figures 1, 3, and 7 lack x and y labels.\nIn Figure 10, the caption does not match what is shown in the figure.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)