Part 1 warned that a longer row of positions is expensive, mostly because of attention. This post measures exactly how expensive, on the modern machine from Part 5, and then fixes it with the trick Mistral's models use: let each hidden state look back only over a fixed window of recent positions.

The code is in part7-sliding-window/, with a follow-along notebook for a Mac, minigpt_follow_along_7.ipynb.

The big picture, in plain English

Why a longer row costs so much

In attention, every hidden state matches its query against the key of every earlier position, and itself. So the number of matches grows much faster than the row.

Every one of those matches is a number the machine has to hold while it trains, in every head of every block. Double the row and you need about four times the memory for attention. On a machine with one fixed pool of memory, that is the wall you hit first.

A sliding window

The fix is to let each hidden state look back only over a fixed window, here the last 256 positions, however long the row is. Then each hidden state makes at most 256 matches, and the total grows only in step with the row: twice the row, twice the matches.

In principle, through the blocks. In block 1, a hidden state gathers information from up to 256 positions back. In block 2, it looks at hidden states that have already gathered from their own windows, so it could reach up to about 512 back, and so on. With 6 blocks, information could travel about 1,500 positions, one window per block, much as Part 1's four blocks let each hidden state reach further back than the one before. But could is not does: the secret-word test below found that my machine never learned to do it.

Writing the window as a mask saves nothing

The obvious way to build the window is the way the earlier-positions rule is built: a grid of "allowed" and "not allowed", here also blocking anything more than 256 positions back. It gives exactly the right behaviour. But it saves no memory at all: the machine still works out every match in the full row-by-row grid, and only then throws most of them away.

To actually save memory, the machine must never build the full grid. So the row is cut into chunks of 256 positions, and each chunk only looks at itself and the chunk before it. That covers every hidden state's window, and the biggest grid ever built is 256 by 512, however long the row.

The memory race

mem_sweep.py trains the 12.6-million-number modern machine for a few steps at each row length, three ways, and records the most memory each used:

• The mask column matches full attention at every length, exactly as the last section predicted.
• At short rows, everything is close. Most of the memory then goes on the MLPs and on scoring against 8,192 rows of lm_head, and both of those grow only in step with the row. Attention takes over at about 2,048 positions.
• At 4,096, full attention needs 35.5 GB, more than half the Mac, and 1.9 seconds a step. The chunked window needs 15.1 GB and 0.8 seconds. At 8,192, full attention would need more memory than the Mac has. The chunked window reaches 16,384 positions in 59.4 GB.

Does the windowed machine still learn?

A window is only useful if the machine still works. I trained the modern machine with a row of 1,024 positions, four times the length used in Parts 3 to 6, once with full attention and once with a 256-position window, for 1,500 steps each:

The windowed machine came out slightly ahead, but only by 0.008, about the size of the luck between two random starts (see Part 5), so the two are level. It also trained faster and in less memory. And its 0.6727 is inside the band of Part 5's 256-position machine (0.672 to 0.678, over three random starts): on these stories, a longer row did not help at all, because most stories are only a few hundred tokens long, so there is nothing further back worth seeing. TinyStories is the wrong text to show what a long row is for, so the next section builds a test that is not. The point here is the memory: on these stories, the windowed machine learns just as well, while its attention cost stays flat as the row grows.

A bird named Bob, a fish that helps him, and "they became good friends": the same shape of story as every machine in this series. The window cost it nothing you can see.

What is a long row for? The secret word

TinyStories cannot show what a long row is for, so I built a test that can. Each test row hides a secret near its start, such as " Lily's secret word is apple.", then carries on with ordinary stories, and ends by asking for it again: " Lily's secret word is". There are 8 names and 24 words, so a machine that cannot see the secret can only guess, right 1 time in 24. Half of each training batch was rows like these, with the secret anywhere from 20 to 1,000 pieces back; the other half was ordinary stories. I trained three machines on exactly this mix, 1,500 steps each, and then asked each one for the word from five distances, 400 times at each:

• All three learned the trick. Whenever the secret was inside what they could see, all three got it right every time.
• The 256-position row cannot see past its row. Beyond 256 pieces, the secret has fallen off the front, and it can only guess.
• Full attention can. With 1,024 positions, it found the word 970 pieces back every time. This is what a long row is for.
• The window did not reach past its window. It behaved exactly like the 256-position row. Through six blocks, the word could have hopped 256 positions at a time to the end, but in 1,500 steps of training the machine never learned to pass it along. Training it four times longer, 6,000 steps, only made it surer of the word close up (96% of the chance on the right word, up from 89%); beyond its window, it still guessed.

My reading of why: to pass the word along, some hidden state in the middle of an unrelated story would have to pick the secret up in one block and hold it for a later block to collect, and no part of that chain is rewarded until the whole chain works. Whatever the reason, the lesson is the same: the reach that the blocks allow in principle is not reach the machine has learned. It is one reason why many models that use windows, such as Google's Gemma 2, also keep some blocks with full attention.

The jargon decoder

The terms for the whole series are collected in one table, the series glossary.

The code, in the order it runs

The windowed attention lives in Part 5's model_llama.py, behind a window setting, and part7-sliding-window/ measures it.

The mask that saves nothing: sliding_window_mask

This builds the full T-by-T grid, and the attention call still works out every match before the mask removes most of them.

The chunks that do save: chunked_swa

The biggest grid is [batch, heads, chunks, w, 2w], which grows in step with the row. With a window as long as the row, it gives exactly the same numbers as full attention.

Choosing between them: Attention in model_llama.py

The secret word: secret_word.py

secret_word.py builds each secret row from Part 3's stories, mixes them half and half with ordinary stories, trains one machine, and then asks for the word from 100 to 970 pieces back. A 256-position machine sees only the last 256 pieces of each row:

run_secret.sh trains the three machines: --block-size 256, --block-size 1024, and --block-size 1024 --window 256, and then the window machine again with --iters 6000.

The race: mem_sweep.py

For each row length from 512 to 16,384, and each of the three ways, mem_sweep.py builds the machine, trains it for 5 steps on random tokens, and records the most memory used and the time per step. It skips full attention and the mask past 4,096 positions, where they would run the Mac out of memory.

The series

Seven parts, from a character-level GPT in a borrowed notebook to a modern small model in MLX:

1. Running it: a trained MiniGPT taken apart while it writes.
2. Growing it: training the same machine from random numbers.
3. Pieces, not letters: three tokenisers, scored fairly in bits per byte.
4. A faster engine: the same machine in Apple's MLX.
5. The modern block: Llama's four changes, one at a time. Only RoPE mattered.
6. Learning from a teacher: distillation, and why the most helpful teacher was not the best storyteller.
7. Reading further: sliding-window attention.

Every machine here is tiny, and none of them is good. That was the point. The tokeniser, the training loop, the engine, and the tricks, distillation, shared keys and values, and windowed attention, are all things you can build and run in an afternoon on one Mac. What separates them from the models I use every day is scale: more text, more numbers, and more computing, applied to substantially the same design.

The series glossary

Every idea in this series in plain words, next to the name the experts use, and the part whose jargon decoder it first appears in. Each term is explained in that part.

Try it yourself

• The follow-along notebook: part7-sliding-window/minigpt_follow_along_7.ipynb, for Jupyter on a Mac with Apple Silicon. It counts the matches, checks that the chunked window gives full attention's numbers when the window covers the whole row, runs a short memory race, and shows the secret-word results. It is saved with the outputs from my own run, so you can read every result on GitHub without a Mac.

• On the command line, after Part 3's prepare_data.py and tokenizers_setup.py:

MLX needs Apple Silicon.