Get sentence embeddings of transformer-based models

Question

I want to get sentence embeddings of transformer-based models (Bert, Roberta, Albert, Electra...).

I plan on doing mean pooling on the hidden states of the second last layer just as what bert-as-service did.

So my questions is that when I do mean pooling, should I include the embeddings related to [PAD] tokens or [CLS] token or [SEP] token?

For example, my sequence is 300 tokens, and are padded into 512 tokens.

The output size is 512 (tokens) * 768 (embeddings).

So should I average the embeddings of first 300 tokens or the embeddings of whole 512 tokens?

Why the embeddings of the last 212 tokens are non-zero?

Answer 1

Usually, padding is excluded from mean pooling.

One approach to derive sentence embeddings by mean pooling excluding padding tokens can be taken from Sentence Transformers. In their pooling source code you can see that they use the attention mask to exclude padding tokens. Here is a simplified implementation of what they do using Huggingface transformers (taken from here):

from transformers import AutoTokenizer, AutoModelForMaskedLM

def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
    return sum_embeddings / sum_mask
  
tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
model = AutoModelForMaskedLM.from_pretrained("xlm-roberta-base")
encoded_input = tokenizer("hello", return_tensors='pt')
model_output = model(**encoded_input)
mean_pooling(model_output, encoded_input['attention_mask'])

Since the input_mask_expanded is 0 for PAD tokens the resulting mean excludes these tokens when calculating sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1). A concrete example is discussed here on the SBERT Github repository.

Answer 2

Ques 1. You decide how you want your padded pooling layer to behave.This is why pytorch's avg pool (e.g., nn.AvgPool2d) has an optional parameter count_include_pad=True: By default (True) Avg pool will first pad the input and then treat all elements the same. On the other hand, if you set count_include_pad=False the pooling layer will ignore the padded elements and the result

When you have too many padding tokens they take part in gradient calculation as part of training and weights updates they gain some weights. Obviously these weights computation are waste of time and should be avoided.

Suggestion 1: Make sure every watch has similar size sentences to avoid a lot of padding column