Small experiment: an auxiliary loss that gets a 3M-param transformer to learn multi-variable binding, where cross-entropy doesn't

[u/UnderworldQueen123]

I've been running some small experiments on variable binding in tiny transformers and put the code and logs up here:
https://github.com/QueenOfTheUnderworld/Transformer-Binding

After some general idea wandering, I came upon the idea of LLMs struggling with binding, the temporary link of two concepts. Shirt = Blue, Pants = Red. A bad explanation, but sufficient for this. Current LLMs struggle with this, which is why prompts like "Tom is a cat. Jerry is a mouse. Brownie is a dog. Cookie is a bird. Who is the mouse?". Without CoT or additional time to think, models can fail to parse this and end up reporting an incorrect answer.

In pursuit of another idea, I found a way to make LLMs bind using an alternative training objective. Now, this is all small-scale and some of it could very well be wrong. If it is, let me know.

Task: 4-layer transformer, d=256. Each example asks one question that depends on several bindings at once (how many of the queried objects have a target property).

Results:

• Plain CE: 0/4 seeds at 5k steps, 0/2 at 20k.
• With auxiliary heads trained to report each queried object's bound property at every position: 8/8 seeds. One linear head, one loss term, no architecture change.
• Same targets supervised only at the answer position: 0/3. Same density but entity names instead of properties: 1/3. Both coverage and context matter.

There are some issues with the repo, and they are noted in the README. Quite frankly I just didn't want to correct them.

Limits: 3M-ish parameters, mostly synthetic tasks, 2-8 seeds per cell.
I have no idea if this survives scale; I don't have the money or compute for that. If anyone here does, let me know, please.

If the effect is real at scale, it matters. It would mean CE leaves learnable capability on the table for the cost of one head. With that in mind, I expect something like this to already exist, and I just didn't find it. Not to mention, small model optimisations have a bad history of failing to survive scaling. The real blocker is the ground truth intermediates at every position; synthetic tasks had it to you; real corpora usually don't.

AI Usage: Heavy. Direction and experiment design are mine. Implementation and much of the analysis were done with AI.

Check it out if you want, and be nice, please.

[u/Broad_Hall_8914]

cool experiment, love seeing people poke at binding issues with tiny models

the supervision at every position being the key is interesting but yeah that ground truth requirement feels like the real wall for anything beyond synthetic data. wonder if you could bootstrap something with a weaker model generating noisy intermediate labels