Softmax

Core Concepts

Softmax converts a vector of arbitrary real numbers (logits) into a probability distribution: all values become non-negative and sum to 1. TokenPrint uses softmax in two distinct places — inside attention and at the output.


Formula

softmax(x_i) = exp(x_i) / sum_j( exp(x_j) )

The exponential amplifies differences between logits non-linearly. A logit 2 units above another receives approximately e^2 ≈ 7.4x more probability mass.


In attention

After computing the scaled dot products Q K^T / sqrt(64), softmax is applied row-wise across the key dimension to produce attention weights. Each row sums to 1.00 — verified independently for the reference model (row sums: 1.0000).


At the output

After the final RMSNorm and the unembedding projection, the model produces a vector of 151,936 logits — one per vocabulary token. Softmax converts these into probabilities. The token with the highest probability is the model's greedy prediction.


Greedy decoding

TokenPrint's Generation mode uses greedy decoding: the token with the highest logit is always selected as the next token. This is deterministic and reproducible. The backend uses real torch.argmax over real logits — no sampling or temperature is applied unless explicitly configured.


TokenPrint visualization

The Softmax & Output chapter in Walkthrough mode shows a bar chart of the top-k predicted tokens with real softmax probabilities for the example sentence. The argmax (the model's actual prediction) is highlighted. The prediction game lets you guess the next token before revealing the real distribution.