8
$\begingroup$

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

enter image description here

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:

enter image description here

$\endgroup$
1
  • 2
    $\begingroup$ Even though here programming questions are generally off-topic, in case the code is necessary, it's preferable to provide it rather than an actual screenshot. So, please, if you can still provide the actual code, edit your post to include it. $\endgroup$
    – nbro
    Commented Jan 16, 2021 at 23:20

4 Answers 4

11
$\begingroup$

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

$\endgroup$
4
  • $\begingroup$ Awesome, thank you this is a great answer. I am still confused on the code part though looking at the docs for numpy the np.dot function should be equivalent to np.matmul in 2 dimensions. $\endgroup$
    – Jacob B
    Commented Dec 8, 2020 at 20:40
  • 2
    $\begingroup$ Ah, you are right! dot is the same as matmul. I'm not sure how to explain the code part in that case. Maybe dimension=2 is too small? The sample variance will itself have a variance that reduces with larger dimension size. $\endgroup$ Commented Dec 8, 2020 at 20:53
  • 1
    $\begingroup$ Hmmm I see I will try with larger shapes, in any case your answer provided great insight to my initial question so thanks for taking the time! $\endgroup$
    – Jacob B
    Commented Dec 8, 2020 at 22:23
  • $\begingroup$ I appreciate it if you answer this related question ai.stackexchange.com/questions/40244/… $\endgroup$
    – Peyman
    Commented Apr 29, 2023 at 21:34
4
$\begingroup$

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.

$\endgroup$
1
  • $\begingroup$ Very good answer! For anyone wondering, we can approximate unit variance for out weights matrix using Xavier-Glorot: $\endgroup$
    – Dude156
    Commented Dec 29, 2023 at 20:54
1
$\begingroup$

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})$.

$\endgroup$
0
$\begingroup$

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.

$\endgroup$
1
  • 1
    $\begingroup$ Hey, I did do this but it does not come out as expected. I have edited the post $\endgroup$
    – Jacob B
    Commented May 21, 2020 at 1:09

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .