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 中：

1. output = scaled_dot_product_attention(query, key, value, attn_mask)

复制代码

此中 attn_mask 是 additive mask，即：

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

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

1. softmax([-inf, -inf, ..., -inf]) → [NaN, NaN, ..., NaN]

复制代码

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

---

产生这种情况的常见缘故起因

这种情况经常发生在任务中存在大量 query（例如图像 patch、token、时间步）本身就不应该 attend 到任何 key，例如背景地区或 padding 地区。
因此，固然逻辑公道，但仍然在数学上不正当。

---

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

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

1. # attention_mask: [B, Q, K]，bool 类型，True 表示“可以 attend”
2. attention_mask_float = torch.full_like(attention_mask, float('-inf'), dtype=query.dtype)
3. attention_mask_float.masked_fill_(attention_mask, 0.0)
5. # fallback：避免某些 query 全为 -inf
6. all_inf_rows = (attention_mask_float == float('-inf')).all(dim=-1, keepdim=True)  # [B, Q, 1]
7. if all_inf_rows.any():
8.     last_key_idx = attention_mask_float.size(-1) - 1
9.     fix_mask = torch.arange(attention_mask_float.size(-1), device=attention_mask.device) == last_key_idx
10.     fix_mask = fix_mask.view(1, 1, -1)  # reshape for broadcast
11.     attention_mask_float = attention_mask_float.masked_fill(all_inf_rows & fix_mask, 0.0)

复制代码

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

---

可视化发起

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

1. # 黑色：可见（0.0），白色：被 mask（-inf）
2. vis_mask = (attn_mask == 0.0).astype(np.uint8)
3. plt.imshow(vis_mask, cmap='Greys', aspect='auto')

复制代码

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

---

总结

条目发起是否允许 query 全被屏蔽语义上允许，数学上不正当（需处理处罚）PyTorch 是否兜底否，需用户自己容错是否应解锁一个 dummy key是，最安全的 fallback 机制可否通过可视化排查是，好坏图可快速识别空行

---