PyTorch Scatter库中scatter_logsumexp函数的输出处理问题分析

问题背景

PyTorch Scatter库是一个用于高效执行散射操作(将数据根据索引分配到不同位置)的扩展库。其中scatter_logsumexp函数用于在log空间执行安全的求和操作，这在深度学习特别是概率模型中非常有用。

问题描述

在最新版本的PyTorch Scatter库中，scatter_logsumexp函数存在一个关键缺陷：当输出张量中有未被索引修改的位置时，这些位置的值会被错误地设置为0，而不是保留其原始值。

问题重现

通过以下代码可以重现该问题：

import torch
from torch_scatter import scatter_logsumexp

src = torch.tensor([-1., -50])  # 输入数据
index = torch.tensor([0, 0])    # 索引，两个元素都映射到位置0

out = torch.full((2,), -10.)    # 初始输出张量，所有位置设为-10

scatter_logsumexp(src=src, index=index, out=out)
# 实际输出: tensor([-0.9999,  0.0000])
# 期望输出: tensor([-0.9999,  -10])

问题分析

1. 预期行为：未被索引修改的输出位置应保持原值不变
2. 实际行为：所有未被修改的位置被强制设为0
3. 影响范围：该问题使得函数只能在所有输出位置都被索引修改的情况下正常工作

技术细节

问题的根源在于函数实现中对输出张量的处理方式。当前实现中有一个nan_to_num_(neginf=0.0)的调用，这会将所有负无穷值转换为0。然而：

1. 在log空间运算中，负无穷(-inf)是表示概率为0的合法值
2. 对于未被索引修改的位置，应该保留其原始值而非强制设为0

解决方案

仓库所有者已经提交了修复该问题的PR。修复方案主要包括：

1. 移除不必要的nan_to_num_转换
2. 确保只修改被索引引用的输出位置
3. 保留未被修改位置的原始值

扩展讨论

这个问题实际上反映了log空间运算中的一个常见陷阱。在概率和深度学习中，log空间运算常用于避免数值下溢，但需要特别注意：

1. 正确处理log(0)的情况(即负无穷)
2. 保持运算的数值稳定性
3. 确保未参与运算的值不被意外修改

总结

PyTorch Scatter库中的scatter_logsumexp函数在处理未被索引修改的输出位置时存在缺陷，这会影响函数在部分场景下的正确性。用户在使用时应注意这个问题，并关注库的更新以获取修复版本。对于需要处理稀疏数据或部分更新的场景，建议暂时验证函数的输出是否符合预期。