为何Python广播在向量差范数计算中比简单循环慢？

首先明确你的测试场景与代码：
你有一个形状为(1000, 10000)的数组M，想要计算每个行向量与第一个行向量v的差的平方范数，分别实现了循环版本和广播版本，却发现广播版本速度反而更慢。

你的测试代码：

import numpy as np

def norm_loop(M, v):
    n = M.shape[0]
    d = np.zeros(n)
    for i in range(n):
        d[i] = np.sum((M[i] - v)**2)
    return d

def norm_bcast(M, v):
    n = M.shape[0]
    d = np.zeros(n)
    d = np.sum((M - v)**2, axis=1)
    return d

M = np.random.random_sample((1000, 10000))
v = M[0]
%timeit norm_loop(M, v)  # 25.9 ms
%timeit norm_bcast(M, v) # 速度慢于循环

核心原因：内存带宽与缓存利用率的瓶颈 #

广播看起来是“向量化”操作，直觉上应该更快，但在这个场景下反而慢，主要有这几个关键点：

• 临时数组的内存开销：广播版本中M - v会创建一个完整的(1000, 10000)临时数组（约80MB的float64数据）。这个数组的大小远超CPU的L2/L3缓存（通常仅几MB到几十MB），导致CPU需要频繁从主存读写数据，而主存的访问速度比缓存慢几个数量级，这直接成为了性能瓶颈。
• 循环版本的缓存友好性：循环版本每次只处理一行数据（M[i] - v是一个(10000,)的数组，约80KB），这个大小完全可以被CPU缓存容纳。每次计算时，数据都在缓存中，内存访问延迟极低，而且np.sum的底层是C实现的高效操作，抵消了Python循环的少量开销。
• 临时数组的创建销毁成本：广播版本需要一次性分配大内存块，计算完成后再释放，这部分额外的内存管理开销也会拖慢速度；而循环版本每次只分配小内存块，开销几乎可以忽略。

优化方案：用数学公式避免临时数组 #

我们可以利用范数的代数性质来规避大临时数组的创建，从而大幅提升性能。对于两个向量a和b，差的平方范数满足：
$$||a - b||^2 = ||a||^2 + ||b||^2 - 2a·b$$

基于这个公式，我们可以实现一个更高效的版本：

def norm_opt(M, v):
    # 计算每个行向量的平方和
    row_norms = np.sum(M ** 2, axis=1)
    # 计算v的平方和
    v_norm = np.sum(v ** 2)
    # 计算M与v的点积（每行与v的点积）
    dot_products = M @ v
    # 代入公式计算
    return row_norms + v_norm - 2 * dot_products

这个版本不需要创建任何大的临时数组，所有操作都是基于原数组的逐元素运算或矩阵乘法，内存占用极小，缓存利用率极高，速度会远快于循环和广播版本。你可以测试一下，这个版本的耗时通常会降到几毫秒级别。

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