Research notes

October 8, 2026

I was helping with one paper of a collegue and working on one project for SCIO (company which does a lot of educational projects in Czech Republic) on a project with LLM tutor. So I did not have much time to work on other projects. However, I have tried all the 4 methods with the two-layer transformer and besides diffusion model, which sucks, all were able to reach 7 digit number multiplciation. Even the version which did not do any recurrence nor memory. However, the variance in the training runs is crazy. Some runs finished without learning 4 digit numbers, wheras some went to 8 digits and were still improving (even with the same configuration). I do not like that diffusion models suck, because I like the concept of them, but I guess with these token based tasks, they are just too hard to learn. So this is similar to the langauge diffusion I have tried. But I will try to make it work, because it seems that Diffusion is a good fit for multiplication.

Then I tried one-layer transformer and here the difference became much mor clear. Both the recurent and memory model were able to learn up to 5 digit multiplication, while the one-layer transformer failed beyond 2 digit number. I used 4 repetitions, so this could be why we cannot go beyond 5. I will try 6 and 8 repetition runs now.

September 30, 2026

I ran the profiling, one of the things that I found out is that linen.Embbed is the bottleneck in small transformers. I had not found any jax issue about this, so it surprised me. It was weird, because it bottlenecks the backward pass and even weirder is that increasing the batch size increased the time, but super-linearly. That was really surprising to me, and at one point, when the batch was big enough it was only linear increase. I replaced the Embbed with custom layer, which is just one-hot encoding with matrix multiplication. This scales much better in these small models, it is sometimes even 8-times faster.

Because of that I ran profiling job for all the implemented techniques. I would like to see whether there are not some other problems like this. It is much better to find them sooner then later. I will hopefully run the experiments that compare the 4 techniques (or even their combination) over night.

September 28, 2026

I finally had some time to work a bit more on this. From the prevous results, I think that just naively learning multiplication will fail. And I think it makes sense, multiplication is O(n log n) operations, so it seems naive to somehow model it efficiently with just forward passing through the network. I know that the network does matrix multplication which is O(N^3) (or O(N^2.3717) with Coppersmith-Winograd), but still, I do not think it can hide behind that. Adding is a linear operations and you can do it with linear passes of single-layer transformer. I am now trying to find a setup where the multplication is learnt without additional supervision. I just don’t like when someone gives the model the answer on how to reason and then they feel like it is smart. In other words, I am a full believer in Bitter Lesson. I understand that it is eonugh for tech companies, because in the large models you do not care about the path to the product, you care about the product. That is the reason why RLHF was so good, even if I think it is lazy research idea. Rant over, for the learning without supervision I have these 4 ideas:

1. Give the model CoT, but do not supervise it, in other words, let the model output 40 tokens and supervise only the ones at that are answer. I think you could use and as tokens and supervise only the part after . I do not like this approach, because it requires you to generate data with different lengths of thinking (such that the model can generalize and does not learn fixed thinking length without any effect).
2. Use Repeated transformer, which instead of passing through the attention once, it passes several times. So the transformer itself learns to make the computation inside the latent representation. I like this approach, because conceptually it is like keeping intermediate results in the memory and workig on them. The problem is that each token would start the memory anew, which seems off
3. Use transformer-RNN hybrid -> You pass through the transformer several times, but you create both memory and new output. This would allow to use the memory for future tokens (as you can just pass the memory onward). I think this is the most computationally demanding approach, so I would assume it would work the best. I mean, if this does not work, the others also won’t, but if it does, there is hope for other.
4. Diffusion Transformer. I like diffusion, so this is an obvious choice for me. I will definitely use it as one of the baselines.

I tried a different approach to experiments now. I just told Claude what it can do on the cluster and I let it experiment freely. I started with 3. and it seems tha it works. I still need to check what exactly was going on, because Claude ran a lot of things, and in usual Claude fashion left a report that is unnecessarily verbose. From the results, it seems that 3. is capable of learning 6 digit multiplication, which is great and better than I’d expect. I will work on this a bit more in the following days.

August 24, 2026

I have trained a transformer for multiplication. The final goal was multiplying 10-digit number with 1-digit one. The goal of this was just to verify whether the transformer can do at least this simple task before moving to more digits as a second operand. I have used following schedule in format (first operand digits, second operand digts, grad_steps): (2, 1, 2000), (4, 1, 6000), (6, 1, 12000), (8, 1, 30000), (10, 1, 50000). I used big model, because this is still just checking out whether the transformer can or cannot do it, so it was exactly 120 000 parameters. Unsurprisingly it takes forever, but at the end, the model gets there and it can do it. So for now I ran another experiment with schedule (2, 1, 2000), (2, 2, 6000), (4, 2, 12000), (4, 3, 30000), (6, 3, 50000).

I am quite interested in these results. For the 1-digit multiplication I have used single-layer transformer with 4 heads. I wonder whether it is even possible to learn efficient multiplication in a single layer, because the naive multiplication approach often requires n+1 steps (where n is the size of the larger operand). Also some papers I have seen argue that without the CoT, the transformers just cannot learn multiplication. We will see about this!

The run did not finish yet, but I am quite pessimistic of the results. When learning the first part of multiplying 2 and 1 digit numbers, the loss goes down to 10^-5, which is unsurprising, it is easy task, when you go to the 2 and 2 digit numbers, the loss does not seem to go below 0.05. Maybe it just needs more training, but I do not know how much. I have tried to just learn the 2,2 with the same architecture. Let’s see how it goes, if it does not work, I’ll try increasing the depth.

If you run training for purely 2-digits (still depth 1), it takes roughly 100k gradient steps to get to 100% accuracy and it is unstable as hell (between checkpoints it sometimes drops below 90%). That is not what I expected. When you increase the depth of the transformer to 2, it converges in roughly 4000 steps and it is more stable, but still not the most stable thing in the world. Moving to transformer with depth 3 for multiplying 3-digit numbers, it again works, but the training slows down significantly. Even after 100k steps, you are not able to get clear 100% accuracy. Mostly it sits at around 98%.

I am thinking that maybe the weight decay from AdamW could be an issue in these kinds of tasks, because you may be at 100%, but it just pushes you out of it. Tried to remove the weight decay, and it is true that there are not these weird spikes, but the whole convergence is much slower.

When making the first operand of the 2-digit multiplication larger, the whole training seems to stop converging. It took quite a lot of iterations even with the 2,2, but when you go up to 10,2, the training becomes tediously slow, which I would take as failure.

August 22, 2026

I was playing with the addition again. I wanted to try some of the tricks used in the small models, whether they would generalize to the multiplication. So I tried to add RMS norm to the Q and K outputs, I tried to create multiple Q heads (but keep single K and V), I tried RoPE. I could not get convergence with any of those, when I decreased the size of the model. Could be unlucky, could be skill issue,

I tried to reduce dimensionality of the embedding from 7 to 6, which would reduce the parameters from 651 to 534, but even when trying couple of samples I could not get it to work. Nice thing is that based on the first part of the training, where you train on the instances with at most 3 digits, based on the loss, you can already tell whether the run will be successful or not. This saves quite a lot of time when running experiments, as it is roughly 2000 iterations. Well, maybe it is not that good of a predictor. I just had run where the 3 digits training worked well, then the 6 digit training just did not converge. I think that the embedding of size 6 cannot be done naively without other tricks I did not incorporate yet (like factorization). I decided to run a sweep through 15 seeds to see whether any of them would work and the best result only got accuracy 50%. A lot of the runs get 0%.

I think the Qwen style transformer (more Q-heads than KV-heads) may be a good idea to try in other projects, especially with the diffusion models.

August 21, 2026

It was interesting to play with the adding of two numbers, but I was ready to take on a bigger fish and that was multiplication. And I thought this is going to be straight-forward, just taking the same codebase and only change the training data. So I decided to use quite large model with ~50k parameters and I limited myself to 3-digit number multiplication. And after few training iterations, I already knew that this is a completely different thing.

Wheras in addition training, the model exhibits the usual grokking behavior, where model does not seem to learn at first, after some time it suddenly jumps in accuracy to nearly 100%, where it stays, or only partially degrades until it comes back to 100%. When you run the multiplication training, it resembles the usual deep learning behavior. The loss gradually decreases, the accuracy steadily rises. This sounds like a good thing at first, but unlike the clear distinction of the addition task, where we can say: The model has learned the algorithm”, it is hard to say the same thing in the multiplication task, because it could just be some overfitting of the model.

I have found this paper (https://arxiv.org/pdf/2510.00184), where the authors claim that the transformers are pretty bad with multiplication. I only skimmed the paper, but if I understand it correctly, the main claim is that the multiplication needs a long-range dependencies, which the difficult fo transformers to learn. They show a CoT approach where they let the model to do the partial sums. I would like to try that.

But besides that I have tried one other thing. That is training the transformer to multiply 7-digit numbers with 1-digit one. Good news is that the model can learn this simpler setting. The addition task starts with 2000 iterations training only on at most 3 digit numbers, then it has 5000 iterations on training of up to 6-digit numbers and only then it uses all 10 digits. Maybe the multiplication needs similar paradigm, but with even more data.

Well, those are the two things to try: 1. Add the intermediate results in the training and find out whether it helps the model to generate the result. 2. Play with a scheduling of the size of training instances to first learn on small ones, which will be gradually increased.

August 20, 2026

While ago I have read about this project of Yeb Havinga (https://github.com/yhavinga/gpt-acc-jax), where he managed to train transformer with 777 parameters to perform addition of 2 10-digit numbers with nearly 100% accuracy. I took this as a proof that the transformer hacan learn the addition algorithm by itself. Still, I wanted to see this myself, so I created a codebase (https://github.com/kubicon/tiny_models), where I reimplemented simple transformer and tried it myself.

When you start with a big model, everything converges quite nicely, but as you are decreasing the model, the whole training is becoming less and less stable. When I got to models at around 1000 parameters I could not get them to converge, no matter how hard I tried, even when I used the exact same hyperparameter and training setting as Havinga.

After a brief discussion with Claude, it identified that my codebase is slightly off because I train different embedding for decoder and encoder, wheras Havinga’s codebase reuses them. When I changed this, it reduced parameter count, but more importantly, it started working. This was surprising to me, becuase I would’ve thought that the model would be capable of learning different embeddings. However, I decided not to pursue this further.

I was toying a bit with this approach, I replaced layer norm with RMS norm and gelu with SwiGLU. I also avoided using the token. I am not sure why Havinga was using it, but it seemed useless to me. After few trials and errors I ended up with transformer that had 651 parameters. These are the final parameters

Layer	Parameters
Token embedding	13 x 7 = 91
Positional embeddings	34 x 7 = 238
RMSNorm pre-attention	7
Attention	7 x 21 = 147
Out projection	7 x 7 = 49
RMSNorm pre-MLP	7
MLP + GLU	7 x (2 x 5) = 70
MLP out	5 * 7 = 35
RMSNorm final	7
Total	651

I was happy about this. Then I found out leaderboards of this challenge (https://github.com/anadim/AdderBoard) and noticed that the record right now is 36 parameters. Well, I am quite far away from that, but I am still happy. The record-holder does few interesting tricks like using only a single matrix in attention, where Q = V and K = rot(Q). All RMSNorms are shared. Instead of the learned positional embeddings, they use rotational embedding.The actually shrink the embedding dimension to 3 and the MLP has just size 2. Well plenty of crazy things, but I decided not to pursue this.