Mathematical Formulation of Dot-Product Attention

Dot-product attention, a widely used attention mechanism, employs the concepts of Query ($Q$), Key ($K$), and Value ($V$). The primary computational step involves calculating a set of scores that represent the relevance or compatibility between each query and every key. These scores are then used to create a weighted sum of the values, effectively allowing a model to focus on the most pertinent information carried by the values based on query-key interactions.

The "dot-product" name comes directly from how these compatibility scores are calculated. For a given query vector $q$ and a set of key vectors $K = \{k_1, k_2, ..., k_m\}$, the score between $q$ and a specific key vector $k_j$ is simply their dot product:

$score(q, k_j) = q \cdot k_j = q^T k_j$

Intuitively, the dot product measures the alignment or similarity between two vectors. If a query vector $q$ is highly similar (points in a similar direction) to a key vector $k_j$, their dot product will be large. Conversely, if they are dissimilar or orthogonal, the dot product will be small or zero.

In practice, we don't compute these scores one by one. Deep learning frameworks excel at parallel processing, especially matrix multiplications. We typically work with matrices of queries, keys, and values:

• $Q$: A matrix containing $n$ query vectors, each of dimension $d_k$. Shape: $[n \times d_k]$.
• $K$: A matrix containing $m$ key vectors, each also of dimension $d_k$. Shape: $[m \times d_k]$.
• $V$: A matrix containing $m$ value vectors, each of dimension $d_v$. Shape: $[m \times d_v]$.

Note that the query and key vectors must have the same dimension ($d_k$) to allow for the dot product calculation. The value vectors can have a different dimension ($d_v$).

To compute all query scores simultaneously, we perform a matrix multiplication between the query matrix $Q$ and the transpose of the matrix $K^T$:

$\text{Scores} = Q K^T$

Let's examine the dimensions: Multiplying $Q$ ($[n \times d_k]$) by $K^T$ ($[d_k \times m]$) results in a Score matrix of shape $[n \times m]$. Each element $(i, j)$ in this Scores matrix represents the dot product between the $i$-th query ($Q_i$) and the $j$-th key ($K_j$), indicating their compatibility.

Matrix multiplication $QK^T$ computes all pairwise dot products between query vectors (rows of Q) and key vectors (columns of $K^T$).

These raw scores ($QK^T$) represent the fundamental compatibility measure in dot-product attention. However, they are not yet ready to be used as weights for the values. The range of dot products can be quite large, potentially leading to issues during training, especially with gradients. Furthermore, we need to convert these scores into a probability distribution that sums to 1 for each query, representing how much attention that query should pay to each value.

The next steps, scaling these scores and applying the softmax function, address these issues and transform the raw dot-product scores into usable attention weights. We will examine these steps in the following sections.