Softmax Function

Why we need it (motivation)

In many ML systems, we compute scores for several options: which class is present, which token to attend to, which action to take. Those scores often live in ℝ: they can be negative, huge, and not constrained to sum to 1.

But downstream we often want a probability distribution:

• •Nonnegative values (so they can represent probabilities)
• •Sum to 1 (so they distribute total mass)
• •Smooth and differentiable (so gradient-based learning works)

Softmax is the standard way to convert a vector of real-valued scores (“logits”) into a probability distribution.

Definition

Let x = (x₁, x₂, …, xₙ) be a vector of real numbers (logits).

The softmax function returns a vector p = softmax(x) where each component is

Intuition: “exponentiate then normalize”

Softmax does two simple things:

1. 1)Exponentiate each logit: $x_i \mapsto e^{x_i}$

• •This makes everything positive.
• •It also makes differences matter: a logit that’s larger by 2 becomes $e^2 \approx 7.39$ times bigger after exponentiation.

2. 2)Normalize by the sum: divide by ∑ⱼ e^{xⱼ}

• •This forces the outputs to sum to 1.

So softmax turns relative score gaps into relative probability mass.

A quick sanity check: it’s a probability distribution

For each i:

• •$e^{x_i} > 0$ ⇒ softmax(x)ᵢ > 0
• •Summation property:

So softmax(x) lies on the probability simplex (the set of all probability vectors).

Terminology you’ll see

• •Logits: the input scores x.
• •Probabilities: the output softmax(x).
• •Attention weights: in attention, softmax converts similarity scores into weights over tokens.

Softmax is simple to write down, but its behavior (and pitfalls) matter a lot in real models—especially numerical stability and temperature scaling, which we’ll build up next.

Why exponentials?

Exponentials have two key effects:

1. 1)Positivity: $e^{x_i}$ is always positive.
2. 2)Multiplicative amplification: differences in logits turn into ratios.

A crucial identity is the ratio form:

This says softmax compares logits via their differences. If $x_i$ exceeds $x_k$ by Δ, then i gets $e^{\Delta}$ times more probability than k.

Two-class case: softmax becomes sigmoid

If n = 2 with logits (a, b), then

That’s exactly a sigmoid in the logit difference (a − b). This is a nice mental model:

• •softmax is the “multi-class sigmoid.”

Invariance to units? Not quite.

Softmax is not invariant to scaling of logits. If you multiply logits by a constant c, softmax typically becomes more or less peaked (we’ll formalize this with temperature later).

Peakedness: how “winner-take-most” emerges

Consider three logits: x = (2, 1, 0).

Compute exponentials:

• •e² ≈ 7.39
• •e¹ ≈ 2.72
• •e⁰ = 1

Sum ≈ 11.11

So probabilities ≈ (0.665, 0.245, 0.090).

A gap of 1 between logits becomes a factor of e ≈ 2.72 in weight; a gap of 2 becomes e² ≈ 7.39. This is why softmax can produce strong preferences even from modest logit gaps.

Geometric view: softmax outputs live on the simplex

For n = 3, the output probabilities (p₁, p₂, p₃) satisfy p₁ + p₂ + p₃ = 1 and each pᵢ ≥ 0. That set is a 2D triangle (a simplex) embedded in 3D.

Here’s an ASCII simplex diagram to orient you:

          p3=1
           ▲
          / \
         /   \
        /  •  \   • interior points: all p_i in (0,1)
       /       \
      /         \
     /___________\
 p1=1             p2=1

• •Vertices correspond to “certain” distributions like (1,0,0).
• •The center corresponds to uniform (1/3,1/3,1/3).

Softmax maps any logits vector x to some point inside this triangle.

Visualization: temperature effect on a 2-option softmax curve

For two options, softmax probability of option 1 depends on the logit difference d = x₁ − x₂:

Below is an inline diagram showing how changing T changes the curve. The horizontal axis is d, vertical is p₁.

p1
1.0 |                         ............  T=0.5 (sharper)
    |                    .....
0.8 |               .....
    |           ....
0.6 |        ...                    _________  T=1 (baseline)
    |     ...                 _____
0.5 |-----+-------------------+----------------------------- d
    |     ...             ____
0.4 |        ...      ____                 - - - - - - - -  T=2 (flatter)
    |           ....__
0.2 |               .....
    |                    .....
0.0 |                         ............
      -6   -4   -2    0    2    4    6

Interpretation:

• •Lower T: transitions faster from 0 to 1 → more “confident.”
• •Higher T: transitions slower → more “uncertain.”

We’ll connect this to attention weights: low temperature makes attention concentrate on a few tokens; high temperature spreads it out.

Practical note: softmax is often applied row-wise

In attention, you’ll see softmax applied to a vector of scores for a given query over all keys. If you have a matrix of scores, softmax is applied per row (or per last dimension), producing a distribution over positions for each query.

This first mechanic—exponentiate then normalize—gives the core behavior. Next we’ll cover the crucial property that makes softmax usable in real systems: shift invariance and numerical stability.

Why this matters

Exponentials can overflow or underflow:

• •$e^{1000}$ is astronomically large (overflow in float32/float64).
• •$e^{-1000}$ is essentially 0 (underflow).

Yet logits in neural nets can easily reach magnitudes where naive exp() is unsafe. So we need a stable way to compute softmax.

Key property: shift-invariance

Softmax is unchanged if you add the same constant c to every logit:

Derivation (showing work):

Let $y_i = x_i + c$.

So adding a constant doesn’t change the output probabilities.

The numerical-stability trick: subtract max

Because of shift-invariance, we can choose c conveniently. The most common choice is

Define $m = \max_i x_i$ and $z_i = x_i - m$.

Then $\max_i z_i = 0$, so every $z_i \le 0$.

Now compute softmax using z:

This is stable because:

• •The largest exponent is $e^0 = 1$ (safe).
• •Others are $e^{\text{negative}} \in (0,1]$ (also safe).

Simple example: stability without changing meaning

Suppose x = (1000, 1001, 999).

Naively, e^{1001} overflows.

Use max trick: m = 1001

• •z = (-1, 0, -2)

Exponentials:

• •e^{-1} ≈ 0.3679
• •e^{0} = 1
• •e^{-2} ≈ 0.1353

Sum ≈ 1.5032

Probabilities ≈ (0.2447, 0.6652, 0.0900)

These are perfectly reasonable—no overflow.

Visualization: shifting logits moves nothing on the simplex

Shifting logits by a constant slides x along the direction 1 = (1,1,1,…). Softmax “forgets” that direction completely.

For n=3, imagine two different logit vectors:

• •x = (2, 1, 0)
• •x' = (2+10, 1+10, 0+10) = (12, 11, 10)

They map to the exact same point (p₁, p₂, p₃) on the simplex triangle.

Here’s a conceptual diagram combining both ideas—shift vs. scale:

Simplex (n=3 probabilities)

          (0,0,1)
             ▲
            / \
           /   \
          /  A  \        A = softmax(x)
         /       \       softmax(x + 10·1) = A  (shift: unchanged)
        /    •    \      softmax(x / T) moves toward vertex or center (scale)
       /___________\
 (1,0,0)           (0,1,0)

- Shift logits: stay at the same point A.
- Scale logits (or change T): slide along a path toward a vertex (peaked) or toward center (uniform).

Implementation note (what you should do in code)

Always compute softmax as:

1) $m = \max_i x_i$

2) $z_i = x_i - m$

3) $p_i = \exp(z_i) / \sum_j \exp(z_j)$

This gives identical results in exact math, and far better results in floating-point.

A related stable quantity: log-softmax

Often you want log probabilities (e.g., for cross-entropy). Use:

Stably, compute:

• •$m = \max_j x_j$
• •$\log\sum_j e^{x_j} = m + \log\sum_j e^{x_j - m}$ (this is the log-sum-exp trick)

Even if you don’t implement it now, it’s important conceptually: stability is not optional when exponentials are involved.

Next we’ll look at temperature scaling, which is like a controlled scaling of logits that changes the softness/hardness of the distribution.

Why introduce temperature?

Sometimes you want probabilities that are:

• •Sharper (more peaked) so the model strongly prefers one option.
• •Flatter (more spread out) so the model remains uncertain or explores alternatives.

Temperature scaling gives a single knob T > 0 that controls this.

Definition

Given logits x, temperature-scaled softmax is

Equivalent viewpoint: dividing by T is like multiplying logits by $\alpha = 1/T$.

• •T = 1 → standard softmax.
• •T < 1 → logits are effectively magnified → sharper.
• •T > 1 → logits are effectively shrunk → flatter.

Limiting behavior (important intuition)

Let p(T) = softmax(x/T).

1) As T → 0⁺:

• •The largest logit dominates.
• •p approaches a one-hot distribution at argmax.

2) As T → ∞:

• •All logits become tiny relative to T.
• •Exponentials become similar.
• •p approaches uniform: $p_i → 1/n$.

You can see this via differences: ratios are

• •If T is small, (xᵢ − x_k)/T is large in magnitude ⇒ ratios explode ⇒ one option dominates.
• •If T is large, (xᵢ − x_k)/T ≈ 0 ⇒ ratios near 1 ⇒ uniform-ish.

Temperature in attention

In dot-product attention, scores often look like

Then attention weights are

The $1/\sqrt{d}$ factor plays a temperature-like role: it prevents dot products from growing too large with dimension d (which would make softmax too peaked too early).

Visual: how T moves you on the simplex (n=3)

Take logits x = (2, 1, 0). Consider three temperatures.

Compute probabilities:

• •T = 2: softmax( (1, 0.5, 0) )
• •exp: (2.718, 1.649, 1)
• •sum: 5.367
• •p ≈ (0.506, 0.307, 0.186) (flatter)

• •T = 1: softmax( (2, 1, 0) )
• •p ≈ (0.665, 0.245, 0.090)

• •T = 0.5: softmax( (4, 2, 0) )
• •exp: (54.598, 7.389, 1)
• •sum: 62.987
• •p ≈ (0.867, 0.117, 0.016) (peaked)

On the simplex triangle, these three points lie along a path from the center-ish region toward the vertex (1,0,0) as T decreases.

Calibration note (probabilities vs confidence)

Temperature scaling is also used for calibration: you can adjust T (often on a validation set) so predicted probabilities better match empirical accuracy.

• •If a classifier is overconfident, increasing T (T > 1) can reduce peakiness.
• •If underconfident, decreasing T can sharpen predictions.

This is a big reason softmax is interpreted carefully: the raw logits contain information beyond just the top class.

At this point you know:

• •What softmax is.
• •How to compute it stably.
• •How temperature changes its behavior.

Next we connect it directly to attention mechanisms, masking, and how to interpret attention scores.

Softmax as “attention allocator”

In attention, you compute a score for each key/value relative to a query. These scores are logits s.

Softmax turns them into weights a that sum to 1:

Then the attention output is a weighted sum:

So softmax is the mechanism that converts similarities into a convex combination of values.

How to read attention weights

Because $a_i \ge 0$ and ∑ᵢ aᵢ = 1:

• •Each aᵢ is a fraction of attention mass.
• •The output is inside the convex hull of the value vectors.

But interpret carefully:

• •Attention weights reflect relative importance under the model’s scoring function.
• •Small changes in logits can cause large changes in weights when the distribution is already sharp (especially at low T).

Masking: forcing probabilities to ignore some positions

In sequence models you often must prevent attending to:

• •padding tokens (padding mask)
• •future tokens (causal mask)

The standard technique: add a large negative number (−∞ in math; a big negative constant in practice) to masked logits before softmax.

Let mask mᵢ be 0 for allowed, and −∞ for disallowed. Define

Then

• •if mᵢ = −∞ ⇒ $e^{s'_i} = 0$ ⇒ probability becomes 0.
• •allowed positions renormalize to sum to 1.

This works because softmax only cares about exponentials; setting a logit to −∞ removes it from the sum.

Numerical detail: choose a safe “−∞”

In floating point, you use something like −1e9 (float32) or a framework-provided mask fill value.

• •Too small in magnitude: masked positions may still get nonzero probability.
• •Too large in magnitude: can cause NaNs if combined with other operations (less common if you use stable softmax).

Connection to cross-entropy and learning signals

Softmax is commonly paired with cross-entropy loss.

If the true class is k and predicted probabilities are pᵢ, then

When the model assigns low probability to the correct class, the loss is large.

A key internal quantity is log-sum-exp:

This is one reason stable log-softmax implementations are so common.

Interpreting logits vs probabilities

Logits contain “un-normalized evidence.” Softmax converts them to probabilities, but:

• •Probabilities can saturate near 0 or 1 (especially at low T), hiding meaningful logit differences.
• •Comparing logits across different contexts can be tricky; softmax probabilities are context-dependent because the denominator includes all options.

Summary of when softmax is the right tool

Use softmax when you need:

• •a distribution over mutually exclusive categories, or
• •nonnegative weights summing to 1 (attention, mixture weights).

Avoid or reconsider when you need:

• •independent multi-label probabilities (use sigmoid per label instead), or
• •hard argmax choices during training (softmax gives a smooth proxy).

With these connections, you’re ready to use softmax as a dependable building block for attention mechanisms, masking, and sequence-to-sequence models.