MiniGPT (Part 5)
Neil Haddley • October 5, 2026
The block inside today's small models: four changes Meta made for Llama (RMSNorm, rotary positions, SwiGLU, and shared keys and values), each turned back off in turn to find out which one actually matters, with a follow-along notebook
Every MiniGPT so far has used the block from 2017, the one in "Attention Is All You Need" and in Karpathy's nanoGPT: normalise, attention, normalise, MLP. Today's small models do not. Meta's Llama 3.2 1B and 3B, the models that ship on phones, use four changes to that block. This post makes all four changes to the machine from Part 4, still in MLX, still on Part 3's 8k pieces and stories, and then turns each change back off on its own, to find out which ones actually do the work.
The code is in part5-modern-block/, with a follow-along notebook for a Mac, minigpt_follow_along_5.ipynb.
| This post's machine | |
|---|---|
| What changed | the block |
| Text | TinyStories |
| Pieces | my 8,192 |
| Blocks | 6 of Llama's design: RMSNorm, a gated MLP, and 6 queries sharing 2 key and 2 values |
| Vector size | 384 |
| Positions | 256, with no position embeddings: queries and keys are turned instead |
| Engine | MLX |
| Size | 12.6 million numbers |
| Score | 0.6755 bits per byte, the average of three random starts |
The big picture, in plain English
Here is where the four changes sit, in Part 1's five steps. Only steps 2 and 3 change: step 2 loses its position embeddings, and step 3 gets a new block.
Scroll sideways, or tap the diagram to open it full sizeThe five steps before and after, and one block of each kind. The green numbers match the four changes below
And here is the new machine at work, the same way Part 1 showed goo, on the start of a sentence from a story:
Scroll sideways, or tap the diagram to open it full sizeThe new-block machine choosing the piece after "The sun was". The numbers come from a second training run of the same design, which scored 0.678 bits per byte
Change 1: a simpler normalise
Before attention, and again before the MLP, every block normalises the hidden states: it rescales each hidden state's numbers to a standard range. The 2017 way, LayerNorm, which Part 1 works through with real numbers, first subtracts the hidden state's average from every number, then divides by how spread out the numbers are. The new way, RMSNorm, skips the subtraction and just divides by the numbers' typical size. It also keeps only the stretch parameters from Part 1, not the shift. That is one calculation fewer every time a hidden state is normalised, and fewer parameters to train.
Change 2: turning queries and keys instead of position embeddings
Since Part 1, the machine has known where each token sits from its position embedding: a fixed vector per position, added to the token embedding. The new way, rotary position embeddings or RoPE, has no position embeddings at all. Instead, inside attention, it turns each query and key by an angle that grows with the token's position.
Picture a clock hand. A token in position 1 has its queries and keys turned a little, position 2 a little more, and so on. When a query is matched against a key, what matters is the difference between their turns, so the match depends only on how far apart the two tokens are, not on where they are in the text. "The token just before me" looks the same at position 5 as at position 205.
Scroll sideways, or tap the diagram to open it full sizeThe pencil exercise, drawn as clock hands. Both pairs are two positions apart, so both are turned 20 degrees apart, however far along the row they are
The notebook turned a real query and key with MLX's RoPE. Pairs with the same gap get the same match, wherever they are in the row
Change 3: an MLP with a gate
The 2017 MLP widens each hidden state to four times its size with one table of weights, bends it, and narrows it back with a second: two tables of weights. The new one, SwiGLU, makes two widened copies of the hidden state with two tables of weights, bends one, and multiplies the two together, number by number, before narrowing back with a third: three tables of weights. The bent copy acts as a gate, deciding how much of the other copy gets through. To keep the machine the same size, the widened vector is narrower: 1,024 numbers instead of 1,536, which makes three tables of weights of 384 × 1,024 exactly as big as two of 384 × 1,536.
Change 4: sharing keys and values
In Part 1's attention, every head made its own queries, keys, and values. Grouped-query attention keeps a query for every head, here 6, but shares the keys and values: just 2 of each, each pair shared by 3 heads. That needs smaller key and value weights, so fewer numbers, and when the machine writes, the KV cache of saved keys and values is a third of the size.
Taking one change out at a time
To find out which change matters, I compared six machines, each trained for 3,000 steps on the same stories with the same 8k pieces. Five are new; the sixth is Part 4's:
- the 2017 block, from Part 4
- the new block, with all four changes
- the new block with one change at a time turned back off
Turning one part off to see what it was doing is called an ablation.
| Machine | Numbers | Bits per byte |
|---|---|---|
| 2017 block (Part 4) | 13.9 million | 0.6885 |
| New block, all four changes | 12.6 million | 0.6755, the average of three random starts |
| … but the old MLP | 12.6 million | 0.6768, the average of three random starts |
| … but full keys and values for every head | 13.8 million | 0.6721 |
| … but the old normalise | 12.6 million | 0.6683 |
| … but position embeddings instead of turning | 12.7 million | 0.7023 |
: only the 2017 block and the position embeddings land clearly outside it](/assets/images/minigpt4/ablation-luck.svg)
Scroll sideways, or tap the diagram to open it full sizeBest bits per byte for each machine. The green band is how far the same new block moved between three random starts, explained below: only the 2017 block and the position embeddings land clearly outside it
The follow-along notebook in Jupyter on my Mac, building each machine and counting its numbers. Full keys and values for every head cost 1.2 million more
How much is luck?
Every machine starts from random numbers and practises on randomly chosen snippets, so training the same design twice never gives exactly the same score. Before reading anything into a difference of a few thousandths, I needed to know how big that luck is. So I trained the new block, and the version with the old MLP, twice more each, from different random starts:
| Machine | Three random starts | Average |
|---|---|---|
| New block, all four changes | 0.6717, 0.6783, 0.6764 | 0.6755 |
| … but the old MLP | 0.6756, 0.6798, 0.6749 | 0.6768 |
The same design landed anywhere in a band about 0.007 wide. So a difference smaller than that could just be luck.
The four extra runs, as they finished. The table uses each run's best exam score, the same rule as every other score in this series
What actually matters
- Turning, RoPE, is the whole difference. On average, the new block beats the 2017 block by 0.013 bits per byte, about twice the band of luck. Put the position embeddings back, keeping the other three changes, and the machine scores 0.7023: worse than the 2017 block, and 0.027 worse than the new block's average, about four times the band. At this size, on these stories, every bit of the new block's advantage comes from turning queries and keys instead of adding position embeddings.
- Sharing keys and values is free. Full keys and values for every head scored 0.6721, inside the band, while needing 1.2 million more numbers, about a tenth of the machine.
- The gated MLP makes no difference I can measure. On my first runs it looked like a small win, 0.6717 against 0.6756. Over three random starts each, the averages are 0.6755 against 0.6768: a gap of 0.001, far inside the band. My first run of the new block had simply been the luckiest.
- The simpler normalise is not about accuracy. The old LayerNorm scored 0.6683, slightly better than any of the three new-block runs, though only just outside the band, and the two runs took within 5 seconds of each other. RMSNorm is in Llama's design because it saves work in machines with dozens of blocks and billions of numbers. At 6 blocks and 12.6 million numbers, there is nothing to see.
Bits per byte while training. The new block and three of its variants run together below the grey 2017 block; the one with position embeddings put back (red) lands on top of it
And the writing? It reads like Part 4's machine. Tim has a toy car, then a tank, and the tank stays the subject through to the end:
The new-block machine continuing "Once upon a time"
A better block at this size buys a small, measurable drop in bits per byte, not a jump you can see in the stories.
The jargon decoder
The terms for the whole series are collected in one table, the series glossary.
| What I called it | What the experts call it |
|---|---|
| the simpler normalise | RMSNorm (root mean square normalisation) |
| the 2017 normalise | LayerNorm |
| turning queries and keys by position | rotary position embeddings, or RoPE |
| the MLP with a gate | SwiGLU |
| sharing keys and values between heads | grouped-query attention, or GQA |
| a query, key, and value for every head | multi-head attention, or MHA |
| turning one change off at a time | an ablation study |
The code, in the order it runs
The code is in part5-modern-block/. It reuses Part 4's MLX data loading and Part 3's stories. I checked the layout against mlx-lm's models/llama.py.
The settings: Config in model_llama.py
Every change has a switch, so one script can train all six machines:
PYTHON
1norm: str = "rms" # "rms" | "layer" 2pos: str = "rope" # "rope" | "learned" 3mlp: str = "swiglu" # "swiglu" | "gelu" 4n_kv_heads: int = 2 # 6 would give every head its own keys and values 5hidden: int = 1024 # SwiGLU width: 384 x 1024 x 3 == 384 x 1536 x 2
Change 1: make_norm
PYTHON
1return nn.RMSNorm(cfg.dim) if cfg.norm == "rms" else nn.LayerNorm(cfg.dim)
Changes 2 and 4: Attention
The key and value weights make 2 heads' worth of keys and values, while the query weights make 6 heads' worth of queries. Then RoPE turns the queries and keys, and MLX's fused attention handles the mismatch in head counts by itself:
PYTHON
1self.wq = nn.Linear(cfg.dim, self.n_heads * self.head_dim, bias=False) # 6 heads 2self.wk = nn.Linear(cfg.dim, self.n_kv_heads * self.head_dim, bias=False) # 2 heads 3self.wv = nn.Linear(cfg.dim, self.n_kv_heads * self.head_dim, bias=False) # 2 heads 4self.rope = nn.RoPE(self.head_dim, traditional=False, base=cfg.rope_base) 5... 6q, k = self.rope(q), self.rope(k) 7out = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask="causal")
Change 3: SwiGLU
PYTHON
1def __call__(self, x): 2 return self.w2(nn.silu(self.w1(x)) * self.w3(x)) # gate (w1), up (w3), down (w2)
The six machines: run_ablation.sh and train_llama.py
train_llama.py is Part 4's compiled MLX training loop, with the switches as options. run_ablation.sh trains the new block and then each one-change-off variant:
BASH
1python train_llama.py --tag modern 2python train_llama.py --tag rms_rope_gelu --mlp gelu 3python train_llama.py --tag rms_rope_mha --gqa-off 4python train_llama.py --tag layer_rope --norm layer 5python train_llama.py --tag rms_learned --pos learned
Then, for How much is luck?, it trains the new block and the old-MLP version from two more random starts:
BASH
1python train_llama.py --tag modern_seed1 --seed 1 2python train_llama.py --tag gelu_seed1 --seed 1 --mlp gelu 3python train_llama.py --tag modern_seed2 --seed 2 4python train_llama.py --tag gelu_seed2 --seed 2 --mlp gelu
The five runs. Every machine is close in size except the one with full keys and values, which adds back 1.2 million numbers
The new block trained in the notebook: 3,000 steps, finishing at 0.6717 bits per byte
figures.py then draws the bar and curve charts above, with Part 4's 2017-block run as the grey baseline.
Try it yourself
- The follow-along notebook:
part5-modern-block/minigpt_follow_along_5.ipynb, for Jupyter on a Mac with Apple Silicon. It builds each machine and counts its numbers, shows RoPE's turning in action, trains the new block, and writes. It is saved with the outputs from my own run, so you can read every result on GitHub without a Mac.
The notebook open in Jupyter on my Mac, with the outputs saved from my run
- On the command line, after Part 3's
prepare_data.pyandtokenizers_setup.py:
BASH
1cd minigpt-series/part5-modern-block 2./run_ablation.sh 3python figures.py 4python generate_llama.py --tag modern --prompt "Once upon a time"
MLX needs Apple Silicon.
Part 6 goes back to Part 4's machine and gives it a teacher: a bigger model whose whole wheel of chances, at every position, becomes something extra for the student to copy.