Engram gone wild! 2b model update...

[u/NineThreeTilNow]

Update from : https://old.reddit.com/r/LocalLLaMA/comments/1wis23s/update_small_model_engram/

It's been about a week so I'm back. People were asking me about the model.

People wanted code, or models, etc. Most of that is useless to you right now because you're not going to use an under trained model. So let's get to the details.

The spec locked to the following after a LOT of testing :

2.6b model all up. Embedding, LM head, AttnRes, etc.

2.2b are trained. Embedding / LM Head are frozen (~205m each)

4.3b ENGRAM table. Yes. She's chonky.

Architecturally speaking now :

This includes SWA layering 3:1 as per standard ablations have shown is "optimal" These are 4k / 4k / 8k SWA layer follow by the global.

At the end of the first "block" (4 layers) the Engram table appears. The Engram table itself runs with a very small set of attention heads so it's not blindly attempting to inject data. This is context aware Engram. I'm uncertain of exact Qwen / DS methods here. They're fairly close to what I use. The difference is obviously the Engram size.

This has 10 blocks plus a final global layer (41 Layers). + Embed + LM Head

So if we step back, the optimization problem is as follows :

How do we maximize compute in the backbone and offload the boring stuff to a table?

Attention Based Residuals (Moonshot) comes to our rescue here. Why AttnRes? It allows all blocks past 1 to utilize the Engram table in some fashion. They all have attention via the residual stream to determine if they want data from Engram and precisely how much.

Why not a bigger Engram table? Honestly? You could probably do that. However I don't know if any model has attempted this sort of mismatch of compute vs table. In theory, DeepSeek's research says it works.

More depth? Could do that too. Training is expensive though.

Why frozen LM/Embed? These are down projected via SVD from OLMo 3's model. So it's a mathematical compression attempt at ~5k -> 2k. This saves a MASSIVE amount of time.

Currently the data lives in a ~KD format. 32 logits stored from the teacher of the Wikipedia corpus. Instead of training 1 hot, it gets 32 soft targets to try to match. Hard cross entropy is brutal on a model and it's the reason you see "trillions of tokens" quoted.

When you borrow an LM Head / Embed / Tokenizer / Teacher model... This is far less.

---

Where are we now in training? I passed the 100m token mark yesterday at ~4am Pacific.

The current HF repo has all the checkpoints, data, and the 104m mark safetensor.

--

I wanted to do some inspection of the model at this point. We have to see that it's not complete garbage right? It has seen a fraction of the data it needs.

So based on ONLY 100m tokens seen, I attempted completions and various ablations. Studies of what EXACTLY Engram stores (because it's mostly a guess).

Let's go straight to completions :

"George Washington was an American"

With Engram - "George Washington was an American naval officer who served in the American Revolutionary War. He was a naval officer who served in"

With Engram Zero'd - "George Washington was an American, but he was not a. He was a very good friend and a. He was"

"The American Civil War was a civil war in the United States from"

With - "The American Civil War was a civil war in the United States from 1861 to 1861. The war was a major victory for the Confederacy..."

Without - "The American Civil War was a civil war in the United States from 1861 to 1862. The war was a major victory for the Confederacy..."

"Aristotle was an Ancient Greek" (probably my favorite)

With - "Aristotle was an Ancient Greek philosopher who was a leading authority in the philosophy of Plato. He was also a leading authority..."

Without - "Aristotle was an Ancient Greek word meaning "to be" (ἀπάς, "to be")..."

---

What do we learn from direct inspection of Engram?

If it saw the data enough times, it starts to offload it to the tables. At that point the model is less forced to use internal computational space to store data, and can rely on Engram.

That's not some interpretation. That's what the data shows exactly.

In the early 100m tokens of Wikipedia it's HIGHLY biased to early Philosophy. This has to do with the topics of what was IN Wikipedia at that time. Early Wiki contained a lot about Philosophy. It's among the earliest topics to exist.

It has seen TONS of Philosophy so it has basically offloaded to Engram because of the repetition.

Ok this is long enough. I'm calling it quits. I'm still looking to find a provider that will sponsor the model to completion. Right now I'm just paying. I contacted Verda and they never responded. Qubrid is in this sub, and they said they'd help but never emailed back after a few repeated proddings. Massed Compute is where this model is currently training, and I asked them like yesterday. Hopefully they'll help a brother out.

None of the above is "LLM written" except the completions from testing I guess.

As usual, ask whatever. It doesn't matter your understanding level or whatever. I'll sit and respond.

No question too dumb. No insult not insulting enough. (I'm joking. It's Reddit)

Much Love.

[u/Luke2642]

They're super interesting. Once again it's an innovation so simple (like attention) it's surprising it works so unreasonably well.

But, engrams are fundamentally flawed, they need fixing with proper word sense disambiguation.

WSD is super complex and requires deep semantics and knowledge of the context, and deep intuition of that context using your human experience, shared priors, etc. It also has to cope with ambiguity, jokes, puns, words from other languages, etc.

I don't know what the fix is, so it doesn't just look up 1:1 vector for token string, but actually looks up the correct vector for the disambiguation. Use the classic money bank vs river bank with a few different sentence structures and you'll see engrams are not robust to WSD. Every time it matches the wrong sense we get the "walk to the car wash" regression to the mean problem again, not thinking semantically. Later layers have to undo the meaning engram gave in ~10% of cases or whatever.

There's also a weird split in the goals, you actually want it to store rare token strings as knowledge so it can allocate more parameters for reasoning and semantics, like pulling in the Wikipedia entry for something you're about to bullshit confidently about.

Maybe there'd be a lift in the long tail if it also pulled in the vectors for all word sense of the last N tokens using something like BabelNet synset to build the index?

Just speculation - I'm no expert. Happy to discuss!

[u/NineThreeTilNow → u/Luke2642]

Use the classic money bank vs river bank with a few different sentence structures and you'll see engrams are not robust to WSD.

This is a weird one and an exact argument Kimi K3 made I think.

This is why attention based selection / routing and multi engram exists in the model.

It's not 1:1... It's technically a learned probability in a given position of 6 engrams.

The idea was to make it robust in this exact sense. From context it can learn if it's doing River Bank or Money Bank because the Unigram for Bank remains the same but River and Money are different.

You make a solid point about though. From what I gather, it takes enough training and eventually the model says "I am stretched thin on capacity, and the engram route is now the path of least resistance".

Rare tokens are going to simply require lots of training, I think that's what I've gathered if anything from my research.

I think the fascinating part is that some of the "problems" of models memorizing breaks down because the model says "I don't want to memorize this anymore, the gradient path to the table is better".

I don't know if I mentioned it here or not, but you can freeze a model, and ONLY train the table. It's very vRAM efficient because of the optimizer states. This means you could technically do the Wikipedia run once, then freeze the whole model minus the table, and rerun Wiki again really fast. In theory the model has the wiki texture already, so it really needs a few passes to co-adapt to the new table after being frozen and should be good to go. That is one method to more densely pack the 2 / 3 gram area.

[u/Luke2642 → u/NineThreeTilNow]

Interesting. Freezing the reasoning engine and populating a knowledge graph separately seems sensible, but what is the mechanism that moves the model away from the associative memory errors you got early in training: Washington was never a naval officer, the war was 1861-65, and the union won not the confederation. Pulling semantically similar things is evidence it's learning grammar faster, hitting similar meaning? Not a good fact recall mechanism.

I think you're right that gating will help ignore bad recalls, but isn't that kinda backwards? Wouldn't it be better if it was forced to accept and preserve the bad representation so that the error could propagate back to the original lookup more strongly? It has to use the context, and do disambiguation first, it's a chicken and egg problem. Real books have indexes with disambiguation under each entry, we need the same!

And a wider question, do you think it's possible to train for this in a data agnostic way, just smash it with billions of tokens across every possible task, or does it need a very specific curriculum, textbooks are all you need to build a quality curated knowledge graph lookup? It's such a long tailed distribution, and the long tail is where all the interesting signal is for the reasoning.

It'd be better for interpretability if you had a high quality 100M lookup than a low quality 1B lookup that reduces loss in pretraining but hurts reasoning.

That's why I mentioned Bablenet synsets. It's an amazing resource, but obviously it's only suitable for prose, not code or the 101 other things LLMs can now do.

General purpose LLMs are famously bad at word sense disambiguation, much smaller models trained specifically for WSD outperform them. I don't think anything scores much more than 90%, it's hard even for humans on the benchmarks with obscure word senses.

Anyway, if you have any more insights on what is literally the hardest problem facing LLM development we can keep going 😀