使用PyTorch张量运算高效实现中文文本批次余弦相似度计算

使用PyTorch张量运算高效实现中文文本批次余弦相似度计算

问题背景与核心挑战

在自然语言处理和文本挖掘任务中，计算文本之间的相似度是一项基础且关键的操作。余弦相似度因其对向量大小的不敏感性，成为衡量文本语义相似度的常用指标。然而，当需要处理大规模中文文本数据集时，如何高效地计算批次内或批次间的文本相似度矩阵成为了一个性能瓶颈。传统的循环遍历方法计算复杂度高，无法充分利用现代硬件（如GPU）的并行计算能力。PyTorch作为主流的深度学习框架，其张量运算能够高效地在GPU上执行，为批量计算余弦相似度提供了理想的技术基础。

余弦相似度的数学原理与向量化表达

余弦相似度通过计算两个向量在空间中的夹角的余弦值来衡量其相似性，取值范围为[-1, 1]。对于向量A和B，其计算公式为：cosine_similarity(A, B) = (A · B) / (||A|| ||B||)。在批量计算场景下，我们需要将这一计算过程向量化。假设我们有一个文本向量矩阵X，其形状为[batch_size, embedding_dim]，我们希望计算该批次内所有文本两两之间的相似度矩阵S，其中S[i, j]表示第i个文本与第j个文本的余弦相似度。向量化计算的关键步骤包括：首先计算每行向量的L2范数，得到一个形状为[batch_size]的范数向量；接着计算矩阵X与其自身的点积，得到形状为[batch_size, batch_size]的矩阵；最后通过广播机制，将点积结果除以两个范数向量的外积，从而得到最终的相似度矩阵。

PyTorch高效实现的关键步骤

利用PyTorch实现上述向量化计算，可以避免低效的Python循环。核心代码如下所示：

import torchdef batch_cosine_similarity(X):    # 对嵌入向量进行L2归一化    X_normalized = torch.nn.functional.normalize(X, p=2, dim=1)    # 计算归一化后向量的矩阵乘法，直接得到余弦相似度矩阵    similarity_matrix = torch.mm(X_normalized, X_norma