Replacing the Output Layer for Classification

After text is converted into token IDs and passed to the model, the model returns numbers rather than a direct answer such as spam or not spam. To prepare the model for classification, replace its final output layer so it produces a number of outputs that matches the task. Then select the output row used for the input's classification decision.

Suppose the model output has shape 1 x 4 x 2. Read the dimensions in order: the first dimension is the batch, which contains one input here; the second dimension has one row for each of the four input tokens; and the third dimension gives the number of values produced for each token. In this example, each token row contains two values. These rows are numeric representations, not word labels or class names.

The last dimension reveals which output layer is in use. With the original output layer, the example shape would instead be 1 x 4 x 50257. That layer produces one value for each vocabulary entry for every token position. After replacing it for a two-class task, the final dimension is 2 rather than 50257.

Read the slice from left to right, matching each part to an output dimension. The first colon keeps the batch dimension. The -1 selects the final position on the token axis. The last colon keeps all values in that token's output row. Starting from a tensor shaped 1 x 4 x 2, this selection gives a result shaped 1 x 2: one batch item with two values.

The classification setup uses one decision per input, so it selects one token row rather than asking for a separate classification decision at every token position. The final row is the chosen readout point in this lesson. This selection sets up classification fine-tuning; it does not itself convert the two numbers into the words spam or not spam.

The original output layer maps each token's internal vector to 50,257 values, one for every entry in the model's vocabulary. Classification between two categories does not need a score for every vocabulary entry. Set the replacement layer's output width to the task's number of classes instead.

The assignment replaces model.out_head and leaves the rest of the network unchanged. The Linear layer's in_features argument sets the width it expects from the preceding model layers, so it uses the model's embedding dimension. The out_features argument sets how many values the new layer produces, so it uses num_classes. With num_classes set to 2, each token's output row has two values.

Using BASE_CONFIG["emb_dim"] instead of writing 768 directly connects the replacement layer to the selected model's configuration. The source gives 768 as the embedding dimension for gpt2-small and explains that reading the value from configuration lets the same replacement code work with larger GPT-2 variants whose embedding dimensions are stored in their configurations.

A replacement line may be accepted before a size problem becomes apparent. The problem usually shows up later, when data first passes through the new layer or when a shape is inspected. Trace from the replacement to its first use, then check the output shape. If the full output has the expected dimensions but the selected result does not, investigate the indexing; if the full output is already unexpected, investigate the model output or replacement layer first.