[Paper Review] Engineering Monosemanticity in Toy Models

Adam S. Jermyn, Nicholas Schiefer+1 more|arXiv (Cornell University)|Nov 16, 2022

Explainable Artificial Intelligence (XAI)4 citations

TL;DR

This paper demonstrates that monosemantic neurons—those responding to a single, interpretable input feature—can be engineered in small, tractable neural network models by manipulating training dynamics and architectural choices. By initializing with negative biases or increasing model width, the authors achieve highly monosemantic solutions without sacrificing performance, enabling full mechanistic interpretability of both monosemantic and polysemantic neurons.

ABSTRACT

In some neural networks, individual neurons correspond to natural ``features'' in the input. Such \emph{monosemantic} neurons are of great help in interpretability studies, as they can be cleanly understood. In this work we report preliminary attempts to engineer monosemanticity in toy models. We find that models can be made more monosemantic without increasing the loss by just changing which local minimum the training process finds. More monosemantic loss minima have moderate negative biases, and we are able to use this fact to engineer highly monosemantic models. We are able to mechanistically interpret these models, including the residual polysemantic neurons, and uncover a simple yet surprising algorithm. Finally, we find that providing models with more neurons per layer makes the models more monosemantic, albeit at increased computational cost. These findings point to a number of new questions and avenues for engineering monosemanticity, which we intend to study these in future work.

Motivation & Objective

• To investigate whether monosemantic neurons can be systematically engineered in small neural network models.
• To understand how training dynamics and architectural choices influence monosemanticity in neural networks.
• To develop methods for achieving high monosemanticity without increasing loss or computational cost.
• To enable mechanistic interpretability of both monosemantic and polysemantic neurons in these engineered models.

Proposed method

• Use a two-layer feedforward architecture with a linear transformation (including bias) followed by a nonlinearity, then a second linear layer without bias.
• Train models on three tasks: feature decoding, random re-projection, and absolute value computation, all designed to have known ground-truth features.
• Control training dynamics via learning rate and bias initialization to steer optimization toward different local minima.
• Apply weight decay on biases to systematically tune the model toward more monosemantic solutions.
• Systematically vary the number of neurons per layer to study the effect of model width on monosemanticity.
• Perform mechanistic interpretability by analyzing neuron activations, input/output maps, and interference patterns in recovered features.

Experimental results

Research questions

• RQ1 Can monosemantic neurons be engineered in small, interpretable models without degrading performance?
• RQ2 How does the choice of training hyperparameters, such as learning rate and bias initialization, affect monosemanticity?
• RQ3 Does increasing the number of neurons per layer improve monosemanticity, and if so, at what cost?
• RQ4 What computational algorithms do monosemantic and polysemantic neurons implement in the trained models?
• RQ5 Can bias patterns serve as a reliable signal for identifying and engineering monosemantic minima?

Key findings

• Models trained with higher learning rates or negative bias initialization converge to more monosemantic local minima without performance degradation.
• More monosemantic minima consistently exhibit moderate negative biases across all three tasks.
• Applying bias weight decay enables engineering of highly monosemantic models, achieving performance comparable to L1 regularization.
• Increasing the number of neurons per layer increases monosemanticity, though at the cost of higher computational load.
• In the monosemantic limit, single neurons implement a balancing algorithm between feature recovery and interference suppression.
• Two monosemantic neurons per feature jointly classify features as 'likely real' or 'likely interference' and recover the strength of real features.