Logits — Linear algebra for LLMs

From the last hidden state to scores

\displaystyle \text{logits} = \operatorname{Norm}(h)\, W_U

→

After the last block, a final norm and the unembedding $W_U$ ( $d \times V$ ) give every position one score per vocabulary token. Llama 3’s $V = 128{,}256$ , so this is one of the biggest matrices in the model.

Write

Logits =

\operatorname{Norm}(h)\,W_U

,

W_U\!: d \times V

, output

(B, T, V)

.

Weight tying: reuse the embedding table

\displaystyle W_U = E^\top \;\Rightarrow\; \text{logit}_j = h \cdot E_{j,:}

Some models use the transposed embedding table as the unembedding. Then the logit for token $j$ is the dot product of the hidden state with token $j$ ’s own embedding: “how much does $h$ look like this token?” GPT-2 ties its weights; Llama 3 8B keeps two separate matrices.

Write

Tied:

W_U = E^\top

, logit

_j = h \cdot E_{j,:}

.

✏️ Your turn

Score the toy vocabulary

h = [1, 0, 1, 0]

With tied weights (and skipping the norm), write the five logits $Eh$ .

Work it out on paper, then go to the next card to check.

Answer

Each logit is $h$ dotted with a row: “the” scores $1 + 2 = 3$ , the highest.

Softmax to probabilities

Softmax (Part 5) turns the logits into a distribution over the next token: “the” gets about 69%. Sampling picks from it; greedy decoding takes the top one.

The training loss

\displaystyle \mathcal{L} = -\log p_{\text{target}}

Cross-entropy: the negative log-probability the model gave to the token that actually came next, averaged over every position. If the target here were “the”, the loss would be $-\ln 0.693 \approx 0.37$ .

Write

Cross-entropy:

\mathcal{L} = -\log p_{\text{target}}

, averaged over positions.

What it costs

\displaystyle 2\,d\,V = 2 \times 4096 \times 128{,}256 \approx 1.05 \times 10^9 \text{ FLOPs per token}

In a transformer

At generation time only the last position’s logits are needed, so code computes logits = lm_head(h[:, -1]). In training every position predicts its next token, so the full $(B, T, V)$ tensor is built: at $T = 8192$ and $V = 128{,}256$ that’s over a billion logits per sequence.

📌 Recap

Logits in one card

Logits = $\operatorname{Norm}(h)\,W_U$ , one score per vocabulary token.
Tied weights: $W_U = E^\top$ , so each logit is $h \cdot$ a token embedding.
Softmax gives next-token probabilities; the loss is $-\log p_{\text{target}}$ .
Cost: $2dV$ FLOPs per token.

Write

The four lines above.