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

Assume that the query embeddings $Q$ and key embeddings $K$ have zero mean and unit std. Then for the variance of the attention score between any query and key we get:

$$ \alpha = q_i k_j^T = \sum_{n=1}^{d_k} q_{in} k_{jn} $$ $$ \text{Var}(\alpha) = d_k $$ $$ \text{std}(\alpha) = \sqrt{d_k} $$

But is it safe to assume that our query embeddings and key embeddings have unit variance?

The answer is yes. The embeddings are computed by multiplying the input with the query and key matrices ($W_Q$ and $W_K$). We can assume that the input already has unit variance by using a normalizing layer (e.g. LayerNorm), and the weights of the $W_Q$ and $W_K$ matrices should be initialized so that variance is preserved.

But why do we want unit variance at all?

Applying softmax on the attention scores with such high variance will result in all of the weight being placed on one random element, while all the other elements will have a weight of zero. Thus, in order to have the attention scores with unit std, we scale by $\sqrt{d_k}$.

For more checkout this.

Answer 3

The mathematical explanation is already explained in the another answer. But I feel bothered why OP's code get an output of 1.6 variance, hence I run my own experiment to prove it:

import numpy as np
a = np.random.randn(1111,111)
print('variance of matrix A =', a.var())
b = np.random.randn(111,1111)
print('variance of matrix B =', b.var())
c = a @ b
print('variance of dot-product =', c.var())
d = c / np.sqrt(111)
print('variance of scaled (normalized) dot-product =', d.var())

The output:

variance of matrix A = 0.9998308190814476
variance of matrix B = 0.9957289239999045
variance of dot-product = 110.5603610976103
variance of scaled (normalized) dot-product = 0.9960392891676604

As shown above, the variance is becoming (close to) 1 after dividing by $\sqrt({d_k})$.

Answer 4

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.