大值log(exp(x))高效计算与累积求和的数值稳定方案咨询

DevOps 小满

2026-7-1

解决方案：累积LogSumExp的高效向量化实现

针对大数值tensor计算累积log(sum(exp(x[0:i])))时的溢出问题，我们可以通过构造下三角掩码矩阵，结合torch.logsumexp实现完全向量化的数值稳定计算，无需循环。

核心思路

直接计算exp(x)会因数值过大溢出为inf，而torch.logsumexp通过以下公式保证数值稳定性：

log(sum(exp(x))) = max(x) + log(sum(exp(x - max(x))))

对于累积求和场景，我们可以构造一个下三角矩阵，每行保留当前位置及之前的所有元素，其余位置填充-inf（因为exp(-inf)=0，不影响求和），然后对每行执行logsumexp即可得到对应位置的累积结果。

代码实现

import torch

# 示例输入tensor
x = torch.tensor([1000.0, 2000.0, 3000.0])
n = x.size(0)

# 构造下三角掩码：行索引 >= 列索引的位置为True
mask = torch.arange(n)[None, :] <= torch.arange(n)[:, None]

# 扩展x为n×n矩阵，并将掩码外的位置设为-inf
x_masked = torch.where(mask, x.expand(n, n), torch.tensor(-float('inf'), device=x.device))

# 对每行执行logsumexp，得到累积结果
result = torch.logsumexp(x_masked, dim=1)

print(result)
# 输出：tensor([1000.0000, 2000.0000, 3000.0000])
# 注：实际计算中，当元素差距较大时，较小的exp项会被主导，结果近似等于当前最大值，这是数值上的合理现象

原理说明

掩码构造：torch.arange(n)[None, :]生成行向量[0,1,2]，torch.arange(n)[:, None]生成列向量[[0],[1],[2]]，两者比较得到下三角掩码矩阵：
```
tensor([[ True, False, False],
        [ True,  True, False],
        [ True,  True,  True]])
```
矩阵填充：将x扩展为n×n矩阵后，用掩码过滤掉每行不需要的元素，替换为-inf，确保这些位置在求和时贡献为0。
logsumexp计算：对每行执行logsumexp，自动处理数值稳定问题，避免直接计算大数值exp导致的溢出。