Research map / Architectures and efficiency / Mixture of experts
Wasserstein Distances, Neuronal Entanglement, and Sparsity
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Wasserstein distance to a Gaussian as a neuron entanglement metric and Sparse Expansion, a clustering-based mixture of SparseGPT-sparsified experts, which improves one-shot LLM pruning accuracy.
Problem
Polysemantic, entangled neurons are sensitive to weight sparsification, and it is unclear how neuronal entanglement relates to post-training sparsity performance in LLMs.
Approach
Quantify neuron entanglement by the 1-Wasserstein distance between a normalized neuron output distribution and N(0,1), identify "Wasserstein neurons," and disentangle them by clustering layer inputs into a mixture of SparseGPT-sparsified experts (Sparse Expansion) that specialize to input subsets without retraining.
Key results
- Sparse Expansion outperforms Magnitude, Wanda, and SparseGPT on Llama-2-7B across all tested sparsities, with the gap growing at higher sparsity.
- Sparsifying only the top 3% Wasserstein neurons in Llama-3-8B degrades perplexity and reasoning far more than sparsifying 3% random neurons.
- On Llama-2-7B at 50% unstructured sparsity, Sparse Expansion reaches perplexity 5.5839 versus 5.7082 for SparseGPT and 6.7757 for Wanda; up to 4.8x layer-wise speedup reported with INT4+2:4 kernels.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Sparse Expansion instead of SparseGPT at 50% unstructured sparsity | Llama-2 family · Wikitext-2 | perplexity: 5.7082 -> 5.5839; 5.0521 -> 4.9728; 3.9013 -> 3.8791 | Helped |
| Use Sparse Expansion instead of SparseGPT with 2:4 sparsity and quantization | Llama-2 family · Wikitext-2 | perplexity: 6.9767 -> 6.4456; 7.2759 -> 6.5745; 13.076 -> 7.0757 | Helped |
| Use Sparse Expansion rather than SparseGPT at 50% sparsity | Llama-3.2-1B · ARC-e, ARC-c, Lambada, SciQ, MMLU | score: 55.640 -> 57.713; 24.403 -> 26.962; 31.613 -> 35.807; 86.8 -> 87.5; 25.046 -> 28.729 | Helped |
| Use Sparse Expansion rather than SparseGPT at 50% sparsity | Pythia-1.4B · ARC-e, ARC-c, Lambada, SciQ | score: 56.608 -> 58.449; 24.061 -> 25.720; 44.615 -> 46.424; 85.8 -> 86.3 | Helped |
| Sparsify the top 3% of neurons selected by highest Wasserstein distance rather than random neurons | Llama-3-8B · Wikitext-2 and reasoning benchmarks | performance | Hurt |
| Expand sparsified neurons into 16 input-clustered sparse experts rather than one SparseGPT expert | Wasserstein neurons · Llama-3-8B, Wikitext-2 | normalized perplexity recovery: 98% of neurons; median of 42% | Helped |
| Increase the number of sparse experts | Sparse Expansion · Wasserstein neuron output distribution | relative improvement over SparseGPT: 18%; 77% | Helped |
| Keep Wasserstein neurons dense while pruning other neurons more to maintain overall sparsity (ablation) | Llama-2-7B · Wikitext-2 | perplexity: 6.219 -> 6.259 (x = 3%); 6.219 -> 6.345 (x = 5%); 6.219 -> 6.366 (x = 7%); 6.219 -> 6.522 (x = 10%) | Hurt |
| Increase preserved dense Wasserstein neurons from 3% to 10% (ablation) | Llama-2-7B · Wikitext-2 at 50% overall sparsity | perplexity: 6.259 -> 6.522 | Hurt |
| Increase the number of experts from 2 to 32 (ablation) | Sparse Expansion · Llama-2-7B | perplexity | Helped |
| Use Sparse Expansion instead of the single-expert SparseGPT | Pythia-1.4B neurons · 80% and 90% sparsity | relative improvement | Helped |
| Sparsify the top 3% highest-Wasserstein neurons instead of random neurons | Llama-3-8B · Wikitext-2 test | perplexity | Hurt |
Surprising findings
- Keeping the highest-WD neurons dense while sparsifying the rest more heavily does not improve perplexity, and performance worsens as more neurons are kept dense.
- Gaussian mixture model component count (BIC) shows almost no correlation (R^2 <= 0.001) with relative improvement from Sparse Expansion, while WD correlates better.
Limitations
- Sparse Expansion increases memory footprint and router overhead, so the authors say it is likely not practically implementable without further optimization.
- Speedup measured only at layer level on a single RTX3090 with INT4+2:4 kernels; end-to-end efficiency not reported.
Experimental setup
Datasets: WikiText-2
Benchmarks: WikiText-2 perplexity ARC-e ARC-c Lambada SciQ MMLU SQuAD 2.0 TriviaQA-Wiki GSM8K BBH
Metrics: perplexity accuracy Wasserstein distance relative improvement speedup
Baselines: Magnitude Pruning Wanda SparseGPT GPTQ
Builds on: SparseGPT Wanda GPTQ
Closest papers by meaning
- Expand Neurons, Not Parameters (ICML 2026 · similarity 0.86)
- Sparsing Law: Towards Large Language Models with Greater Activation Sparsity (ICML 2025 · similarity 0.85)
- $\alpha$TC-VAE: On the relationship between Disentanglement and Diversity (ICLR 2024 · similarity 0.84)
- The Sparsity Whisperer (NeurIPS 2026 · similarity 0.84)
- Synergy Between Sufficient Changes and Sparse Mixing Procedure for Disentangled Representation Learning (ICLR 2025 · similarity 0.84)
- Scaling Laws and Tradeoffs in Recurrent Networks of Expressive Neurons (NeurIPS 2026 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.