为何带JIT的矩阵余弦相似度计算比循环向量计算慢？

我正在优化余弦相似度的计算效率，最初是把向量存入矩阵后逐行调用计算；后来改成移除内层循环，直接对整个矩阵做计算——无JIT时处理时间确实减少了，但带JIT的向量循环计算速度还是更快。给矩阵计算函数加上JIT后，处理时间反而大幅飙升，效率变成了最低。

我曾怀疑问题出在matrix ** 2，试过改用np.linalg.norm，但当前Numba版本不支持axis参数。我已经找到一种高效的Numba矩阵余弦计算实现，但还是想弄明白我的代码为什么会出现这种反向优化的问题。

实验代码 #

模块导入 #

import numpy as np
from numba import jit
from timeit import default_timer as timer

计算函数 #

# @jit(nopython=True)
def cosine_sim_matrix(matrix, vector):
    dot_product = np.dot(matrix, vector)
    magnitude_a = np.sqrt((matrix ** 2).sum(axis=1))
    magnitude_b = np.sqrt((vector ** 2).sum())
    similarity_score = dot_product / (magnitude_a * magnitude_b)
    similarity_score[similarity_score > 1] = 1
    similarity_score[similarity_score < -1] = -1
    return similarity_score

# @jit(nopython=True)
def cosine_sim_vector(vector_1, vector_2):
    dot_product = np.dot(vector_1, vector_2)
    magnitude_a = np.sqrt((vector_1** 2).sum())
    magnitude_b = np.sqrt((vector_2 ** 2).sum())
    similarity_score = dot_product / (magnitude_a * magnitude_b)
    if similarity_score > 1:
        similarity_score = 1.0
    elif similarity_score < -1:
        similarity_score = -1.0 
    return similarity_score

测试代码 #

mymatrix = np.random.rand(10000,1000)

# 计时矩阵计算
start = timer()
for idx in range(mymatrix.shape[0]):
    vec = mymatrix[idx,:]
    cosine_sim_matrix(mymatrix, vec)
end = timer()
print(f"process time, matrix={end-start}")

# 计时向量计算
start = timer()
for idx in range(mymatrix.shape[0]):
    vec = mymatrix[idx,:]
    for idx2 in range(mymatrix.shape[0]):
        vec2 = mymatrix[idx2,:]
        cosine_sim_vector(vec, vec2)
end = timer()
print(f"process time, vector={end-start}")

实验结果 #

# 无JIT
process time, matrix=348.5313815
process time, vector=839.2151422999999

# 带JIT
process time, matrix=883.4681065999998
process time, vector=275.0485957000001

补充：高效实现的测试结果 #

process time, no parallel, matrix=109.96504489999916
process time, w/ parallel, matrix=21.85965199999964 

问题原因分析 #

• Numba对Numpy高阶操作的优化局限性：cosine_sim_matrix依赖大量Numpy广播（matrix ** 2）和按轴聚合（sum(axis=1)）操作，Numba的JIT编译器无法像优化纯循环代码那样生成高效机器码，反而会因兼容Numpy API语义引入额外开销，比如临时内存对象创建、不必要的数组拷贝，抵消了原本的性能优势。
• 向量循环的JIT优化更彻底：cosine_sim_vector是纯向量级计算，逻辑简单，Numba能直接将其编译为高度优化的机器码，没有Numpy高阶操作的负担；同时循环逻辑在JIT下可充分优化（如寄存器重用、循环展开），所以带JIT的向量循环反而更快。
• 无JIT时矩阵计算的优势来自Numpy底层：未加JIT时，cosine_sim_matrix的Numpy操作调用的是底层优化过的C实现，比纯Python循环的向量计算快，但加上JIT后，Numba未能有效优化这些Numpy操作，反而导致性能下降。

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