LM6-T2 — Classification Heads and Sequence Representations

Learning objectives

After this lesson, you should be able to:

• Describe the classification head and its input and output shapes.
• Explain why the last real token represents the sequence in a causal model, and why the first token does not.
• Select the hidden state at a different position for each sequence in a padded batch.
• Name alternatives to last-token pooling and when they apply.

The idea: one position that has read everything

After the backbone we have hidden states of shape (B, n, d): a vector per token. To classify, we pick one vector per sequence, shape (B, d), and feed it to a small linear layer that outputs K numbers, one per class. Those are the logits; softmax turns them into class probabilities and cross-entropy compares them with the label.

Why the last token

Under the causal mask, the hidden state at position i depends on tokens 0 to i only. The first position has seen nothing but itself, so it cannot know what the sentence says: classifying from it gives a model that can only guess the base rate (the notebook shows an F1 of 0 for this). The last position has attended to every earlier token, so its state can summarise the sentence. This is the same reason that, during generation, the logits of the last position are used to predict the next token.

Padding

Batches are rectangles: short texts are padded to a common length. The padding sits at the end (right padding), after the real tokens, so the causal mask guarantees that it cannot influence the real positions. But the last position of the tensor is then a padding position for most sequences and holds no summary of the text. For each sequence we must pick its own last real position, length − 1. The book avoids this by padding every example to the same length and using the final position; both work, but the first approach wastes less compute on padding while the second needs every sequence to end in the same place (and ignores the cost of attending to padding).

Alternatives

• Mean pooling: average the hidden states of all real tokens. Every token contributes; common with bidirectional models. With a causal model, early tokens have seen little, so the average dilutes the useful late states.
• A special classification token appended to the text (or prepended in bidirectional models such as BERT): its hidden state is used. In a causal model it must be appended at the end, where it equals the last-token approach with a dedicated token.
• Bidirectional encoders: models that read in both directions make every position informative, so the first token or a mean is natural.

The head

The head is usually a single linear layer: logits = Wh + b. A deeper head (one hidden layer with a nonlinearity) is sometimes used when the backbone is frozen, to give the classifier more capacity. The head is initialised randomly, so before training its output is meaningless.

The mathematics: pooling and the classification loss

Let the backbone output H of shape (B, n, d) and let ℓ_b be the number of real tokens of sequence b (1 ≤ ℓ_b ≤ n). Last-token pooling selects

and the head maps it to class scores:

ℓ_b

the number of real (non-padding) tokens of sequence b; positions ℓ_b … n−1 are padding.

z_b

the pooled vector of sequence b; shape (d,).

W, c

the head’s weight matrix and bias; K·d + K parameters.

o_b

the logits (class scores) of sequence b; shape (K,).

The class probabilities are softmax(o_b) and the loss over a batch of labels y_b is

Mean pooling would instead use z_b = (1/ℓ_b) Σ_{i < ℓb} H[b, i, :], masking the padding positions out of the average.

Shapes. For the notebook: B = 32, n = 32 characters, d = 64, K = 2: H is (32, 32, 64), z is (32, 64) and the logits are (32, 2). The head has 64 · 2 + 2 = 130 parameters. For GPT-2 small, d = 768 and K = 2 give 1,538.

Untrained head. With small random weights the logits are close to 0 and the loss starts near ln K (0.693 for two classes), the analogue of ln V for language modelling.

A worked example: picking the right position

A batch of three sequences padded to length 6 has real lengths 5, 3 and 4. The hidden states are random numbers for illustration; what matters is which rows are selected.

Using the last tensor position (5) for all three would pick padding for sequences 1 and 2.

import torch
import torch.nn as nn

torch.manual_seed(0)
B, n, d, K = 3, 6, 4, 2
H = torch.randn(B, n, d)                       # backbone output
lengths = torch.tensor([5, 3, 4])

pos = lengths - 1                              # tensor([4, 2, 3])
z = H[torch.arange(B), pos]                    # advanced indexing: one row per sequence
print(pos.tolist(), tuple(z.shape))            # [4, 2, 3] (3, 4)
print(torch.equal(z[1], H[1, 2]))              # True: sequence 1 uses its position 2

# the same with gather, and mean pooling with a mask for comparison
z2 = H.gather(1, pos.view(B, 1, 1).expand(B, 1, d)).squeeze(1)
print(torch.equal(z, z2))                      # True
mask = (torch.arange(n)[None, :] < lengths[:, None]).float()      # 1 for real tokens
mean = (H * mask[..., None]).sum(1) / lengths[:, None]
print(tuple(mask.shape), tuple(mean.shape))    # (3, 6) (3, 4)

head = nn.Linear(d, K)
logits = head(z)
print(tuple(logits.shape), sum(p.numel() for p in head.parameters()))   # (3, 2) 10  = d*K + K

# padding after the end cannot change the pooled state of a causal model: here a causal toy "backbone"
cum = lambda x: x.cumsum(1)                    # position i depends only on positions <= i
x = torch.randn(B, n, d); x2 = x.clone(); x2[1, 3:] = 99.0
print(torch.equal(cum(x)[1, 2], cum(x2)[1, 2]))   # True: position 2 is unaffected by changes after it

The final lines use a cumulative sum as a stand-in for a causal layer: each output depends only on the inputs up to its own position, so changing the padding at the end cannot change an earlier output. The notebook verifies the same property on the real GPT.

Engineering implications and limitations

Right padding with explicit lengths. Carry the lengths alongside the batch. Compute them from the token IDs or from the attention mask; never assume the last tensor position is the last real token.

Left padding changes the answer. With left padding and absolute position embeddings, the same text gets different position vectors, and the last tensor position is the last real token. If you switch conventions between training and inference, accuracy drops without any error.

Test with an altered padding. Change the padding tokens and check that the output is unchanged (Self-check 1 of the notebook). It is a cheap way to catch the wrong position and attention to padding.

Truncation. Texts longer than the context length must be truncated. For last-token pooling, truncate from the left (keep the end) or the right according to where the information is; document the choice. Cutting the end removes the position that matters.

Class order. Fix the mapping from class names to indices once and store it with the model. A swapped mapping gives a confidently wrong classifier.

Dropout on the pooled vector. A small dropout before the head helps when data are few; keep it off at inference.

For robots. The same head turns a language model into a command-intent classifier (stop, go, report) or a log-message triage tool. For safety-relevant intents, thresholds and a rejection class matter more than the architecture (Lesson LM6-T4).

Common misconceptions

Engineering takeaways

1. A classifier maps one vector per sequence to K logits with a linear layer: o = Wz + c, K·d + K parameters.
2. In a causal model use the hidden state of the last real token, position ℓ − 1; the first token has seen only itself.
3. With right padding, index each row separately: H[torch.arange(B), lengths − 1].
4. Mean pooling and special tokens are alternatives; mean pooling suits bidirectional models better.
5. Check that changing the padding does not change the output.

References and further reading

See Chapter 6 of the book (adding a classification head, and calculating the classification loss and accuracy).

• Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762.
• Build a Large Language Model (From Scratch), Chapter 6: Finetuning for Classification (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: which parameters to train and which to freeze, and what each choice costs.

Related content

Nothing below is required to finish this lesson.

• Before
  LM6-T1: Pretraining Versus Supervised Fine-Tuning

  Contrast the objectives, data and costs of pretraining and supervised fine-tuning.

• Lab
  LM6-L2: Fine-Tune a Language Model for Classification

  Add a classification head to a small GPT, compare frozen and fine-tuned variants and analyse the errors.