The Softmax operation takes a vector of raw scores which are proven to be unnormalized log-probabilities (often called logits) and transforms them into a valid probability distribution. [1, 2]

Why It Works

To make any set of numbers look like a probability distribution, they must meet two rules: [1]

1. Positive values: Every number must be between 0 and 1.
2. Sum to 1: All numbers added together must equal exactly 1.0. [1, 2, 3]

How It Works (The Formula)

For any given score $z_{i}$ in a row:

$\text{Softmax}(z_{i})=\frac{e^{z_{i}}}{\sum _{j=1}^{n}e^{z_{j}}}$

- Step 1: Turn everything positive. It applies the exponential function ($e^{z}$) to every score. Even negative scores become positive numbers.
- Step 2: Normalize. It divides each exponential value by the sum of all the exponential values. This forces the final numbers to add up to 1. [1, 2, 3, 4, 5]

Quick Example

Imagine your model outputs three raw scores: [2.0, 1.0, 0.1]

1. Apply $e^{z}$:
  - - $e^{2.0} \approx 7.39$
    - $e^{1.0} \approx 2.72$
    - $e^{0.1} \approx 1.11$
    - Sum of these values = 11.22
2. Divide by the sum:
  - - $7.39 / 11.22 = \mathbf{0.66}$ (66% chance)
    - $2.72 / 11.22 = \mathbf{0.24}$ (24% chance)
    - $1.11 / 11.22 = \mathbf{0.10}$ (10% chance)

The final distribution is [0.66, 0.24, 0.10], which adds up perfectly to 1.

PROOF (if input to softmax is $z_i = \ln(p_i)$)

Using Multinomial Logits ($z_i = \ln(p_i)$)

Let our input probabilities be $p_1, p_2, p_3$ such that $p_1 + p_2 + p_3 = 1$. Let the input logits to the softmax function be defined as $z_i = \ln(p_i)$ for $i \in \{1, 2, 3\}$.

1. Apply the Softmax Formula [1]

The softmax output $q_{i}$ for each logit $z_{i}$ is defined as:

$q_{i}=\frac{e^{z_{i}}}{e^{z_{1}}+e^{z_{2}}+e^{z_{3}}}$

2. Substitute $z_i = \ln(p_i)$

Substitute the logit values into the exponential terms. Because $e^{\ln(p_i)} = p_i$, we get:

$q_{i}=\frac{e^{\ln (p_{i})}}{e^{\ln (p_{1})}+e^{\ln (p_{2})}+e^{\ln (p_{3})}}=\frac{p_{i}}{p_{1}+p_{2}+p_{3}}$

3. Simplify Using the Constraint

We are given that the probabilities sum to 1 ($p_1 + p_2 + p_3 = 1$). Substitute this into the denominator:

$q_{i}=\frac{p_{i}}{1}=p_{i}$

Thus, $q_1 = p_1$, $q_2 = p_2$, and $q_3 = p_3$.

Here is the structural proof showing why a linear layer’s output inherently models unnormalized log-probabilities. [1]

Step 1: Define the Generative Model (Bayes’ Theorem)

Let $x$ be the hidden vector from the last Transformer block, and let $y \in \{1, \dots, K\}$ be the target token class.

According to Bayes’ Theorem, the posterior probability that the token belongs to class $i$ given the representation $x$ is:

$$P(y=i \vert{} x) = \frac{P(x \vert{} y=i) \cdot P(y=i)}{P(x)}$$

Step 2: Apply the Exponential Family Assumption

In machine learning, we assume the conditional data distribution $P(x \vert{} y=i)$ belongs to the exponential family (e.g., a multivariate Gaussian distribution with a shared covariance matrix across classes).

This means we can write the conditional probability in its general exponential form:

$$P(x \vert{} y=i) = \exp\left( f_i(x) \right)$$

Where $f_i(x)$ is a function determined by the parameters of class $i$.

Step 3: Expand the Posterior Probability

Substitute this exponential form back into Bayes’ Theorem:

$$P(y=i \vert{} x) = \frac{\exp(f_i(x)) \cdot P(y=i)}{P(x)}$$

Using the identity $A = \exp(\ln(A))$, we can rewrite the prior probability $P(y=i)$ as $\exp(\ln(P(y=i)))$:

$$P(y=i \vert{} x) = \frac{\exp(f_i(x)) \cdot \exp(\ln(P(y=i)))}{P(x)}$$

Combine the exponents in the numerator:

$$P(y=i \vert{} x) = \frac{\exp\Big( f_i(x) + \ln(P(y=i)) \Big)}{P(x)}$$

Step 4: Map to the Final Linear Layer (Logits)

In a Transformer, the final layer is a linear projection (matrix multiplication plus bias) applied to the hidden state $x$. Let’s define the output of this linear layer for class $i$ as $z_i$:

$$z_i = w_i^T x + b_i$$

For a Gaussian exponential distribution, the term $f_i(x) + \ln(P(y=i))$ simplifies exactly into a linear equation matching this form, where the weights $w_i$ and bias $b_i$ absorb the distribution parameters and class priors.

Therefore, we can substitute $z_i$ directly into our numerator:

$$P(y=i \vert{} x) = \frac{\exp(z_i)}{P(x)}$$

Step 5: Eliminate the Total Probability

By the Law of Total Probability, the denominator $P(x)$ must equal the sum of all possible conditional outcomes in the numerator:

$$P(x) = \sum_{j=1}^{K} \exp(z_j)$$

Substituting this back gives us the standard Softmax function:

$$P(y=i \vert{} x) = \frac{\exp(z_i)}{\sum_{j=1}^{K} \exp(z_j)}$$

Step 6: Take the Logarithm to Extract $z_i$

Now, take the natural logarithm of the true probability $P(y=i\vert{}x)$:

$$\ln P(y=i \vert{} x) = \ln\left( \frac{\exp(z_i)}{\sum_{j=1}^{K} \exp(z_j)} \right)$$

$$\ln P(y=i \vert{} x) = z_i – \ln\left(\sum_{j=1}^{K} \exp(z_j)\right)$$

Rearranging to solve for the raw transformer output $z_i$:

$$z_i = \ln P(y=i \vert{} x) + \ln\left(\sum_{j=1}^{K} \exp(z_j)\right)$$

Conclusion

Because $\ln\left(\sum_{j=1}^{K} \exp(z_j)\right)$ is a constant scalar value that is identical for every single class token, it acts strictly as a normalization constant.

Therefore, $z_i$ is proven to be the true log-probability $\ln P(y=i\vert{}x)$ shifted only by an unnormalized global constant.

Does this generative perspective using Bayes’ Theorem make it clearer why the linear mathematical layer produces log-probabilities?

Using Binary Log-Odds Logits ($z_i = \ln\frac{p_i}{1-p_i}$)

If your definition of the logit for each class is the binary log-odds relative to its complement, then the input logits are:

$z_{1}=\ln \left(\frac{p_{1}}{1-p_{1}}\right),\quad z_{2}=\ln \left(\frac{p_{2}}{1-p_{2}}\right),\quad z_{3}=\ln \left(\frac{p_{3}}{1-p_{3}}\right)$

1. Evaluate the Exponential of the Logits [1]

First, compute the exponential of any logit $z_{i}$:

$e^{z_{i}}=e^{\ln \left(\frac{p_{i}}{1-p_{i}}\right)}=\frac{p_{i}}{1-p_{i}}$

2. Apply Softmax to these Logits [1]

The softmax output $q_{1}$ for the first class is:

$q_{1}=\frac{\frac{p_{1}}{1-p_{1}}}{\frac{p_{1}}{1-p_{1}}+\frac{p_{2}}{1-p_{2}}+\frac{p_{3}}{1-p_{3}}}$

3. Is this exactly equal to $p_{1}$?

In the multinomial case, using binary log-odds directly inside a standard softmax function does not generally map back to $p_{1}$ unless $K=2$. For $K > 2$ classes, the binary log-odds do not inherently preserve the multinomial structure inside a joint softmax without adjusting for the other classes.

There for it would not work if the output of transformer was what statisticians call logit

There are two places that SoftMax is used

Variable definitions

claude

Symbol	Meaning	Shape
$n$	sequence length (number of tokens)	scalar
$d_{\text{model}}$	hidden/embedding dimension	scalar
$h$	number of attention heads	scalar
$i$	head index, $i = 1,\dots,h$	—
$d_k$	dimension per head, $d_k = d_{\text{model}}/h$	scalar
$	V	$
$X$	input to an attention block (e.g., $\text{LN}(H^{l-1})$)	$n \times d_{\text{model}}$
$W_Q^{(i)}, W_K^{(i)}, W_V^{(i)}$	learned projection matrices for head $i$ (each head has its own)	$d_{\text{model}} \times d_k$
$Q_i, K_i, V_i$	query, key, value matrices for head $i$	$n \times d_k$
$S_i$	raw attention scores for head $i$	$n \times n$
$P_i$	normalized attention weights for head $i$	$n \times n$
$O_i$	output of head $i$	$n \times d_k$
$Z$	concatenation of all $h$ heads’ outputs	$n \times d_{\text{model}}$
$W_o$	attention output projection	$d_{\text{model}} \times d_{\text{model}}$
$A$	attention block output (before residual add)	$n \times d_{\text{model}}$
$H^L$	final hidden state, after all $L$ layers	$n \times d_{\text{model}}$
$W_{\text{vocab}}$	unembedding/output projection	$d_{\text{model}} \times
$\hat y$	predicted probability distribution over vocabulary	$n \times

Your statement is correct: there are two structurally different softmax operations in the forward pass — one repeated $h$ times per layer (attention weights, shape $n\times n$ each), and one applied a single time at the very end (vocabulary logits, shape $n\times|V|$). Full derivation of each below.

Softmax #1 — inside each attention head, on attention scores

This computation runs independently and in parallel for every head $i=1,\dots,h$, each with its own weight matrices. Nothing below is shared across heads until Step 5.

Step 1 — projections, for head $i$: $$Q_i = XW_Q^{(i)}, \qquad K_i = XW_K^{(i)}, \qquad V_i = XW_V^{(i)}$$

Purpose: project the shared input $X$ into three different learned subspaces — “what this token is looking for” ($Q_i$), “what this token offers” ($K_i$), and “what this token contributes if attended to” ($V_i$) — separately for each head, so each head can learn a different notion of relevance.

Step 2 — raw scores, for head $i$: $$S_i = \frac{Q_iK_i^\top}{\sqrt{d_k}} \qquad (\text{shape } n\times n)$$

Purpose: $S_{i,jk}$ measures the raw compatibility (dot-product similarity) between query token $j$ and key token $k$, within head $i$. Division by $\sqrt{d_k}$ rescales the dot product so its variance stays roughly constant as $d_k$ grows, preventing softmax from saturating into near-one-hot outputs (which would kill gradients).

Step 3 — softmax, row-wise, for head $i$: $$P_{i,,jk} = \frac{e^{,S_{i,,jk}}}{\displaystyle\sum_{k’=1}^{n} e^{,S_{i,,jk’}}} \qquad (\text{shape } n\times n)$$

Purpose: converts row $j$ of $S_i$ (raw, unnormalized compatibility scores against every token in the sequence) into a valid probability distribution — $P_{i,j,:}$ sums to 1, and each entry says “how much of token $k$’s value should token $j$ absorb, within head $i$.” This is the softmax your question refers to; its input is $S_i = Q_iK_i^\top/\sqrt{d_k}$, and its output shape is $n\times n$, never touching vocabulary size $|V|$.

Step 4 — weighted sum of values, for head $i$: $$O_i = P_i V_i \qquad (\text{shape } n \times d_k)$$

Purpose: each output token $j$ becomes a weighted average of all value vectors $V_i$, weighted by how much attention $P_{i,j,:}$ token $j$ pays to each other token. This produces one $O_i$ per head — $h$ separate matrices total.

Step 5 — recombine heads: $$Z = \text{Concat}(O_1, \dots, O_h) \qquad (\text{shape } n \times d_{\text{model}})$$

Purpose: stitch the $h$ independent per-head outputs back into a single $d_{\text{model}}$-wide representation ($h$ matrices of width $d_k$ concatenated $\rightarrow$ width $h\cdot d_k = d_{\text{model}}$). This is the only step where the per-head computations merge.

$$A = Z W_o \qquad (\text{shape } n \times d_{\text{model}})$$

Purpose: $W_o$ linearly mixes information across heads (since concatenation alone keeps each head’s contribution in disjoint column blocks — $W_o$ lets the model combine what different heads learned). $A$ then enters the residual stream of the current layer $l$: $H^l = H^{l-1} + A + F^l$, where $F^l$ is the FFN sub-block’s output.

Key property: this softmax happens $L \times h$ times per forward pass (once per head, per layer) — for a 4-layer, 8-head model, that’s 32 separate instances, each operating on an $n\times n$ matrix.

Softmax #2 — on vocabulary logits, once, at the very end

Step 1 — final hidden state. After all $L$ layers have run: $$H^L \qquad (\text{shape } n \times d_{\text{model}})$$

Step 2 — unembedding projection: $$\text{logits} = H^L W_{\text{vocab}} \qquad (\text{shape } n \times |V|)$$

Purpose: $W_{\text{vocab}}$ maps each token position’s $d_{\text{model}}$-dimensional hidden representation into a length-$|V|$ vector of raw scores — one score per vocabulary token, indicating (before normalization) how likely that token is to be the correct next token at that position.

Step 3 — softmax, row-wise, over vocabulary: $$\hat y_{i,v} = \frac{e^{,\text{logits}{i,v}}}{\displaystyle\sum{v’=1}^{|V|} e^{,\text{logits}_{i,v’}}} \qquad (\text{shape } n\times|V|)$$

Purpose: normalizes each position $i$’s raw logits into a valid probability distribution over the entire vocabulary — $\hat y_{i,:}$ sums to 1. This is the second, distinct softmax: its input is $\text{logits} = H^LW_{\text{vocab}}$ (shape $n\times|V|$), completely different in both source and shape from Softmax #1’s input.

Purpose of the full pipeline at this stage: $\hat y_{i,v}$ becomes the model’s predicted probability that vocabulary token $v$ is the correct token at position $i$, which feeds directly into the training loss: $$\mathcal{L} = -\frac{1}{n}\sum_{i=1}^n \log \hat y_{i,,y_i}$$ where $y_i$ denotes the true (ground-truth) token index at position $i$.

Key property: this softmax happens exactly once per forward pass — after the last layer’s output ($H^L$) is projected into vocabulary space by $W_{\text{vocab}}$.

Side-by-side contrast

	Softmax #1 (attention)	Softmax #2 (vocabulary)
Input	$S_i = Q_iK_i^\top/\sqrt{d_k}$	$\text{logits} = H^LW_{\text{vocab}}$
Input shape	$n \times n$	$n \times
Count per forward pass	$L \times h$	$1$
Row meaning	token $j$’s attention distribution over the $n$ sequence tokens, within head $i$	position $i$’s probability distribution over the $
Feeds into	$O_i = P_iV_i$, continuing within the same layer/head	$\mathcal{L}$ directly

Demosophy.org

Softmax and so Called Logits

PROOF (if input to softmax is \(z_i = \ln(p_i)\))