在学习bert的时候,我们知道bert是输出每个token的embeding。但在使用hugging face的bert模型时,发现除了last_hidden_state
还多了一个pooler_output
输出。
例如:
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased")
inputs = tokenizer("I'm caixunkun. I like singing, dancing, rap and basketball.", return_tensors="pt")
outputs = model(**inputs)
print("last_hidden_state shape:", outputs.last_hidden_state.size())
print("pooler_output shape:", outputs.pooler_output.size())
last_hidden_state shape: torch.Size([1, 20, 768])
pooler_output shape: torch.Size([1, 768])
许多人可能以为pooler_output
是[CLS]
token的embedding,但使用last_hidden_state shape[:, 0]
比较后,发现又不是,然后就很奇怪。
先说一下结论: pooler_output可以理解成该句子语义的特征向量表示。
那它是怎么来的?和[CLS]
token的embedding区别在哪?
我们将Bert模型打印一下,会发现最后还有一个BertPooler层,pooler_output就是从这来的。如下所示:
BertModel(
(embedding): BertEmbeddings(
....
)
(encoder): BertEncoder(
... # 12层TransformerEncoder
)
(pooler): BertPooler(
(dense): Linear(in_features=768, out_features=768, bias=True)
(activation): Tanh()
)
)
其中encoder就是将BERT的所有token经过12个TransformerEncoder进行embedding。而pooler
就是将[CLS]
这个token再过一下全连接层+Tanh激活函数,作为该句子的特征向量。
我们可以从Bert源码中验证以上结论。在transformers.models.bert.modeling_bert.BertModel.forward
方法中这么一行代码:
# sequence_output就是last_hidden_state
# self.pooler就是上面的BertPooler
pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
我们再来看看transformers.models.bert.modeling_bert.BertPooler
的源码:
class BertPooler(nn.Module):
def __init__(self, config):
super().__init__()
self.dense = nn.Linear(config.hidden_size, config.hidden_size)
self.activation = nn.Tanh()
def forward(self, hidden_states):
# hidden_states的第一个维度是batch_size。所以用[:, 0]取所有句子的[CLS]的embedding
first_token_tensor = hidden_states[:, 0]
pooled_output = self.dense(first_token_tensor)
pooled_output = self.activation(pooled_output)
return pooled_output
从上面的源码可以看出,pooler_output
就是[CLS]
embedding又经历了一次全连接层的输出。我们可以通过以下代码进行验证:
print("pooler:", model.pooler)
my_pooler_output = model.pooler(outputs.last_hidden_state)
print(my_pooler_output[0, :5])
print(outputs.pooler_output[0, :5])
pooler: BertPooler(
(dense): Linear(in_features=768, out_features=768, bias=True)
(activation): Tanh()
)
tensor([-0.8129, -0.6216, -0.9810, 0.8090, 0.9032], grad_fn=)
tensor([-0.8129, -0.6216, -0.9810, 0.8090, 0.9032], grad_fn=)
我们知道,BERT的训练包含两个任务:MLM和NSP任务(Next Sentence Prediction)。 对这两个任务不熟悉的朋友可以参考:BERT源码实现与解读(Pytorch) 和 【论文阅读】BERT 两篇文章。
其中MLM就是挖空,然后让bert预测这个空是什么。做该任务是使用token embedding进行预测。
而Next Sentence Prediction就是预测bert接受的两句话是否为一对。例如:窗前明月光,疑是地上霜
为 True,窗前明月光,李白打开窗
为False。
所以,NSP任务需要句子的语义信息来预测,但是我们看下源码是怎么做的。transformers.models.bert.modeling_bert.BertForNextSentencePrediction
的部分源码如下:
class BertForNextSentencePrediction(BertPreTrainedModel):
def __init__(self, config):
super().__init__(config)
self.bert = BertModel(config)
self.cls = BertOnlyNSPHead(config) # 这个就是一个 nn.Linear(config.hidden_size, 2)
...
def forward(...):
...
outputs = self.bert(...)
pooled_output = outputs[1] # 取pooler_output
seq_relationship_scores = self.cls(pooled_output) # 使用pooler_ouput送给后续的全连接层进行预测
...
从上面的源码可以看出,在NSP任务训练时,并不是直接使用[CLS]
token的embedding作为句子特征传给后续分类头的,而是使用的是pooler_output。个人原因可能是因为直接使用[CLS]
的embedding效果不够好。
但在MLM任务时,是直接使用的是last_hidden_state
,有兴趣可以看一下transformers.models.bert.modeling_bert.BertForMaskedLM
的源码。