MultiHeadAttention attn_mask argument doesn't seem to mask nan values

jagandecapri: <p>I am experimenting with MultiHeadAttention and tried the code below. However, I couldn’t get the <code>attn_mask</code> argument to work as what I’m thinking.</p>
<pre><code class="lang-python">import torch
import torch.nn as nn


tgt = torch.rand(2,3,4)
src = torch.rand(2,3,4)
src[:, 0, :] = torch.nan
print(f"tgt: {tgt.shape}, {tgt}")
print(f"src: {src.shape}, {src}")

print("First example")

mask = torch.zeros(2, 3, 3, dtype=torch.bool)
mask[:, :, 0] = True

print(f"mask: {mask.shape}, {mask.dtype}, {mask}")

multihead_attn = nn.MultiheadAttention(embed_dim=4, num_heads=1, batch_first=True)
attn_out, attn_weights = multihead_attn(tgt, src, src, attn_mask=mask)

print(f"multihead_attn attn_weights: {attn_weights.shape}, {attn_weights}")
print(f"multihead_attn out: {attn_out.shape}, {attn_out}")

print("Second example")

mask = torch.zeros(2, 3, 3, dtype=torch.bool)
mask[:, 0, :] = True

print(f"mask: {mask.shape}, {mask.dtype}, {mask}")

multihead_attn = nn.MultiheadAttention(embed_dim=4, num_heads=1, batch_first=True)
attn_out, attn_weights = multihead_attn(tgt, src, src, attn_mask=mask)

print(f"multihead_attn attn_weights: {attn_weights.shape}, {attn_weights}")
print(f"multihead_attn out: {attn_out.shape}, {attn_out}")
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">tgt: torch.Size([2, 3, 4]), tensor([[[0.5204, 0.8147, 0.7542, 0.5528],
         [0.9538, 0.2621, 0.3353, 0.3032],
         [0.3012, 0.1474, 0.5201, 0.7111]],

        [[0.4666, 0.8606, 0.8641, 0.0374],
         [0.3397, 0.3314, 0.7925, 0.1301],
         [0.5559, 0.0732, 0.4313, 0.4095]]])
src: torch.Size([2, 3, 4]), tensor([[[   nan,    nan,    nan,    nan],
         [0.6643, 0.2720, 0.9359, 0.5026],
         [0.5256, 0.6496, 0.8701, 0.9192]],

        [[   nan,    nan,    nan,    nan],
         [0.9559, 0.8751, 0.0662, 0.0216],
         [0.2466, 0.2948, 0.5645, 0.5404]]])
First example
mask: torch.Size([2, 3, 3]), torch.bool, tensor([[[ True, False, False],
         [ True, False, False],
         [ True, False, False]],

        [[ True, False, False],
         [ True, False, False],
         [ True, False, False]]])
multihead_attn attn_weights: torch.Size([2, 3, 3]), tensor([[[nan, nan, nan],
         [nan, nan, nan],
         [nan, nan, nan]],

        [[nan, nan, nan],
         [nan, nan, nan],
         [nan, nan, nan]]], grad_fn=&lt;MeanBackward1&gt;)
multihead_attn out: torch.Size([2, 3, 4]), tensor([[[nan, nan, nan, nan],
         [nan, nan, nan, nan],
         [nan, nan, nan, nan]],

        [[nan, nan, nan, nan],
         [nan, nan, nan, nan],
         [nan, nan, nan, nan]]], grad_fn=&lt;TransposeBackward0&gt;)

Second example
mask: torch.Size([2, 3, 3]), torch.bool, tensor([[[ True,  True,  True],
         [False, False, False],
         [False, False, False]],

        [[ True,  True,  True],
         [False, False, False],
         [False, False, False]]])
multihead_attn attn_weights: torch.Size([2, 3, 3]), tensor([[[nan, nan, nan],
         [nan, nan, nan],
         [nan, nan, nan]],

        [[nan, nan, nan],
         [nan, nan, nan],
         [nan, nan, nan]]], grad_fn=&lt;MeanBackward1&gt;)
multihead_attn out: torch.Size([2, 3, 4]), tensor([[[nan, nan, nan, nan],
         [nan, nan, nan, nan],
         [nan, nan, nan, nan]],

        [[nan, nan, nan, nan],
         [nan, nan, nan, nan],
         [nan, nan, nan, nan]]], grad_fn=&lt;TransposeBackward0&gt;)
</code></pre>
<p>If I understood the code correctly on MultiHeadAttention at <a href="https://github.com/pytorch/pytorch/blob/735e6ae801d69721893eca2481077168019f949a/torch/nn/functional.py#L5415-L5418" rel="noopener nofollow ugc">here</a>, the <code>nan</code> values should be masked before a Softmax is computed. I am not sure why in my case the <code>nan</code> values in <code>src</code> is not masked in both examples by the MultiHeadAttention?</p>

nairbv: <p>Be aware that: “For a binary mask, a <code>True</code> value indicates that the corresponding position is <em>not</em> allowed to attend.” Your masks are mostly False, which signifies not masking in most cases.</p>

jagandecapri: <p>I installed PyTorch from source to debug what’s going on. It looks like when there are <code>nan</code> values in the tensor, the values produced by <code>torch.baddbmm</code> at <a href="https://github.com/pytorch/pytorch/blob/735e6ae801d69721893eca2481077168019f949a/torch/nn/functional.py#L5415" rel="noopener nofollow ugc">here</a> produces <code>nan</code> values in the tensor causing the softmax to produce <code>nan</code> everywhere after that. Setting the <code>nan</code> values to <code>0</code> before feeding it to the multi head attention seems to work. Example as below. I’m not sure whether this is an intended behaviour or a bug.</p>
<pre><code class="lang-python">import torch
import torch.nn as nn

torch.manual_seed(42)

tgt = torch.rand(2,3,4)
src = torch.rand(2,3,4)
src[:, 0, :] = torch.nan
print(f"tgt: {tgt.shape}, {tgt}")
print(f"src: {src.shape}, {src}")

mask = torch.zeros(2, 3, 3, dtype=torch.bool)
mask[:, :, 0] = True
print(f"mask: {mask.shape}, {mask.dtype}, {mask}")

print("First example")

multihead_attn = nn.MultiheadAttention(embed_dim=4, num_heads=1, batch_first=True)
attn_out, attn_weights = multihead_attn(tgt, src, src, attn_mask=mask)

print(f"multihead_attn attn_weights: {attn_weights.shape}, {attn_weights}")
print(f"multihead_attn out: {attn_out.shape}, {attn_out}")

print("Second example")

src_replace_nan = torch.where(torch.isnan(src), 0, src)
print(f"src_replace_nan: {src_replace_nan.shape}, {src_replace_nan}")

multihead_attn = nn.MultiheadAttention(embed_dim=4, num_heads=1, batch_first=True)
attn_out, attn_weights = multihead_attn(tgt, src_replace_nan, src_replace_nan, attn_mask=mask)

print(f"multihead_attn attn_weights: {attn_weights.shape}, {attn_weights}")
print(f"multihead_attn out: {attn_out.shape}, {attn_out}")
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">tgt: torch.Size([2, 3, 4]), tensor([[[0.8823, 0.9150, 0.3829, 0.9593],
         [0.3904, 0.6009, 0.2566, 0.7936],
         [0.9408, 0.1332, 0.9346, 0.5936]],

        [[0.8694, 0.5677, 0.7411, 0.4294],
         [0.8854, 0.5739, 0.2666, 0.6274],
         [0.2696, 0.4414, 0.2969, 0.8317]]])
src: torch.Size([2, 3, 4]), tensor([[[   nan,    nan,    nan,    nan],
         [0.5472, 0.0062, 0.9516, 0.0753],
         [0.8860, 0.5832, 0.3376, 0.8090]],

        [[   nan,    nan,    nan,    nan],
         [0.6343, 0.3644, 0.7104, 0.9464],
         [0.7890, 0.2814, 0.7886, 0.5895]]])
mask: torch.Size([2, 3, 3]), torch.bool, tensor([[[ True, False, False],
         [ True, False, False],
         [ True, False, False]],

        [[ True, False, False],
         [ True, False, False],
         [ True, False, False]]])
First example
multihead_attn attn_weights: torch.Size([2, 3, 3]), tensor([[[nan, nan, nan],
         [nan, nan, nan],
         [nan, nan, nan]],

        [[nan, nan, nan],
         [nan, nan, nan],
         [nan, nan, nan]]], grad_fn=&lt;MeanBackward1&gt;)
multihead_attn out: torch.Size([2, 3, 4]), tensor([[[nan, nan, nan, nan],
         [nan, nan, nan, nan],
         [nan, nan, nan, nan]],

        [[nan, nan, nan, nan],
         [nan, nan, nan, nan],
         [nan, nan, nan, nan]]], grad_fn=&lt;TransposeBackward0&gt;)
Second example
src_replace_nan: torch.Size([2, 3, 4]), tensor([[[0.0000, 0.0000, 0.0000, 0.0000],
         [0.5472, 0.0062, 0.9516, 0.0753],
         [0.8860, 0.5832, 0.3376, 0.8090]],

        [[0.0000, 0.0000, 0.0000, 0.0000],
         [0.6343, 0.3644, 0.7104, 0.9464],
         [0.7890, 0.2814, 0.7886, 0.5895]]])
multihead_attn attn_weights: torch.Size([2, 3, 3]), tensor([[[0.0000, 0.3902, 0.6098],
         [0.0000, 0.4472, 0.5528],
         [0.0000, 0.4268, 0.5732]],

        [[0.0000, 0.5090, 0.4910],
         [0.0000, 0.5182, 0.4818],
         [0.0000, 0.4987, 0.5013]]], grad_fn=&lt;MeanBackward1&gt;)
multihead_attn out: torch.Size([2, 3, 4]), tensor([[[ 0.0439, -0.0275,  0.0880, -0.1783],
         [ 0.0369, -0.0280,  0.0826, -0.1601],
         [ 0.0394, -0.0278,  0.0845, -0.1666]],

        [[ 0.1194,  0.0144,  0.1713, -0.2186],
         [ 0.1204,  0.0149,  0.1723, -0.2194],
         [ 0.1184,  0.0139,  0.1702, -0.2178]]], grad_fn=&lt;TransposeBackward0&gt;)
</code></pre>