I’m implementing adding new tokens to a pretrained Transformer model. For that I need to reinitialise the embedding vectors (expand them to include new tokens). However, I noticed that my expand_embeddings method, when run, breaks something in the model (loss goes up, predictions go rubbish).
Here’s what it looks like:
def expand_embeddings(self, new_vocab_size, device):
# Clone weights from embedding layer
old_embeddings = self.embed.token_embed.weight.data.clone().transpose(0,1).to(device)
# TEST: Don't include new tokens, instead simply reinitialise the layers as you would
# if new tokens were to be added
new_vocab_size = 32004
# Initialise the embedding layer to the old weights
self.embed = EmbeddingLayer(new_vocab_size, self.d_model, self.max_length, old_embeddings).to(device)
# Tie weights with the output layer
self.decoder.generator = Generator(self.d_model, new_vocab_size, self.embed.token_embed.weight.data).to(device)
Here are the Generator and EmbeddingLayer init’s for reference:
class Generator(nn.Module):
def __init__(self, d_model, vocab_size, tied_weights):
super(Generator, self).__init__()
self.proj = nn.Linear(d_model, vocab_size)
self.proj.weight.data = tied_weights
def forward(self, x):
(...)
class EmbeddingLayer(nn.Module):
def __init__(self, vocab_size, d_model, max_length, tied_weights=None):
super(EmbeddingLayer, self).__init__()
self.token_embed = Embedding(vocab_size, d_model)
if tied_weights is not None:
# Must transpose weights since Linear transposes by default but Embedding does not
self.token_embed.weight.data = tied_weights.clone().transpose(0, 1)
self.pos_embed = Embedding(max_length, d_model)
self.vocab_size = vocab_size
def forward(self, x, pos):
(...)
When expand_embeddings is not run, the training loss starts at 15.12, when it is run it starts at 93.06. However, I’d expect these values to be identical.