From attention to a full language model

You know what attention does. This section fills in everything around it: how text becomes vectors, what one block contains, what stops a model reading ahead when it is trained to predict the next word, and how the final vector becomes a token. Having the whole path in your head is worth the hours because most of the practical questions, such as why generation is slow, why streaming exists, or why long context costs what it does, are answered at some specific point along it.

token ids        [9906, 1917, 0]
embedding        one learned vector per id, width d_model
position         rotate or add, because attention cannot see order

repeat N times:
  x = x + attention(norm(x))        mix information across positions
  x = x + feed-forward(norm(x))     transform each position on its own

final norm
unembedding      d_model -> one score per vocabulary entry
sample           scores become probabilities, one is chosen

Embeddings#

Tokenization turns text into integers, and each integer is a row index into a learned table. Look up the row and you have a vector of width d_model, which for a model of this size is 4,096 numbers. Nothing about that lookup is clever: the table is a plain matrix of vocabulary size by d_model, trained like everything else, and the row for a token is the only thing the model knows about it before any layer has run.

Two consequences follow and both come up in interviews. The embedding table is large because vocabulary is large, so on a small model it is a noticeable share of the parameters and on a large one it is not. And because every token becomes the same width, the rest of the network never has to know how many tokens there are; d_model is fixed and the sequence length is whatever it is.

Why position has to be supplied at all#

This is the part that sounds like a detail and is not. Attention compares every position with every other and blends the results, and none of those operations knows where anything sits. Shuffle the input tokens and you get the same outputs back in a different order. That is not an approximation or a weakness of the explanation; it is exactly true, and you can watch it happen in twelve lines.

import numpy as np

rng = np.random.default_rng(0)
d = 8
Wq, Wk, Wv = (rng.standard_normal((d, d)) / np.sqrt(d) for _ in range(3))

def attention(X):
    Q, K, V = X @ Wq, X @ Wk, X @ Wv
    s = Q @ K.T / np.sqrt(d)
    p = np.exp(s - s.max(axis=1, keepdims=True))
    return (p / p.sum(axis=1, keepdims=True)) @ V

X = rng.standard_normal((4, d))        # four token embeddings, no position anywhere
order = [3, 0, 2, 1]                   # shuffle the tokens

same = np.abs(attention(X)[order] - attention(X[order])).max()
print(f"no position:  shuffling the input just shuffles the output, gap {same:.1e}")

P = rng.standard_normal((4, d)) * 0.5  # one learned vector per slot in the sequence
moved = np.abs(attention(X + P)[order] - attention(X[order] + P)).max()
print(f"with position: the same shuffle changes every value, gap {moved:.3f}")

# no position:  shuffling the input just shuffles the output, gap 4.4e-16
# with position: the same shuffle changes every value, gap 1.378

So position is injected from outside. The original models added a learned vector per slot, which is what the second half of the code above does. Current models rotate the query and key vectors by an angle proportional to position instead, which makes the comparison between two positions depend on the gap between them rather than on where each one sits, and removes the fixed-size table that capped how far a model could read. That is rotary position embedding, and the modern-variants section in this topic covers it and the tricks used to stretch it.

Residual connections and normalisation#

Each block adds its output to its input rather than replacing it. Written out, x becomes x plus attention of x, then x plus feed-forward of x. The running vector that every block reads from and writes back into is usually called the residual stream, and it is the closest thing the architecture has to a working memory: a layer that has nothing to contribute at a position can contribute nearly nothing and the information already there passes through untouched.

def cos(a, b):
    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b)))

def rms_norm(v):
    return v / np.sqrt(np.mean(v * v) + 1e-6)

D = 256
g = np.random.default_rng(1)
A = [g.standard_normal((D, D)) / np.sqrt(D) for _ in range(32)]   # 32 blocks
h0 = g.standard_normal(D)
plain = res = h0.copy()

print("layer    no residual          with residual")
for layer, Ai in enumerate(A, start=1):
    plain = np.tanh(plain @ Ai)
    res = res + np.tanh(rms_norm(res) @ Ai)
    if layer in (1, 2, 4, 8, 16, 32):
        print(f"{layer:5d}    cos {cos(plain, h0):+.3f}  |h| {np.linalg.norm(plain):6.3f}"
              f"    cos {cos(res, h0):+.3f}  |h| {np.linalg.norm(res):6.2f}")

# layer    no residual          with residual
#     1    cos -0.026  |h| 10.339    cos +0.843  |h|  19.04
#     2    cos -0.033  |h|  7.728    cos +0.760  |h|  22.35
#     4    cos -0.100  |h|  5.283    cos +0.653  |h|  25.06
#     8    cos +0.003  |h|  3.820    cos +0.479  |h|  33.14
#    16    cos +0.017  |h|  2.723    cos +0.409  |h|  39.62
#    32    cos -0.027  |h|  2.209    cos +0.256  |h|  59.01

That is the case for residuals in one table, and it is the same argument in the forward direction as the usual one about gradients flowing backwards. The cost is visible in the right-hand column: the norm of the residual stream grows as every layer adds to it, reaching 59 after thirty-two layers from a start of about 16. Normalisation is what keeps that under control. Each sub-layer normalises its input before using it, so a block sees a vector of predictable scale whatever the stream has accumulated, and a final normalisation runs before the unembedding for the same reason.

Layer normalisation subtracts the mean, divides by the standard deviation, then applies a learned scale and shift. RMSNorm, which almost everything current uses, drops the mean and the shift: divide by the root mean square and apply a scale. Two learned vectors of width d_model per block, so a few thousandths of one per cent of the parameters, which is worth knowing because it is a common guess that normalisation is expensive. It is not; it is a rounding error in the parameter count and a real cost only in memory traffic.

The feed-forward block, where the parameters actually are#

After attention has moved information between positions, the feed-forward block transforms each position independently: project up to a wider hidden dimension, apply a non-linearity, project back down. Older models used two matrices with a GELU between them. Current ones use a gated variant, SwiGLU, which does the up-projection twice and multiplies one copy through a swish by the other, so there are three matrices instead of two and the hidden width is cut to compensate.

People consistently expect attention to dominate the parameter count, because attention is the part everyone talks about. Work it out for a real configuration and it is not close.

# Shapes taken from a published 8B open-weights config.
V, d_model, layers = 128256, 4096, 32
heads, kv_heads, head_dim, ffn = 32, 8, 128, 14336

parts = [
    ("token embedding", V * d_model),
    ("attention Q", layers * d_model * heads * head_dim),
    ("attention K and V", layers * 2 * d_model * kv_heads * head_dim),
    ("attention output", layers * heads * head_dim * d_model),
    ("feed-forward (3 matrices)", layers * 3 * d_model * ffn),
    ("RMSNorm scales", layers * 2 * d_model + d_model),
    ("unembedding", V * d_model),
]
total = sum(n for _, n in parts)
print(f"{'component':26}{'parameters':>14}{'share':>8}")
for name, n in parts:
    print(f"{name:26}{n:14,}{n/total:8.1%}")
print(f"{'TOTAL':26}{total:14,}{1.0:8.1%}")

attn = sum(n for name, n in parts if name.startswith("attention"))
ffnp = dict(parts)["feed-forward (3 matrices)"]
print(f"\nfeed-forward is {ffnp/attn:.2f}x the whole attention mechanism")
print(f"per layer: attention {attn//layers:,}  feed-forward {ffnp//layers:,}")
for bits in (16, 8, 4):
    print(f"weights at {bits:2d}-bit: {total * bits / 8 / 1024**3:6.2f} GiB")

# component                     parameters   share
# token embedding              525,336,576    6.5%
# attention Q                  536,870,912    6.7%
# attention K and V            268,435,456    3.3%
# attention output             536,870,912    6.7%
# feed-forward (3 matrices)  5,637,144,576   70.2%
# RMSNorm scales                   266,240    0.0%
# unembedding                  525,336,576    6.5%
# TOTAL                      8,030,261,248  100.0%
#
# feed-forward is 4.20x the whole attention mechanism
# per layer: attention 41,943,040  feed-forward 176,160,768
# weights at 16-bit:  14.96 GiB
# weights at  8-bit:   7.48 GiB
# weights at  4-bit:   3.74 GiB

Seventy per cent in the feed-forward layers, seventeen across all of attention, thirteen in the two vocabulary matrices, and effectively nothing in the norms. The gap is wider still per layer: 176 million parameters of feed-forward against 42 million of attention, a ratio of 4.2 to 1. Attention gets the attention because it is the interesting mechanism and because its cost grows with the square of sequence length, but in sheer weight the model is mostly a stack of position-wise transformations, which is the usual reading of where a model keeps what it knows.

The same table answers a practical question. Weights at 16-bit are 14.96 GiB, which does not leave much of a 24 GiB consumer card for anything else, and at 4-bit they are 3.74 GiB, which does. The inference-cost section in the quantization topic takes that further, including what the cache on top of the weights costs.

Causal masking#

Language models predict the next token, so each position may only attend to earlier positions. A mask blocks attention to the future. This is what lets the whole sequence train in parallel while still learning left-to-right prediction. Drop the mask and you have an encoder-only model, such as BERT, where every position sees the whole input at once: useful for classification and embedding, and unable to generate text left to right.

The mask is also why the same architecture serves two different jobs. With it you have a decoder-only model that generates. Without it you have an encoder, where every position is conditioned on the whole input at once, which is what you want for classification and for producing a single embedding of a document and useless for writing the next word. Nothing else about the block changes.

Output: from the last vector to a token#

The final vector at the last position is projected to a score for every token in the vocabulary. Sampling from those scores produces the next token, which is appended to the input, and the process repeats. That autoregressive loop is generation.

The projection is the unembedding, a matrix of d_model by vocabulary size, and it is the largest single matrix in the model at 525 million parameters in the table above. Some models tie it to the embedding table, reusing the same weights transposed, which saves those parameters and costs a little quality; the 8B configuration above does not, which is why the two rows are counted separately. Its output is one raw number per vocabulary entry, called logits, and logits are not probabilities. Turning them into probabilities is a softmax, and everything a sampling parameter does happens between those two steps.

vocab = ["the", "a", "this", "my", "an", "our", "some", "no"]
logits = np.array([8.2, 6.9, 5.4, 4.8, 4.1, 3.7, 2.2, 1.0])

def softmax(z, temperature=1.0):
    z = z / temperature
    e = np.exp(z - z.max())
    return e / e.sum()

print(f"{'token':8}" + "".join(f"{f'T={t}':>9}" for t in (0.5, 1.0, 1.5)))
cols = [softmax(logits, t) for t in (0.5, 1.0, 1.5)]
for i, w in enumerate(vocab):
    print(f"{w:8}" + "".join(f"{c[i]:9.4f}" for c in cols))

p = softmax(logits)
order = np.argsort(-p)
cum = np.cumsum(p[order])
keep = order[: int(np.searchsorted(cum, 0.9)) + 1]
print(f"\ntop-p 0.9 at T=1 keeps {len(keep)} of {len(vocab)} tokens: "
      f"{[vocab[i] for i in keep]}")
print(f"the other {len(vocab) - len(keep)} share {1 - cum[len(keep)-1]:.4f} of the mass "
      f"and are dropped before sampling")
print(f"greedy would always pick {vocab[order[0]]!r} at p={p[order[0]]:.4f}")

# token       T=0.5    T=1.0    T=1.5
# the        0.9264   0.7155   0.5495
# a          0.0688   0.1950   0.2310
# this       0.0034   0.0435   0.0850
# my         0.0010   0.0239   0.0570
# an         0.0003   0.0119   0.0357
# our        0.0001   0.0079   0.0274
# some       0.0000   0.0018   0.0101
# no         0.0000   0.0005   0.0045
#
# top-p 0.9 at T=1 keeps 2 of 8 tokens: ['the', 'a']
# the other 6 share 0.0895 of the mass and are dropped before sampling
# greedy would always pick 'the' at p=0.7155

Three things are worth taking from that. Temperature is a division applied before the softmax, not a scaling of probabilities afterwards, which is why its effect is not linear and why the difference between 0.7 and 1.0 is larger than it looks. Top-p truncates the sorted distribution at a cumulative mass and renormalises, so on a confident step it keeps two tokens and on an uncertain one it keeps dozens; that is the point of it, and the reason it is usually preferred to a fixed top-k. And temperature zero is not a temperature at all, since dividing by zero is undefined: implementations treat it as take the highest logit, which is why it is deterministic in a way no small temperature is.

Then the loop. The chosen token is appended and the whole thing runs again, which is the fact that explains most of what you notice about using a model: output is produced strictly one token at a time so latency scales with answer length, streaming is possible because each token is finished when it appears, and a model cannot revise a token it has already emitted because that token is now part of its own input.