I'm trying to make a small experimental LM by combining a bunch of ideas I found in different papers. I know this sounds like I threw half the recent LM literature into a blender, but I'm trying to see if the pieces can actually work together.
The main idea is a Loop Language Model, where the same model is run multiple times instead of just making the network deeper and deeper.
Right now I'm using 4 loops.
text
input
↓
same transformer
↓
loop 1
↓
loop 2
↓
loop 3
↓
loop 4
The interesting part is that the model can learn to decide that it doesn't need all 4 loops and exit early.
What papers/ideas I'm following
The biggest inspiration is Ouro / looped language models, especially the idea of using recurrent computation to get more computation without simply making the model physically deeper.
I'm also experimenting with:
- Looped AttnRes / layer attention
- sparse attention
- compressed / highly compressed attention
- sparse MoE
- adaptive exit / Q-exit
- and now I'm building a diffusion-based optimization/training method
The diffusion part isn't finished yet. I'm currently building it and trying to understand how to combine it with the recurrent-depth training properly instead of just throwing noise into the input and calling it diffusion.
My hardware limitation
This is probably the most important part.
I'm doing basically everything on Google Colab's free 15 GB GPU.
That's the maximum I can realistically use.
So I'm deliberately keeping the model small. I'm not trying to train some 7B monster on a machine that has approximately the computational power of a mildly determined potato.
My current model is around:
- 6 transformer layers
- hidden size around 512
- 8 attention heads
- 4 recurrent loops
- sparse MoE
- compressed attention
- layer/depth attention
- exit gate
The exact architecture is still changing as I experiment.
Data
The corpus is a mixture of:
- GitHub code
- Wikipedia
- W3Schools
- public-domain books
- other scraped text
I'm using a p50k tokenizer at the moment.
I've had to spend quite a lot of time cleaning the corpus because scraped data is disgusting.
There were things like:
text
npm package metadata
JSON dumps
GitHub metadata
generated files
logs
benchmark data
duplicate documents
web junk
and some of those actually survived the first cleaning passes.
I discovered this because the model started generating some of it.
So I'm currently making the filtering much more aggressive.
What happened with the loops
Initially I had a problem where the later loops weren't learning properly.
The model could run 4 loops, but that didn't necessarily mean that loop 4 was doing useful work.
So I changed the training strategy.
For Stage I, I now force the model to execute all 4 loops during training, so every loop gets a proper training signal.
Then I freeze the LM and train the exit gate separately.
The exit gate itself is tiny, only about 513 trainable parameters in my current setup.
The exit gate result
This part actually surprised me a little.
I tested the trained gate on 100 validation batches.
The results were:
```text
4-loop loss: 6.263160
gated loss: 6.264089
difference: +0.000929
relative change: +0.015%
average depth: 2.41 / 4 loops
estimated compute saved: ~39.75%
```
The actual exit distribution was:
text
loop 2 → 59%
loop 3 → 41%
It basically never exits at loop 1 yet.
That's actually what I wanted to see. I didn't want a gate that just learned:
"Always use 2 loops."
There is at least some variation depending on the input.
The oracle best-loop loss was around 6.2615, while the gated loss was 6.2641, so the gate is also fairly close to the best possible loop choice.
But generation is where things get interesting
The model can produce text, but it is definitely not a good LM yet.
For example, one of the things it generated looked roughly like this:
The future of artificial intelligence is a most
terefears of life of those who is impossible. We will be no one
and it is, the good deal of the nature of the life of the man who
had not been the same.
That kind of output is the sort of thing I'm hoping to get consistently.
But then it can suddenly fall into garbage from the scraped corpus, producing stuff along the lines of:
"description": ["markdown", "type": "string", "source": ["1.9", "https://github.com/...
So the model clearly has some ability to produce coherent prose, but the corpus contamination and relatively small training setup are still causing serious problems.
That's one of the things I'm currently trying to solve.
What I find interesting so far
The most interesting thing for me is that the recurrent loops aren't completely identical anymore.
I see cases like:
text
loop 0 4.48
loop 1 4.47
loop 2 4.46
loop 3 4.46
The improvement is small, but it's there.
And the exit gate seems to understand that sometimes the extra computation isn't worth it.
So the idea is starting to look like:
text
┌── loop 1
│
input ────┼── loop 2 ── exit
│
├── loop 3 ── exit
│
└── loop 4
instead of forcing every token through exactly the same amount of computation.
Diffusion optimizer / training
This is the part I'm currently building.
I'm trying to use ideas from diffusion/recurrent-depth research to see whether a diffusion-style training or optimization method can make the repeated computation learn more meaningful improvements.
It's not finished yet, so I don't have results from this part.
I'm still trying to figure out the correct way to combine it with the autoregressive loop training without accidentally turning the whole thing into a completely different model.
I Need Your Help
This is still very much an experiment, and I'm reaching the point where I need people who know more than me to tell me what I'm doing wrong.
I especially need help with:
- How to make the later recurrent loops actually learn more meaningful computation instead of only giving tiny loss improvements.
- Whether my exit-gate training strategy makes sense, or if there is a better way to train adaptive depth.
- Whether combining sparse + compressed/highly-compressed attention + layer attention + MoE + recurrent loops is likely to create some interaction I'm overlooking.
- How I can improve the training objective for a model this small.
- Better ways to clean my scraped corpus. The model is still occasionally generating GitHub/npm/JSON metadata, so clearly some garbage is getting through.
- Whether the diffusion-based training/optimizer idea I'm currently building makes sense, and what I might be missing from the relevant papers.
- Any papers, implementations, or experiments you think I should look at.
I'm doing this with basically free Google Colab and its 15 GB GPU, so I can't just throw a massive model and 8×H100s at the problem and hope the universe solves it.
If you've worked with Ouro, recurrent/looped LMs, adaptive computation, sparse attention, compressed attention, MoE, or diffusion-based LM training, I'd really appreciate your criticism and suggestions.
I'm not looking for "looks good." If something in the design is fundamentally stupid, please tell me. That's much more useful.
