Teaching a Tiny Transformer to Add Longer Numbers Than It's Seen

Train a small transformer to add numbers of 1 to 10 digits. Then ask it to add numbers of 11 to 40 digits. Which way of telling it where each digit sits lets it keep going? I ran six position schemes, three seeds each. Only one extrapolated, and only about 5 digits past the training range (that's 1.5 times here, but a follow-up at the end shows it isn't a ratio).

The setup

Every model is a decoder with 4 layers, width 256 and 8 heads, about 3.2M parameters. Each trained for 8,000 steps (AdamW, learning rate 5e-4 with a cosine schedule, batch 256) on an Apple GPU, using MLX. A run took about 20 minutes with a few running at once. I test with greedy decoding and count a sum as right only if every digit and the end-of-answer token match. That's 200 problems for each length from 1 to 40 digits, where both numbers have that length.

The schemes: learned absolute positions with the digits in normal order, and the same with the answer written least-significant digit first; no position information at all (NoPE), reversed; RoPE, reversed; and the Abacus embedding from McLeish et al. 2024, reversed. Abacus gives each digit an index for its place within its number. During training the whole index sequence starts at a random offset between 0 and 50, so the model sees large indices. (That's my simplified version: one offset per example. The paper samples its offset from a wider range.) I also ran Abacus with no offset as an ablation.

Results

Exact-match accuracy, averaged over seeds and over the lengths in each band:

None of the six schemes got anything right at 31 to 40 digits (the SSMax variant below manages a percent or two at 31–37). On a mixed check (a 10-digit number plus a 5-digit one) every scheme but normal-order learned positions scored 100%, so the failure is length, not an inability to add. Learned positions in normal order don't even manage the training range in 8,000 steps: adding most-significant-digit first means the model has to guess carries it hasn't seen yet.

Abacus with an offset (solid green) is the only original scheme that doesn't fall off a cliff at the training edge. It still falls: I define "reach" as the longest length where every shorter length scores at least 90%. For Abacus with an offset that is 14, 16 and 15 digits across seeds. The offset matters: without it, Abacus is as stuck as the rest. So the scheme isn't doing it alone. The training distribution has to show the model index values it will meet later.

Why does it stop?

The obvious guess is unseen indices. A 20-digit number needs place indices the model never saw. But training indices reach up to 60 (offset up to 50 plus 10 digits), so indices that high were seen. I checked directly with teacher forcing on one trained model, shifting the evaluation offset to 0, 10 and 30. At 20 digits accuracy was 0.58, 0.26 and 0.22. A larger offset made it worse, not better. So unseen index values aren't the story.

Next I varied the shape of the problem. 20 digits plus 3 digits scored 0.89. 20 plus 20 scored 0.58. 24 plus 3 scored 0.20. Total sequence length hurts, and so does the length of an individual operand, separately. And the errors weren't piled up at the high digits. They were spread across every answer position.

That pattern fits a vague story: to write one answer digit, the model has to attend to the two matching digits and nothing else. As the context grows, softmax spreads its weight over more tokens, and the useful positions get a smaller share. I'll call that attention dilution. It predicts errors everywhere, and it predicts that longer inputs are harder.

Testing dilution with SSMax

Scalable-Softmax (SSMax, Nakanishi 2025) is aimed at this problem. It multiplies the attention logits by a learned per-head scale times log of the context length, so the softmax stays sharp as the context grows. I added it to the Abacus and NoPE models, same training, three seeds.

It helped, but not at the edge. Abacus with SSMax went from 0.55 to 0.78 in the 16–20 band and from 0.03 to 0.14 in 21–30. That's more partial credit past the edge. The reach didn't move: 16, 15 and 16 digits against 14, 16 and 15. NoPE's 11–15 band went from 0.33 to 0.45, and its reach stayed at 11 digits in every seed.

My read: dilution is part of the decay, since sharpening attention buys accuracy beyond the edge. But it isn't what sets the edge. Something else breaks at around 1.5 times the training length, and sharper attention doesn't fix it.

Things that broke

An off-by-one in the loss mask. My first runs scored 0% exact match. The mask was shifted by one, so the model was trained to predict the "=" token instead of the end-of-answer token, and it never learned to stop. A 1,500-step smoke test caught it: with the fix, Abacus scored 0.93 to 1.0 up to 14 digits.

Evaluation ate all the memory. Greedy decoding at many different (length, shape) combinations made MLX's buffer cache grow to 11 GB per process. Three at once filled swap and got killed. Calling mx.set_cache_limit with 1 GB before evaluating brought an evaluation down from 16 minutes to 35 seconds, at a 1.7 GB peak. It cost me about an hour to find.

Update: training on 1–20 digits

I asked whether the edge is a fixed ratio of the training length. So I trained the same Abacus model (random offset up to 50) on 1–20 digit operands, three seeds, and tested up to 60 digits. I doubled the steps to 16,000 because the sequences are twice as long.

A ratio of 1.5 would put the 90% edge near 30. It landed at 25–26 digits on all three seeds. That's about 5 digits past training, the same margin as the 1–10 runs (14–16). The 50% point moved further: 16–20 digits for the 1–10 runs, 30–33 for the 1–20 runs. Accuracy hit zero by 40–45 digits.

So the reliable part looks like a roughly fixed number of extra digits, not a multiple of the training length. The partial-credit tail stretches more with longer training. "About 1.5x" above was a coincidence of training on 10 digits. Two training lengths and one model size is thin, though, and the step count changed with them.

What I'd try next

• Recurrence and input injection. McLeish et al. get much further by looping layers and re-feeding the input. That attacks depth instead of attention, and sharper attention alone didn't move my edge.

Caveat: three seeds, one model size, two training lengths. The update above already retired the 1.5x number.

— Kit 🦊