Implementing Softmax and Argmax for Class Labels

In this classification task the only difference is the size of the vector: there is one value per class, so the row is 2-dimensional, one value for not spam and one for spam, rather than 50,257 values.

Programming

From Raw Scores to a Spam Decision

Imagine a model that reads a text message and must answer one question: is this message spam or not spam? The model does not write a sentence back. It picks one of two labels. That narrow task changes what the model's outputs mean and how they are turned into a decision. This article follows the path from the model's raw numbers to a final label, and then asks a sharper question: when the model is adapted for this task, does it matter which token's output is used for the decision?

The exercise this article is built around sits at the end of a chapter that adapts a pretrained language model to sort text messages into spam and not spam. The main setup is already in place. The exercise changes one thing: instead of fine-tuning the output for the last token in the input, you fine-tune the output for the first token, and then you observe how the predictions change.

Why the Output Row Is a Design Choice

A language model does not produce one answer for a whole message. It produces a row of outputs for each token it reads. For classification, the chapter takes one of those rows and treats it as the summary that the final decision depends on. The figure that accompanies the exercise uses the row for the last token. The open question is whether a different row, the one for the first token, gives a different quality of prediction once the model has been fine-tuned on it.

This is why the choice of token is a real design decision. The model is not being asked to produce new text, so the only thing that determines the label is which row feeds the decision. Changing that row can change the labels that come out, even when the dataset and the rest of the pipeline stay the same.

Turning Scores into Probabilities with Softmax

The model's raw outputs are not probabilities. They are scores that can be positive or negative. The softmax function converts a row of such scores into probabilities. The chapter recalls that the same method was used earlier, when the model picked its next token from 50,257 possible outputs. In this classification task the only difference is the size of the vector: there is one value per class, so the row is 2-dimensional, one value for not spam and one for spam, rather than 50,257 values.

Formula

Read the formula as follows. Each raw score z_i is exponentiated, so every value becomes positive. Each result is then divided by the sum of all the exponentiated scores in the row. Because every output is divided by the same total, the resulting probabilities add up to 1. Larger raw scores receive larger shares of that total.

Worked example

Softmax on a Small Two-Class Row

Suppose a row of raw scores is [2.0, 1.0] for the classes not spam and spam. Convert it to probabilities with softmax.

1
Exponentiate each score
e raised to 2.0 is about 7.389, and e raised to 1.0 is about 2.718.
2
Sum the exponentiated values
7.389 plus 2.718 is about 10.107.
3
Divide each value by the sum
7.389 divided by 10.107 is about 0.731, and 2.718 divided by 10.107 is about 0.269.
4
Check the total
0.731 plus 0.269 is 1.0, so the outputs form a valid probability distribution.

Result

Probabilities of about [0.731, 0.269]. The first class has the larger share, so it is the more likely label for this row.

Picking the Winner with Argmax

After the probabilities are computed, the chapter needs one decision for each input. The argmax function provides it. Argmax looks at the row of probabilities and returns the index position that holds the highest value. That index is the predicted class. Because the classes are stored by position, finding the largest value tells you which class wins. Once the index is known, the probability value itself is no longer needed to produce the label.

In the diagram, the row holds one probability per class. Argmax reports index 1, the position with the highest value, and that position is labeled spam. The row is read by position, so the meaning of each index must be fixed by the setup, not recalculated for each message.

Tracing Two Messages Through the Pipeline

Follow a message through the whole path. The second message in the chapter is 'Do you have time'. Its raw output row is [-3.9846, 5.2940]. The first number is negative, which is allowed for raw scores. After softmax, this row becomes approximately [0.01, 0.99]. Argmax then returns index 1, which the chapter labels as spam.

Code walkthrough

import numpy as np

def softmax(z):
    z = np.asarray(z, dtype=float)
    e = np.exp(z - np.max(z))
    return e / e.sum()

raw = [2.0, 1.0]
probs = softmax(raw)
label_index = int(np.argmax(probs))
print(probs, label_index)

1
Define softmax so it takes any row of raw scores and returns probabilities that sum to 1.
2
Pass the row [2.0, 1.0] through softmax to get about [0.731, 0.269].
3
Call argmax on the probabilities; index 0 holds the largest value, so the predicted class is index 0.

softmax subtracts the maximum score before exponentiating; this does not change the result but avoids very large exponentials
np.argmax returns the index of the largest probability

Output

[0.73105858 0.26894142] 0

Both messages go through identical steps. Only the numbers differ. This sameness is important for the experiment later: the code path is fixed, so any difference in predictions has to come from the data the model produces, not from changing the decision logic.

Comparing First and Last Token Targets

The exercise asks you to change one variable: which output token is fine-tuned. The baseline setup reads the decision from the output row for the last token. The experiment reads it from the output row for the first token instead. The dataset, the data loaders, the softmax and argmax steps, and the evaluation utilities all stay fixed. Holding these constant means any difference in performance can be attributed to the choice of token.

Aspect	Last output token	First output token
Role in the setup	Baseline target shown in the main figure	Alternative target tested in the exercise
Which row feeds the decision	Output row for the last token	Output row for the first token
What stays fixed	Dataset, loaders, softmax, argmax, evaluation utilities	Dataset, loaders, softmax, argmax, evaluation utilities
Known result from the source	Used as the reference point	Not stated in the source; must be measured

The only intended difference between the two runs is which output token is fine-tuned.

Mistakes When Reading Softmax and Argmax Output

Common mistake
Treating raw model outputs as probabilities
What it looks like
Reading a raw score of 5.2940 as a 529 percent chance of spam.
Why it fails
Raw scores can be negative and do not sum to 1. Softmax is the step that produces probabilities.
Do this instead
Apply softmax first, then interpret the values as probabilities.
Common mistake
Using the probability value as the label
What it looks like
Reporting 0.99 as the class instead of spam.
Why it fails
The label comes from the position of the largest value, not from the value itself.
Do this instead
Use argmax to get the index, then map that index to its class name.
Common mistake
Assuming the index order is arbitrary
What it looks like
Mapping index 1 to not spam because it seems like a neutral choice.
Why it fails
The index-to-label mapping is fixed by the setup. Changing it flips every prediction.
Do this instead
Confirm the mapping used in the chapter before reading any predictions.
Common mistake
Skipping the evaluation utilities and judging the token choice by eye

Designing the First-Token Experiment

What do you think happens?

Given the probability row [0.20, 0.80] for the classes not spam and spam, which label does argmax return?

not spam (index 0)
spam (index 1)

Reveal answer

Answer:

spam (index 1)

Argmax returns the index holding the highest value. 0.80 is larger than 0.20 and sits at index 1, which the chapter labels spam.

MEDIUM

Design the experiment that compares the first and last output tokens as fine-tuning targets. List what you will keep fixed, what you will change, and which evaluation utilities you will use to measure the difference in predictive performance.

Hints

Only one variable should change between the two runs.
The evaluation utilities must be the same for both runs so the comparison is fair.
Do not state the expected result before you have measured it.

Your attempt (not saved)

Key Takeaways

A pretrained language model outputs a row of values for each token, and classification reads one chosen row as the basis for its decision.
Softmax converts raw, possibly negative scores into probabilities that sum to 1 while keeping the largest score the largest value.
Argmax returns the index of the highest probability, and that index maps to the predicted class label.
The exercise compares fine-tuning the first output token against fine-tuning the last output token, holding everything else fixed.
The source does not report the result of that comparison, so it must be measured with the same evaluation utilities for both runs.

Related concepts

More concepts starting with I