414k 参数 Transformer 学会 Boids 群飞，线性探针读出的规则是否真被使用？

I wrote a small boid simulator (12 birds), recorded it flying, 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).
• for comparison, a small network that only gets a bird's own move plus the flock's average move (no attention, no positions) gets R² 0.40. one seed, same held-out clips.
• the three rules were easy to read out of its hidden states with a linear probe. cohesion (stay close to your neighbours) could be read before any training (probe R² about 0.85).
• alignment (line up with your neighbours) was readable after the first attention layer, but if I make the second layer attend only to itself, it is gone in all 4 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 to 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