Overview

One of the active research areas for interpretability involves distilling neural network activations into clean, labeled features. This is made difficult because of superposition, where a neuron may fire in response to multiple, disparate signals making that neuron polysemantic. To date, research has focused on one type of such superposition: compressive superposition where a network can represent more features than it has neurons. I report on another type of superposition that can arise when a network has more neurons than features: “symmetric mixtures”. Essentially, this is a form of “favored basis” that allows a network to reinforce the magnitude of its logits via parallelism. I believe understanding this concept can help flesh out the conceptual foundations of how DNNs represent features and how we may interpret them.

Contents

Example of feature mixing that arose studying a toy model
Demonstration of how feature mixing operates to reduce loss
Theoretical framework for why this happens
Predictions for where this may arise in popular models

Toy Model

I was studying a small model on the modular addition problem, as in Nanda+23 and Chughtai+23. The former used an attention-only 1-layer transformer, while that latter used a simpler "embedded MLP" with ReLU layer. The attention mechanism of the transformer allows for products of input signals, which makes it suitable for the learned algorithm, while the ReLU layer struggles to approximate multiplication. Meanwhile, I had also come across this note from Lee Sharkey which makes a case for use of bilinear layers. Inserting a bilinear layer into the embedded MLP model seemed a natural thing to do! The goal was a compact, interpretable model for modular addition.

After experiments with training the model, I found results that mostly aligned with expectations. The model still learns the same solution: Fourier embedding space, trigonometric angle addition, etc. Also this architecture is indeed more parameter efficient. But there remained a few interesting and confusing aspects:

The embedding contains two mutually prime frequencies, and it superposes the frequencies together in each of the four columns of the embedding matrix.
The learned bilinear layer weights only take values in roughly {-1, 1}. Rounding these values improves loss.

So why does the network learn superposed frequencies when there's enough capacity ...

Posts

Wikitag Contributions

Comments