{"slug": "challenge-hand-coding-weights-for-efficient-sequence-memorisation", "title": "Challenge: Hand coding weights for efficient sequence memorisation", "summary": "A research team hand-coded weights for one-layer MLPs that memorize labels for input token sequences of length two, finding that the number of facts memorized with 90% accuracy scales roughly linearly with parameter count. The researchers challenge the community to find better constructions that close the gap to trained solutions, aiming to clarify how look-ups are encoded in large language models.", "body_md": "We hand coded weights for one layer MLPs that memorises labels for input token sequences of length two. The number of facts our hand-coded models can memorise with 90% accuracy [1]scales roughly linearly with the models' parameter count\n\n**We pose a challenge to the community:** Find better constructions than ours that can memorise more facts in the same number of weights and close the gap to the trained solution further.[[4]](https://www.lesswrong.com/feed.xml#fn-jXdzaFrLEj5wYD5bh-4)\n\nThe goal of this research is to get a clearer understanding of how look-ups are encoded in LLMs.\n\nWe expect that a lot of the information stored in the weights of an LLM is memorized facts, rather than general circuits. We don't assume a clean separation between what is a \"general circuit\" vs a \"memorized fact\", but a clear example of the former is this [addition circuit](https://arxiv.org/abs/2605.01148), and a clear example of the latter is [knowing what sport some specific athlete is playing](https://www.lesswrong.com/s/hpWHhjvjn67LJ4xXX/p/iGuwZTHWb6DFY3sKB).\n\nThe goal of mech-interp is to be able to take a model (possibly together with its training data), and pick it apart into different components that do human-understandable tasks. Since we expect that many of these tasks are factual look-ups, it would be useful to know what we should expect look-up to look like in a transformer model.\n\nPrior work already points at the MLP layers as the main storage site: [Geva et al.](https://arxiv.org/abs/2012.14913) describe transformer feed-forward layers as key-value memories, [Dai et al.](https://arxiv.org/abs/2104.08696) identify individual MLP neurons that control specific facts, [ROME](https://arxiv.org/abs/2202.05262) locates and edits factual associations in mid-layer MLPs, [MEMIT](https://arxiv.org/abs/2210.07229) inserts thousands of facts into mid-layer MLPs with closed-form weight updates, and [Allen-Zhu & Li](https://arxiv.org/abs/2404.05405) measure how many bits of knowledge a transformer can store per parameter. On the toy-model side, a recent theory line models transformer weight matrices as associative memories that store facts as sums of outer products: [Bietti et al.](https://arxiv.org/abs/2306.00802) for bigram lookups, [Cabannes et al.](https://arxiv.org/abs/2310.02984) for the capacity of linear memories over random embeddings (only facts), and [Nichani et al.](https://arxiv.org/abs/2412.06538), who prove that one-layer transformers can store a number of facts linear in their parameter count, up to log factors.\n\nNone of this yet amounts to a weight-level account of *how* the storage in real models actually works. That's the level of understanding we're after.\n\nIn this post, we study sequence memorization as a toy model for any lookup where some combination of input signals carries a meaning that is substantially different from a linear combination of the individual signals.[[5]](https://www.lesswrong.com/feed.xml#fn-jXdzaFrLEj5wYD5bh-5)\n\nNote that the kinds of sequences we consider here are deliberately structureless. Random input pairs are mapped to random labels, so there is nothing to compress. The model can do nothing except memorize.\n\nWe mostly focused for now on creating algorithms that memorise as much as possible, not yet on trying to understand how exactly the solutions trained models use work. The hope is that once we have any solution at all that performs somewhat similarly to the trained models', figuring out what exactly the trained models are doing might become a lot easier. That does mean that we focus on algorithms that qualitatively seem to us like they're not too far removed from something gradient descent methods might be able to learn.\n\nThe training data are sequences of three tokens, two input tokens and one output token. Given an input of two tokens, the network is trained to predict the next token (i.e. the output token).\n\nThe code for generating the training data is not very long, and is quoted below in full, if you prefer to read code.\n\nHyperparameters for the data generation:\n\n`n_facts`\n\n-- Number of facts.`input_vocab_size`\n\n`output_vocab_size`\n\n`seed`\n\n-- Random seed. The code first generates a list of every possible input combination. Then this list is shuffled, and the first `n_facts`\n\npairs from the shuffled list are used as the inputs for the `n_facts`\n\nfacts. These facts are then divided as equally as possible among the `output_vocab_size`\n\ntarget labels.\n\nCode\n\n``` python\ndef generate_facts(n_facts: int, # of facts to generate,                    input_vocab_size: int, # of unique tokens in the vocabulary                   output_vocab_size: int, # of unique targets                   seed: int = 42                  ) -> dict[str, torch.Tensor]:        if n_facts > input_vocab_size ** 2:        raise ValueError(f\"Cannot generate {n_facts} unique facts with a vocabulary of size {input_vocab_size}. Maximum unique facts: {input_vocab_size ** 2}\")        device = torch.tensor(0).device  # respect default device    generator = torch.Generator(device=device).manual_seed(seed)    all_possible_inputs = torch.cartesian_prod(torch.arange(input_vocab_size),                                               torch.arange(input_vocab_size))    inputs = all_possible_inputs[torch.randperm(all_possible_inputs.size(0),                                                generator=generator)[:n_facts]]        targets = torch.arange(n_facts) % output_vocab_size    sorted_indices = torch.argsort(targets)        return {\"inputs\": inputs[sorted_indices],             \"targets\": targets[sorted_indices]}\n```\n\nWe want to find out how these facts can be encoded in a one layer transformer. However, that turned out to be hard. But if we know in what part of the model the main action is, then maybe we can simplify the toy model to only that part and start with understanding that.\n\nTo test what parts of the model are important for the sequence memorization task, we made a transformer model, where every part of the model can be turned on or off. Then we trained all variants of this model and compared their performance.\n\nThe full toy model consists of:\n\nAfter the attention, the model only continues its computation in the second token position. This is because the model is only trying to predict the third token, and not the second token.\n\nWe want to be able to simplify the model by removing the attention. However, the problem with doing so is that the information from the first token has to reach the second token position somehow. Therefore, we can't just remove the attention, but will have to replace it with something else.\n\n*Mixing* is the part of the model that combines the first token and second token information into the same residual stream. We have three different variants for this.\n\nThere are a number of variants regarding the MLP. Firstly the MLP can either be present or be missing. Secondly if there is an MLP layer, each of the following can be varied\n\nThe norms can also be turned on and off. Each of the norms for the read-in to the attention and MLP only exists if both that part of the network is present (Unif Attn or Lrn Attn for the attention), and Norms are turned on. The last norm, just before the unembedding only depends on the norm setting, and is there if norms are turned on and not there if norms are turned off.\n\nThe following is just a summary of the results. To see the full results, including the details of the experimental setup, see Appendix B.\n\nTo find out which parts of the network matter for the memorization task, we trained every combination of the architectural variants described above, and measured the maximum number of facts each one could learn. All models in this experiment used the same size:, , , .[[10]](https://www.lesswrong.com/feed.xml#fn-jXdzaFrLEj5wYD5bh-10)\n\nThe networks are trained using gradient descent [11], on a cross entropy loss. We say that they have successfully learned some number of facts if, at the end of training, taking argmax over the logits always gives the correct label. To find the maximum learnable facts for a specific architecture variant, we do a binary search over the number of facts.\n\nNote that this is measuring the maximum number of facts *that Adam can find weights for*, not how many facts an architecture can represent in theory. So, some of the observed effects below (definitely the learned-attention results, possibly also the norm results) are likely about trainability rather than architectural capacity.\n\nIn the bullets below, percentage ranges like \"X% - Y% more facts\" give the range of the effect of one setting across all combinations of the other settings. For absolute numbers see table B.1 in Appendix B.\n\nThe **MLP** is by far the most important part.\n\n**Mixing** is the second or third most influential setting.[[12]](https://www.lesswrong.com/feed.xml#fn-jXdzaFrLEj5wYD5bh-12)\n\nUniform attention being better than learned attention has to be due to learned attention having training difficulties, since learned attention is strictly more expressive. Consistent with this, the learned attention results are also by far the least stable across repeated runs. It's not surprising that dual embedding does better than uniform attention, since it's both strictly more expressive [13]and should be no harder to train.\n\nThese results suggest that the attention probably isn't doing anything importantly different from just linearly adding together the embeddings of the two tokens.[[14]](https://www.lesswrong.com/feed.xml#fn-jXdzaFrLEj5wYD5bh-14)\n\n**Norms** is the second or third most influential setting.[[12]](https://www.lesswrong.com/feed.xml#fn-jXdzaFrLEj5wYD5bh-12)\n\n**The remaining settings** matter less.\n\nFinally, there is one notable outlier: the combination **MLP**=✅, **Norms**=✅, **Res**=❌, **Bias**=❌, **ReLU** (combined with **any Mixing** option) does far worse than the individual settings would predict, almost as bad as having no MLP at all. [15]We don't know why.\n\nHow does the number of learnable facts grow with model size? To find this out, we picked two of the very many model architecture varieties and scaled them up.\n\nThese models are:\n\nWe test how many facts each of these models can learn for a range of model dimensions with\n\nAnd the result is\n\nThe number of facts each model can learn scales roughly as . This is basically what we'd expect a priori from an information theory perspective, because the number of bits in the model's parameters scales with. The number of bits required to store a label scales with .\n\nCan we write down weights for the sequence memorization toy model, either by hand or with some algorithm that isn't raw numeric optimisation à la gradient descent, such that the resulting model matches the performance of a trained model?\n\nUltimately we also want this algorithm to yield weights and neuron activations qualitatively similar to those real models produce, but for now we're mostly focusing on raw performance.\n\nThere are two reasons why this is a useful framing.\n\nWe think that our current best attempt (which is presented further down) is some non-zero progress on this challenge, but there is still far to go. We encourage all readers to give it a try.\n\nMost of the sequence memorization capacity is in the MLP. We therefore propose focusing on a toy model with only this part and everything else cut out. This would be something like the *simple* model in the previous scaling experiment (Figure 2). However, that architecture still has unnecessarily many weights, a legacy of being a cut-down version of the full model in Figure 1: the MLP's input matrix sits directly after the embeddings, and its output matrix directly before the unembedding. Two consecutive linear maps can always be folded into one, so we absorb into the two embedding matrices and into the unembedding, with no loss of expressivity. Doing so gives us this architecture:\n\nUse this architecture for the challenge if you want to be able to compare your results with ours. However, if you want to go for something slightly different, or even the full 1-layer transformer, we'd still be interested in what you can do.\n\nNote that in our model architecture, **the MLP's hidden dimension ** is always scaled to stay smaller than the input vocabulary **, and this seems somewhat load-bearing. For the case , we already have a more efficient construction which we think probably gets pretty close to the theoretical capacity limit, though it doesn't seem like something a model is likely to learn in training. See Appendix A for that construction.\n\nFor each of the four criteria, [18]see how the maximum number of facts scales with model size.\n\nHand-coded memorization has a long history in learning theory: [Baum (1988)](https://www.sciencedirect.com/science/article/pii/0885064X88900209) memorizes binary-labeled points in *general position* [19]in with hidden threshold units,\n\nJust before publishing this, we actually found one construction in the literature that seems promising: [Dugan et al.](https://arxiv.org/abs/2512.00207). They give a weight construction for gated one-hidden-layer MLPs. The gating weights are chosen at random. The requirement that every key hit its target output (margin-optimal label directions, compressed to roughly dimensions with a random projection that the output matrix then undoes) then becomes a linear system in the entries of the up-projection matrix , which they solve exactly.\n\nWe have some preliminary results for adapting this construction to our setting that indicate it performs better than our current hand-coded solution or even our hybrid solution, though it still falls well short of the trained one. The reason we haven't switched over to their construction yet is that it seems to us very unlike something models trained with gradient descent would or could learn. A bit like our own construction in Appendix A. We do ultimately want to reverse engineer trained models, getting storage efficiency up is just a stepping stone for that. Nevertheless, adapting this construction would be a valid challenge entry and we'd probably have included an adaptation of their algorithm in this post if we'd found it earlier.\n\nThis is our [6]attempt at solving the challenge we proposed above. Below we present first our algorithm, and then how it performed relative to trained models.\n\nYou can find the code for our construction [here](https://github.com/LindaLinsefors/Memory-Toy-Models/blob/master/hand_coded_models/hc2.py).[[20]](https://www.lesswrong.com/feed.xml#fn-jXdzaFrLEj5wYD5bh-21)\n\nFirst we associate each label with a unique set of neurons. [21]These neurons will somehow identify facts with this label.\n\nIt would be great if, given a fact with label, all the neurons associated with label were active, and no other neurons are active. However for most sets of facts, this will not be possible.[[22]](https://www.lesswrong.com/feed.xml#fn-jXdzaFrLEj5wYD5bh-23)\n\nAn alternative, which is less ideal, but has the significant advantage of being possible, is to set up the embedding weights such that, for any input fact with any label, all neurons associated with will be active. Some additional neurons will also be active, but not all of them. With some skill and some luck, none of the wrong labels will have all their neurons activated.\n\nHowever, there is one more problem. \"Active\" isn't a single number. Even if the correct label has all its neurons active, and none of the incorrect labels has all of their neurons active, an incorrect label can still have higher overall activation, if that label's neurons activate more strongly. One way to solve this is to make sure all active neurons have the exact same value, but that would add an unnecessary constraint. A better solution is to flip the script, and make sure that all neurons associated with label are *inactive* for any fact with label . This exploits the one place where a ReLU network gives us exact equality for free: every negative pre-activation is mapped to exactly zero. \"All of this label's neurons are inactive\" is therefore a condition the readout can check exactly, by giving those neurons negative weight into the label's logit and checking for a logit of exactly zero.\n\nIn summary:\n\nOur algorithm for assigning weight matrix values has these steps:\n\nThere are ReLU neurons, and labels. Each label gets assigned neurons. In most of our experiments , which means for any , the assignments will overlap.\n\nOne problem our network needs to solve is that there will likely be some pattern of facts\n\nAny weight allocation, on this model architecture, where the logit for some label only depends on a single ReLU neuron, will fail at encoding this pattern. Therefore, the network either needs more ReLU neurons than labels (not realistic) or the labels will have to somehow share neurons, i.e. some sort of [superposition](https://transformer-circuits.pub/2022/toy_model/index.html) encoding.\n\nWe don't know a priori what the best value of is, so is given as a hyperparameter, to search over later. Given , and , we want to find an allocation where the assignments are spread out nicely. I.e. we want all neurons to be used by approximately the same number of labels, and we want to minimize the max neuron overlap between any pair of labels.\n\nWe had Claude Code write the script that does this, and verified that the outputs look good. There are probably many ways to achieve similar allocations, and we can't think of any reason why the exact method matters, so we will not go into this further. If you want to look into the details, see [the code](https://github.com/LindaLinsefors/Memory-Toy-Models/blob/master/hand_coded_models/hc2.py).[[20]](https://www.lesswrong.com/feed.xml#fn-jXdzaFrLEj5wYD5bh-21)\n\nThe next step is to make sure that every ReLU neuron always output zero on all facts with a label assigned to that neuron. Recall that in this architecture (Figure 4), the embedding weights map input tokens directly to neuron pre-activations, so \"embedding weight\" below means the weight from a token (at position one or two) to a neuron.\n\nIn broad strokes, our algorithm for each neuron is:\n\nWe set the unembedding weights between each label and its assigned neurons to . We set all other unembedding weights to .\n\nWe did also try assigning positive values everywhere else, but for the success criterion we use (looking at argmax of the logits), adding these positive values makes no difference. We first noticed this empirically, but it's also a mathematical fact.\n\nSame as above except the unembedding weights are randomly initialized and then trained.\n\nWhen testing the capacity of this model design, we always do a hyperparameter search over and .\n\nFor the hand-coded models the winning is typically where is the model dimension. For hybrid models, this number is a bit smaller.\n\nFor the hand-coded models the winning is typically in the range . For the hybrid models the number is a bit larger.\n\nSee Appendix C and D for details.\n\nThe plot below shows our data from binary search to find the maximum number of facts a model can learn.\n\nThe best-fit lines from Figure 5:[[24]](https://www.lesswrong.com/feed.xml#fn-jXdzaFrLEj5wYD5bh-26)\n\nCondition | acc = 1 | acc ≥ 0.9 | |\n|---|---|---|---|\ntrained | |||\nhybrid | |||\nhand-coded | |||\nrand-emb |\n\nThere might be multiple different algorithms to achieve high memorisation capacity, so even if we match the trained model's performance that doesn't necessarily mean we understand how the trained model works yet, though we think it'd be a good intermediate result. That said, how close do the weight matrices our algorithm constructs look to the trained weight matrices?\n\nThe first row shows our hand coded weights. The second and third rows are weights for models trained on the same number of facts as the hand-coded ones, and as many as they can fit respectively.\n\nSome surface-level observations: Our embedding weights only take three different values, 1, 0 or -1. The network's weights are more smeared out, though the ones in the second row are quite trimodal, with one narrow cluster at zero and two broader clusters of positive and negative weights. The weights in the third row are much more spread out. They could maybe be three broad peaks grown together?[[25]](https://www.lesswrong.com/feed.xml#fn-jXdzaFrLEj5wYD5bh-27)\n\nThe unembedding is pretty different even in the second row, with our hand-coded weights being bimodal whereas the trained weights are trimodal again, with a small peak at zero and big positive and negative peaks.\n\nFor the neurons, while our hand-coded algorithm stores facts as patterns of inactive neurons, trained models seem to use patterns of active neurons instead. Maybe as a result of that, the proportions of positive vs. zero activations appear to be somewhat flipped. The trained models' neuron activations are also more smeared out and gradual than ours, as one would expect from their more smeared out embedding weights.\n\nLinda did most of the work. Lucius gave advice and encouragement, came up with the alternative construction in Appendix A, and helped write and edit the post.\n\nThanks to Nico Penttilä and Mikhail Doroshenko for suggesting the Uniform Attention variant.\n\nLinda is supported by a grant from Coefficient Giving. Lucius works at Goodfire AI.\n\nThroughout the main text the hidden width is scaled to stay *b*elow* the input vocabulary size (we use ). This appendix treats the boundary case just past that constraint. The hidden layer is given one neuron more than the input vocabulary size,. In this case a hand-coded construction can store all possible facts with 100% accuracy. It uses the same architecture as Figure 4 (two embedding matrices, one ReLU layer, an unembedding; no norms, residuals, or bias terms [26]).\n\nThis algorithm does seem qualitatively very unlike the ones trained models learn, with only two active neurons per fact. So it might be a bit of a dead end. We nevertheless include it in case there are ways to adapt ideas from it for future constructions.\n\nLet be the label of the fact with input .\n\nThe first neurons will be **selector neurons**, with neuron switching on only if first-position token is active, and the final neuron is an always-on bias. The weights are:\n\nThe overall scale of is free: multiplying it by any positive constant leaves the argmax unchanged but makes the softmax arbitrarily confident in the correct label.\n\nIn words: the first token uses the large negative weight to switch every selector neuron except the one it is assigned to off. The second token's embedding then writes the fact's label (shifted up by one) into that surviving selector neuron as its activation value. The shift keeps the value positive, so that a label of is distinguishable from a silent neuron. The bias neuron is wired to be permanently on and supplies the per-class thresholds that are required to read the labels out linearly (which is why we cannot fold it away in a bias-free architecture, and why we need the extra neuron).\n\n**How it works.** Take the fact. On input , selector neuron has pre-activation and fires with value ; every other selector neuron has pre-activation and stays silent; the bias neuron is . In general, writing ,\n\nThe bias neuron's output weights make the activations for different labels linearly separable. The labels with closest to the neuron's activation value then always gets the highest logit. This holds for every one of the pairs, so any set of up to possible facts is memorized perfectly.\n\nThe point of this experiment is to figure out which parts of the network are important for the sequence memorization task, so that we know which parts are safe to ignore or even remove, in order to make understanding the model easier.\n\nWhat we found was that having an MLP is the most important part (approximately responsible for learning half of the facts), and everything else only matters a bit.\n\nWe trained all different versions of the toy model, to see how many facts each of them could learn. There are some patterns, but unfortunately, for most of them, we can't separate what is a result of changing the expressivity of the model architecture and what is a result of making it easier or harder for Adam to learn good solutions.\n\nFor this experiment, we used a single model size across all architectures:\n\nWe say that a model has \"learned a fact\" if, when the model is given the first two tokens of this sequence, it correctly predicts the third token. And by \"correctly predicts\" we mean that argmaxing over the logits locates the correct output token.\n\nTo find the maximum number of facts a model can learn, we performed a binary search over the number of facts, to find the highest number of facts such that the model learned all of them.\n\nInside the binary search, for each number of facts tested, we trained 11 models in parallel, with the exact same facts but differently randomly initialized weights. We used three different success criteria — \"Any\", \"Most\" and \"All\" — meaning that we counted the model as having succeeded at learning all the facts if it succeeded in any, most, or all of the 11 trials.[[27]](https://www.lesswrong.com/feed.xml#fn-jXdzaFrLEj5wYD5bh-29)\n\n**All tables (below) show all of \"Any\", \"Most\" and \"All\". However, the analysis in the text, and in the main post, is based on the results for \"Any\", since that's the most stable setting.**\n\nEach individual training run was done as follows. We used Adam with no weight decay, and with all facts included in every batch. [28]Each model was initially trained with learning rate\n\n`lr=1e-2`\n\n, for up to 50,000 epochs, or until early stopping due to reaching full accuracy, or accuracy not improving for the last 5,000 epochs. If after this step the accuracy is less than 100% but above 95%, then training continues with `lr=3e-3`\n\n, with the same stopping criterion. If still not at 100% accuracy, training continues with `lr=1e-3`\n\n, with the same stopping criterion.For each architecture and each of Any/Most/All we ran a binary search to find the maximum number of facts it could learn. Furthermore, we repeated each such binary search 4 times, to check for stability. The \"Any\" setting had the highest stability (similar max number of facts over all 4 duplicate experiments), and \"All\" had the worst stability.[[29]](https://www.lesswrong.com/feed.xml#fn-jXdzaFrLEj5wYD5bh-31)\n\nAll the results for all the experiments are shown in the table below.\n\nNote that 1024 is the dataset ceiling, so configurations reaching it have saturated the data rather than the model.\n\nTo see the effect of each of Mixing, Norms, Res, Bias and Act, we'll look at pairs (or triples for Mixing) of model architectures that are the same except for that variable. E.g. for Norms, we look at every pair that differs only in whether it has norms or not, to see how much models with norms typically outperform the ones without norms. We're doing this analysis on the \"Any\" runs, since these have the most stable outcomes.\n\nHowever, before doing all that, it's worth noting one major outlier.\n\nFor any form of mixing (learned attention, uniform attention, or dual embedding), networks with **MLP**=✅, **Norm**=✅, **Res**=❌, **Bias**=❌ and **Act**=** ReLU** do really badly. Almost as badly as (and in one case slightly worse than) removing the MLP.\n\nThis combination is extra bad for some reason that isn't just the sum of its parts. We don't know why. Specifically, we don't know if the limitation is due to training dynamics or due to what is possible for this architecture.\n\nInstead of writing \"except for **MLP**=✅, **Norm**=✅, **Res**=❌, **Bias**=❌ **Act**=** ReLU**\" in every subsection below, we'll just point this out here. Having pointed this out, this data will be excluded from the triple or pairwise comparisons below.\n\nNow back to triple and pairwise comparisons. I.e., we compare outcomes (number of learned facts) for pairs (or triples, when varying Mixing) of model architectures where the only difference is a single setting.\n\nIt is notable that Mixing (the setting that determines the embedding and attention) has the least effect precisely when everything else is turned off, i.e. when there is only the embedding, possibly attention, and the unembedding.\n\nIt's not surprising that dual embedding does the best, since this architecture is (arguably) the most powerful. Because attention is non-linear and the dual embedding is linear, there are things that the attention can express that the dual embedding can't. But on the other hand, the dual embedding gets to encode the input for each token position entirely separately, which gives the network more freedom. Additionally, this no-attention setup should be easier to train, since it's simpler.\n\nMore surprising is that uniform attention outperforms learned attention, given that learned attention is strictly more powerful. Therefore, this has to be because of ease of training. This interpretation is also supported by the observation that the number of facts networks with learned attention manage to learn is unstable. You can see this in the number of outliers in Table B.1, and also in how much the number of facts drops from Any to Most to All.\n\nNot surprisingly, adding the MLP makes the biggest difference to the number of learnable facts out of any of the settings.\n\nNorms are generally useful for learning more facts. Norms make a bigger difference if there is attention, and if there is no MLP. Possibly this means that the norms in front of the attention and unembedding are helpful, while the norm in front of the MLP is anti-helpful. Or possibly the MLP and norms overlap somewhat in function, such that the MLP makes the norms less useful.\n\nThe fact that the MLP has its biggest effect when there are no other non-linearities in the network points to the overlapping-function hypothesis. However, the unusually bad performance of **MLP**=✅, **Norms**=✅, **Res**=❌, **Bias**=❌, **Act**=** ReLU** might be a sign that adding norms can be bad for the performance of the MLP.\n\nAdding a bias ought to be strictly helpful, but for some reason it's anti-helpful in a few cases.\n\n**GELU** is typically better than **ReLU**, but the difference is small.\n\nThe data points in Figures 3 and 5 in the main post are for the \"Any\" condition. Below are the same two plots with \"Most\" and \"All\" included.\n\nis (at least for the hand-coded model) the number of neurons used by any label. We're interested in how this scales with various model parameters, since it might give us some clue about what we should expect superposition to look like over ReLU and ReLU-like neurons. To be clear, anything we see here is at best a small hint, with no guarantee of having anything to do with how computations are distributed in fully trained models. But it's still a little bit of Bayesian evidence, and maybe if it meets up with other evidence later on, it will tell us something. This is why we think it's worth recording.\n\nIn the experiments in the main post, in order to get the best version of our hand-coded model, we did a hyperparameter sweep over and . From this we can extract the optimal for different model sizes by looking at from the winning (, ) pair.\n\nIn the scaling experiments (see main post) we investigated models with dimensions, , and for . Looking at the best-performing from these runs, we find that .\n\nHowever, in these experiments the relations between, , and are locked together, i.e. we can't tell from the data which of them influences the ideal value of .\n\nIn the next experiment we did a binary search for the max number of facts for every combination of the following parameters:\n\nBelow you can see how the optimal depends on all of them.\n\nNote that we swept over . A small number of runs may have hit the ceiling, i.e. the optimal value of is actually something above .\n\nAs you can see, the picture is less clean when the different model dimensions are varied independently. Just scaling up the hidden layer increases the optimal faster than . But also decreases with , just enough to add up to the pattern we see in the first figure of this appendix when they increase together.\n\nIn the hybrid model, also plays a role.\n\nThe other hyperparameter (besides ) used when creating the embedding matrix for the hand-coded and hybrid models is a variable we dubbed . See the main post for the definition.\n\n**We don't expect that looking into this variable will tell you anything interesting. We don't recommend paying attention to this section unless you know something we don't.**\n\nBut for completeness, and because it cost us almost no extra work to add this, here are the same plots as in Appendix D, but for instead of .\n\nNote that we only swept over . Some runs with the hybrid model seem to have hit the ceiling, i.e. the real best is something above .\n\nMeaning the model memorises the correct label for of the facts it is trained on. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-1)\n\nUp to a log factor that comes from larger output label dictionaries needing more bits to index. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-2)\n\nNot a trivial one though. We couldn't find an existing analytic solution for linear classification under softmax cross entropy loss, or linear classification under argmax. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-3)\n\nImportantly, the MLP's hidden layer size here is smaller than the input token dictionary . For the case (or if we get output biases), we already have a more efficient construction, see Appendix A. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-4)\n\nAn even simpler example of lookup, not studied in this post, is [bi-gram statistics, which is (at least sometimes) encoded in the embedding + unembedding matrices](https://transformer-circuits.pub/2021/framework/index.html#zero-layer-transformers). Bi-grams don't need any token-combination machinery, which is why they can live entirely in the embedding/unembedding. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-5)\n\nWe used a fixed random seed when generating facts in order to avoid some runs getting lucky and getting easier facts, and to specifically have the same facts for trained networks and hand-coded networks. The last part kind of failed because torch random functions give different results when run on CPU vs GPU, even when the seed is the same. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-7)\n\nThanks to Nico Penttilä and Mikhail Doroshenko for suggesting this variant. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-8)\n\nIf the bias is included, that means both the linear projection into and the linear projection out of the ReLU/GELU neurons get a bias. (Making them actually not linear functions but affine functions, in strict math terminology.) If the bias is not included, this means neither of them get a bias. All other linear connections in the rest of the network (e.g. embeddings, etc) are always bias-free. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-9)\n\nInitially was the same as all the other values, but too many networks maxed out the number of possible facts, so we doubled [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-10) [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-10:1) [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-10:2)\n\nAdam to be specific, with the following learning rate schedule: We start out with lr=1e-2, and train for either 50000 epochs or until accuracy [defined as ] has not improved for 5000 epochs. If at the end of this accuracy is below 1 but above 0.95, we continue the training with lr dropped to 3e-3 and same stopping criteria, and finally repeat for lr=1e-3. Training also ends if the accuracy=1 is reached. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-11)\n\nIf there is an MLP block then Mixing is generally more important than norms, and the other way round without the MLP. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-12) [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-12:1)\n\nStrictly more expressive because: with uniform attention, the residual stream at the second position is plus positional terms: Two per-position linear maps that are constrained to share the same embedding matrix . 2Emb replaces these with two fully independent matrices, so it can represent anything uniform attention can, and more. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-13)\n\nWith some rotation added to the first token embedding so as to be able to tell apart inputs with the two tokens swapped, e.g. to differentiate the input \"1,2\" from the input \"2,1\". [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-14)\n\nJust turning off the norms or turning on the MLP bias or switching from ReLU to GELU (changes that normally have small effects) is sufficient to restore this architecture to the normal capabilities range. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-15)\n\nThat is, if you want to be able to compare your results with ours. But if you make progress on this challenge using some other architecture, we'd be interested in that too. Just make sure to include something like an MLP layer, since that is where most of the sequence memorization capacity lives in the trained models. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-16)\n\nExcept for the hybrid condition where you can use gradient descent-esque methods for the unembedding weights. See \"Evaluation Criteria\". [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-17)\n\nThere are four combinations:\n\nMeaning no k+1 of the points lie on a common affine hyperplane. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-19) [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-19:1)\n\nI (Linda) cleaned this one file in the repo. Everything else there is a mess. I don't recommend trying to read any other file. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-21) [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-21:1)\n\nTo clarify: Typically each neuron will be used by multiple labels, but no two labels share all their neurons. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-22)\n\nGiven some set of facts (i.e. all the facts with label , or all the facts associated with any label that is associated with neuron ), you typically can't decide that some neuron will be active (or inactive) for only that set of facts, and no others. That's only possible if the inputs to those facts are linearly separable, which is typically not the case. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-23)\n\nThe embedding matrix weights are drawn from a uniform distribution centered on zero. Specifically it's uniform over . However, the scale should not matter, since that can be compensated for by the unembedding weights. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-25)\n\nNote that the prefactors here are a little different from the ones in Figure 0 because we're fitting both the prefactors and the exponents to confirm which exponents seem close to the desired . Any exponents also shouldn't be taken too seriously. They're either noise in the fit or ought to vanish at large enough . [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-26)\n\nThe absolute scale of the weights doesn't matter that much because ReLUs are scale invariant. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-27)\n\nIf we have output biases, we can adjust the construction to only need . [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-28)\n\n\"Most\" means that there are more successes than failures, i.e. at least 6 out of 11. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-29)\n\nI.e. batches and epochs are the same thing. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-30)\n\nIt's not surprising that \"All\" had bad stability, since it only takes one bad run to throw off the entire batch. But it was not a priori obvious to us that \"Any\" would be more stable than \"Most\". [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-31)\n\n**2Emb** and **Unif Attn** learn **208** facts in each of the four repeated experiments. Lrn Attn learns **208** in two of the replications, slightly more in one, and slightly less in one. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-32)\n\nSee the appendix for our much less 'realistic' but much more efficient storage construction in the easier setting. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-20)\n\nThis could have been any negative number, and it would work equally well. It just happens to be -2 in the code. [↩︎](https://www.lesswrong.com/feed.xml#fnref-jXdzaFrLEj5wYD5bh-24)", "url": "https://wpnews.pro/news/challenge-hand-coding-weights-for-efficient-sequence-memorisation", "canonical_source": "https://www.lesswrong.com/posts/KWtchKwwnJkd4bwCi/challenge-hand-coding-weights-for-efficient-sequence-1", "published_at": "2026-07-23 18:05:43+00:00", "updated_at": "2026-07-23 18:30:16.738949+00:00", "lang": "en", "topics": ["machine-learning", "neural-networks", "ai-research"], "entities": ["Geva et al.", "Dai et al.", "ROME", "MEMIT", "Allen-Zhu & Li", "Bietti et al.", "Cabannes et al.", "Nichani et al."], "alternates": {"html": "https://wpnews.pro/news/challenge-hand-coding-weights-for-efficient-sequence-memorisation", "markdown": "https://wpnews.pro/news/challenge-hand-coding-weights-for-efficient-sequence-memorisation.md", "text": "https://wpnews.pro/news/challenge-hand-coding-weights-for-efficient-sequence-memorisation.txt", "jsonld": "https://wpnews.pro/news/challenge-hand-coding-weights-for-efficient-sequence-memorisation.jsonld"}}