Skip to content

7/22 Dot-Product Attention

Representing Local Context

Suppose given an input data entry or a hidden representation \(X^{(l)}\in\mathbf R^{s\times d}\), where \(s\) is a sequence dimension. When producing the next hidden representation \(X^{(l+1)}\), at the sequence index position \(1\le i\le s\), we would like to have a representation of the local context. To this end, we have seen the following option so far:

In a translation-equivariant, that is convolution layer, we can aggregate context at different sequence elements with the same kernel:

\[ X^{(l+1)}_i \leftarrow \sum_{j=1}^sX^{(l)}_{\overline{i + j}}K_j \]

We saw that to get local context, we can make \(K_j=0\) for \(|j|>C\), for some constant \(C\).

One may wonder: can we use the features \(X^{(l)}_i\) at element \(i\) to decide how much the features \(X^{(l)}_j\) at different element position matter for the local context features \(X^{(l+1)}_i\) at the next hidden representation level?

Dot-Product Attention

We want

\[ X^{(l+1)}_i=\sum_{j=1}^sw_{ij}X^{(l)}_jW_V^{(l)}, \]

where:

  1. the value weight matrix \(W_V^{(l)}\in\mathbf R^{d_v\times d}\) processes element-level features, and
  2. for element indices \(1\le i,j\le s\), the weight \(w_{ij}\) determines how much the level \(l\) representation \(X^{(l)}_j\) of element \(j\) matters for the level-(\(l+1\)) representation of element \(i\).

For stability, we want \(\sum_{j=1}^sw_{ij}=1\). Thus, we take the Data Scientist's default choice:

\[ w_{ij}=\mathop{\mathrm{softmax}}(w'_{ij'}:1\le j'\le s)_j. \]

The idea was first proposed in [1], for machine translation. We follow [2], whence the success of the method stems.

Self-Attention

We want to get the logits \(w'_{ij}\) from the same data entry \(X^{(l)}\in\mathbf R^{s\times d}\). One can get them as scalar products

\[ w'_{ij} \leftarrow\left(X^{(l)}_iW^{(l)}_Q\right) \left(X^{(l)}_jW^{(l)}_K\right)^T =X^{(l)}_i \left(W^{(l)}_Q(W^{(l)}_K)^T\right) (X^{(l)}_j)^T :=\mathbf q_i^T\mathbf k_j. \]

We call

  1. \(W^{(l)}_Q,W^{(l)}_K\in\mathbf R^{d\times d_k}\) the query and key weight matrices, and
  2. \(\mathbf q_i,\mathbf k_j\in\mathbf R^d\) query and key vectors.

Scaled Dot-Product Attention

Suppose that we initialize our parameters by the book, and thus the values in \(\mathbf q_i\) and \(\mathbf k_j\) are independent, of mean 0 and variance 1. Then the logit \(w_{ij}'\) will have mean 0 and variance \(d\). To counteract this effect, we divide the logits by \(\sqrt d\).

Multi-Head Attention

It may help training if we aggregate contextual data along not one, but several attention patterns. To this end:

  1. We introduce a new hyperparameter: the number of attention heads \(h\).

  2. We stack \(h\) query, key and value weight matrices into 3-tensors of shape (h, d, d_k) in the first two cases and (h, d, d_v) in the last.

  3. Accordingly, we get a stack of aggregated values, of shape (h, s, d_v)

  4. By transposing the first two dimensions, then flattening the last two, we get the aggregated values in the shape (s, h * d_v).

  5. We can multiply this matrix on the right by the output weight matrix \(W_O^{(l)}\in\mathbf R^{hd_v\times d}\). We denote the result of this by \(\mathop{\mathrm{MultiHead(Q,K,V)}}\) or \(\mathop{\mathrm{MHSA}}^{(l)}(X^{(l)})\).

In most transformer architectures, the following formulas hold:

  1. \(h\mid d\)
  2. \(d_k, d_v = d/h\).

References

[1] Dzmitry Bahdanau, Kyunghyun Cho and Yoshua Bengio: Neural Machine Translation by Jointly Learning to Align and Translate. International Conference on Learning Representations (ICLR), 2015. link

[2] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin: Attention Is All You Need. link