如何高效为数组A中每个词嵌入查找B中的Top10余弦相似词？

一、用向量化计算替代两层循环（精确匹配） #

两层循环完全浪费了张量框架的并行计算能力，直接通过归一化+矩阵点积就能批量完成余弦相似度计算：

1. 对A、B的嵌入向量做L2归一化（余弦相似度 = (A·B)/(||A||*||B||)，归一化后向量模长为1，点积结果等价于余弦相似度）
2. 计算A与B转置的矩阵乘积，得到形状为[100, 400000]的相似度矩阵
3. 对矩阵每行取前10个最大值及其对应索引

示例代码：

import torch
import torch.nn.functional as F

# 假设A: [100,50], B: [400000,50]
A_normalized = F.normalize(A, p=2, dim=1)
B_normalized = F.normalize(B, p=2, dim=1)

# 批量计算相似度矩阵
similarity_matrix = torch.matmul(A_normalized, B_normalized.T)

# 获取每个A中词的前10相似结果
top10_scores, top10_indices = torch.topk(similarity_matrix, k=10, dim=1)

该方法比循环快数倍至数十倍，若将张量移至GPU（.cuda()），速度还能再提升一个数量级。

二、近似最近邻(ANN)库（超大规模B场景） #

当B规模达百万级或A样本量大幅增长时，精确矩阵乘法会占用大量内存（比如10000个A样本对应40亿次运算），此时用近似最近邻库能在极小精度损失下，将查询速度提升10-100倍。

常用库推荐： #

• FAISS（Facebook开源）：支持CPU/GPU，精度与速度平衡最优，适合大规模数据。提供多种索引类型，IndexFlatIP用于精确内积查询，IndexHNSWFlat用于快速近似查询。
• Annoy（Spotify开源）：轻量级，内存占用低，适合资源有限的部署环境，查询速度快。
• HNSWLib：单独实现HNSW算法，比FAISS的同类型索引更轻量，速度相当。

FAISS示例代码：

import faiss
import numpy as np

# 若用torch张量，先转numpy
A_np = A_normalized.cpu().numpy()
B_np = B_normalized.cpu().numpy()

# 构建内积索引（归一化后内积等价于余弦相似度）
index = faiss.IndexFlatIP(50)
index.add(B_np)

# 查询前10相似词
top10_distances, top10_indices = index.search(A_np, k=10)
# 返回的distance即为余弦相似度值

追求更快速度可换用近似索引：

index = faiss.IndexHNSWFlat(50, 32)  # 32为图连接数，值越大精度越高、速度越慢
index.add(B_np)
top10_distances, top10_indices = index.search(A_np, k=10)

三、额外优化技巧 #

• 预归一化B：若B固定不变，提前归一化并保存，避免重复计算。
• 混合精度计算：PyTorch中用torch.float16计算，减少内存占用并提升GPU运算速度，精度损失可忽略。
• 分块计算：内存不足时，将B分成多块，逐块计算相似度后合并topk结果，避免一次性加载全部B。

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