[Paper Review] Polysemanticity and Capacity in Neural Networks

Adam Scherlis, Kshitij Sachan+3 more|arXiv (Cornell University)|Oct 4, 2022

Neural Networks and Applications4 citations

TL;DR

This paper introduces feature capacity as a framework to understand polysemanticity in neural networks, showing that optimal capacity allocation leads to monosemantic representation of high-importance features, polysemantic representation of medium-importance features, and complete omission of low-importance features. The key finding is that polysemanticity emerges naturally when marginal loss reduction per feature is balanced, with geometry revealing block-semi-orthogonal structures dependent on model architecture.

ABSTRACT

Individual neurons in neural networks often represent a mixture of unrelated features. This phenomenon, called polysemanticity, can make interpreting neural networks more difficult and so we aim to understand its causes. We propose doing so through the lens of feature \emph{capacity}, which is the fractional dimension each feature consumes in the embedding space. We show that in a toy model the optimal capacity allocation tends to monosemantically represent the most important features, polysemantically represent less important features (in proportion to their impact on the loss), and entirely ignore the least important features. Polysemanticity is more prevalent when the inputs have higher kurtosis or sparsity and more prevalent in some architectures than others. Given an optimal allocation of capacity, we go on to study the geometry of the embedding space. We find a block-semi-orthogonal structure, with differing block sizes in different models, highlighting the impact of model architecture on the interpretability of its neurons.

Motivation & Objective

• To understand the causes of polysemanticity in neural network neurons, where a single neuron represents multiple unrelated input features.
• To formalize the concept of feature capacity as the fractional dimension a feature consumes in the embedding space, bounded between 0 and 1.
• To investigate how loss functions incentivize polysemantic representations through capacity allocation, particularly in relation to feature importance and sparsity.
• To analyze the geometric structure of the embedding space under optimal capacity allocation, revealing block-semi-orthogonal configurations.
• To compare theoretical predictions with empirical results across different model architectures and input distributions.

Proposed method

• Propose a toy model with a one-layer network using quadratic activations to analytically compute feature capacity allocation as a function of feature importance and sparsity.
• Define feature capacity as $ C_i = \frac{(w_i \cdot w_i)^2}{\sum_{j=1}^N (w_i \cdot w_j)^2} $, representing the fraction of embedding dimension consumed by feature $ i $.
• Use unitary matrix constructions to prove that any capacity tuple $ (C_1, \dots, C_N) $ with $ 0 \leq C_i \leq 1 $ and $ \sum C_i = D $ is achievable in the embedding space.
• Construct a phase diagram showing regions where features are monosemantic ($ C_i = 1 $), ignored ($ C_i = 0 $), or polysemantic ($ 0 < C_i < 1 $) based on their importance and sparsity.
• Analyze the embedding geometry and show it exhibits a block-semi-orthogonal structure, with block sizes varying by model architecture.
• Demonstrate that polysemanticity increases with input kurtosis and sparsity, and is more prevalent in certain architectures due to structural constraints on capacity allocation.

Experimental results

Research questions

• RQ1 Under what conditions does a neural network assign full, partial, or zero capacity to individual features in the embedding space?
• RQ2 How does feature importance and input sparsity jointly influence the emergence of polysemantic neurons?
• RQ3 What geometric structure emerges in the embedding space when capacity is optimally allocated across features?
• RQ4 Why do some models exhibit sharp transitions between monosemantic and polysemantic representations, while others do not?
• RQ5 How does model architecture constrain the capacity allocation and thus affect the interpretability of neurons?

Key findings

• Features with high importance relative to the loss are always represented monosematically, consuming a full embedding dimension ($ C_i = 1 $).
• Features with low importance are entirely ignored, receiving zero capacity ($ C_i = 0 $), even when they contribute to the loss.
• Medium-importance features are represented polysemantically, sharing embedding dimensions with others, with capacity allocated such that marginal loss reduction is equal across them.
• The embedding space geometry exhibits a block-semi-orthogonal structure, where blocks correspond to groups of features sharing embedding dimensions, and block size varies with model architecture.
• Polysemanticity increases with input sparsity and kurtosis, and is more prevalent in architectures with constrained or non-orthogonal embedding structures.
• All feasible capacity tuples $ (C_1, \dots, C_N) $ with $ 0 \leq C_i \leq 1 $ and $ \sum C_i = D $ are mathematically achievable via unitary matrix constructions, confirming the theoretical framework's generality.