Introduction

Softmax converts a vector of scores into a probability distribution whose values sum to 1. The implementation is mathematically short, but a numerically stable version is important because the naive formula can overflow when logits are large.

Basic Softmax Formula

Given logits z, softmax is:

• exponentiate each score
• divide by the sum of all exponentials

In Python with NumPy, the naive version looks like this:

1import numpy as np
2
3
4def softmax_naive(x):
5    exp_x = np.exp(x)
6    return exp_x / np.sum(exp_x)
7
8
9print(softmax_naive(np.array([2.0, 1.0, 0.1])))

This works for small values, but it is unsafe for real models because np.exp(1000) overflows.

Use a Numerically Stable Implementation

The standard fix is to subtract the maximum logit before exponentiation. This does not change the final probabilities, because softmax is invariant to adding or subtracting the same constant from every element.

1import numpy as np
2
3
4def softmax(x):
5    shifted = x - np.max(x)
6    exp_x = np.exp(shifted)
7    return exp_x / np.sum(exp_x)
8
9
10logits = np.array([2.0, 1.0, 0.1])
11print(softmax(logits))

[0.65900114 0.24243297 0.09856589]

Those values sum to 1, which is exactly what you want from a probability distribution.

Handle Batches Correctly

In machine learning code, you often apply softmax to a matrix where each row is one sample. That means the implementation must work along an axis.

1import numpy as np
2
3
4def softmax_batch(x):
5    shifted = x - np.max(x, axis=1, keepdims=True)
6    exp_x = np.exp(shifted)
7    return exp_x / np.sum(exp_x, axis=1, keepdims=True)
8
9
10logits = np.array([
11    [2.0, 1.0, 0.1],
12    [1.0, 3.0, 2.0],
13])
14
15print(softmax_batch(logits))

Using keepdims=True keeps the array shapes aligned for broadcasting and avoids shape bugs.

Why Softmax Is Used

Softmax is common in multiclass classification because it turns arbitrary real-valued scores into comparable class probabilities. The largest logit still maps to the largest probability, but the full output now has a probabilistic interpretation.

That said, softmax does not magically calibrate a model. A model can still be confidently wrong. The function only normalizes scores into a distribution.

Relationship to Cross-Entropy

In practice, softmax is often paired with cross-entropy loss. Many machine learning frameworks combine the two into one optimized operation for numerical stability. That is why you often use a built-in loss function instead of manually computing softmax during training.

For learning and debugging, though, implementing it manually is valuable because it makes the role of logits and normalization explicit.

Sanity Checks for an Implementation

A correct softmax implementation should satisfy a few simple checks:

• every output value is between 0 and 1
• the values sum to 1
• larger logits produce larger probabilities

You can verify that quickly:

result = softmax(np.array([1.0, 2.0, 3.0]))
print(result)
print(result.sum())

Those checks are simple, but they catch axis mistakes and normalization bugs faster than staring at formulas.

Common Pitfalls

• Implementing softmax directly with np.exp(x) and ignoring overflow risk.
• Applying the normalization over the wrong axis in batched input.
• Forgetting keepdims=True, which can cause subtle broadcasting errors.
• Interpreting softmax output as proof that the model is well calibrated.
• Reimplementing training-time softmax and cross-entropy manually when the framework already provides a more stable fused version.

Summary

• Softmax transforms logits into probabilities that sum to 1.
• The numerically stable version subtracts the maximum value before exponentiation.
• Batch implementations must normalize along the correct axis.
• Softmax is common in multiclass classification but does not guarantee calibrated confidence.
• Built-in framework losses often combine softmax with cross-entropy more safely for training.