Transformers: how does the decoder final layer output the desired token?

Question

In the paper Attention Is All You Need, this section confuses me:

In our model, we share the same weight matrix between the two embedding layers [in the encoding section] and the pre-softmax linear transformation [output of the decoding section]

Shouldn't the weights be different, and not the same? Here is my understanding:

For simplicity, let us use the English-to-French translation task where we have $n^e$ number of English words in our dictionary and $n^f$ number of French words.

• In the encoding layer, the input tokens are $1$ x $n^e$ one-hot vectors, and are embedded with a $n^e$ x $d^{model}$ learned embedding matrix.

• In the output of the decoding layer, the final step is a linear transformation with weight matrix $d^{model}$ x $n^f$, and then applying softmax to get the probability of each french word, and choosing the french word with the highest probability.

How is it that the $n^e$ x $n^{model}$ input embedding matrix share the same weights as the $d^{model}$ x $n^f$ decoding output linear matrix? To me, it seems more natural for both these matrices to be learned independently from each other via the training data, right? Or am I misinterpreting the paper?

Answer 1

I found the answer by reading the paper referenced by that section, Using the output embedding to improve language models

Based on this observation, we propose threeway weight tying (TWWT), where the input embedding of the decoder, the output embedding of the decoder and the input embedding of the encoder are all tied. The single source/target vocabulary of this model is the union of both the source and target vocabularies. In this model, both in the encoder and decoder, all subwords are embedded in the same duo-lingual space.

It seems like they learned a single embedding matrix ($n^e + n^f$) x $d^{model}$ in dimension.

Answer 2

Well actually the dictionary does not contain entire words but rather byte pair encodings. At the end of chapter 3 of the paper that you listed they briefly mention that around 85-90% of the subwords are shared between english and french (see also Table 1). So in this case sharing the weights between the source embedding layer and the target embedding layer makes a lot of sense.

I think that if you are in a different setting (e.g. english to chinese), where your final layer contains $n^e + n^f$ outputs, then this doesn't make much sense because during decoding you actually want to generate words only from the target language, and training this huge softmax layer might be really tricky and hard to do.

Now regarding sharing the weights between the decoder embedding layer and the decoder output layer. I think this is a very nice idea, and you can always do this regardless of the source language or the task altoghether. However, you might want to be careful with the initialization of these weights. You can checkout this if you want to read more.