The Softmax Function for Attention Weights

The calculation of scaled dot-product scores between Queries ($Q$) and Keys ($K$) produces a matrix of raw alignment scores. This computation yields $\frac{QK^T}{\sqrt{d_k}}$. Although these scores reflect the compatibility between query and key vectors, they are unnormalized and can span any range of values, making them difficult to interpret directly as contribution weights.

To transform these raw scores into a usable set of weights that represent a distribution of attention, we apply the softmax function independently to each row of the score matrix. For a specific query $q_i$ (corresponding to the $i$-th row of $Q$), the raw score for its alignment with key $k_j$ (corresponding to the $j$-th column of $K^T$) is denoted as $s_{ij} = \frac{q_i k_j^T}{\sqrt{d_k}}$. The softmax function converts the vector of scores $s_i = [s_{i1}, s_{i2}, ..., s_{iN}]$ for query $q_i$ across all $N$ keys into a vector of attention weights $\alpha_i = [\alpha_{i1}, \alpha_{i2}, ..., \alpha_{iN}]$ where each weight $\alpha_{ij}$ is calculated as:

$$\alpha_{ij} = \text{softmax}(s_{ij}) = \frac{\exp(s_{ij})}{\sum_{l=1}^{N} \exp(s_{il})}$$

Here, $N$ represents the sequence length of the Key/Value pairs.

Properties and Interpretation

Applying the softmax function yields several important properties for the resulting attention weights $\alpha_{ij}$:

1. Normalization: The denominator $\sum_{l=1}^{N} \exp(s_{il})$ ensures that the sum of all attention weights for a given query $q_i$ across all keys equals 1. That is, $\sum_{j=1}^{N} \alpha_{ij} = 1$. This allows us to interpret the weights as a probability distribution.
2. Non-negativity: Since the exponential function $\exp(x)$ is always positive for any real input $x$, each individual attention weight $\alpha_{ij}$ is guaranteed to be positive.
3. Probabilistic Interpretation: The set of weights $\{\alpha_{i1}, \alpha_{i2}, ..., \alpha_{iN}\}$ for a query $q_i$ represents the probability distribution of attention over the input sequence. The weight $\alpha_{ij}$ indicates the proportion of attention the model assigns to the $j$-th input element (represented by its value vector $v_j$) when computing the output representation for the $i$-th position.
4. Highlighting Importance: The exponential function inherently amplifies larger scores more than smaller ones. If one score $s_{ik}$ is significantly larger than others in the row $s_i$, its corresponding weight $\alpha_{ik}$ will be close to 1, while others will be close to 0. This mechanism allows the model to sharply focus on the most relevant input elements identified by the dot-product scoring.

Role in Computing the Final Output

These calculated attention weights $\alpha_{ij}$ are the coefficients used to compute a weighted sum of the Value vectors ($V$). The attention mechanism's output for the $i$-th query is obtained by multiplying the attention weight distribution $\alpha_i$ with the Value matrix $V$:

$$\text{output}_i = \sum_{j=1}^{N} \alpha_{ij} v_j$$

In matrix notation, this corresponds directly to the final step in the scaled dot-product attention formula:

$$\text{Attention}(Q, K, V) = A V = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Where $A$ is the attention weight matrix whose elements are $A_{ij} = \alpha_{ij}$. The softmax function is thus fundamental for converting the raw similarity scores into a normalized distribution that dictates how information from different parts of the input sequence (represented by the Value vectors) is aggregated to form the output.

Visualization Example

Consider a simplified case with 4 raw alignment scores for a single query: $s = [1.0, 0.5, 2.5, -0.1]$. Applying the softmax function transforms these scores:

• $\exp(1.0) \approx 2.718$
• $\exp(0.5) \approx 1.649$
• $\exp(2.5) \approx 12.182$
• $\exp(-0.1) \approx 0.905$

The sum of exponentials is $2.718 + 1.649 + 12.182 + 0.905 = 17.454$.

The resulting softmax weights are:

• $\alpha_1 = 2.718 / 17.454 \approx 0.156$
• $\alpha_2 = 1.649 / 17.454 \approx 0.094$
• $\alpha_3 = 12.182 / 17.454 \approx 0.698$
• $\alpha_4 = 0.905 / 17.454 \approx 0.052$

Notice how the highest raw score (2.5) corresponds to the dominant attention weight (0.698), effectively focusing the attention mechanism on the third input element in this example. The weights sum to approximately 1 ($0.156 + 0.094 + 0.698 + 0.052 = 1.000$).

Comparison of raw alignment scores and the resulting attention weights after applying the softmax function for a single query. Softmax converts scores into a probability distribution, highlighting the most relevant positions.

In summary, the softmax function acts as an important normalization step within the attention mechanism. It transforms raw alignment scores into a probability distribution, enabling the model to selectively weight and combine information from the Value vectors based on Query-Key similarities, forming the foundation for how Transformers process sequential information.