machine-learning - 如何使用注意掩码计算 HuggingFace Transformers BERT token 嵌入的平均值/最大值？

Question

我正在使用 HuggingFace Transformers BERT 模型，并且我想使用 mean 计算句子中标记的摘要向量(也称为嵌入)。或 max功能。复杂的是一些 token 是 [PAD] ，所以我想在计算平均值或最大值时忽略这些标记的向量。
这是一个例子。我最初实例化一个 BertTokenizer和一个 BertModel :

import torch
import transformers
from transformers import AutoTokenizer, AutoModel

transformer_name = 'bert-base-uncased'

tokenizer = AutoTokenizer.from_pretrained(transformer_name, use_fast=True)

model = AutoModel.from_pretrained(transformer_name)

然后我在分词器中输入一些句子并退出 input_ids和 attention_mask .值得注意的是，一个 attention_mask值 0 表示 token 是 [PAD]我可以忽略。

sentences = ['Deep learning is difficult yet very rewarding.',
             'Deep learning is not easy.',
             'But is rewarding if done right.']
tokenizer_result = tokenizer(sentences, max_length=32, padding=True, return_attention_mask=True, return_tensors='pt')

input_ids = tokenizer_result.input_ids
attention_mask = tokenizer_result.attention_mask

print(input_ids.shape) # torch.Size([3, 11])

print(input_ids)
# tensor([[  101,  2784,  4083,  2003,  3697,  2664,  2200, 10377,  2075,  1012,  102],
#         [  101,  2784,  4083,  2003,  2025,  3733,  1012,   102,     0,     0,    0],
#         [  101,  2021,  2003, 10377,  2075,  2065,  2589,  2157,  1012,   102,   0]])

print(attention_mask.shape) # torch.Size([3, 11])

print(attention_mask)
# tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
#         [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
#         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0]])

现在，我调用 BERT 模型来获取 768-D token 嵌入(顶层隐藏状态)。

model_result = model(input_ids, attention_mask=attention_mask, return_dict=True)

token_embeddings = model_result.last_hidden_state
print(token_embeddings.shape) # torch.Size([3, 11, 768])

所以在这一点上，我有:

[3, 11, 768] 矩阵中的标记嵌入:3 个句子、11 个标记、每个标记的 768 维向量。

[3, 11] 矩阵中的注意力掩码:3 个句子，11 个标记。 1 值表示非 [PAD] .

我如何计算 mean/ max在有效的向量上，非 [PAD] token ？
我尝试使用注意力掩码作为掩码，然后调用 torch.max() ，但我没有得到正确的尺寸:

masked_token_embeddings = token_embeddings[attention_mask==1]
print(masked_token_embeddings.shape) # torch.Size([29, 768] <-- WRONG. SHOULD BE [3, 11, 768]

pooled = torch.max(masked_token_embeddings, 1)
print(pooled.values.shape) # torch.Size([29]) <-- WRONG. SHOULD BE [3, 768]

我真正想要的是形状 [3, 768] 的张量。也就是说，3 个句子中的每一个都有一个 768 维向量。

Answer 1

对于 max ，你可以乘以 attention_mask :

pooled = torch.max((token_embeddings * attention_mask.unsqueeze(-1)), axis=1)

对于 mean ，您可以沿轴求和并除以 attention_mask沿着那个轴:

mean_pooled = token_embeddings.sum(axis=1) / attention_mask.sum(axis=-1).unsqueeze(-1)