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

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

其中 attn_mask 是 additive mask，即：

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

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

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

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

产生这种情况的常见原因

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

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

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

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

# 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 分布：

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

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

总结