How to use an embedding layer as a linear layer in PyTorch?

user_7766024 0
Viewed 1232

I'm currently working on a personal reimplementation of the Transformer paper and had a question.

On page 5 in section "3.4 Embeddings and Softmax," it states:

In our model, we share the same weight matrix between the two embedding layers and the pre-softmax linear transformation.

I've currently implemented my model to use just one embedding layer for both source and target tensors, but I'm wondering if there would be a way that I could use the weights of the embedding layer as a linear layer. What I've currently done is something like:

output = previous_layer(previous_input)
final_output = torch.matmul(output, embedding_layer.embedding.weight.transpose(1, 0))

I've transposed the weight matrix before matrix multiplication because it's of shape (vocab_size, embedding_dim) and the shape of output is (batch_size, seq_len, embedding_dim). Is this the proper way to use an embedding layer as a linear layer? If not, I'd like some tips on what I should be doing.

Thanks.

Edit

Specific line for weight sharing from link in answer.

pytorch
1 Answers

You could define a nn.Linear layer and replace its weights by copying the weights from the nn.Embedding:

trg_emb = nn.Embedding(trg_enc_dim, embedding_dim)
src_emb = nn.Embedding(src_enc_dim, embedding_dim)

trg_projection = nn.Linear(embedding_dim, trg_enc_dim, bias=False)
trg_projection.weight = trg_emb.weight # copy to nn.Linear

You can also copy to the source embedding so your two embedding layers share the same weights as well.

src_emb.weight = trg_emb.weight

This would mean the source embedding would end up having the same size as the target embedding i.e. trg_enc_dim x embedding_dim.

Here's a possible PyTorch implementation for inspiration.

user_63313690
Related