np.einsum与np.sum求和的差异及精度差异原因探究

问题背景 #

测试发现，用np.einsum直接完成乘法+归约的结果，和先做乘法再用np.sum归约的结果存在精度差异，但单独的乘法步骤结果完全一致。已知np.sum是基于np.add的成对求和实现，想了解np.einsum的求和机制有何不同。

测试代码与结果 #

import numpy as np

test_data = np.random.random((100, 90, 180)).astype(np.float32)
test_weights = np.random.random((90, 180)).astype(np.float32)

# einsum直接完成乘法+归约
test1 = np.einsum('...abc,...bc->...a', test_data, test_weights)
# 先einsum乘法再sum归约
test2 = np.sum(np.einsum('...abc,...bc->...abc', test_data, test_weights), axis=(1, 2))
print(np.array_equal(test1, test2)) 
# 输出：False

# 先multiply再sum归约
test3 = np.sum(np.multiply(test_data, test_weights), axis=(1, 2))
print(np.array_equal(test1, test3)) 
# 输出：False

# 两种乘法方式结果一致
print(np.array_equal(test2, test3)) 
# 输出：True
print(np.array_equal(np.einsum('...abc,...bc->abc', test_data, test_weights), np.multiply(test_data, test_weights)))
# 输出：True

核心原因分析 #

1. np.sum的求和机制 #

np.sum本质是调用np.add进行成对迭代累加：每次从数组中取两个元素相加，将结果存入临时变量，再继续和下一个元素相加。对于float32这种低精度类型，每次加法的舍入误差会不断累积，当累加元素数量较多时（比如本例中每个a对应90*180=16200个元素），误差会被放大。

2. np.einsum的归约机制 #

当einsum处理类似...abc,...bc->...a的收缩操作时，底层并非简单调用np.sum，而是做了以下优化：

• 调用BLAS库的优化实现：这类收缩本质是对每个a维度的切片，执行与test_weights的逐元素相乘后求和，等价于矩阵-向量的点积操作。einsum会自动识别这种模式，调用BLAS中的gemv（矩阵向量乘法）或类似优化函数，这类函数通常会用**更高精度的中间存储（比如float64）**来累加结果，大幅减少舍入误差。
• 更优的累加顺序：部分BLAS实现会采用分块累加或树状累加的方式，相比成对累加能更有效地降低误差累积。

3. 精度差异的本质验证 #

用np.allclose替代np.array_equal可以看到，两者的差异是微小的舍入误差，而非逻辑错误：

print(np.allclose(test1, test3, atol=1e-6))
# 输出：True

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