Building a Minimal Transformer for 10-digit Addition

The Challenge

I saw a post on Twitter/X about training a model with fewer than 1,000 parameters to do 10-digit addition, and how Claude and Codex were able to do that on their own — that this would be a very impressive thing. I, for one, did not find it that impressive, because I've seen the impressive things these models models can do so it wasn't necessarily surprising to me, but also because I felt I could just do a lot better. Just squinting at it and guessing, I figured I could probably do it with 100 parameters or fewer. So I thought about what the best way to do it would be...

The Idea

Knowing about transformers and their architecture, you'd obviously use attention to accumulate over things. You have queries, keys, and values, and ideally you'd want them to be just one dimension each, so they don't add extra parameters. The natural choice, if you're only going to have one value, would be for that value to be the running sum.

So what do you need for that? You need the digits, and you need the powers of 10, and you need the powers of 10 in a format that lets you accumulate them. You need to multiply the digits by the powers of 10. It doesn't have to be exactly the powers of 10 — you can have the right relative scale and then pick out a arbitrary absolute scale later. (The code is can be found here and a leaderboard of submissions here)

Design Principles: What Counts as a "Plausible Transformer"?

I wanted to have a plausible transformer — something where, if somebody saw the ONNX file, they'd say, "Yep, that's a transformer." Maybe there are some funky things, as might happen when you're really pushing for low parameter counts, but you shouldn't just be making stuff up. As far as I know, nobody is applying RoPE to values. Maybe there's some weird ablation in some super niche paper, but it's not really a thing. So I nixed that.

More broadly, here are the criteria I'm working with:

What counts as a parameter? This is obviously opinion. Zeros aren't really parameters — if everything in a tensor is fixed zeros, that shouldn't count, no weight and no bias can be specified by omission. Zeros are not parameters in much the same way infinity is not a number. The identity matrix is pretty natural too; I'm not sure that should count as parameters either.

What's in-bounds for the architecture? It might make sense to go slightly out-of-bounds with things like having a different dimension for each layer. But we shouldn't be doing ad hoc things like using different activations between layers, or different forms of feedforward or attention between layers. We want to stay within bounds that wouldn't totally creep someone out looking at the ONNX file. I use ReGLU, alibi -- only on one head equivalent to other heads having zero slope, doubles, a different model dim and n_heads per layer (could be the same with implicit zeros and/or sparse tensors), no/implicit identity lm_head.

And of course, we want as few parameters as possible. If we count with dense matrices all the non-zero params it would be 95, ignoring identities 36. Reusing a dim of the input embeddings used for digits for the candidates, (scaling, shifting), 28. If I were to push the level of standardness and use rope instead of k values to select. If we push "Standard Architecture" and use rope for some heads and alibi for others we can get down to ~22, and if we don't count embedding parameters then we could get down to ~12. You could probably push it even further with some combination of technicalities and hacking it. My intuition is for a structured sequence like this using both alibi, rope, softmax1 you could probably get rid of a few more attention parameters by just having them be a fixed pattern and not count it as params. Counting only unique parameters and allowing use of arange without counting as 10 parameters gets things almost arbitrarily small. Arguable the albi slope should count as a param.

Format

The codex solution reversed the order which makes sense for making carry logic easy, but it is less clean. So I prefer to just use 0 padding in the front. So we have two ten digit digit inputs and an 11 digit output. Example:

$$7650676663+0149460439=07800137102$$

The RoPE Detour (and Why I Abandoned It)

For getting the powers of 10, I was thinking about rotary positional embeddings. You can rotate, and if you think about very small rotations, they're basically linearly spaced relative to each other. With linearly spaced numbers, you can rescale or add biases, and if you're using a ReGLU, the sigmoid has regions — mainly the very negative regions — that are approximately exponential, so you could get there that way. But this approach requires applying RoPE to the values, which isn't natural at all. And that brings me to a key design principle. You can also use rope to effectively mask positions for attention, but I wanted to focus on parameters not hyperparameters.

ALiBi for Exponential Decay

Having nixed RoPE on values, the thing I'm doing instead is using ALiBi. ALiBi is very nice for having exponential decay, which is exactly what we want. We want descending powers of 10: 1e9 for the first digit, all the way down to 1e0. Then we want a reset after we see the plus sign, and another reset when we get to the equals sign, at which point we start generating. We want to sum up the values from the first two numbers, appropriately normalized relative to each other. So our attention output is directly the running sum over the sequence length.

Embedding Strategy: Keeping the Model Dimension Small

If we want an easy time getting the actual values of the digits, we could do a one-hot embedding — but that's already 100 parameters, more than we want. Even if the identity matrix "shouldn't really count," it feels kind of wasteful having a big model dimension. So instead, we're going to have the first dimension just be the digit value, and for everything else (operators, special tokens), that dimension will be zero. This is a convenient representation if we want our attention values to just grab that. Mentally note: we'll have a 1 somewhere in our value projection that's basically grabbing that dimension of the embedding.

Using Additional Dimensions for Control Tokens

We have the liberty to use other dimensions, so we're going to use them to have a 1 for BOS, for plus, and for equals. Theoretically we could use negative one on the same dimension as a flag value, but it's easier to just work with separate dimensions. We're also going to have another dimension that is 1 only for =, and why that is will become apparent soon.

The Mean vs. Sum Problem

If you want to do operations with attention, you can do things like means, but not really sums. Because of softmax: it's normalized to be within the range of 0 to 1. You can do different weightings within that range, but going beyond it is more difficult. We want the sum, but what we've got is the mean. That's a problem.

However, there's something really simple we can do. If the queries and keys are equal for all positions, and you have a 1 in the BOS position, then each position gets equal weight. You can't directly get N (the position in the sequence), but you can get $\frac{1}{N}$ pretty easily — or anything-else-over-N as well.

So if we have the mean, we aren't really able to multiply by N (we don't have a good facility for multiplication), but if we're operating on something else, we can divide that by N, and that helps us along.

The Scale Factor and Padded Format

What we're going to do is have a scale that's just descending powers of 10, with special values associated with the BOS, plus, and equals tokens. Conveniently, if we're using a padded format, all of these are 11 positions away from each other. So we don't even have to use our 1/N — we know exactly the scaling we want applied.

With ALiBi providing exponential decay, we can use those values to do the reset and re-normalize the powers of 10 appropriately.

Computing the Output: Running Sum and Digit Selection

For the equals portion, we effectively want a running sum of how much we have and how much we have left. We'll actually count the digits of what it equals as negative, because we want the task to remain the same: "write out the remaining digits."

Here's what we're going to do to accomplish that: we have our proposed values (the digits — we have ready access to them), and we adjust them with the normalization factor (the scale). Then we look at the difference between the value we have and the digit plus 0.5.

Why plus 0.5? If we're trying to minimize the difference, we effectively want to do a floor operation for each digit (except the last one, which we'll return to later). We want to find the greatest digit that would not exceed the number. Adding 0.5 to the digit lets us reframe the task as finding the least difference from those values, which simplifies things.

Using ReLU for Computing Differences

If we can compute these differences with appropriate scales, there are multiple ways to proceed. With a ReLU network, one approach is to have fixed up-projections (just set them to 1 or any non-zero constant) and have gate projections that correspond to the error and negative the error (weights of 1 and -1), then sum those values to get the absolute error.

You can also do something similar with squaring by using the linear region of ReLU. With the squaring aspect, you can do a sort of half-space N-squared, and various other things as well.

From Differences to Output Logits

This is the overall strategy. Once we have these differences, we want to select the item with the least difference. The simplest way to do that: instead of having the absolute value of the difference, have the negative absolute value of the difference, and treat those as logits directly.

You don't need a separate output projection, because you already have it such that the greatest logit will correspond to the digit you want — except for the last position, which needs to be rounded up. To minimize the effect of the residual I scaled these up a lot. Your loss function may care about big magnitude logits but I just care about the right ordering.

The Remaining Pieces: Normalization and Cumulative Sum

From here we just need the normalization factors and the cumulative sum of the digits. As discussed, the cumulative sum can't be done directly — we get the cumulative mean instead. To use that as the sum, we get $\frac{1}{N}$ and use it as a scale factor for whatever we're comparing to (in this case, the proposed digits for computing our error).

Why Softmax1: Solving the 1/N Problem

With ALiBi, it's easier because you have exponential decay. But with fully standard attention, it doesn't behave nicely — you're adding additional things and doing exponential decay on that, so you don't get an exponential decay since you are either adding more things that are effecting it in different ways in different sequence positions or you are not and the output is constant attention weight to the position. This is actually something that's very well and cleanly solved by Softmax1. The key observation is that standard Softmax alone is not properly normalized, and Softmax1 is. $$\frac{e^{x_i}}{1 + \sum_j{e^{x_j}}}$$ We have a great illustration of this with our case: we have a value and key associated with the first element, and we'd like its weight to be exponentially decaying. But it needs something to be relative to, and functionally that has to be the attention scores of the other elements, which are also subject to the decay. So it ends up looking like something entirely different — a much less reasonable basis to get exponential decay from.

On the topic of $\frac{1}{N}$, there's some difficulty, you just don't get a clean $\frac{1}{N}$, because of the 1 in Softmax1. We are dividing by N + 1, however since we are doing this both for the mean and the scale factor it doesn't actually matter.

So what we're going to do is just go ahead and use Softmax1. This gets us $\frac{1}{N}$ by having one token have weight 1 and all the others, 0, and having all keys and values equal to zero.

However for averaging the special tokens for the rescaling of the powers we want behavior like 1->decay, 1->decay, -1->decay uniformly. However if we started with a score around 0 then our decay will start off approximately linearly, leading to bad values for our powers of 10, and since the contributions of the early ones will be very large we can't afford even small errors so we use a large negative bias of -30 to keep the decay very close to perfectly exponential.

Using Doubles

Since we want to be able to represent the sum directly in one activation we will need at least 11 decimal digits of precision in the mantissa, 32 bit floats can only handle about 7, so with double precision we go.

Reflections

This was a fun nerd snipe. Not quite ready to hang up my hat and let the AIs have all the fun with these sorts of things. I wrote a complete implementation myself with minimal AI involvement, basically just boilerplate from copilot. At some point I set claude code on some debugging to my surprise I don't recall it actually solving any of the bugs, it seemed much more concerned with "correcting" the funky things I was intentionally doing. Naturally transformers are quite expressive. I don't think this is something a transformer could really feasibly learn, in part because the logits would almost certainly make it unstable in training, not to mention the alibi head and the fact that effectively nobody is training in double precision. There are probably analogous things that could be learned, that fact alone is interesting but so is the gap, only certain algorithms for doing it would be amenable to learning by gradient descent. I can hand code this one but not the one for writing poetry. Is there another way to go from one to the other (the learned vs. the coded).