Is it necessary to use attention mask for mean pooling for BERT?

Question

I am working on a project involving the analysis of clinical texts using the "emilyalsentzer/Bio_ClinicalBERT" model from Hugging Face's transformers library. My goal is to extract meaningful sentence embeddings from the model for downstream tasks. I am aware of several pooling strategies for BERT embeddings, but I'm uncertain which would be most effective for my specific use case.

The prevalent method is to use the CLS token for embeddings. However, inspired by a recent paper using mean pooling and an algorithm for optimal layer selection, I've been considering mean pooling of the last hidden state layer. A common practice in many repositories is to directly calculate the mean of this last layer. I believe, though, that this might overlook a critical aspect: the attention masks.

Intuitively, it seems more accurate to first multiply the last hidden state with the attention masks, effectively filtering out padded sequences, before pooling. This, in my understanding, would ensure that only the significant token embeddings contribute to the mean.

Is this modified approach of incorporating attention masks into mean pooling more effective than the typical direct mean calculation? I am seeking insights on whether this method enhances the quality of the sentence embedding for downstream applications.

A code example

import torch.nn as nn
from transformers import AutoModel, AutoTokenizer

proj_dim = 512
bert_model = AutoModel.from_pretrained("emilyalsentzer/Bio_ClinicalBERT", output_hidden_states=True)
tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
tokenizer.model_max_length = 256
projection_head = nn.Linear(768, proj_dim)

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()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

# Example usage
output = bert_model(input_ids=input_ids, attention_mask=attention_mask)
embed = mean_pooling(output, attention_mask)