Mechanistic interpretability aims to reverse-engineer neural networks, to understand the actual algorithms they learn. A key part of this is understanding how a model stores information. A significant portion of a large language model's weights is dedicated to memorized facts, rather than general circuits. To make progress in interpretability, we need to know what a look-up table looks like inside a transformer. That's the question at the heart of a recent challenge: can you hand-code the weights of a neural network to memorize random sequence pairs as efficiently as a model trained with gradient descent?
The Sequence Memorization Task
The research focuses on a deliberately simple task to act as a toy model for any look-up where a combination of input signals carries a meaning that is different from a linear combination of the individual signals. The training data consists of sequences of three tokens: two input tokens and one output token. The model sees two input tokens and must predict the output label. The input-to-output mappings, or "facts," are generated uniformly at random, so there's nothing to compress. The model can do nothing but memorize them.
The architecture is a simplified one-layer transformer, but to understand how facts are encoded, the researchers first needed to identify which parts of the model are doing the heavy lifting.
Finding the Key Component: The MLP
To find out which parts of the network matter for this memorization task, the researchers made every architectural component optional: attention, MLP, norms, residual connections, and biases. They trained every possible variant and measured the maximum number of facts each one could learn.
The main finding was decisive: the MLP is by far the most important component. Replacing the attention with a simple sum of per-position token embeddings worked just as well or better, suggesting that attention wasn't doing much beyond linearly mixing the two tokens' information. The experiment also showed that an MLP with ReLU or GELU neurons, where those neurons are the only non-linearity in the network, was the most crucial setting. Adding the MLP made the biggest difference to the number of learnable facts out of any of the settings.
Armed with this knowledge, the researchers scaled up two variants: the full one-layer transformer and a heavily stripped-down version consisting only of embeddings, a ReLU layer, and an unembedding. Both stored a similar number of facts, scaling roughly with the model dimension, close to the information-theoretic expectation.
The Hand-Coding Challenge
This led to the central question: Can you write down weights for the toy model by hand, or with any algorithm that isn't gradient descent, such that it matches the performance of a trained model? The spirit of the rule is that the algorithm should embody an explanation of how the facts are embedded. If we understand how facts are stored, we should be able to replicate this without gradient descent.
The researchers proposed a challenge with a specific architecture: a two-layer embedding matrix, a single ReLU layer, and an unembedding layer, with no biases or norms. The goal was to find a construction that could memorize a maximum number of facts, with two evaluation criteria: 100% accuracy, and a more lenient 90% accuracy.
The Proposed Construction: Inactive Neurons as a Code
Their hand-coding algorithm works on a clever principle: storing facts as patterns of inactive neurons. The idea is to ensure that for any input fact, all pre-activations for neurons assigned to that fact's label are exactly zero, while other labels have at least one active neuron. The readout can then check for an exact zero.
The algorithm works in stages. First, it assigns each of the n ReLU neurons to one or more labels. The assignment is designed to be "spread out" so that all neurons are used by approximately the same number of labels, and the maximum neuron overlap between any pair of labels is minimized.
Next, it sets the embedding weights to ensure that for every fact, the ReLU neurons assigned to its label output zero. This is done by:
- Setting weights from the first and second input tokens to -1 for their assigned neurons.
- Setting a baseline -1 weight for the most common tokens.
- Assigning a +1 weight from a "safety token" (an unseen token) to a neuron, so that if the neuron isn't inhibited, it fires.
This ensures that the correct label's neurons are all inactive (output zero), while some neurons in incorrect labels will be active (output positive). The unembedding matrix then simply gives a negative weight from each neuron to its assigned labels, so that an inactive neuron contributes nothing to the logit, while active neurons penalize their labels. The correct label, with all its neurons silent, wins by default.
Results and the Gap to Trainability
The results showed some success but also a significant gap. For the 90% accuracy threshold, the fully hand-coded models scaled almost as well as the trained models, with similar exponents, but with a worse prefactor of about 9.7x fewer facts. At 100% accuracy, the hand-coded models fell further behind. A hybrid solution, where only the embedding matrix is hand-coded and the unembedding is learned, boosted performance to roughly 3.5x fewer facts than the fully trained model.
The Interpretability Angle
This challenge is more than just a puzzle. It's a concrete test of our understanding. If our hand-coded solution qualitatively resembles what a trained network learns, it supports our hypotheses about its internal mechanisms. However, the initial analysis shows some stark differences. The hand-coded embedding weights take only the values 1, 0, or -1, while trained weights are more "smeared out." More importantly, the hand-coded model stores facts as patterns of inactive neurons, while trained models seem to use patterns of active neurons instead.
This challenge is a call to action for the mechanistic interpretability community. The gap between our best hand-coded solution and the trained model shows that we have more to learn about how neural networks memorize. But every step toward closing that gap is a step toward truly understanding what's happening under the hood of the AI systems we're building. The hope is that once we have any solution at all that performs similarly to trained models, figuring out exactly what the trained models are doing might become a lot easier.