SBERT无论输入何种句子输出结果均一致，求问题排查

问题描述 #

我编写了一个测试SBERT的脚本，代码如下：

import torch
from transformers import BertTokenizer, BertModel
from sklearn.cluster import KMeans

# 1. Use SBERT to compare two sentences for semantic similarity.
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

input_ids_1 = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
input_ids_2 = torch.tensor(tokenizer.encode("Hello, my cat is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
outputs_1 = model(input_ids_1)
outputs_2 = model(input_ids_2)
last_hidden_states_1 = outputs_1[0]  # The last hidden-state is the first element of the output tuple
last_hidden_states_2 = outputs_2[0]  # The last hidden-state is the first element of the output tuple

# 2. Take SBERT embeddings for both sentences and cluster them.
kmeans = KMeans(n_clusters=2, random_state=0).fit(last_hidden_states_1.detach().numpy()[0], last_hidden_states_2.detach().numpy()[0])

# 3. Print the clusters.
print(kmeans.labels_)
print(kmeans.cluster_centers_)

输出结果为：

[0 0 0 0 0 0 0 1]
[[-0.2281394   0.29968688  0.3390873  ... -0.40648264  0.2930719
   0.41721284]
 [ 0.6079925   0.26097086 -0.3130729  ...  0.03109726 -0.6282735
  -0.19942412]]

无论将第二个句子更换为任何内容（例如"The capital of France is Paris"），输出结果始终相同。显然在数据传递或转换环节存在错误，请问问题出在哪里？

问题原因及解决方法 #

1. KMeans参数传递错误（核心问题） #

KMeans.fit()的参数定义为fit(X, y=None)：

• X是必需的特征矩阵，形状为[n_samples, n_features]
• y是可选的样本标签，仅用于监督学习变种，KMeans作为无监督算法会直接忽略该参数

你将第二句的token embedding作为y传入，导致KMeans完全没有用到第二句的数据，仅对第一句的8个token embedding进行聚类，因此更换第二句内容不会改变结果。

2. 混淆了BERT与SBERT #

当前使用的是原生BertModel，并非SBERT。SBERT是针对语义相似度任务微调后的BERT变种，通常需通过sentence-transformers库加载，原生BERT的输出并不适合直接用于语义相似度或聚类任务。

修正方案 #

方案一：对两个句子的所有token embedding进行聚类

将两个句子的token embedding合并为一个特征矩阵后传入KMeans：

import torch
import numpy as np
from transformers import BertTokenizer, BertModel
from sklearn.cluster import KMeans

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

input_ids_1 = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)
input_ids_2 = torch.tensor(tokenizer.encode("Hello, my cat is cute", add_special_tokens=True)).unsqueeze(0)

outputs_1 = model(input_ids_1)
outputs_2 = model(input_ids_2)

last_hidden_states_1 = outputs_1[0]
last_hidden_states_2 = outputs_2[0]

# 提取并合并两个句子的token embedding，形状为[总token数, 隐藏层维度]
emb1 = last_hidden_states_1.detach().numpy()[0]
emb2 = last_hidden_states_2.detach().numpy()[0]
combined_embeddings = np.vstack([emb1, emb2])

# 对合并后的embedding聚类
kmeans = KMeans(n_clusters=2, random_state=0).fit(combined_embeddings)
print(kmeans.labels_)
print(kmeans.cluster_centers_)

方案二：使用SBERT生成句子级embedding并聚类

这是更贴合语义相似度需求的做法，直接生成句子向量后聚类：

from sentence_transformers import SentenceTransformer
from sklearn.cluster import KMeans

# 加载预训练SBERT模型
model = SentenceTransformer('all-MiniLM-L6-v2')
sentences = ["Hello, my dog is cute", "Hello, my cat is cute"]

# 生成句子级embedding
sentence_embeddings = model.encode(sentences)

# 对句子向量聚类
kmeans = KMeans(n_clusters=2, random_state=0).fit(sentence_embeddings)
print(kmeans.labels_)
print(kmeans.cluster_centers_)

内容的提问来源于stack exchange，提问作者reallymemorable