I am trying to implement the Attention is All You Need paper from scratch. The authors mentioned in section 3.4 that "In our model, we share the same weight matrix between the two embedding layers and the pre-softmax linear transformation, similar to [30]" [30] here is the paper Using the output embedding to improve language models by Ofir Press and Lior Wolf.

Excuse me but

  1. could you explain why we need output embedding between encoder stack and the decoder stack? Isn't the output of encoder d_model dimension vectors (d_model = 512 in the original paper) so they don't need embedding layer in-between? Can I treat this as a linear layer?

  2. could you explain how the weights can be shared between the pre-softmax linear layer and the input embedding? Since in the original paper, it was about the translation between German and English, shouldn't the weight matrices be different for pre-softmax linear layer and the input embedding since the number of vocabularies for two languages aren't the same?

Thank you for reading my question!


1 Answer 1


There are three weight matrices that can be shared:

  1. The weights of the embedding layer for the encoder -- that is the layer that embeds the source sequence before forwarding through the transformer encoders. If your source vocab has size $V_\text{src}$ then this linear layer will have weight matrix of shape $ V_\text{src} \times d_\text{model}$.
  2. The weights of the embedding layer for the decoder -- If your target vocab has size $V_\text{tgt}$ then this linear layer will have weight matrix of shape $ V_\text{tgt} \times d_\text{model}$.
  3. The weights of the final linear layer that maps from the decoder embeddings into your target vocabulary space. The weight matrix of this layer has shape $ d_\text{model} \times V_\text{tgt}$

So regarding your first question:
There is nothing that you need to do with the output of the encoder.

Regarding your second question:
Since matrix 2 and matrix 3 have the same shapes you can easily share the weights. Matrix 3 is just the transpose of matrix 2.
Sharing the same weights between matrix 1 and matrix 2 and 3 is only possible if your source vocabulary matches the target vocabulary. So this makes perfect sense for tasks like text summarization or question answering, where both sequences share the same vocabulary.

In the case of machine translation where you have different vocabularies, this doesn't make much sense at first glance. But note that most of the times your are not using a word vocabulary, but a sub-word vocabulary (think character-level encodings or byte pair encodings). Now if you take a look at Table.1 in the paper from Ofir Press nad Lior Wolf, you will see that for translating between French and English (or German and English) almost 80% of the sub-words appear in both dictionaries. So what they do is they use a single vocabulary $V_\text{total}$ that contains all subwords that appear in both languages.

I really don't think that this idea makes sense for languages where there is very little overlapping of words / sub-words / characters, e.g. translating between English and Chinese.

If you want to see how this could be implemented you can check out the code snippets that I wrote here:
In this implementation not only the weight matrices are shared, but also the weights for the positional encodings.


You must log in to answer this question.

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