Transformers: how to get the output (keys and values) of the encoder?

Question

I was reading the paper Attention Is All You Need.

It seems like the last step of the encoder is a LayerNorm(relu(WX + B) + X), i.e. an add + normalization. This should result in a $n$ x $d^{model}$ matrix, where $n$ is the length of the input to the encoder.

How do we convert this $n$ x $d^{model}$ matrix into the keys $K$ and values $V$ that are fed into the decoder's encoder-decoder attention step?

Note that, if $h$ is the number of attention heads in the model, the dimensions of $K$ and $V$ should both be $n$ x $\frac{d^{model}}{h}$. For $h=8$, this means we need a $n$ x $\frac{d^{model}}{4}$ matrix.

Do we simply add an extra linear layer that learns a $d^{model}$ x $\frac{d^{model}}{4}$ weight matrix?

Or do we use the output of the final Add & Norm layer, and simply use the first $\frac{d^{model}}{4}$ columns of the matrix and discard the rest?

Answer 1

I have read the OpenNMT source code (https://github.com/OpenNMT/OpenNMT-py/blob/cd29c1dbfb35f4a2701ff52a1bf4e5bdcf02802e/onmt/modules/multi_headed_attn.py).

It seems like an extra linear layer learns the weights $W^{key}$ and $W^{value}$ (plus biases), so to get the output (keys and values), you multiply the output of the encoder's final add + norm layer by $W^{key}$ to get the keys, and by $W^{value}$ to get the values.

Additionally, these weights and biases seem to be independent across each of the decoding layer. So you feed the same encoder output (add + norm layer output), but multiply by different $W^{key}$ and $W^{value}$ matrices and add by different biases for each of the decoding layer, resulting in different keys and values for each layer