After The Bubble - Snake II

Second part of the snake on a transformer explanation, you can find the first part here.

Snake, for real

So far we have a moving point on a grid, that’s not nothing, but snake it is not: we need for it to remember its past (to draw the whole tail), food position, and possible death, and the game logic must all live in the transformer, since the frontend will mostly only send direction orders.

This time the grid is 6x6, and the snake length goes from 3 to 8.

Protocol

Here the protocol is more complex than the point in a grid one, it has 4 different token roles: <Movement>, <SnakePosition>, <SnakeState>, <FoodPosition,SnakeLength>.

The conversation is composed of repeated blocks of those 4 parts, with the frontend adding each time the <Movement> part and sending the whole conversation to ollama, asking it to generate 3 tokens that will complete the block.

The game though can’t cold start: the transformer is stateless, and unable to produce random values on its own, so the frontend sends the starting state before the game begins, composed of 2 consecutive blocks:

<NOOP><P:13><OK><F:f0,3><NOOP><P:7><OK><F:f0,3>

We can read those bootstrap blocks as such: the first one has no movement, tells us the position of the snake tail, that the snake is OK, that food is at position f0 and the snake is long 3. The second one is similar, but shows us the position of the next part of the snake, the head.

Here f0 is a placeholder: the frontend randomly chooses an initial food cell, excluding the two occupied cells, and puts its number in both tokens.

There is a small shortcut in the bootstrap: the declared length is 3, but we only provide 2 body positions. The first normal move adds the third.

The model generates the 3 output tokens one after another: the new position is available when generating the state, and the state is available when generating food and length.

Snake State

The SnakeState token has 3 possible values: OK, when the snake survived without eating, DEAD, when it died with the latest movement, and ATE, when the snake just ate.

The SnakeState is fundamental to actually draw the whole snake, since the model reads its value and an AND gate network in the MLP is used to update the snake length: if SnakeState is ATE, the length increases by 1, up to the maximum of 8, otherwise it stays the same.

The frontend draws the body using the recent head positions, keeping as many as the returned length allows. The whole tail is therefore already in the conversation: there is no need for a token containing all its coordinates.

Something similar happens for the food position: if the last SnakeState is OK or DEAD, nothing happens, otherwise it updates its position. To do that, since it’s unable to produce random values on its own¹, the transformer uses a deterministic lookup table that uses the current token position in the conversation as key. This means that, given the same initial food and moves, the game is perfectly deterministic.

The food lookup table is stored in the position embeddings. At each state-token position, the embedding adds a one hot food candidate to a dedicated block of the residual stream. An AND gate checks that candidate together with ATE, and writes the new food position only when the snake has eaten.

There’s another issue: the food might appear under the snake body, and thus only becomes visible when the snake moves away. I don’t actually care too much.

Collision detection

This is likely the hardest part to build, but there is a cool shortcut I never thought of before: self-collisions during normal movement can only happen between positions separated by an even number of moves.

This becomes clear if you imagine the snake moving on a chessboard: each move gets the head to a different color, so if now the head is on a white cell, I know that 1, 3, 5, etc. moves ago it must have been on a black one. Since I know at the very least that odd moves ago the head (and now the body) was on a black cell, there is no need to check if that black cell is the same as the white one the head is in now.

This means that the network needs only to attend to half the body coordinates of the snake when checking if the head is going against something. With a maximum length of 8, the relevant offsets are 2, 4 and 6 moves into the past. Reversing into the snake’s neck is the case with offset 2.

There is one exception: the walls.

As in the point-on-a-grid model, the movement table clamps coordinates at the borders. Trying to move outside the grid leaves the head in its current cell, so we also compare the new head with the position from 1 move ago.

A match there means that we hit a wall. This kills the snake immediately, so the move that breaks the chessboard colour alternation never becomes part of a surviving game’s history.

We therefore need 4 attention heads for these checks, fetching positions from 1, 2, 4 and 6 moves ago.

Since each turn occupies 4 tokens, those are token offsets -4, -8, -16 and -24, measured from the new position token. Each head copies its cell into a separate residual block, and the MLP uses AND gates to compare it with the new head position.

The tail has already moved

Matching an old position is not enough to declare a collision: that position must still belong to the body.

For example, if the snake is only 3 segments long, returning to a cell from 6 moves ago must not kill it.

In this implementation, the collision check considers the previous L - 1 positions, where L is the current length. A match at offset k only counts when:

k <= L - 1

The length is also read through attention, and the collision neurons receive negative weights from the length dimensions that make the match irrelevant. If the snake is too short, those inputs prevent the neuron from firing.

Food detection is simpler: another group of AND gates checks whether the new head and the food occupy the same cell.

The unembedding gives the state outputs a priority: DEAD over ATE over OK. If the snake collides and reaches food at the same time, it still dies.

Attention across the turn

The model has 6 attention heads in total: the 4 for collisions, and 2 for the remaining information.

The first reads the previous head position while processing the movement token. This is the same mechanism as in the first post, followed by the MLP’s cell/movement table, now with 36 * 4 = 144 neurons.

That head is also used while processing the state token, this time to fetch the previous food and length. In both cases the required token is at offset -3.

The other head reads food and length at offset -2 while processing the new head position. Those values are needed to check for food and inhibit collisions with positions that no longer belong to the body.

Writing food and length together

Food and length are computed in separate residual blocks, even though they end up in the same token.

For each possible <F:c,l> token, the unembedding adds the score for food cell c to the score for length l. The correct pair gets both contributions, while every other pair misses at least one.

So the MLP can update the two values separately, without needing a neuron for every food/length combination.

A nice schema

graph TD
  M["Movement token"]
  MOVE["Attention: previous head"]
  STEP["MLP: cell + movement"]
  P["New head token"]
  READ["Attention: food, length, body positions"]
  CHECK["MLP: meal and collision checks"]
  STATE["State token: OK, ATE or DEAD"]
  OLD["Attention: previous food and length"]
  HASH["Position embedding: food candidate"]
  UPDATE["MLP: copy or update food and length"]
  F["Food and length token"]

  M --> MOVE
  MOVE --> STEP
  M --> STEP
  STEP --> P
  P --> READ
  READ --> CHECK
  P --> CHECK
  CHECK --> STATE
  STATE --> OLD
  OLD --> UPDATE
  STATE --> UPDATE
  HASH --> UPDATE
  UPDATE --> F

Conclusion

The whole thing still fits in a single transformer layer, with 6 attention heads and 409 MLP neurons.

There is one more limit: the context is 236 tokens. After the 8-token bootstrap, that leaves room for 57 turns.

So even if you avoid the walls and your own tail, eventually you die of context exhaustion. Or boredom: this is a really limited version of snake.

Next step

If by any chance you are still here, I’ll post in a few days/weeks the last post of this series on snake on a transformer, where I detail what I found out when I tried to train transformers to replicate the capabilities of this “handcrafted” one.