Scaled Dot-Product Attention

Scaled Dot-Product Attention is an essential refinement for stable and effective training. It uses dot products between Query ($Q$) and Key ($K$) vectors to measure alignment.

The core dot-product attention mechanism computes $QK^T$. Consider the dimensions involved. If the query and key vectors have dimension $d_k$, the dot product between a single query $q$ and a single key $k$ is $\sum_{i=1}^{d_k} q_i k_i$. If we assume the components of $q$ and $k$ are independent random variables with zero mean and unit variance, the expected value of the dot product is 0, but its variance is $d_k$.

In practice, especially for large values of $d_k$ (e.g., 64, 128, or higher, common in Transformers), the magnitude of these dot products can become significantly large. Why is this problematic? These dot products are the inputs to the softmax function, which calculates the attention weights.

$\text{softmax}(z)_i = \frac{e^{z_i}}{\sum_j e^{z_j}}$

The softmax function exhibits saturation behavior. If its inputs (the dot product scores) have large magnitudes (either large positive or large negative), the function enters regions where its gradient is extremely close to zero. Visualize the derivative of the softmax function; it approaches zero as the input values move away from zero.

The gradient of the softmax function is largest near input values of zero and diminishes rapidly for large positive or negative inputs. Large dot products push the computation into these low-gradient regions.

Small gradients drastically slow down or even halt the learning process during backpropagation, as weight updates become vanishingly small. This makes training deep models very difficult.

To counteract this effect, the "Attention Is All You Need" paper introduced a scaling factor of $1/\sqrt{d_k}$. The complete formula for Scaled Dot-Product Attention is:

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

Here, $Q$, $K$, and $V$ are matrices representing the packed sets of queries, keys, and values, respectively. $d_k$ is the dimension of the key (and query) vectors.

By dividing the dot products $QK^T$ by $\sqrt{d_k}$, we effectively moderate their magnitude, keeping the inputs to the softmax function closer to a regime where gradients are more substantial. This scaling ensures that the variance of the inputs to the softmax remains approximately constant regardless of the choice of $d_k$, contributing significantly to more stable training dynamics. It prevents the attention weights from becoming overly peaky (concentrated on a single input) or overly diffuse too early in training simply due to the dimensionality of the representations.

This scaling is not just a minor implementation detail; it's an important component that enables the successful training of deep Transformer models by mitigating potential gradient issues stemming directly from the core attention calculation. It allows the model to learn effectively even when using high-dimensional vector representations for queries and keys.