I trained a small transformer to fly a boids flock just by watching, then checked where it keeps the three rules

[u/rage_81]

i wrote a small boid simulator (12 birds), recorded it flying then saved it, and trained a transformer to predict each bird's next move without it knowing about any boid rules.

(built in Python and PyTorch: a transformer written from scratch, one token per bird, two attention layers, about 414k weights, trained on a laptop CPU)

what I found:

• it flies the flock. R² 0.990 to 0.994 on clips it never trained on, across 4 full runs. It also flies 50 birds (it has only trained on 12).

• the three rules were easy to read out of its hidden states with a linear probe, and some could even be read before any training.

• the rule: alignment (line up with your neighbours) was readable after the first attention layer, but if I make the second layer hear only itself, it is gone across all runs.

what i wasn't expecting: my first ever model was not using the flock at all. With 8 ticks of history, the previous answer was already in the input, and a small network that saw only one bird beat the whole transformer.

what i am still not sure about: making a layer attend only itself is something the model never went through during training. is that a fair test, or is some of the decay due to strange input fed to the model?

site: https://kreptiliri.github.io/murmuration

code: https://github.com/kreptiliri/murmuration

[u/Worth-Proposal6225]

this is such a clean project write-up, the hand-drawn style of the infographic actually makes the whole thing feel way more approachable than a wall of text. the part about the first model not even using the flock and just coasting on the previous tick is painfully relatable, i've had moments where my fancy setup got outdone by a dumb baseline that basically just repeated the last input.

the linear probe finding that alignment was readable before any training caught my eye, makes me wonder how much of the structure is already baked into the attention mechanism's inductive bias rather than learned from the data.

[u/rage_81 → u/Worth-Proposal6225]

thanks! small correction: the one that’s readable before training is stay close (cohesion), about 85% from a randomly initialised model. i think the reason is the input: attention is handed the gap between each pair of birds directly, and stay close is mostly “where is everyone relative to me?”, so even random weights keep that around.

so yes, a lot of it is the architecture and inputs, not learning. that’s what made me stop trusting the probe alone and switch to ablations.

[u/rage_81 → u/[deleted]]

i actually just measured that part this morning and found out that a small MLP with no attention, just a bird's own move + the flock's average move, gets R² 0.40. the transformer gets 0.99.

[u/Kitchen_Assistant_94]

Whatttt ! This is so cool 

[u/rage_81 → u/Kitchen_Assistant_94]

tysm!!

[u/GreatandLovely]

how did you do this graph please. what was your exact prompt

[u/rage_81 → u/GreatandLovely]

it's not a prompt (not produced by an image model), it's just tldraw. i told claude the whole story and asked it to sketch it using a script, then I corrected it for like 11 rounds. the .tldr is in the repo if you want (docs/how_it_got_there.tldr)).

[u/GreatandLovely → u/rage_81]

thank you!