如何使用注意掩码计算 HuggingFace Transformers BERT 令牌嵌入的均值/最大值?
Posted
技术标签:
【中文标题】如何使用注意掩码计算 HuggingFace Transformers BERT 令牌嵌入的均值/最大值?【英文标题】:How to compute mean/max of HuggingFace Transformers BERT token embeddings with attention mask? 【发布时间】:2021-03-12 23:11:11 【问题描述】:我正在使用 HuggingFace Transformers BERT 模型,我想使用 mean
或 max
函数计算句子中标记的摘要向量(也称为嵌入)。复杂之处在于某些标记是[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 意味着令牌是我可以忽略的 [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 令牌嵌入(顶层隐藏状态)。
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]
。
如何针对有效的非[PAD]
标记计算向量上的mean
/ max
?
我尝试使用注意掩码作为掩码,然后调用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 维向量。
【问题讨论】:
【参考方案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)
【讨论】:
谢谢。mean_pooled
不能与最大池类似地实现,除非您使用 torch.mean
?像这样:mean_pooled = torch.mean((token_embeddings * attention_mask.unsqueeze(-1)), axis=1)
?
使用token_embeddings * attention_mask.unsqueeze(-1)
非常巧妙。我没想到。
这似乎不正确 - PAD 标记没有嵌入为 0 向量;它们代表上下文和位置,因此具有真正有价值的元素。因此,将使用 PAD 向量计算令牌轴上的总和,这可能不是我们想要的。我刚想出的一项解决方法是,如果注意掩码为 0,则将嵌入设置为 0,然后求和,然后除以 # of tokens。在代码中,state[inputs["attention_mask"] == 0] = 0
。不过这感觉效率低下,所以我正在寻找更优雅的解决方案
@AlexL:你写了This seems incorrect
。你指的是什么This
?【参考方案2】:
亚历克斯是对的。 Look on hidden states for strings that go into tokenizer. For different strings, padding will have different embeddings.
因此,为了正确池化嵌入,您需要忽略那些填充向量。
假设您想从 BERT 的最后 4 层中获取嵌入(因为它会产生最好的分类结果):
#iterate over the last 4 layers and get embeddings for
#strings without having embeddings from PAD tokens
m = []
for i in range(len(hidden_states[0])):
m.append([hidden_states[j+9][i,:,:][tokens["attention_mask"][i] !=0] for j in range(4)])
#average over all tokens embeddings
means = []
for i in range(len(hidden_states[0])):
means.append(torch.stack(m[i]).mean(dim=1))
#stack embeddings for all strings
pooled = torch.stack(means).reshape(-1,1,3072)
【讨论】:
以上是关于如何使用注意掩码计算 HuggingFace Transformers BERT 令牌嵌入的均值/最大值?的主要内容,如果未能解决你的问题,请参考以下文章
huggingface-transformers:训练 BERT 并使用不同的注意力对其进行评估
如何在 Huggingface Trainer 课程中恢复训练时避免迭代 Dataloader?