Output Layers in Neural Networks

A large language model is built to continue text, one token at a time. Its ordinary output is a choice among tens of thousands of possible next tokens. That is a sensible shape for writing a sentence, but it is the wrong shape for a simple decision such as whether a text message is spam or not spam. Classification fine-tuning is the process that adapts the model to this kind of task. In this article we follow the model's output side step by step, so you can see what state it is in before, during, and after the change.

The output layer is the part of the model that turns its internal computation into a final result. In the original setup, this layer has one output node for every unique token in the vocabulary. The vocabulary size used here is 50,256, so the original layer has 50,256 nodes.

Notice what this means for the state of the model before any fine-tuning. Its output is not a verdict. It is a score over tens of thousands of candidate tokens, and whichever candidate scores highest becomes the next word the model produces.

Classification fine-tuning swaps the large original output layer for a new, much smaller one. The new layer is sized to the task rather than to the vocabulary. The rest of the model is the part being adapted for the task, while the output layer is the part that changes so that answers take the form of class labels.

Read the diagram as a before and after of one component. The body of the model is carried over, and the output layer is the piece that changes. In the before state the output layer has 50,256 nodes. In the after state, for a two-class task, it has 2 nodes. The number of nodes is the only size change this section describes; the source does not specify other changes to the layer.

The number of output nodes in the new classification layer equals the number of classes in the task. Each node corresponds to one label the model is allowed to output. The class count is the deciding factor, so the first step in designing the layer is to list the labels.

The diagram shows the spam case from the source. There are two labels, so there are two output nodes, and each node is tied to one label. Which node takes which label is a convention chosen when the model is set up; what the source fixes is that each node corresponds to exactly one class.

The change in the output layer also changes what the model is trained to produce. During pretraining, the model predicts the next token in the text. Classification fine-tuning instead trains the model to output a correct class label, such as spam or not spam. The question the model answers moves from what word comes next to which category this input belongs to.

Trace a single text message through the fine-tuned model. The text enters the model body, where it is processed into an internal computation. That computation reaches the classification layer, whose nodes correspond to the labels. The label tied to the highest-scoring node is the model's predicted label. The output is a class, not a continuation of the text.