How to get gradient wrt to a word embedding layer pytorch?
Author: zeyuyun1
Pull request: no

# ❓ Questions & Help

<!-- The GitHub issue tracker is primarly intended for bugs, feature requests,
     new models and benchmarks, and migration questions. For all other questions,
     we direct you to Stack Overflow (SO) where a whole community of PyTorch and
     Tensorflow enthusiast can help you out. Make sure to tag your question with the
     right deep learning framework as well as the huggingface-transformers tag: 
     https://stackoverflow.com/questions/tagged/huggingface-transformers 
     
     If your question wasn't answered after a period of time on Stack Overflow, you
     can always open a question on GitHub. You should then link to the SO question 
     that you posted.
     -->

## Details
<!-- Description of your issue -->'
The idea behind this is I want to try some old school gradient ascent style visualization with Bert Model.

I want to know the effect of changes of embedding layer on a specific layer's specific dimension. Thus, I took the gradient of the output of a specific layer's specific dimension wrt the first word embedding layer's output.

The best thing I can do here is the following:


```
from transformers import BertTokenizer, BertModel
model = BertModel.from_pretrained('bert-base-uncased', output_attentions=True,output_hidden_states=True)
tokenizer = BertTokenizer.from_pretrained(model_version, do_lower_case=True)
s = 'I want to sleep'
inputs = tokenizer.encode_plus(s,return_tensors='pt', add_special_tokens=False,is_pretokenized=True)
input_ids = inputs['input_ids']
output = model(input_ids)
hidden_states = output[-2]

X = hidden_states[0] #embedding space, shape: [1,4,768] (batch_size,sentence_length,embedding dimension)
y = hidden_states[3][0][0][0] ##the 0th position and 0th dimension of output of 3rd hidden layer. Dimension should just be [1], a scalar.

torch.autograd.grad(y,X,retain_graph=True, create_graph=True) #I take the gradient of y wrt. Since y is scalar. The dimension of the gradient is just the dimension of X.
```

This is, however, not good enough. I want the gradient wrt the actual word embedding layer. However, Transformer's embedding contains "position_embedding" and "token_type_embedding". Here's the code for the first layer embedding:

```
class BertEmbeddings(nn.Module):
    """Construct the embeddings from word, position and token_type embeddings.
    """
    def __init__(self, config):
        super(BertEmbeddings, self).__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=0)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
        # any TensorFlow checkpoint file
        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, input_ids, token_type_ids=None, position_ids=None):
        seq_length = input_ids.size(1)
        if position_ids is None:
            position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
        if token_type_ids is None:
            token_type_ids = torch.zeros_like(input_ids)

        words_embeddings = self.word_embeddings(input_ids)
        position_embeddings = self.position_embeddings(position_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

        embeddings = words_embeddings + position_embeddings + token_type_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings
```

Ideally, I want the gradient wrt JUST “words_embeddings" Rather than, wrt "words_embeddings + position_embeddings + token_type_embeddings" and follows by layerNorm and dropout.

I think I can do this by modifying changing the model. Is there a way to it without changing the model?



<!-- You should first ask your question on SO, and only if
     you didn't get an answer ask it here on GitHub. -->
**A link to original question on Stack Overflow**:  https://stackoverflow.com/questions/62750798/how-to-get-gradient-wrt-to-a-specific-layers-output-pytorch


Comments:
Hi, the `word_embeddings` variable is not returned so you won't obtain these except if you modify the file.

Thanks!