Wasserstein Distances, Neuronal Entanglement, and Sparsity

2024-05-24

AI Summary

This work addresses performance degradation in large language models (LLMs) during weight sparsification, caused by strong entanglement among polysemous neurons. We propose a novel paradigm for quantifying neuronal entanglement based on the Wasserstein distance and introduce the concept of “Wasserstein neurons”—neurons whose outputs significantly deviate from Gaussianity and exhibit high sensitivity to model accuracy. To enable efficient disentanglement, we design a layer-wise input-separated Mixture-of-Experts (MoE) architecture that achieves post-training sparse disentanglement without retraining. Experiments demonstrate that a small subset of highly entangled Wasserstein neurons dominates sparsification error; our method substantially improves sparse-model accuracy while enhancing neuronal semantic interpretability. This work provides a new pathway toward efficient, interpretable LLM compression and inference.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: (Large) Language Models
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Web Mining and Content Analysis: Large pretrained models with web data

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.

Problem

• Disentangling polysemantic neurons in LLMs
• Measuring neuronal entanglement using Wasserstein distance
• Impact of weight sparsity on model accuracy

Innovation

• Wasserstein distance measures entanglement
• Mixture of experts sparsifies neurons
• Disentangling polysemantic neurons improves accuracy