如何在 BERT 中获得“正确”的嵌入层？

Question

我想利用 BERT 来评估两段文本之间的相似度：

from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn.functional as F
import numpy as np

tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")
model = AutoModel.from_pretrained("bert-classifier")

def calc_similarity(s1, s2):
    inputs = tokenizer(s1, s2, return_tensors='pt', padding=True, truncation=True)

    with torch.no_grad():
        outputs = model(**inputs)
        embeddings = outputs.last_hidden_state[:, 0, :].cpu().numpy()

    cosine_similarity = F.cosine_similarity(embeddings[0], embeddings[1])
    return cosine_similarity

这里提出的相似性源自 BERT 情感分类器，这是一个基于 BERT 架构微调的模型。

我的询问主要围绕这行代码：

embeddings = outputs.last_hidden_state[:, 0, :].cpu().numpy()

我观察到关于这条线的至少三种不同的实现；除了上述检索第一行的版本之外，还有其他两种变体：

embeddings = outputs.last_hidden_state.mean(axis=1).cpu().numpy()

和

embeddings = model.bert.pooler(outputs.last_hidden_state.cpu().numpy())

事实上，向量

outputs.last_hidden_state

是一个9*768的张量，前面提到的三种方法可以将其转化为1*768的向量，从而为后续的相似度计算提供基础。从我的角度来看，第一种方法在分类任务定义的语义空间内不合适，因为我们的目标不是预测下一个单词。让我困惑的是第二种和第三种方法之间的选择，特别是是采用简单平均还是利用模型本身的池化层。

如有任何帮助，我们将不胜感激！

Answer 1

第一种方法不是一个好的选择，因为如果 BERT 针对相似性匹配之外的任务进行了微调，那么直接利用 [CLS] 令牌嵌入可能不是最佳方法。

• 特定于任务的嵌入：[CLS] 令牌嵌入受 bert 模型训练任务的影响。
• 平均：取所有标记嵌入的平均值，我们可以得到输入的更一般的表示。该方法通过考虑所有标记的上下文嵌入来平衡表示。

考虑取平均值或池化（穿过另一个密集层）会起作用。