Problem when using Autograd with nn.Embedding in Pytorch

Question

I am in trouble with taking derivatives of outputs logits with respect to the inputs input_ids. Here is an example of my input:

# input_ids is a list of token ids got from BERT tokenizer
input_ids = torch.tensor([101., 1996., 2833., 2003., 2779., 1024., 6350., 102.], requires_grad=True)

content_outputs = self.bert(input_ids,
                        position_ids=position_ids, token_type_ids=token_type_ids,
                        attention_mask=attention_mask, head_mask=head_mask)

# content_outputs[1] is sentence embedding
logits = F.linear(content_outputs[1], self.W_s, self.b_s)

# Notes:
# - input_ids.dtype = torch.float32
# - input_ids.is_leaf = True
# - input_ids.requires_grad = True
# - Torch version: 1.0.1

To compute the gradients for input_ids, input_ids.dtype must be float; otherwise, I will get the following error:

RuntimeError: Only Tensors of floating point dtype can require gradients

However, my model is using an embedding layer which requires the input with dtype=long causing another problem if I use the input_ids initialized with float type above:

RuntimeError: Expected tensor for argument #1 'indices' to have scalar type Long; but got torch.FloatTensor instead (while checking arguments for embedding)

I have also searched on Stack Overflow as well as Stack Exchange but I found nothing related to this problem. Please help me with this. Any comments would be appreciated!

Answer 1

So, this is just half answer, I write it here to be able to format the text clearly. You are facing troubles because you are trying to do something that you shouldn't, which is applying gradient to indices instead of embeddings. When using embeddings (all kinds, not only BERT), before feeding them to a model, sentences must be represented with embedding indices, which are just number associated with specific embedding vectors. Those indices are just values mapping words to vectors, you will never apply any operation on them, especially you will never train them, cause they are not parameters.

['Just an example...'] # text
       |
       | Words are turned into indices 
       v
[1., 2., 3., 4., 4., 4.] # indices, tensor type long, no sense in applying gradient here
       |
       | Indices are used to retrieve the real embedding vectors
       v
[0.3242, 0.2354, 0.8763, 0.4325, 0.4325, 0.4325] # embeddings, tensor type float, this is what you want to train

If you want to fine tune BERT for a specific task, I suggest you to take a look to this tutorial BERT-fine-tuning