手动近似计算log-sum-exp为何比NumPy/SciPy内置函数更快？

在处理大数组时，我发现手动计算对数和指数（log-sum-exp）的速度居然比SciPy的scipy.special.logsumexp甚至NumPy的numpy.logaddexp.reduce内置函数更快。原本以为库实现会更优化，想知道是不是操作有误或者忽略了某些关键点？

测试代码 #

# test_logsumexp.py

from functools import partial
import timeit
import numpy as np
from scipy.special import logsumexp


def logsumexp_manual(a, axis=-1):
    max_vals = np.max(a, axis=axis)
    exp_terms = np.exp(a - max_vals[..., np.newaxis])
    sum_exp_terms = np.sum(exp_terms, axis=axis)
    return max_vals + np.log(sum_exp_terms)

# 有效性校验
arr = np.random.rand(10_000, 30, 3)
scipy_result = logsumexp(arr, axis=-1)
numpy_result = np.logaddexp.reduce(arr, axis=-1)
manual_result = logsumexp_manual(arr)
np.testing.assert_allclose(scipy_result, manual_result)
np.testing.assert_allclose(scipy_result, numpy_result)


n_loops = 10
setup = ("import numpy as np; "
         "from scipy.special import logsumexp, softmax; "
         "arr = np.random.rand(10_000, 30, 3)")

total_scipy = timeit.timeit(
    "logsumexp(arr, axis=-1)",
    setup=setup,
    number=n_loops,
)

total_numpy = timeit.timeit(
    "np.logaddexp.reduce(arr, axis=-1)",
    setup=setup,
    number=n_loops,
)

logsumexp_partial = partial(logsumexp_manual, a=arr)
total_manual = timeit.timeit(
    logsumexp_partial,
    setup=setup,
    number=n_loops,
)

print(f"Scipy  logsumexp:  {total_scipy  / n_loops:.6f} seconds per loop")
print(f"Numpy  logaddexp:  {total_numpy  / n_loops:.6f} seconds per loop")
print(f"Manual logsumexp:  {total_manual / n_loops:.6f} seconds per loop")
smallest = min(total_scipy, total_numpy, total_manual)

控制台输出 #

Scipy  logsumexp:   0.070824 seconds per loop
Numpy  logaddexp:   0.044226 seconds per loop
Manual logsumexp:   0.031754 seconds per loop

原因分析 #

• SciPy的logsumexp有额外通用开销：它支持权重参数b、nan/inf异常处理、多轴灵活解析等通用场景，这些适配逻辑会带来额外计算成本，而手动实现只针对当前需求做了最精简的操作。
• NumPy的logaddexp.reduce是累积式计算：reduce采用逐元素迭代累积结果的方式，而手动实现用的是全量向量化的max、exp、sum操作。在你测试的场景中，最后一维长度仅为3，向量化批量操作的开销远低于迭代累积的循环开销。
• 场景匹配度差异：手动实现完全适配当前数组的维度和需求，没有多余分支判断；而库函数要兼容各种维度、数据类型和异常情况，在特定小维度场景下反而不如定制化的手动实现高效。

如果换用更大的最后一维（比如长度1000），库函数的底层BLAS加速求和等优化会体现出优势，此时手动实现的速度可能会被反超。

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