Mistral模型对不同输入生成相同嵌入的问题排查

使用预训练的Mistral-7B-Instruct-v0.2模型生成输入文本的嵌入时出现异常：无论输入文本如何不同，输出的嵌入完全一致。相关代码及输出如下：

代码 #

from transformers import pipeline, AutoTokenizer, AutoModel
import numpy as np
import torch  # 原代码遗漏torch导入，运行会报错

PRETRAIN_MODEL = 'mistralai/Mistral-7B-Instruct-v0.2'
tokenizer = AutoTokenizer.from_pretrained(PRETRAIN_MODEL)
model = AutoModel.from_pretrained(PRETRAIN_MODEL)

def generate_embedding(document):
    inputs = tokenizer(document, return_tensors='pt')
    print("Tokenized inputs:", inputs)
    with torch.no_grad():
        outputs = model(**inputs)
    embedding = outputs.last_hidden_state[0, 0, :].numpy()
    print("Generated embedding:", embedding)
    return embedding

text1 = "this is a test"
text2 = "this is another test"
text3 = "there are other tests"

embedding1 = generate_embedding(text1)
embedding2 = generate_embedding(text2)
embedding3 = generate_embedding(text3)

are_equal = np.array_equal(embedding1, embedding2) and np.array_equal(embedding2, embedding3)

if are_equal:
    print("The embeddings are the same.")
else:
    print("The embeddings are not the same.")

输出结果 #

Tokenized inputs: {'input_ids': tensor([[   1,  456,  349,  264, 1369]]), 'attention_mask': tensor([[1, 1, 1, 1, 1]])}
Generated embedding: [-1.7762679  1.9293272 -2.2413437 ...  2.6379988 -3.104867   4.806004 ]
Tokenized inputs: {'input_ids': tensor([[   1,  456,  349, 1698, 1369]]), 'attention_mask': tensor([[1, 1, 1, 1, 1]])}
Generated embedding: [-1.7762679  1.9293272 -2.2413437 ...  2.6379988 -3.104867   4.806004 ]
Tokenized inputs: {'input_ids': tensor([[   1,  736,  460,  799, 8079]]), 'attention_mask': tensor([[1, 1, 1, 1, 1]])}
Generated embedding: [-1.7762679  1.9293272 -2.2413437 ...  2.6379988 -3.104867   4.806004 ]
The embeddings are the same.

问题原因 #

你取嵌入的位置是outputs.last_hidden_state[0, 0, :]，也就是序列第一个token（，即bos token）的隐藏状态。对于Mistral这类大语言模型，bos token作为句子的起始标记，其隐藏状态不会随输入文本内容变化，因此不同文本生成的嵌入完全一致。

def generate_embedding(document):
    inputs = tokenizer(document, return_tensors='pt')
    with torch.no_grad():
        outputs = model(**inputs)
    # 取最后一个token的嵌入
    embedding = outputs.last_hidden_state[0, -1, :].numpy()
    return embedding

def generate_embedding(document):
    inputs = tokenizer(document, return_tensors='pt')
    with torch.no_grad():
        outputs = model(**inputs)
    
    # 基于attention mask计算平均池化嵌入
    mask = inputs['attention_mask'].unsqueeze(-1).expand(outputs.last_hidden_state.shape)
    masked_embeddings = outputs.last_hidden_state * mask
    summed_embeddings = torch.sum(masked_embeddings, dim=1)
    # 避免除以0，加极小值
    token_counts = torch.clamp(mask.sum(dim=1), min=1e-9)
    embedding = (summed_embeddings / token_counts).squeeze(0).numpy()
    return embedding