Overview:

We train a tiny ReLU network to output sparse top- distributions over a vocabulary much larger than its residual dimension. The trained network seems to converge to a mechanism closely resembling a Bloom filter: tokens are assigned sparse binary hashes, the hidden layer computes an approximate union indicator, and the output logits are linearly read from this union.

Here’s what a small network trained on a toy version of the sparse top- distribution task learns to use:

That’s the input weight matrix of the trained network. Every entry is either or . The network has effectively encoded a binary hash for each token—and as we’ll show, this seems to enable the network to approximately simulate a Bloom filter, and so output the correct set of top- tokens with high probability.

We provide a theoretical construction showing how to set the weights to exactly implement a Bloom filter. The real network seems to learn to do something similar, and seems to behaviourally act like a Bloom filter, but while we provide a fair bit of mechanistic evidence, we don’t yet have a complete mechanistic explanation of the trained network.

Additionally, this is just a toy network, so it doesn’t directly tell us about what larger models might do. But I found the discreteness of the learned algorithm to be very interesting.

It seems to learn a probabilistic solution which scales like in width, where is the MSE loss. In contrast, a standard JL style superposition solution would have error scaling polynomially in with respect to width.

The Task:

As a result of reading The Softmax Bottleneck, I have gotten quite interested in what mechanisms toy MLP networks learn to output sparse probability distributions. I have tried to devise a toy problem to address this question.

We sample token indices uniformly at random without replacement in , and uniformly from for .

We give as input to a -layer ReLU neural network (where represents an embedding of the th token, which we are allowed to choose).

Let be the probability distribution over which assigns probability mass to each of the tokens , and probability mass to the remaining tokens.

The task is to minimise the expected KL divergence between and the distribution , where is the final residual stream, and . Where the expectation is taken over random samples of token indices and token weights .

We vary because we are trying to simulate a scenario where the uniform distribution is being produced in parallel with a distribution specific solution giving logit adjustments . It turns out that once we’ve got a solution that can produce the uniform distribution with no error, we can scale up the uniform distribution an outrageous amount, and then we can make the logit adjustments using a width that doesn’t depend on .

can be relaxed to if we want, it just makes it less clean to discuss the network and there are weird edge cases that would never actually happen that have to be discussed, so we avoid it.

Construction:

The construction below is my hypothesis for what the -layer ReLU network I trained on the task above is approximately doing.

I had Claude prepare an animation to visualise the construction.

The initial residual stream is given as a weighted sum of the hashes of the top- tokens . The MLP learns to turn this residual stream into a discrete indicator set for the union of the hashes. The output logits fall into discrete categories, depending on the overlap of their hash with the union of the hashes of the top- tokens.

The discreteness of the logits allows exactly uniform probability to be assigned to the top- tokens, with rare Bloom filter false positives when the hash of a non top- token happens to lie inside the union of the hashes of the top- tokens.

Formal construction:

We use a -layer neural network with ReLU activation, , and a residual connection.

Fix with , with , and d .

For each , sample a uniformly random subset of size of . From the Bloom filter perspective, represents the independent hashes of token .

For each set if , else .

Set , . Add a uniform negative bias to each hidden layer neuron.

Set .

Analysis of a single forward pass:

The residual stream initially is given by an input vector , which has support (here we use the assumption that ).

By construction, the th component of the output from the MLP layer to the residual stream is given by . This is given by if , and otherwise.

Then the final residual stream is given by if , and otherwise.

is set such that the th logit is given by , with equality if and only if . The indices therefore each attain the logit . If no other tokens attain this logit, we will have assigned uniform probability to the top- tokens, and lower probability to the remaining tokens.

So we want to have with high probability over uniform random choice of that ” if and only if .”

This is precisely the same situation as that of a Bloom filter (Youtube video that explains them better than I could), also see the wikipedia article.

The standard bloom filter analysis applies and we get scaling of for the residual stream width, i.e: to get success with probability, and then failure (false positive bloom filter match) otherwise.

Interestingly the false positive case just looks like another token being included in the top- set. So it’s a quite robust solution. Bloom filters never have false negatives, so if a token is in the top- then it will definitely be among the tokens given the maximum tier of logit.

Training:

I trained a -layer residual ReLU network on the top- distribution task using online random sampling of token sets.

The model used vocabulary size , residual width , hidden width , and .

For each training example, a uniformly random subset of distinct tokens was sampled together with random positive coefficients . The target distribution assigned probability mass uniformly across the selected tokens and zero elsewhere.

The network was trained for optimisation steps with batch size using AdamW and a cosine learning-rate schedule beginning at . Training was performed in bfloat precision.

The final model achieved total variation distance on held-out random samples, assigning nearly all probability mass to the target top- set while maintaining low leakage onto non-target tokens.

Behavioural analysis of the trained network:

We draw sets of tokens uniformly without replacement from , and use these as our test top- token sets as input to our trained network.

So we have a distinct rarely occurring cluster where the top- probability mass falls below .

A typical top- output distribution for a token set in the typical region looks as follows:

After samples of random token sets, had false positives, as defined by having probability mass below on the top- tokens (boundary of the bottom probability mass cluster above). Below is a random sample from that set of false positives:

You can see how the network fairly robustly handles the false positive. It doesn’t significantly disrupt the top- token set. If the false positive weren’t highlighted, we wouldn’t be able to tell which token among the top- it was, which is plausibly due to the discrete nature of the logits.

Interestingly, we can also provide tokens in the input, and the model will handle it and give uniform probability among the set of given tokens. So is more like a soft limit to the number of tokens the hashes are designed to be able to handle without breaking down.

Conditioning on there being even a single false positive, we expect the union of hashes to be large, which means we should expect other tokens to have large intersection with the union. We see this in the above figure, with significant probability mass going to tokens , , and as well.

As a baseline, here is Token on the same token set:

In a token set iteration I performed, none of them had probability above assigned to a token where that token was not fully covered by the union of the hashes of the top- tokens. It’s quite interesting how similar in probability the near misses above are, suggesting that there is a discrete probability assigned at each intersection size with the union of the hashes (indeed on other samples we get similar probabilities for near misses.)

Mechanistic analysis of the trained network:

Now we have seen that the model seems to behave like a Bloom filter, we give a partial mechanistic analysis of how the trained weights help to implement the filter.

The fact that the positive values vary in magnitude (in contrast to our construction, which uses a constant value for the positive values) doesn’t change the story very much. Remember we are primarily interested in extracting an indicator set of a union of token hashes, and so as long as each hidden layer neuron fires if and only if it is contained in the union of the token hashes, this story is ok.