Text Classification with Language Models
Classification fine-tuning replaces the output layer of a large language model with a smaller classification layer.
A Text Generator Asked a Yes or No Question
A large language model is built to continue text, one token at a time. Suppose you want that same model to answer a different kind of question: is this text message spam or not spam? The model's ordinary output is a choice among tens of thousands of possible next tokens. That is the wrong shape for a two-way decision, so the model needs to be adapted before it can answer this kind of question well.
What do you think happens?
If you want a model to decide between two labels, spam and not spam, how many output nodes should its final layer have?
Reveal answer
Two
The new classification layer has one output node per class, and a spam versus not spam task has two classes, so it uses two nodes.
This article follows the model's output side step by step. The goal is to see what state the model is in before, during, and after the change, so that the idea of classification fine-tuning becomes concrete rather than a label you memorize.
The Original Output Layer Before Fine-Tuning
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 source gives the vocabulary size as 50,256, so the original layer has 50,256 nodes. Each node stands for one possible next token, which is why the model can continue any piece of text.
Replacing the Output Layer
Classification fine-tuning swaps the large 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 piece that changes so that the answers take the form of class labels.
Counting Output Nodes for a Task
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. Nothing else about the layer is fixed by the source, so the class count is the deciding factor.
Sizing the layer for a four-class task
A team wants a model that sorts news articles into four topics: sports, business, science, and arts. How many output nodes does the new classification layer need?
- 1Identify the classes
The labels the model may output are sports, business, science, and arts. That is four distinct classes.
- 2Apply the rule
The new classification layer has one output node per class, so the number of nodes equals the number of classes.
- 3Compare with the vocabulary
The vocabulary size does not decide the count here. The original layer would have 50,256 nodes for the vocabulary, but the classification layer is sized to the four topics.
The classification layer needs 4 output nodes, one for each topic label.
Next Token Versus Class Label
The change in 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.
| Aspect | Next-token prediction | Class-label output |
|---|---|---|
| Question answered | Which token comes next? | Which category does this input belong to? |
| Output layer size | One node per vocabulary token (50,256 in the source) | One node per class |
| Example output | A choice among possible next tokens | spam or not spam |
| Training target | The next token in the text | A correct class label |
Contrasting what the model predicts before and after classification fine-tuning
Consider an email filter built on a language model. Asked to continue a sentence, the model chooses among tens of thousands of possible next words. Once adapted, the same kind of model only needs to answer one of two labels for each message, spam or not spam, so its output is a small set of class choices instead of a word choice.
Mistakes When Sizing and Interpreting the Output
Common mistake Keeping the 50,256-node output layer for a classification task
What it looks likeUsing the original vocabulary-sized layer to decide between spam and not spam
Why it failsThat layer's outputs are next-token choices, which is the wrong shape for a two-way decision.
Do this insteadReplace the output layer with a classification layer sized to the task.
Common mistake Sizing the classification layer by the vocabulary
What it looks likeGiving a spam task 50,256 output nodes because the vocabulary has that many tokens
Why it failsThe source states that the new layer's node count equals the number of classes, not the vocabulary size.
Do this insteadCount the labels the model may output and use one node for each.
Common mistake Treating the classification output as a next-token prediction
What it looks likeReading the model's output for a message as the next word it would write
Why it failsAfter fine-tuning, the model is trained to output a correct class label.
Do this insteadInterpret the output as a choice among the defined class labels.
Practice Sizing a Classification Layer
A company wants a model to label customer feedback as positive, neutral, or negative. Before you answer, decide what the new classification layer replaces, then state how many output nodes it needs and what each node stands for.
Hints
- Start by counting the class labels the model is allowed to output.
- The new layer has one output node per class, not one per vocabulary token.
- Each node corresponds to exactly one label.
Key Takeaways
- Classification fine-tuning replaces a large language model's output layer with a smaller classification layer.
- The original output layer has one node per unique vocabulary token, which the source gives as 50,256.
- A classification layer has one output node per class, so a spam versus not spam task uses two nodes.
- Pretraining teaches the model to predict the next token, while classification fine-tuning trains it to output a correct class label.