Old spear still good, but Ugg sharpen every part

Long ago wise tribe make the Transformer — big brain-machine that read words and see pictures. Made year 2017. Same spear shape today! But clever cavemen quiet-swap almost every part. New sharper stone, new tighter binding. Ugg count nine swaps. Know all nine, Ugg pass the test.

Why care? Feed brain-machine a trillion berries, one bad part make it fall over dead; good parts make it grow strong like GPT-4. These not shiny extras. These life-or-death.

Keep the river clean (stability swaps)

Pre-Norm — put the rock-cleaner before the helper, not after. New way:

Old way, Post-Norm, put cleaner after: $x_{\text{out}}=\text{LN}(x+\text{Sublayer}(x))$.

The straight line $x$ — Ugg call it the residual stream — is a clean river running front to back of cave, no squeezer touch it. Teach-signals swim up this river fast to reach early layers. Post-Norm squeeze the river every layer; deep cave (24+ layers) shake and fall without slow careful start. So Pre-Norm give signals a clean unbroken path. All big machines use it.

LayerScale — tiny volume-dial on each helper's branch:

Start dial almost all-off ($\lambda$ tiny). At start every helper say nothing, machine act like plain pass-through — calm, no fall. As hunt go on, useful dials grow loud. For very deep cave only.

Make all rocks same size (normalise swaps)

Old LayerNorm do two things: slide all rocks to middle (take away mean $\mu$), then shrink to same length (divide by spread $\sigma$), plus two learn-dials $\gamma ,\beta$:

RMSNorm skip the slide-to-middle. Just shrink by root-mean-square, one dial $\gamma$, no $\beta$:

Fewer moves, run faster, work same or better. Ugg warn: no do this! RMSNorm not just "drop the $\beta$" — it ALSO drop the mean-slide! Both gone.

Fire-picture (the test-trap): scatter dots on cave wall. LayerNorm slide-then-shrink squish every dot onto one slant line ($x_1+x_2=0$, the anti-diagonal) — dots land on two matching spots. RMSNorm no slide, so dots spread all round the full circle (ball of radius $\sqrt{D}$). Chant it: LayerNorm to small circle on slant line (two constraints); RMSNorm to full ball (one constraint). Dropping mean no hurt — machine got other ways to center (the dials, the linear-layer bias).

QK-Norm — the sniff-score $Q{K}^{\top }$ can grow monster-big, then softmax choke (one choice eat all, no learning). Old fix divide by $\sqrt{d_k}$ — not enough at big size. QK-Norm clean $Q$ and $K$ each before the sniff:

Now each pair-dot bounded $\pm 1$, softmax stay calm. No choke, no blow-up.

Give tribe scratch-stones (attention hygiene)

Registers — garbage-catcher tokens. Big picture-brain do weird thing: it dump global notes onto blank sky-patches, corrupt them, attention map go crazy. Fix: glue on 4–8 blank scratch-stone tokens with no picture behind them:

Now brain scribble on scratch-stones, not on the good hunting map. Attention clean, dense tasks better. Ugg warn: no do this! Scratch-stones get NO position-mark — that is the whole point. Throw them away at the end.

Hunt faster (efficiency swaps)

Flash Attention — same answer, way faster. Not a lazy guess — exact, same numbers as before. Slow part is not thinking; it is hauling the giant $N\times N$ sniff-rock between big-slow store (HBM) and tiny-fast pocket (SRAM, ~100× faster). Trick: cut $Q,K,V$ into small tiles that fit the pocket, sniff tile-by-tile with a running-biggest $m$ and running-total $\ell$ (the online-softmax), never build the whole giant rock. Memory drop $\mathcal{O}(N^2)\to \mathcal{O}(N)$, 2–4× faster. Two exam-words: exact and IO-aware — speed come from less hauling, not less thinking.

KV-cache — talk-time trick. When machine make words one-by-one, the past-token keys $K$ and values $V$ never change! So store them in a memory-pouch. Each new step: make $Q,K,V$ for the new token only, drop new $K,V$ in pouch, sniff new $Q$ against all pouch-keys. Work per token $\mathcal{O}(N^2)\to \mathcal{O}(N)$; whole talk $\mathcal{O}(N^3)\to \mathcal{O}(N^2)$. This why chatbot stream fast. Cost: pouch grow big with long talk. Ugg warn: no do this! Flash Attention and KV-cache NOT same thing — one is memory-hauling, one is skip-recompute.

GQA (Grouped-Query Attention) — share the key-heads. In MHA (Multi-Head), $H$ query-heads each get own $K,V$ head → huge pouch. In MQA (Multi-Query), all share one $K,V$ → tiny pouch but worse hunt. GQA sit in middle: split $H$ query-heads into $G$ groups, each group share one $K,V$. Pouch shrink by:

Typical $H=64,G=8$ → 8× smaller pouch, near-same quality. Ugg warn: no do this! GQA is not a broken MQA — it the sweet-spot between MHA and MQA.

The footprint trick (positional swap)

RoPE (Rotary Position Embedding) — mark position by spinning the arrow, not adding to it. For token at spot $m$:

Spin $q$ by $m$, spin $k$ by $n$, then the sniff-score only care about the gap:

Because two spins add to one spin of angle $n-m$. Good things: reach longer than trained (no fixed lookup-table), speak relative distance, and it twist not add. Ugg warn: no do this! RoPE is multiply-spin, NOT add.

Ugg remember

• Nine swaps, 2017 → now: Pre-Norm, RMSNorm, LayerScale, QK-Norm, Registers, Flash Attention, RoPE, KV-cache, GQA. Same spear, all-new parts.
• Clean river beats squeezed river: Pre-Norm is the unbroken residual highway for teach-signals; LayerScale start each helper near-silent ($\lambda \approx 10^{-4}$) so deep cave begin calm.
• RMSNorm = LayerNorm minus mean minus bias (cheaper, same-good). Fire-picture: LayerNorm squish dots to slant-line circle; RMSNorm spread on full ball. QK-Norm keep softmax from choking.
• **Flash Attention is EXACT + IO-aware; KV-cache turn $\mathcal{O}(N^3)\to \mathcal{O}(N^2)$ talk** — different tricks, no mix them. GQA shrink key-pouch by $H/G$ (about 8×) with near-MHA quality.
• **RoPE spins for relative position** — multiplicative, reaches past training length. Registers are position-free scratch-stones, thrown away at the end.