Scaled Dot-Product Attention Formula

Learning Objectives

By the end of this section, you will:

1. Derive the dot-product attention formula step by step
2. Understand why scaling by √d is necessary for stable training
3. Apply softmax normalization to obtain attention weights
4. Master the complete attention formula
5. Trace through a worked numerical example

Why This Matters: Scaled dot-product attention is the core building block of modern attention mechanisms, from Transformers to our AMNL model. Understanding its mathematical formulation reveals why it works and how to debug issues like vanishing gradients or attention collapse.

Dot-Product Attention

The simplest way to measure query-key similarity is the dot product.

Similarity via Dot Product

For a query $q$ and key $k$, both vectors in $\mathbb{R}^{d_k}$:

Properties of the dot product:

• Positive: Similar directions → high value
• Negative: Opposite directions → low value
• Zero: Orthogonal → no similarity

All Pairs at Once

For all queries and keys, compute similarities via matrix multiplication:

Where $Q \in \mathbb{R}^{T \times d_k}$ and $K \in \mathbb{R}^{T \times d_k}$. Each element $(i, j)$ gives the similarity between query i and key j.

1Score matrix (T=30 timesteps):
2
3         Key positions (j)
4         1    2    3   ...   30
5       ┌─────────────────────────┐
6Q   1  │ s₁₁  s₁₂  s₁₃  ...  s₁₃₀ │  ← Query 1's similarity to all keys
7u   2  │ s₂₁  s₂₂  s₂₃  ...  s₂₃₀ │
8e   3  │ s₃₁  s₃₂  s₃₃  ...  s₃₃₀ │
9r   :  │  :    :    :         :   │
10y  30  │ s₃₀₁ s₃₀₂ s₃₀₃ ... s₃₀₃₀ │
11       └─────────────────────────┘
12
13Each row becomes attention weights after softmax

Why Scale by √d?

Raw dot products can become very large, causing problems with softmax. Scaling fixes this.

The Variance Problem

If $q$ and $k$ have independent components with mean 0 and variance 1:

The variance of the dot product grows with dimension!

The Scaling Solution

Divide by $\sqrt{d_k}$ to restore unit variance:

Now scores have consistent magnitude regardless of dimension.

Impact on Softmax

Softmax Normalization

Softmax converts raw scores into a probability distribution over keys.

Softmax Formula

For score vector $\mathbf{s} = [s_1, ..., s_T]$:

Properties of softmax outputs:

• $\alpha_t > 0$ for all t (positive)
• $\sum_t \alpha_t = 1$ (normalized)
• Higher scores → higher weights (monotonic)
• Smooth and differentiable (trainable)

Row-wise Application

Softmax is applied independently to each row of the score matrix:

Row i contains query i's attention distribution over all keys.

Complete Attention Formula

Combining all components, the scaled dot-product attention formula is:

Breakdown

1. Compute scores: $S = QK^T$ — all pairwise similarities
2. Scale: $S' = S / \sqrt{d_k}$ — stabilize variance
3. Normalize: $A = \text{softmax}(S')$ — row-wise probabilities
4. Aggregate: $\text{Output} = AV$ — weighted sum of values

Dimensions

For our model: T = 30, dₖ = dᵥ = 32 (per head).

Worked Example

Let us trace through attention with concrete numbers.

Summary

In this section, we derived scaled dot-product attention:

1. Dot-product similarity: $QK^T$ measures query-key alignment
2. Scaling: Divide by $\sqrt{d_k}$ to prevent variance explosion
3. Softmax: Converts scores to probability distribution
4. Complete formula: $\text{softmax}(QK^T/\sqrt{d_k}) \cdot V$
5. Output: Weighted combination of values based on query-key similarity

Looking Ahead: A single attention head learns one type of query-key relationship. To capture multiple types of relationships (different aspects of degradation), we use multi-head attention—running multiple attention mechanisms in parallel and combining their outputs.

With single-head attention understood, we now extend to multi-head attention with 8 parallel heads.