Why does this multiplication of $Q$ and $K$ have a variance of $d_k$, in scaled dot product attention?

Question

In scaled dot product attention, we scale our outputs by dividing the dot product by the square root of the dimensionality of the matrix:

The reason why is stated that this constrains the distribution of the weights of the output to have a standard deviation of 1.

Quoted from Transformer model for language understanding | TensorFlow:

For example, consider that $Q$ and $K$ have a mean of 0 and variance of 1. Their matrix multiplication will have a mean of 0 and variance of $d_k$. Hence, square root of $d_k$ is used for scaling (and not any other number) because the matmul of $Q$ and $K$ should have a mean of 0 and variance of 1, and you get a gentler softmax.

Why does this multiplication have a variance of $d_k$?

If I understand this, I will then understand why dividing by $\sqrt({d_k})$ would normalize to 1.

Trying this experiment on 2x2 arrays I get an output of 1.6 variance:

Answer 1

In statistics, if $X$ and $Y$ are independent and randomly distributed variables:

$\mathbb{E}[X + Y] = \mathbb{E}[X] + \mathbb{E}[Y] \\ Var(X + Y) = Var(X) + Var(Y) \\ \mathbb{E}[XY] = \mathbb{E}[X]\mathbb{E}[Y] \\ Var(XY) = (Var(X) + \mathbb{E}[X]^2)(Var(Y) + \mathbb{E}[Y]^2) - \mathbb{E}[X]^2\mathbb{E}[Y]^2$

Let $Q$ and $K$ be random $d_k$ x $d_k$ matrices, where each entry is a some random distribution with $0$ mean and $1$ variance. Every entry is independent from each other.

Since each entry of $Q$ and $K$ have identical distribution, we can focus only on the top-left-most element of $QK$ without loss of generality. The same applies to every other element.

The top-left-most element of $QK$ is $\displaystyle \sum_{i=0}^{d_k} Q_{1,i} K_{i, 1}$.

Since $Q$ and $K$ are independent:

$\mathbb{E}[Q_{1, i} K_{i, 1}] = \mathbb{E}[Q_{1, i}] \mathbb{E}[K_{i, 1}] = 0 \\ Var(Q_{1, i} K_{i, 1}) = (Var(Q_{1, i}) + \mathbb{E}[Q_{1, i}]^2)(Var(K_{i, 1}) + \mathbb{E}[K_{i, 1}]^2) - \mathbb{E}[Q_{1, i}]^2\mathbb{E}[K_{i, 1}]^2 = 1$

And so summing up $d_k$ of them:

$\mathbb{E} \left[\displaystyle \sum_{i=0}^{d_k} Q_{1,i} K_{i, 1} \right] = \displaystyle \sum_{i=0}^{d_k} \mathbb{E} \left[ Q_{1,i} K_{i, 1} \right] = 0 \\ Var\left(\displaystyle \sum_{i=0}^{d_k} Q_{1,i} K_{i, 1} \right) = \displaystyle \sum_{i=0}^{d_k} Var\left( Q_{1,i} K_{i, 1} \right) = d_k$

For your code block, you are computing the dot product of matrices $a$ and $b$, when you should be doing a matrix multiplication (the attention function multiplies $Q$ by $K$ after all, which is the vectorized form of dot-product -- it doesn't actually do dot product). It should work out to unit variance.

edit: the last paragraph is incorrect, as dot is the same as matrix multiply in the above case

Answer 2

It might help to take two small matrices that match the assumptions (mean of zero and variance of one) and just do the matrix multiplication. The dimensionality of K scales Q in the multiplication, scaling the variance simultaneously.