Reproduce “Disentangling MLP Neuron Weights in Vocabulary Space” on OpenResearch

About this paper

Asaf Avrahamy, Yoav Gur-Arieh, Mor Geva

ROTATE separates MLP-neuron weights into vocabulary-aligned directions called channels by optimizing vocabulary-projection kurtosis, using weights and the unembedding without activation data or forward passes for discovery. In Llama-3.1-8B-Instruct layer 18, it scores 0.71 on faithfulness and 0.55 on completeness, above the best SAE baselines; evaluations cover selected layers in two architectures, and vocabulary-unaligned functions remain outside scope.

Key insights

• Vocabulary-space kurtosis tends to be higher for known concept-specific neurons and increases during pretraining, particularly in middle and later layers where model representations are more aligned with the vocabulary basis.
• A single neuron’s behavior can be separated into distinct channels, such as temporal and polarity-related input mechanisms, while the input weights determine activation magnitude and sign.
• Causal channel ablations selectively suppress the associated input examples or output tokens, indicating that the decomposition captures functional subcomponents rather than merely correlated token groups.

Summary by alphaXiv.