63695×384维矩阵余弦相似度矩阵的向量化计算方法

核心数学原理 #

两个行向量$x_i$、$x_j$的余弦相似度公式为：
$$sim(x_i,x_j) = \frac{x_i \cdot x_j}{|x_i|_2 \times |x_j|_2}$$
其中$|x_i|_2$是向量$x_i$的L2范数。如果提前把所有向量归一化为L2范数等于1的单位向量，余弦相似度就直接等于两个向量的点积，整个计算可以完全转化为底层高度优化的矩阵运算，不需要任何手写Python层循环。

具体实现流程 #

• 第一步：类型转换压缩内存
  先把原始63695×384的矩阵转为float32类型，相比默认的float64可以直接省一半内存，384维向量的相似度计算用float32精度完全足够，不会有可感知的精度损失。
• 第二步：计算行向量L2范数
  沿矩阵的行方向（axis=1）计算每个行向量的L2范数，得到形状为63695×1的范数矩阵，计算时给所有范数值加一个极小值1e-8，避免原始矩阵存在全零行时触发除零错误产生NaN值。

  注意：计算范数时必须保留行维度（设置keepdims=True），否则后续无法通过广播机制完成逐行归一化。

• 第三步：生成单位向量矩阵
  将原始矩阵的每一行除以对应行的L2范数，得到归一化后的矩阵$X_{norm}$，形状仍为63695×384，此时矩阵每一行的L2范数均为1。
• 第四步：矩阵乘法得到相似度矩阵
  用归一化矩阵乘以自身的转置，即$sim = X_{norm} @ X_{norm}.T$，得到的63695×63695的矩阵就是全量余弦相似度矩阵，其中sim[i][j]就是第i个行向量和第j个行向量的余弦相似度，取值范围为[-1,1]，对角线元素恒为1（向量与自身的余弦相似度为1）。

关键注意事项 #

• 内存占用提前评估：全量相似度矩阵如果用float32存储，总大小约15GB，计算过程中还需要额外存储原始矩阵、归一化矩阵，建议可用内存至少20G再跑全量计算；如果内存不足，不要一次性加载全量结果，采用分块计算策略：每次取归一化矩阵的k行（比如k=5000~10000），和归一化矩阵的转置做乘法得到k×63695的相似度块，块计算完成后直接写入磁盘或做下游处理，逐块跑完所有行即可。
• 硬件加速选择：如果有NVIDIA显卡，可将矩阵放到GPU显存中计算，矩阵乘法速度比CPU快10~100倍，代码逻辑和CPU版本完全一致，只需把数组替换为PyTorch、CuPy支持的GPU张量即可。
• 避坑：不要手动写Python层的二重/一重循环遍历向量计算，哪怕逻辑正确，运行速度也会比底层优化的矩阵乘法慢几百上千倍，完全达不到向量化计算的效率。

参考代码 #

NumPy CPU版本 #

import numpy as np

# 原始矩阵X，形状为(63695, 384)
X = X.astype(np.float32)
# 计算行L2范数，加极小值防除零
row_norm = np.linalg.norm(X, axis=1, keepdims=True) + 1e-8
# 逐行归一化
X_norm = X / row_norm
# 矩阵乘法算相似度
sim_matrix = X_norm @ X_norm.T

PyTorch GPU版本（速度更快） #

import torch

# 自动选择可用GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 原始矩阵转torch张量，放到GPU，转float32
X = torch.tensor(X, dtype=torch.float32, device=device)
# 计算行范数
row_norm = torch.norm(X, p=2, dim=1, keepdim=True) + 1e-8
# 归一化
X_norm = X / row_norm
# 矩阵乘法算相似度
sim_matrix = X_norm @ X_norm.T
# 如果需要转回numpy数组，先移回CPU再转
# sim_matrix = sim_matrix.cpu().numpy()

内容的提问来源于stack exchange，提问作者An old man in the sea.