如何从Hugging Face预训练模型中通过embedding向量反向查找对应单词

1 获取BERT嵌入层权重 #

你需要的embedding_matrix可以直接从加载完成的BERT模型对象中提取，对应属性为model.embeddings.word_embeddings.weight，代码示例如下：

# 获取的嵌入矩阵形状为 [词典大小, 隐藏层维度]，bert-base-uncased对应维度为[30522, 768]
embedding_matrix = model.embeddings.word_embeddings.weight

2 注意事项 #

你当前的实现思路存在两个需要留意的问题：

• 你计算得到的last_hidden_states.mean(1)是BERT最后一层上下文隐状态的均值向量，和输入层的静态词嵌入不在同一个向量空间，直接用该向量和输入嵌入矩阵计算相似度匹配单词，结果会存在较大偏差，没有实际语义参考价值。
• 检索最相似单词时建议使用余弦相似度而非直接点积，避免向量长度对匹配结果的干扰。

3 完整检索示例代码（仅供测试参考） #

如果只是做原理验证，可以参考以下实现：

import torch.nn.functional as F

# 1. 获取嵌入矩阵
embedding_matrix = model.embeddings.word_embeddings.weight

# 2. 对池化向量和嵌入矩阵做归一化，方便计算余弦相似度
pooled_vec_norm = F.normalize(last_hidden_states, p=2, dim=1)
embedding_norm = F.normalize(embedding_matrix, p=2, dim=1)

# 3. 计算余弦相似度，取最高的前5个索引
cos_sim = pooled_vec_norm @ embedding_norm.T
topk_indices = torch.topk(cos_sim, k=5).indices.squeeze().tolist()

# 4. 用tokenizer转成对应单词
for idx in topk_indices:
    print(f"索引{idx}对应单词：{tokenizer.decode(idx)}，相似度：{cos_sim[0][idx].item():.4f}")

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