From Attention to a Working Language Model

Parallelizing the Transformer, Masking the Future, and the Language Modeling Head

Last post built self-attention and the transformer block, but left two promises unkept: I showed the computation one token at a time (so where's the famous parallelism?), and I never showed how any of it actually predicts a word. This post closes both gaps. By the end, you'll know how the whole attention computation collapses into a few big matrix multiplies that run in parallel, why a causal language model has to mask out the future and how a triangle of negative infinity does it, where the input vectors actually come from (token embeddings plus a second thing we've been ignoring: position), and how the language modeling head turns the top-of-stack vector into a probability distribution over the next word. That last piece is the literal proof of the claim that an LLM "just predicts the next token." We'll finish with how these models scale and how you fine-tune one without melting your GPU.

Two ideas separate a transformer from an RNN, and this post is about cashing them in. The first is attention, which we built last time; it replaces recurrence. The second is parallel computation: processing every token at once rather than one at a time. Last post showed attention but still walked through it serially. Today we make it actually parallel, then bolt on the machinery that turns a stack of blocks into something that generates text.

Recap: Where We Left Off

Quick reset. A transformer block takes a token's d-dimensional vector, runs it through layer norm, multi-head attention, another layer norm, and a feedforward layer (with residual connections threading the original vector through), and emits a d-dimensional vector. Same shape in, same shape out, so you stack these blocks, 12 to 96 or more. Only the attention step looks at other tokens; everything else refines a token in place.

We did all of that one token, one residual stream at a time. But here's what we need to emphasize: the computation for each token is independent of the others (except inside attention). That independence is what lets us stop looping and start multiplying matrices.

Doing It All at Once: The Matrix X

Instead of processing token i, then token i+1, then i+2, we pack the entire input sequence into one matrix. Call it $X$ , of shape $[N \times d]$ — N rows (one per token), each row the d-dimensional embedding of that token. For real models N (the context window) runs from 1K to 32K tokens and beyond.

Now the per-token query/key/value projections become single matrix multiplies over the whole sequence at once:

Q = X W^{Q}; K = X W^{K}; V = X W^{V}

With $W^{Q}, W^{K}$ of shape $[d \times d_{k}]$ and $W^{V}$ of shape $[d \times d_{v}]$ , this gives $Q, K$ of shape $[N \times d_{k}]$ and $V$ of shape $[N \times d_{v}]$ — every token's query, key, and value computed in one shot.

Now the satisfying part. To get all the query-key comparisons (every token scored against every other token), you just multiply $Q$ by $K^{T}$ :

Q K^{T} has shape [N \times N]

Cell $(i, j)$ of this matrix is $q_{i} \cdot k_{j}$ , the relevance of token j to token i. The entire attention score table, in a single matrix multiplication. Scale it, softmax it, multiply by $V$ , and you have the attention output for every token at once:

head = softmax (mask (\frac{Q K ^{T}}{d _{k}})) V

A = head W^{O}

That mask in there is new. It's the price of doing everything at once, and it's worth understanding properly.

Masking the Future

The $Q K^{T}$ matrix computes a score for every query against every key, including keys that come after the query in the sequence. For a causal language model, that's cheating.

Think about what the model is trained to do: predict the next word. If, while computing the representation of token 3, you let it attend to tokens 4, 5, 6… you've handed it the answer. Guessing the next word is pretty easy if you already know it.

So we mask out the future. In the $N \times N$ score matrix, every cell where the key is ahead of the query (the upper triangle) gets set to $- \infty$ :

M_{ij} = - \infty \forall j > i, M_{ij} = 0 otherwise

Why negative infinity? Because the very next step is a softmax, and $e^{- \infty} = 0$ . Those future positions get exactly zero attention weight. The model is mathematically blindfolded to anything it hasn't reached yet.

The mask is what makes "parallel" and "causal" coexist. Computing every token at once would normally let each token peek at all the others, including future ones. The triangular

- \infty

mask surgically removes the future from each token's view, so you get the speed of parallel computation and the honesty of left-to-right prediction. (Bidirectional models like BERT use a different mask; they're allowed to look both ways. That's for another day.)

One Catch: Attention Is Quadratic

That $Q K^{T}$ matrix is $N \times N$ . Which means the cost of attention grows with the square of the sequence length; you're computing a dot product between every pair of tokens.

Double the context window, and you quadruple the attention computation. This is why long context is expensive, why running attention over an entire novel is hard, and why "rooms full of GPUs" is not a figure of speech. Modern models still push contexts into the tens of thousands of tokens, but the quadratic cost is the wall everyone is pushing against.

The Parallel Block, and the Two Big Differences

With attention parallelized, the whole transformer block has a clean matrix form (prenorm — layer norm before each sublayer):

O = X + MultiHeadAttention (LayerNorm (X))

H = O + FFN (LayerNorm (O))

Both $X$ and $H$ are $[N \times d]$ : the full sequence in, the full sequence out, so blocks stack exactly as before. For the first block, $X$ is the input embeddings; for block k, it's the output $H$ of block k−1.

Look at how far we've come from RNNs. The transformer introduces exactly two new ideas relative to a recurrent net:

Attention replaces recurrence. No hidden state is passed step to step; instead, every token directly attends to every (prior) token.
Parallel computation replaces iteration. No waiting for token i−1 to finish before starting token i; the whole sequence goes through the matrix multiplications together.

Everything else (embeddings, softmax, predicting the next word) we've seen since the n-gram days. Those two ideas are all that's actually new.

Quick reference: the shapes

X (input sequence) → [N × d], one row per token.

W^Q, W^K → [d × dₖ]; W^V — [d × dᵥ].

Q, K → [N × dₖ]; V → [N × dᵥ].

QKᵀ (all query-key scores) → [N × N].

Attention output A → [N × d]; same shape as X, so blocks stack.

Logits u (one score per vocab word) → [1 × |V|].

What Is X, Really? Token and Position Embeddings

I've been waving my hands about "the input embeddings." Time to be honest about what $X$ actually contains, because it's not just token embeddings. Each row is the sum of two embeddings.

Token embeddings

The familiar part. There's an embedding matrix $E$ of shape $[∣ V ∣ \times d]$ — one row per vocabulary token. To get the input, you tokenize the string (with byte-pair encoding, which can split into subword pieces), convert each token to its vocabulary index, and select the matching rows of $E$ . So "Thanks for all the" might become indices $[5, 4000, 10532, 2224]$ , and you pull rows 5, 4000, 10532, 2224.

Position embeddings

One thing here is easy to miss: Because we process all tokens in parallel, the model has no inherent sense of order. Nothing in the matrix multiplication says token 3 came before token 4. We have to inject that information explicitly.

So alongside the token embedding, we add a position embedding: a vector that says "this is position 3." The simplest scheme, absolute position, keeps a learned embedding matrix $E_{p os}$ with one randomly initialized vector per position, learned during training just like word embeddings. Just as the embedding for fish captures what fish means, the embedding for position 3 captures something about what tends to appear in position 3.

The final input is their sum:

x_{i} = E [token i] + E p os [i]

Both are $[1 \times d]$ , so the sum is too. That composite vector (word meaning + position) is what actually enters the first transformer block.

A near miss: simple absolute positions (the quiz question). Learned absolute position embeddings almost work, which is what makes the flaw worth seeing. You learn them from training data — but how many training examples land at position 25,000 versus positions 1 through 5? Far fewer. So the embeddings for late positions are undertrained and generalize badly. The fixes pick up exactly where absolute positions fall short: sinusoidal embeddings (sine/cosine functions of the position, which extend to any length and naturally encode that position 4 is closer to 5 than to 17), or relative schemes baked into attention at each layer instead of added once at the input.

The Language Modeling Head

We have a stack of transformer blocks turning input embeddings into rich contextual representations. But a pile of d-dimensional vectors isn't a prediction. The language modeling head is the final circuit that turns the top of the stack into an actual next-word distribution.

The recipe is short. Take the output of the last token from the last layer, call it $h_{N L}$ , a single $[1 \times d]$ vector. Don't underestimate this vector: it has been through every layer of attention and feedforward, so it's a deep contextual summary of everything seen so far. If the model has done its job, it encodes exactly what's needed to guess what comes next.

Now project it to a score for every word in the vocabulary. We multiply by an unembedding matrix, and there's a nice trick here called weight tying: instead of learning a fresh matrix, we reuse the transpose of the input embedding matrix, $E^{T}$ , of shape $[d \times ∣ V ∣]$ . At the input, $E$ maps a word to an embedding; at the output, $E^{T}$ maps an embedding back to scores over words. Same weights, run in reverse, hence "unembedding."

u = h_{N L} E^{T}

y = softmax (u)

The vector $u$ holds the logits — raw, unnormalized scores, one per vocabulary word ( $[1 \times ∣ V ∣]$ ). The softmax turns them into a probability distribution $y$ over the whole vocabulary. Then you sample a word from $y$ — greedily (take the most probable), or with top-k / top-p sampling for more diversity (which we covered back in the LLM post).

This is the proof. When people say "all ChatGPT does is predict the next token," this is the mechanism they're describing. The last token's representation, after all that computation, gets multiplied by

E^{T}

and softmaxed into a distribution over words — and one word is sampled. That's it. Everything impressive about an LLM is that distribution being startlingly good, because the vector feeding it carries so much context. The architecture is humble; the representation is not.

One bit of terminology: a transformer used this way (left-to-right, masked, predicting the next token) is called a decoder-only model. It's the decoder half of the original encoder-decoder transformer, repurposed.

The Full Picture

Stand back and trace one token through the finished machine:

Input encoding — token embedding + position embedding → $x_{i}$ .
L transformer blocks — each applying masked multi-head attention and a feedforward layer over the residual stream, all N tokens in parallel, the representation getting richer at every layer.
Language modeling head — take the final layer's last-token vector, unembed with $E^{T}$ , softmax into a distribution, sample the next word.

That's a transformer language model, end to end. We started this course counting n-grams; we're now generating text with a 96-layer attention network. Same task the whole way (predict the next word), just a vastly better predictor.

What a transformer is NOT. It is not a reasoning engine with hidden logic, a database it looks things up in, or a little homunculus deciding what to say. Strip away the scale and it's this: matrix mul...

Source：https://dev.to/kambleakash0/parallelizing-the-transformer-p9b

My favorites