Scaled Dot-Product Attention
Introduction
Let three matrices $\mathbf{Q} \in M^{d_{k} \times n}$, $\mathbf{K} \in M^{d_{k} \times m}$, $\mathbf{V} \in M^{d_{v} \times m}$, called the query, key, and value respectively, and a score function $f : M^{d_{k} \times n} \times M^{d_{k} \times m} \to M^{m \times n}$ for the query and the key be given. Attention refers to the function defined as follows.
$$ \begin{align*} \operatorname{Attention}: M^{d_{k} \times n} \times M^{d_{k} \times m} \times M^{d_{v} \times m} &\to M^{d_{v} \times n} \\ (\mathbf{Q},\mathbf{K},\mathbf{V}) &\mapsto \mathbf{V} \operatorname{Softmax}\left(f(\mathbf{Q},\mathbf{K})\right) \tag{1} \end{align*} $$
Here $\operatorname{Softmax}$ denotes the function that, for a given matrix $\mathbf{X} = \begin{bmatrix} \mathbf{x}_{1} & \cdots & \mathbf{x}_{N}\end{bmatrix}$, applies the softmax $\operatorname{softmax}$ to each column vector so that each column sums to $1$.
$$ \operatorname{Softmax}(\mathbf{X}) := \begin{bmatrix} \underset{\vert}{\overset{\vert}{\operatorname{softmax}(\mathbf{x}_{1})}} & \cdots & \underset{\vert}{\overset{\vert}{\operatorname{softmax}(\mathbf{x}_{N})}} \end{bmatrix} $$
The type of attention is determined by how $f$ is defined, and the form of attention proposed in 『Attention Is All You Need』1, the paper in which Ashish Vaswani and seven others first proposed the Transformer architecture, is as follows.
Definition
The attention function of the following form is called scaled dot-product attention.
$$ \operatorname{Attention} (\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \mathbf{V} \operatorname{Softmax} \left( \frac{\mathbf{K}^{\mathsf{T}} \mathbf{Q}}{\sqrt{d_{k}}} \right) $$
Explanation
Just as the name says, the score function of the query vector $\mathbf{q}$ and the key vector $\mathbf{k}$ is taken to be the inner product together with a scaling factor.
$$ \operatorname{score}(\mathbf{q}, \mathbf{k}) = \dfrac{1}{\sqrt{d_{k}}} \mathbf{k}^{\mathsf{T}} \mathbf{q} $$
$$ f(\mathbf{Q}, \mathbf{K}) = \frac{\mathbf{K}^{\mathsf{T}} \mathbf{Q}}{\sqrt{d_{k}}} = \frac{1}{\sqrt{d_{k}}} \begin{bmatrix} \mathbf{k}_{1} \cdot \mathbf{q}_{1} & \mathbf{k}_{1} \cdot \mathbf{q}_{2} & \cdots & \mathbf{k}_{1} \cdot \mathbf{q}_{n} \\ \mathbf{k}_{2} \cdot \mathbf{q}_{1} & \mathbf{k}_{2} \cdot \mathbf{q}_{2} & \cdots & \mathbf{k}_{2} \cdot \mathbf{q}_{n} \\ \vdots & \vdots & \ddots & \vdots \\ \mathbf{k}_{m} \cdot \mathbf{q}_{1} & \mathbf{k}_{m} \cdot \mathbf{q}_{2} & \cdots & \mathbf{k}_{m} \cdot \mathbf{q}_{n} \end{bmatrix} $$
Why Divide by $\sqrt{d_{k}}$
The reason for going out of the way to divide by $\sqrt{d_{k}}$ instead of using the inner product itself as the score is that, as the dimension $d_{k}$ grows, the magnitude of the inner product grows on the scale of $\sqrt{d_{k}}$ and thus saturates the softmax1. Suppose that the components $q_{i}$, $k_{i}$ of the query $\mathbf{q}$ and the key $\mathbf{k}$ are all random variables with mean $0$ and variance $1$ that are independent of one another. Then, by the linearity of expectation and by independence, the mean of the inner product $\mathbf{k}^{\mathsf{T}} \mathbf{q} = \sum_{i=1}^{d_{k}} k_{i} q_{i}$ is $0$ as follows.
$$ \mathbb{E} \left[ \mathbf{k}^{\mathsf{T}} \mathbf{q} \right] = \mathbb{E} \left[ \sum_{i=1}^{d_{k}} k_{i} q_{i} \right] = \sum_{i=1}^{d_{k}} \mathbb{E} [ k_{i} q_{i} ] = \sum_{i=1}^{d_{k}} \mathbb{E} [ k_{i} ] \mathbb{E} [ q_{i} ] = 0 $$
Since all of the components are independent of one another, the products $k_{i} q_{i}$ are also independent of one another, and the variance of a sum of independent random variables is the sum of the individual variances. Moreover, since the mean is $0$, we have $\mathbb{E} [ k_{i}^{2} ] = \Var [ k_{i} ] = 1$, so the variance of the inner product is exactly $d_{k}$ as follows.
$$ \begin{align*} \Var \left[ \mathbf{k}^{\mathsf{T}} \mathbf{q} \right] &= \sum_{i=1}^{d_{k}} \Var [ k_{i} q_{i} ] = \sum_{i=1}^{d_{k}} \left( \mathbb{E} [ k_{i}^{2} q_{i}^{2} ] - \left( \mathbb{E} [ k_{i} q_{i} ] \right)^{2} \right) \\ &= \sum_{i=1}^{d_{k}} \mathbb{E} [ k_{i}^{2} ] \mathbb{E} [ q_{i}^{2} ] = \sum_{i=1}^{d_{k}} \Var [ k_{i} ] \Var [ q_{i} ] = d_{k} \end{align*} $$
In other words, the standard deviation of the score $\mathbf{k}^{\mathsf{T}} \mathbf{q}$ is $\sqrt{d_{k}}$, and if the dimension is $d_{k} = 512$, the standard deviation is $\sqrt{512} \approx 22.6$, so scores whose absolute values reach into the tens appear. When values with large absolute value enter the softmax, the exponential function blows up the largest component overwhelmingly, so the output saturates close to a one-hot vector; in the saturated region the output barely changes no matter how much the input changes, so the gradient becomes extremely small and learning is hindered by the vanishing gradient. On the other hand, if the score is divided by $\sqrt{d_{k}}$, the variance is normalized back to $1$, so the magnitude of the score stays constant no matter how large $d_{k}$ becomes.
$$ \Var \left[ \frac{\mathbf{k}^{\mathsf{T}} \mathbf{q}}{\sqrt{d_{k}}} \right] = \frac{1}{d_{k}} \Var \left[ \mathbf{k}^{\mathsf{T}} \mathbf{q} \right] = 1 $$
The figure below shows this difference. Sampling the components from the standard normal distribution, one query $\mathbf{q}$ and ten keys $\mathbf{k}_{1}, \dots, \mathbf{k}_{10}$ of dimension $d_{k} = 512$ were generated, and the ten scores $\mathbf{k}_{j}^{\mathsf{T}} \mathbf{q}$ were passed through the softmax. If they are passed through as they are without scaling, as on the left, the probability is in effect concentrated on a single key and saturates into a one-hot vector. On the other hand, if the same scores are divided by $\sqrt{d_{k}} = \sqrt{512}$ before being passed through, then as on the right the probability is spread evenly over several keys, so that attention can refer broadly to many keys instead of looking at only one particular key.
