如何不使用循环 一次性计算矩阵A与多个B矩阵的Cosine similarity

优化方案 #

有更高效的实现方式，核心思路是把所有B类矩阵按样本维度拼接成一个大矩阵，一次性完成相似度计算后再按原来的B矩阵长度拆分即可，时间复杂度比循环低得多，避免了重复计算A的归一化项以及多次函数调用的开销。

具体步骤 #

• 第一步：先把所有B1、B2…Bn在样本维度（axis=0）拼接为一个整体矩阵B_total，假设B1有m1个样本，B2有m2个样本…Bn有mn个样本，拼接后的B_total形状为(m1+m2+…+mn, 向量维度)，和单个Bi的维度完全匹配。
• 第二步：直接调用一次cosine_similarity(A, B_total)，得到形状为(A的样本数, m1+m2+…+mn)的整体相似度矩阵。
• 第三步：提前记录每个Bi的样本长度，按长度拆分整体相似度矩阵的列维度，就能得到每个A和Bi对应的相似度结果。

代码示例 #

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

# 你的原始编码数据
A = sbert.encode([sentence1, sentence2, sentence3])
B_list = [
    sbert.encode([sentence4, sentence5]),  # B1
    sbert.encode([sentence6, sentence7, sentence8]),  # B2
    # 剩余所有Bn都放入该列表
]

# 记录每个B的样本数，用于后续拆分
B_lengths = [b.shape[0] for b in B_list]
# 按行拼接所有B矩阵
B_total = np.concatenate(B_list, axis=0)
# 一次性计算所有相似度
sim_total = cosine_similarity(A, B_total)
# 拆分得到A与每个Bi的相似度矩阵
sim_results = np.split(sim_total, np.cumsum(B_lengths)[:-1], axis=1)

# 输出结果和原循环逻辑完全一致
for sim in sim_results:
    print(sim)

额外性能优化 #

如果数据量较大，还可以提前手动对A和B_total做L2归一化，直接用矩阵乘法计算点积得到余弦相似度，比调用sklearn的通用函数速度更快：

def l2_normalize(mat):
    return mat / np.linalg.norm(mat, axis=1, keepdims=True)

A_norm = l2_normalize(A)
B_total_norm = l2_normalize(B_total)
# 矩阵乘法直接得到余弦相似度，结果和cosine_similarity完全等价
sim_total = A_norm @ B_total_norm.T

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