BERT Self-Attention layer

Question

I am trying to use the first individual BertSelfAttention layer for the BERT-base model, but the model I am loading from torch.hub seems to be different then the one used in hugginface transformers.models.bert.modeling_bert:

import torch, transformers

tokenizer = transformers.BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)
torch_model = torch.hub.load('huggingface/pytorch-transformers', 'model', 'bert-base-uncased')

inputs = tokenizer.encode_plus("Hello", "World", return_tensors='pt')

output_embedding = torch_model.embeddings(inputs['input_ids'], inputs['token_type_ids'])

output_self_attention = torch_model.encoder.layer[0].attention.self(output_embedding)[0]

# compare output with using the huggingface model directly
bert_self_attn = transformers.models.bert.modeling_bert.BertSelfAttention(torch_model.config) 

# transfer all parameters
bert_self_attn.load_state_dict(torch_model.encoder.layer[0].attention.self.state_dict())
# <All keys matched successfully>

output_self_attention2 = bert_self_attn(output_embedding)[0]
output_self_attention != output_self_attention2 # tensors are not equal?

Why is output_self_attention2 different from output_self_attention? I thought they would give the same output given the same input.

Answer 1

The problem is that you are using the models in training mode (the default mode), and the stochastic elements like dropout are active in training mode, therefore you obtain different results, not only between different models, but also on different runs of the same model.

You should invoke model.eval(), which makes some elements like dropout or batch normalization behave in inference mode, i.e. dropout is disabled altogether, batch-normalization uses the remembered statistics.

Also, it is better to avoid gradients being computed with torch.no_grad(), as you do not need them. The same effect can be achieved with torch.set_grad_enabled(False).

import torch, transformers


torch.set_grad_enabled(False) # avoid wasting computation with gradients


tokenizer = transformers.BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)
torch_model = torch.hub.load('huggingface/pytorch-transformers', 'model', 'bert-base-uncased')

torch_model.eval() # <------ set model in inference mode  

inputs = tokenizer.encode_plus("Hello", "World", return_tensors='pt')
output_embedding = torch_model.embeddings(inputs['input_ids'], inputs['token_type_ids'])
output_self_attention = torch_model.encoder.layer[0].attention.self(output_embedding)[0]


# compare output with using the huggingface model directly
bert_self_attn = transformers.models.bert.modeling_bert.BertSelfAttention(torch_model.config) 
bert_self_attn.load_state_dict(torch_model.encoder.layer[0].attention.self.state_dict())

bert_self_attn.eval()  # <------ set model in inference mode  

output_self_attention2 = bert_self_attn(output_embedding)[0]
output_self_attention != output_self_attention2 # tensors are equal

Answer 2

Setting dropout = 0.0 solves this issue:

torch_model.encoder.layer[0].attention.self.dropout.p = 0.0
bert_self_attn.dropout.p = 0.0

I thought that dropout was only used during the training process. Is the dropout in BERT different from conventional dropout layers?