I am experimenting with the code that accompanied the paper “Regularizing and Optimizing LSTM Language Models” (Merity et al. 2018): https://github.com/salesforce/awd-lstm-lm/ 56. In it, they apply dropout on the embedding. I was wondering if it was possible to make the implementation a bit more conservative on memory (it creates an embedding-sized tensor in the process). I managed to do that, but when I put my implementation in place of theirs, I didn’t get the same numbers as them from the second iteration on.
(A simplified version of) their implementation:
def embedded_dropout(embed, words, dropout=0.1):
    mask = embed.weight.data.new_empty((embed.weight.size(0), 1)).bernoulli_(1 - dropout).expand_as(embed.weight) / (1 - dropout)
    masked_embed_weight = mask * embed.weight
    return torch.nn.functional.embedding(words, masked_embed_weight)

Mine:
def embedded_dropout(embed, words, dropout=0.1):
    emb = embed(words)
    mask = embed.weight.data.new_empty((embed.weight.size(0), 1)).bernoulli_(1 - dropout) / (1 - dropout)
    m = torch.gather(mask.expand(embed.weight.size(0), words.size(1)), 0, words)
    return (emb * m.unsqueeze(2).expand_as(emb)).squeeze()

The outputs of the functions are identical. However, the gradients of embed (after going through a few LSTM layers + softmax) differ by 2.0763528e-19 after the first iteration, and then it gets progressively worse, to about 0.0037 after the 200th.
My questions are:

Am I using something that is not stable w.r.t. the gradient computation?
Is it possible to come up with a solution that obtains the same gradients, but only requires a words-sized mask (as in my code) as opposed to an embedding-sized one (as in the original)?