PyTorch 的 F.scaled_dot_product_attention 返回Nan


"为什么 PyTorch 的 scaled_dot_product_attention 会输出 NaN?如何正确构造 Attention Mask"


引言:看似正常的 mask,为什么会引发 NaN?

在使用 F.scaled_dot_product_attention 构建跨模态或多源注意力时,我们常通过 attention_mask 控制每个 query 位置能看到哪些 key。但如果不小心构造出某些 query 对所有 key 都不可见的情况,就会在 softmax 中触发 NaN,进而让模型 loss 崩溃。

这个问题隐蔽却常见,且 PyTorch 不会自动容错,需要我们显式处理。


问题复现:全 -inf 行将导致 NaN

在 PyTorch 的 scaled attention 中:

python 复制代码
output = scaled_dot_product_attention(query, key, value, attn_mask)

其中 attn_maskadditive mask,即:

  • 0.0: 表示该位置可见;
  • -inf: 表示该位置被屏蔽,不可 attend。

当某个 query 行的 mask 全为 -inf 时,softmax 输入类似于:

python 复制代码
softmax([-inf, -inf, ..., -inf]) → [NaN, NaN, ..., NaN]

这将污染整个计算图,最终导致 loss 为 NaN。


产生这种情况的常见原因

这种情况经常发生在任务中存在大量 query(例如图像 patch、token、时间步)本身就不应该 attend 到任何 key,例如背景区域或 padding 区域。

因此,虽然逻辑合理,但仍然在数学上不合法


解决方案:fallback 解锁最后一个 key

为避免 NaN,可在转换 bool mask → float mask 时引入一个 fallback:

python 复制代码
# attention_mask: [B, Q, K],bool 类型,True 表示"可以 attend"
attention_mask_float = torch.full_like(attention_mask, float('-inf'), dtype=query.dtype)
attention_mask_float.masked_fill_(attention_mask, 0.0)

# fallback:避免某些 query 全为 -inf
all_inf_rows = (attention_mask_float == float('-inf')).all(dim=-1, keepdim=True)  # [B, Q, 1]
if all_inf_rows.any():
    last_key_idx = attention_mask_float.size(-1) - 1
    fix_mask = torch.arange(attention_mask_float.size(-1), device=attention_mask.device) == last_key_idx
    fix_mask = fix_mask.view(1, 1, -1)  # reshape for broadcast
    attention_mask_float = attention_mask_float.masked_fill(all_inf_rows & fix_mask, 0.0)

这样即便某个 query 原本完全不可见,也能保证 softmax 至少有一个有效分布。


可视化建议

可以使用 matplotlib.imshow 直接可视化 [Q, K] 的 mask 分布:

python 复制代码
# 黑色:可见(0.0),白色:被 mask(-inf)
vis_mask = (attn_mask == 0.0).astype(np.uint8)
plt.imshow(vis_mask, cmap='Greys', aspect='auto')

可视化能帮助你快速定位全白 query 行,即潜在 NaN 风险点。


总结

条目 建议
是否允许 query 全被屏蔽 语义上允许,数学上不合法(需处理)
PyTorch 是否兜底 否,需用户自己容错
是否应解锁一个 dummy key 是,最安全的 fallback 机制
可否通过可视化排查 是,黑白图可快速识别空行

相关推荐
一只鹿鹿鹿26 分钟前
【网络安全】信息网络安全建设方案(WORD)
人工智能·安全·spring·web安全·低代码
小拇指~27 分钟前
梯度下降的基本原理
人工智能·算法·计算机视觉
AndrewHZ1 小时前
【图像处理基石】如何对遥感图像进行实例分割?
图像处理·人工智能·python·大模型·实例分割·detectron2·遥感图像分割
No0d1es1 小时前
第13届蓝桥杯Python青少组中/高级组选拔赛(STEMA)2021年11月27日真题
python·青少年编程·蓝桥杯·选拔赛
CodeShare1 小时前
某中心将举办机器学习峰会
人工智能·机器学习·数据科学
天天找自己1 小时前
精通分类:解析Scikit-learn中的KNN、朴素贝叶斯与决策树(含随机森林)
python·决策树·机器学习·分类·scikit-learn
那就摆吧2 小时前
U-Net vs. 传统CNN:为什么医学图像分割需要跳过连接?
人工智能·神经网络·cnn·u-net·医学图像
深度学习实战训练营2 小时前
中英混合的语音识别XPhoneBERT 监督的音频到音素的编码器结合 f0 特征LID
人工智能·音视频·语音识别
WADesk---瓜子2 小时前
用 AI 自动生成口型同步视频,短视频内容也能一人完成
人工智能·音视频·语音识别·流量运营·用户运营
赵英英俊2 小时前
Python day31
开发语言·python