How a Transformer Plays Tic-Tac-Toe

The goal is to get three in a row before your moves fade. Only the last 6 moves stay on the board, older ones vanish. Moves about to disappear turn light gray, the latest played move is green.

Each square is numbered 0 to 8, so a game can be encoded as a sequence of moves:X5 O6 X4 O3 X0 O8 X2 O1 X6 O4 X7 …Because old moves disappear, the game can keep going for as long as it takes for someone to win. It also makes move order matter. In standard Tic-Tac-Toe only position counts. Here, when a move was played matters too.

The Transformer receives this sequence and predicts the next move. Six models are available. Two are complete architectures, four are ablations with one component removed:

• 1-head — single-head attention, 32-dimensional.
• 2-head — two-head attention, 16 dimensions each.
• no-pos — positional encoding removed. The model sees tokens but not their order.
• no-causal — causal mask removed. The model can attend to future positions during training.
• no-MLP — MLP block removed. Attention is the only transformation.
• no-res — residual stream removed. Each block overwrites rather than accumulates.

The Transformer consists of six main blocks: tokenizer, embedding, positional encoding, attention, MLP, and unembedding.

It processes a sequence of moves in parallel — one vector per token. Each vector flows through the blocks via the residual stream, which carries information forward so that every block adds to the representation instead of overwriting it. At the end, each vector produces a prediction. Only the last one is used to choose the next move (the others are used for training). It gets appended to the sequence, and the process repeats.

Attention is the only block where tokens interact with each other. In every other block they are processed in total isolation.

Blocks are duplicated to show they operate independently on each token. Only attention uses all tokens at once.

• Tokenizer — Converts moves into integers.
• Embedding — Maps each integer to a 32-dimensional vector.
• Positional encoding — Marks when each move was played in the sequence.
• Attention — The communication layer where tokens read information from each other.
• MLP — A reasoning step that refines the representation of each move.
• Unembedding — Turns each vector into scores over possible next moves.

The model works with vectors stored in a matrix, so it needs a numerical index to fetch them. The tokenizer maps each move to that integer index.

The 20-token vocabulary includes 9 squares for X and 9 squares for O plus special tokens START and PAD

The embedding layer maps each token to a point in a 32-dimensional vector space the model can do math with.

These vectors are learned. They shift during training until the model finds a useful arrangement.

The math in the attention block is order-agnostic. Without help, the model cannot distinguish between sequences like X1 O2 and O2 X1.

Positional encoding fixes this by adding a unique vector to each move based on its position in the sequence. This "stamps" every token with its timing to keep track of move order.

The model uses learned positional embeddings (as in GPT-2). These parameters shift during training to represent the game's timeline.

Attention is the communication layer. Each move looks back at the sequence to gather context about the game state. The entire exchange is described by a single formula:

This formula says: compare every move to every other move (QKᵀ), restrict to the past (M), turn scores into weights (softmax), and use them to mix information (V).

Before computing attention, the vectors are normalized so all values are on a comparable scale for stable training.

Note: Layer normalization is not necessary for this small model, but it is shown here as it is a standard part of the Transformer architecture.

The normalized vectors are transformed into three learned projections: Queries (Q), Keys (K), and Values (V). Because all three come from the same sequence, this is called self-attention.

Note: Matrix multiplications visualised as a row–column diagram. Rows are separated to show each token is processed independently.

Multiplying Q and K^⊤ is the first time tokens communicate. It produces a score for every pair of moves, measuring how relevant each move is to every other.

A causal mask (M) sets future positions to -∞ so the model can only look at the past. Scores are divided by $\sqrt{d}$ for stability, where d = 32 is the attention head dimension. Row-wise softmax then turns each row into attention weights that sum to 1.

The model uses these weights to pick information from each move's Value (V). This is the final exchange of information. Weighted rows of V are combined into context vectors summarizing the game state.

The context vectors are added back to the residual stream, updating each token with information from the sequence.

Multi-head attention runs several attention operations in parallel, each on a lower-dimensional slice of the residual stream. The 2-head model uses two heads of 16 dimensions each. The benefit is not capacity but diversity, as each head can specialise on different patterns.

Both heads compute Q, K, and V independently from the same input. The total Q, K, V capacity remains the same as single-head attention, just split equally between the two heads.

Each head computes its own score matrix. The two are entirely independent and can weight the sequence differently.

The causal mask and softmax are applied per head. Because the space is split, the head dimension is now d = 16.

Each set of attention weights is multiplied by its own V to produce 16-dimensional context vectors.

Since each output is only 16-dimensional, the two are concatenated and mixed through a learned projection back into 32 dimensions. The result is then added to the residual stream. This projection is the only addition of learned weights compared to single-head attention.

While attention moves information between tokens, the MLP (Multi-Layer Perceptron) processes each move individually. It acts as a reasoning step that updates the representation of each move.

The MLP also begins with normalization, using its own independent learned parameters.

The first linear layer expands the vector to 128 dimensions. A GELU activation provides the nonlinearity needed to recognize strategic patterns like winning lines, blocks, and traps.

The second linear layer projects the 128 dimensions back down to 32. This compresses the reasoning so it fits back into the residual stream.

The final result is added back to the stream. This update is much larger than the values in the residual stream and reshapes how the model sees each move.

The architecture can stack more attention-MLP layers to handle harder tasks. For Tic-Tac-Toe, one layer is enough to produce an accurate prediction.

Unembedding is the final step. The model projects every vector in the residual stream into scores (logits) for the 20 possible tokens in the vocabulary.

These logits are then converted into a probability distribution via row-wise, temperature-scaled softmax.

The temperature controls the sharpness of the distribution. Low temperature concentrates probability on the top choice. High temperature flattens it, making weaker moves more likely.

To predict the next move, the model takes the last row of probabilities and builds a cumulative distribution. It then draws a random horizontal line and selects the move where it crosses the staircase.

The selected token is appended to the sequence, and the process repeats for the next turn.

Attention

Every model learns a different attention pattern.

• 1-head — learns to ignore moves that have faded from the board.
• 2-head — does the same, but splits attention across two heads: one tracks the opponent's moves, the other its own.
• no-res — unlike the 1-head and 2-head models, the diagonal is activated because there is no residual stream to carry the current move forward.
• no-pos — fails to identify which moves have disappeared due to the lack of a timeline.
• no-causal — cheats by looking at the next move to predict it, so it has no strategy for the final token when the next move is unknown.
• no-MLP — tracks player turns and basic lines but lacks the logic to calculate blocks or winning moves.

Order

Without order, the model cannot tell which X move is about to fade. It plays O5 in both games regardless of the order. The 1-head model blocks correctly on O4 in the second game.

Note: in the no-pos model, the attention scores for each move in the last row are identical regardless of its position in the sequence.

Swapping the last move does change the output, since its attention scores are what drive the prediction. No-pos now plays O1 in both cases instead of O5, but still misses the block on O4 in the second game.

Temperature

Fading Tic-Tac-Toe is deterministic. X has a forced winning strategy (try playing as O against the 1-head or 2-head model in argmax mode). Most positions have one correct move, so temperature near 0 works best. A small amount of randomness helps vary the opening, where several moves are equally good. Around temperature 1 the model starts blundering, and by temperature 2 it may predict illegal moves like placing on an occupied square or playing twice for the same player.

Data

In practice, training a neural network to play Tic-Tac-Toe would be overkill. The game is small enough to solve with the minimax algorithm. An adaptation of this approach is used to simulate 50000 fading Tic-Tac-Toe games, with 20% of moves chosen at random so the model sees weaker lines and learns to respond when the opponent makes mistakes.

START marks the beginning of every game, providing context for the first move. PAD is used for training only. It fills the remaining slots to a fixed length of 17 so the model can be trained in batches.

Model

The 1-head and 2-head models use a single attention-MLP layer. The four ablations are created from the 1-head model by removing one component at a time.

Width D = 32 provides the best balance between performance and visualisation. D = 16 is too narrow to learn reliable play. D = 64 produces a stronger model but the visualisations become cumbersome.

The context length is 16. In fading Tic-Tac-Toe only 6 moves can be on the board at once, so a context of 6 would suffice. The longer window lets the model learn to ignore fading moves, one of the key behaviours explored in the Analysis section.

If a game runs longer than 16 moves, the sequence is truncated from the left. The START token is always kept, followed by the next X move to maintain a consistent sequence structure. This does not impact the game since moves older than 6 turns have already faded from the board.

Training

Both models were trained for 30000 steps with a batch size of 512 using AdamW. Play uses only the last prediction, but training uses every row. This provides 16 training signals per game.

The loss is cross-entropy, calculated for all positions in parallel. This makes training much faster than sequence models like RNNs or LSTMs, which process tokens one by one.

For scale, here is how the 1-head model compares to larger models.

1-head model compared to nanoGPT, GPT-2, and Gemma 2.

The full training code is available in the Jupyter notebook.

The Transformer consists of six blocks. The tokenizer, embedding, and positional encoding prepare the vectors. Attention and the MLP process the information. Unembedding produces the final prediction. A residual stream ties these components together.

The same mechanism is used in much larger models. Large language models predict the next word from a vast vocabulary using the exact same matrix operations. Every prediction is the result of these blocks working in sequence.

The following resources were helpful while making this article:

• Architecture diagram inspired by the Gemma visualisation in this Welch Labs video.
• Transformer implementation based on nanoGPT and the Let's build GPT video by Andrej Karpathy.
• Tic-Tac-Toe Minimax algorithm explained step by step in this Coding Train video and covered in depth in Artificial Intelligence: A Modern Approach by Stuart Russell and Peter Norvig.