Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Scaled Dot-Product Attention › Softmax Over Attention Logits

Michael BrenndoerferMay 25, 20254 min read

Language AI HandbookAll chapters

Part 13 of 67Self-Attention

Chapter 92Scaled Dot-Product Attention › Softmax Over Attention Logits

Reading section 4 of 15 · Softmax Over Attention Logits · 4 minSection 4 of 15 · 4 min

Section 4 of 15 · 4 min read

Scaled Dot-Product Attention › Softmax Over Attention LogitsOne section · Sections in this chapter4 / 15 ↓

Sections in this chapter15

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

With scaled scores in hand, we apply softmax row-wise to convert each token's scores into a probability distribution over positions. The purpose of this step is to ensure the attention weights are non-negative and sum to one, which is necessary for the subsequent weighted sum of values to be a proper convex combination. Without softmax, the weights could be negative or unbounded, making the output an arbitrary linear combination rather than a meaningful weighted average.

The softmax is applied independently to each row of the scaled score matrix:

A = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right)

where the softmax is applied to each row independently. Entry $A_{ij}$ gives the attention weight: the fraction of position $j$ 's value that position $i$ incorporates into its output. The attention weights satisfy:

\sum_{j=1}^{n} A_{ij} = 1 \quad \text{for all } i

and all entries are non-negative ( $A_{ij} \geq 0$ ), making each row a proper probability distribution.

Attention Weights

Attention weights are the output of softmax applied to scaled dot-product scores. They form a probability distribution over all positions for each query position. A weight $A_{ij}$ close to 1 means position $i$ almost exclusively attends to position $j$ ; uniform weights indicate the output will be an even blend of all positions. The row-sum constraint ensures the output is a convex combination of value vectors, which keeps output magnitudes in a predictable range and aids in training stability.

The softmax has a critical secondary effect: it sharpens the distribution. Positions with high raw scores receive exponentially more attention than positions with low scores. After training, the model learns to assign high scores to relevant positions, producing peaked attention distributions that focus on key information while still being smooth enough to allow gradient-based learning. This sharpening effect is what gives attention its ability to be selective: a model that has learned to identify the antecedent of a pronoun will have a very peaked attention distribution when that pronoun is the current query, with nearly all weight on the antecedent token.

Softmax is applied row-wise, not to the entire matrix. Each query independently produces its own probability distribution over keys. This means different tokens can attend to entirely different parts of the sequence simultaneously, which is exactly the flexibility we need. Token $i$ 's attention distribution is completely independent of token $j$ 's, even though they are computed from the same key matrix.

In[10]:

Code

# Apply softmax to scaled scores
A = softmax(S_scaled, axis=1)  # row-wise softmax

Out[11]:

Console

Attention weight matrix A (each row sums to 1):
[[0.1234 0.2426 0.2145 0.0651 0.3545]
 [0.2248 0.1815 0.158  0.2645 0.1712]
 [0.2517 0.1531 0.1539 0.2708 0.1704]
 [0.1666 0.222  0.1869 0.1256 0.299 ]
 [0.1575 0.2351 0.2459 0.2124 0.1491]]

Row sums: [1. 1. 1. 1. 1.]
All weights non-negative: True

Each row is now a valid probability distribution. Row $i$ tells us how much each position contributes to token $i$ 's output representation. Looking at these numbers, you can read off how each position distributes its attention: does it focus on one position heavily, or spread its attention broadly? Both patterns are useful in different linguistic contexts, and the model learns which pattern to produce based on the content of the sequence.

Out[12]:

Visualization

Heatmap of 5x5 attention weights showing which positions each query token attends to. — Attention weight heatmap for a 5-token sequence after applying scaled dot-product attention. Each row shows how one query token distributes its attention across all key tokens, with darker cells indicating stronger attention. The weights in every row sum to 1, making each row a valid probability distribution over key positions. Notice how some tokens attend broadly while others show concentrated attention on specific positions.