# Wasserstein Distances, Neuronal Entanglement, and Sparsity

arXiv: https://arxiv.org/abs/2405.15756 · HTML: https://mlsanity.com/paper/2405.15756/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in Circuits and mechanistic analysis, one of 13 directions in Mechanistic interpretability.
- Builds on 33 of its 51 references that are in the map.
- First posted on arXiv 24 May 2024.

## Placed in

- Mechanistic interpretability / [Circuits and mechanistic analysis](https://mlsanity.com/mechinterp/circuits-and-mechanistic-analysis/)

## Builds on (most cited first, 30 of 33)

- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [Learning both Weights and Connections for Efficient Neural Networks](https://mlsanity.com/paper/1506.02626/) (2015)
- [GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers](https://mlsanity.com/paper/2210.17323/) (2022)
- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them](https://mlsanity.com/paper/2210.09261/) (2022)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [Know What You Don't Know: Unanswerable Questions for SQuAD](https://mlsanity.com/paper/1806.03822/) (2018)
- [Sparse Autoencoders Find Highly Interpretable Features in Language Models](https://mlsanity.com/paper/2309.08600/) (2023)
- [AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration](https://mlsanity.com/paper/2306.00978/) (2023)
- [SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot](https://mlsanity.com/paper/2301.00774/) (2023)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [Toy Models of Superposition](https://mlsanity.com/paper/2209.10652/) (2022)
- [A Simple and Effective Pruning Approach for Large Language Models](https://mlsanity.com/paper/2306.11695/) (2023)
- [Crowdsourcing Multiple Choice Science Questions](https://mlsanity.com/paper/1707.06209/) (2017)
- [QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs](https://mlsanity.com/paper/2404.00456/) (2024)
- [Sparsity in Deep Learning: Pruning and growth for efficient inference and training in neural networks](https://mlsanity.com/paper/2102.00554/) (2021)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [SpQR: A Sparse-Quantized Representation for Near-Lossless LLM Weight Compression](https://mlsanity.com/paper/2306.03078/) (2023)
- [Optimal Brain Compression: A Framework for Accurate Post-Training Quantization and Pruning](https://mlsanity.com/paper/2208.11580/) (2022)
- [SqueezeLLM: Dense-and-Sparse Quantization](https://mlsanity.com/paper/2306.07629/) (2023)
- [Atom: Low-bit Quantization for Efficient and Accurate LLM Serving](https://mlsanity.com/paper/2310.19102/) (2023)
- [Extreme Compression of Large Language Models via Additive Quantization](https://mlsanity.com/paper/2401.06118/) (2024)
- [Gradient Descent Happens in a Tiny Subspace](https://mlsanity.com/paper/1812.04754/) (2018)
- [Compositional Explanations of Neurons](https://mlsanity.com/paper/2006.14032/) (2020)
- [Linear Algebraic Structure of Word Senses, with Applications to Polysemy](https://mlsanity.com/paper/1601.03764/) (2016)
- [Universal Neurons in GPT2 Language Models](https://mlsanity.com/paper/2401.12181/) (2024)

## Built on by (most cited first, 6 of 6)

- [Adversarial Attacks Leverage Interference Between Features in Superposition](https://mlsanity.com/paper/2510.11709/) (2025)
- [Negative Pre-activations Differentiate Syntax](https://mlsanity.com/paper/2509.24198/) (2025)
- [Sparse CLIP: Co-Optimizing Interpretability and Performance in Contrastive Learning](https://mlsanity.com/paper/2601.20075/) (2026)
- [NeuroADDA: Active Discriminative Domain Adaptation in Connectomic](https://mlsanity.com/paper/2503.06196/) (2025)
- [Resting Neurons, Active Insights: Robustifying Activation Sparsity in LLMs via Spontaneity](https://mlsanity.com/paper/2512.12744/) (2025)
- [The Sparsity Whisperer](https://mlsanity.com/paper/2608.06630/) (2026)

## Abstract (arXiv)

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.