MultiheadAttention after LSTM returns the same output for all input

Vu_Nguyen: <p>Hi,</p>
<p>I would like to use MultiheadAttention as self-attention after applying LSTM on a single sequence.</p>
<pre><code class="lang-auto">import torch

# shape: (sequence length, batch size, embedding dimension)
inp = torch.randn(5, 3, 10)
lstm = torch.nn.LSTM(input_size=10, hidden_size=10, num_layers=2)
self_attn = torch.nn.MultiheadAttention(embed_dim=10, num_heads=2)

x, _ = lstm(input=inp)
# x is query, key, value at the same time
out, weight = self_attn(query=x, key=x, value=x)

out
tensor([[[ 0.0137,  0.0457, -0.0169, -0.0393, -0.0214,  0.0162,  0.0534,
           0.0202,  0.0519, -0.0192],
         [ 0.0116,  0.0359, -0.0044, -0.0320, -0.0294,  0.0175,  0.0573,
           0.0373,  0.0316, -0.0242],
         [ 0.0128,  0.0410, -0.0083, -0.0347, -0.0234,  0.0136,  0.0642,
           0.0148,  0.0425, -0.0273]],

        [[ 0.0137,  0.0457, -0.0169, -0.0393, -0.0214,  0.0162,  0.0534,
           0.0202,  0.0519, -0.0192],
         [ 0.0116,  0.0359, -0.0044, -0.0320, -0.0294,  0.0175,  0.0573,
           0.0373,  0.0316, -0.0242],
         [ 0.0128,  0.0410, -0.0083, -0.0347, -0.0234,  0.0136,  0.0642,
           0.0148,  0.0425, -0.0273]],

        [[ 0.0137,  0.0457, -0.0169, -0.0393, -0.0214,  0.0162,  0.0534,
           0.0202,  0.0519, -0.0192],
         [ 0.0116,  0.0359, -0.0044, -0.0320, -0.0294,  0.0175,  0.0573,
           0.0373,  0.0316, -0.0242],
         [ 0.0128,  0.0410, -0.0083, -0.0347, -0.0234,  0.0136,  0.0642,
           0.0148,  0.0425, -0.0273]],

        [[ 0.0137,  0.0457, -0.0169, -0.0393, -0.0214,  0.0162,  0.0534,
           0.0202,  0.0519, -0.0192],
         [ 0.0116,  0.0359, -0.0044, -0.0320, -0.0294,  0.0175,  0.0573,
           0.0373,  0.0316, -0.0242],
         [ 0.0128,  0.0410, -0.0083, -0.0347, -0.0234,  0.0136,  0.0642,
           0.0148,  0.0425, -0.0273]],

        [[ 0.0137,  0.0457, -0.0169, -0.0393, -0.0214,  0.0162,  0.0534,
           0.0202,  0.0519, -0.0192],
         [ 0.0116,  0.0359, -0.0044, -0.0320, -0.0294,  0.0175,  0.0573,
           0.0373,  0.0316, -0.0242],
         [ 0.0128,  0.0410, -0.0083, -0.0347, -0.0234,  0.0136,  0.0642,
           0.0148,  0.0425, -0.0273]]], grad_fn=&lt;AddBackward0&gt;)

# shape: (batch size, sequence length, embedding dimension)
out.transpose(0, 1)
tensor([[[ 0.0137,  0.0457, -0.0169, -0.0393, -0.0214,  0.0162,  0.0534,
           0.0202,  0.0519, -0.0192],
         [ 0.0137,  0.0457, -0.0169, -0.0393, -0.0214,  0.0162,  0.0534,
           0.0202,  0.0519, -0.0192],
         [ 0.0137,  0.0457, -0.0169, -0.0393, -0.0214,  0.0162,  0.0534,
           0.0202,  0.0519, -0.0192],
         [ 0.0137,  0.0457, -0.0169, -0.0393, -0.0214,  0.0162,  0.0534,
           0.0202,  0.0519, -0.0192],
         [ 0.0137,  0.0457, -0.0169, -0.0393, -0.0214,  0.0162,  0.0534,
           0.0202,  0.0519, -0.0192]],

        [[ 0.0116,  0.0359, -0.0044, -0.0320, -0.0294,  0.0175,  0.0573,
           0.0373,  0.0316, -0.0242],
         [ 0.0116,  0.0359, -0.0044, -0.0320, -0.0294,  0.0175,  0.0573,
           0.0373,  0.0316, -0.0242],
         [ 0.0116,  0.0359, -0.0044, -0.0320, -0.0294,  0.0175,  0.0573,
           0.0373,  0.0316, -0.0242],
         [ 0.0116,  0.0359, -0.0044, -0.0320, -0.0294,  0.0175,  0.0573,
           0.0373,  0.0316, -0.0242],
         [ 0.0116,  0.0359, -0.0044, -0.0320, -0.0294,  0.0175,  0.0573,
           0.0373,  0.0316, -0.0242]],

        [[ 0.0128,  0.0410, -0.0083, -0.0347, -0.0234,  0.0136,  0.0642,
           0.0148,  0.0425, -0.0273],
         [ 0.0128,  0.0410, -0.0083, -0.0347, -0.0234,  0.0136,  0.0642,
           0.0148,  0.0425, -0.0273],
         [ 0.0128,  0.0410, -0.0083, -0.0347, -0.0234,  0.0136,  0.0642,
           0.0148,  0.0425, -0.0273],
         [ 0.0128,  0.0410, -0.0083, -0.0347, -0.0234,  0.0136,  0.0642,
           0.0148,  0.0425, -0.0273],
         [ 0.0128,  0.0410, -0.0083, -0.0347, -0.0234,  0.0136,  0.0642,
           0.0148,  0.0425, -0.0273]]], grad_fn=&lt;TransposeBackward0&gt;)
</code></pre>
<p>As we can see, every token in the input sequences has the same attention output. That turns out no token was actually attended to.</p>
<p>I have tried using attn_mask and key_padding_mask but it’s no use.</p>
<p>Can someone explain what happened? And please correct me if I am applying them the wrong way.</p>
<p>Thanks!</p>

pascal_notsawo: <p>I tried to execute the steps as shown in the <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L4817" rel="noopener nofollow ugc">source code</a>, but I don’t have the same results as them (I must have forgotten something):</p>
<p>Let’s start by doing it directly:</p>
<pre><code class="lang-auto">import torch
import torch.nn.functional as F

torch.manual_seed(0)

seq_len, bsz, embedding_dim = 3, 1, 6
src_len = seq_len
tgt_len = seq_len

num_heads = 1
hidden_size = num_heads*4

inp = torch.randn(seq_len, bsz, embedding_dim)
lstm = torch.nn.LSTM(input_size=embedding_dim, hidden_size=hidden_size, num_layers=2)

embed_dim = hidden_size
self_attn = torch.nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads)

x, _ = lstm(input=inp) # seq_len x bsz x hidden_size
"""
tensor([[[-0.0383,  0.2014,  0.0948,  0.0756]],

        [[ 0.0603,  0.2379,  0.1403,  0.2036]],

        [[ 0.1172,  0.2602,  0.1392,  0.2943]]], grad_fn=&lt;StackBackward&gt;)
"""

out, weight = self_attn(query=x, key=x, value=x) # seq_len x bsz x hidden_size, ...
"""
tensor([[[ 0.0262, -0.0770,  0.1162,  0.1179]],

        [[ 0.0262, -0.0770,  0.1162,  0.1179]],

        [[ 0.0262, -0.0770,  0.1162,  0.1179]]], grad_fn=&lt;AddBackward0&gt;)

and

tensor([[[0.3330, 0.3334, 0.3336],
         [0.3327, 0.3334, 0.3339],
         [0.3325, 0.3335, 0.3340]]], grad_fn=&lt;DivBackward0&gt;)
"""
</code></pre>
<p>Now, if we want to do it ourselves, we can start by noticing that :</p>
<pre><code class="lang-auto">self_attn.q_proj_weight, self_attn.k_proj_weight, self_attn.v_proj_weight
"""
(None, None, None)
"""

self_attn.in_proj_weight
"""
tensor([[ 0.5861,  0.4650, -0.2232,  0.3442],
        [-0.3479, -0.0960,  0.5200,  0.0253],
        [-0.4331, -0.2047, -0.1662, -0.1181],
        [ 0.0586,  0.5663,  0.0328, -0.3781],
        [ 0.0314,  0.2936,  0.3038, -0.5597],
        [-0.1096, -0.4551, -0.2613,  0.2206],
        [-0.4349,  0.2276,  0.5198,  0.0402],
        [-0.4081, -0.2194,  0.1337, -0.4668],
        [ 0.3042, -0.5560, -0.5887, -0.5950],
        [-0.1242,  0.4118, -0.5796,  0.5090],
        [-0.2450,  0.1794,  0.0279, -0.5522],
        [ 0.5079,  0.3297,  0.6087,  0.3094]], requires_grad=True)
"""
</code></pre>
<p>So, as <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L4949" rel="noopener nofollow ugc">here</a>, we have (also, in the forward method of <a href="https://pytorch.org/docs/stable/_modules/torch/nn/modules/activation.html#MultiheadAttention" rel="noopener nofollow ugc">MultiheadAttention</a>):</p>
<pre><code class="lang-auto">use_separate_proj_weight = False
if not use_separate_proj_weight:
        #q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)
        q, k, v = F.linear(input=x, weight=self_attn.in_proj_weight, bias=self_attn.in_proj_bias).chunk(3, dim=-1)
"""
(tensor([[[ 0.0761,  0.0452, -0.0493,  0.0864]],
 
         [[ 0.1847,  0.0343, -0.1222,  0.0659]],
 
         [[ 0.2600,  0.0141, -0.1620,  0.0475]]], grad_fn=&lt;SplitBackward&gt;),
 tensor([[[ 0.0444, -0.0956,  0.1148, -0.0512]],
 
         [[ 0.0004, -0.1066,  0.1091, -0.1531]],
 
         [[-0.0423, -0.1027,  0.0925, -0.2237]]], grad_fn=&lt;SplitBackward&gt;),
 tensor([[[-0.2244,  0.0712,  0.0064,  0.1280]],
 
         [[-0.3177,  0.1128, -0.0806,  0.2575]],
 
         [[-0.3661,  0.1617, -0.1407,  0.3212]]], grad_fn=&lt;SplitBackward&gt;))
"""
</code></pre>
<p>In our case, <code>attn_mask is None</code> (see <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L4962" rel="noopener nofollow ugc">here</a>), <code> key_padding_mask is None</code> (see <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L4983" rel="noopener nofollow ugc">here</a>) et <code>bias_k is None and bias_v is None</code> (see <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L4988" rel="noopener nofollow ugc">here</a>), so we have :</p>
<pre><code class="lang-auto">head_dim = embed_dim // num_heads
scaling = float(head_dim) ** -0.5

q = q * scaling

q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
if k is not None:
  k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
if v is not None:
  v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)


attn_output_weights = torch.bmm(q, k.transpose(1, 2))
"""
tensor([[[0.3373, 0.3288, 0.3339],
         [0.3328, 0.3401, 0.3271],
         [0.3250, 0.3258, 0.3491]]], grad_fn=&lt;SoftmaxBackward&gt;)

"""

attn_output_weights = F.softmax(attn_output_weights, dim=-1)
attn_output_weights = F.dropout(attn_output_weights, p=0, training=False)
"""
tensor([[[0.3352, 0.3331, 0.3317],
         [0.3355, 0.3331, 0.3314],
         [0.3357, 0.3331, 0.3312]]], grad_fn=&lt;SoftmaxBackward&gt;)
"""

attn_output = torch.bmm(attn_output_weights, v)
"""
tensor([[[-0.3025,  0.1151, -0.0714,  0.2352],
         [-0.3024,  0.1150, -0.0713,  0.2351],
         [-0.3024,  0.1150, -0.0713,  0.2351]]], grad_fn=&lt;BmmBackward0&gt;)
"""

attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
attn_output = F.linear(attn_output, self_attn.out_proj.weight, self_attn.out_proj.bias)
"""
tensor([[[ 0.0262, -0.0770,  0.1162,  0.1179]],

        [[ 0.0262, -0.0770,  0.1162,  0.1179]],

        [[ 0.0262, -0.0770,  0.1162,  0.1179]]], grad_fn=&lt;AddBackward0&gt;)

"""

need_weights = True
if need_weights:
    # average attention weights over heads
    attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len) 
    weights = attn_output_weights.sum(dim=1) / num_heads
"""
tensor([[[0.3352, 0.3331, 0.3317],
          [0.3355, 0.3331, 0.3314],
          [0.3357, 0.3331, 0.3312]]], grad_fn=&lt;DivBackward0&gt;))
"""
</code></pre>
<p>It seems to me that there is a problem.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/3/13016958f4d5ac02c8824fea9340fe77feddac7f.png" data-download-href="https://discuss.pytorch.org/uploads/default/13016958f4d5ac02c8824fea9340fe77feddac7f" title="mha"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/3/13016958f4d5ac02c8824fea9340fe77feddac7f_2_690x328.png" alt="mha" data-base62-sha1="2I85Xnxc18CdDTul8DIC4iiyyWr" width="690" height="328" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/3/13016958f4d5ac02c8824fea9340fe77feddac7f_2_690x328.png, https://discuss.pytorch.org/uploads/default/original/3X/1/3/13016958f4d5ac02c8824fea9340fe77feddac7f.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/3/13016958f4d5ac02c8824fea9340fe77feddac7f.png 2x" data-dominant-color="3E3E3E"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">mha</span><span class="informations">989×471 40.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Here is an implementation that could also help you (readapted from <a href="https://github.com/facebookresearch/XLM/blob/master/xlm/model/transformer.py#L149" class="inline-onebox" rel="noopener nofollow ugc">XLM/xlm/model/transformer.py at main · facebookresearch/XLM · GitHub</a>), it’s always what I use in my projects :</p>
<pre><code class="lang-auto">import torch.nn as nn
class MultiHeadAttention(nn.Module):

    def __init__(self, n_heads, dim, dropout, temperature = None):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.dropout = dropout
        assert self.dim % self.n_heads == 0

        self.q_lin = nn.Linear(dim, dim)
        self.k_lin = nn.Linear(dim, dim)
        self.v_lin = nn.Linear(dim, dim)
        self.out_lin = nn.Linear(dim, dim)

        self.dim_per_head = dim // n_heads # d_k = d_v = d_k
        if temperature is None :
            self.temperature = float(self.dim_per_head) ** 0.5 # math.sqrt(self.dim_per_head) 
        else :
            self.temperature = temperature

    def forward(self, input, mask, kv=None, cache=None, need_weights = True):
        """
        Self-attention (if kv is None) or attention over source sentence (provided by kv).
        """
        # Input is (bs, qlen, dim)
        # Mask is (bs, klen) (non-causal) or (bs, klen, klen)
        bs, qlen, dim = input.size()
        if kv is None:
            klen = qlen if cache is None else cache['slen'] + qlen
        else:
            klen = kv.size(1)
        assert dim == self.dim, 'Dimensions do not match: %s input vs %s configured' % (dim, self.dim)
        n_heads = self.n_heads
        dim_per_head = self.dim_per_head
        mask_reshape = (bs, 1, qlen, klen) if mask.dim() == 3 else (bs, 1, 1, klen)

        def shape(x):
            """  projection """
            return x.view(bs, -1, self.n_heads, dim_per_head).transpose(1, 2)

        def unshape(x):
            """  compute context """
            return x.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * dim_per_head)

        q = shape(self.q_lin(input))                                          # (bs, n_heads, qlen, dim_per_head)
        if kv is None:
            k = shape(self.k_lin(input))                                      # (bs, n_heads, qlen, dim_per_head)
            v = shape(self.v_lin(input))                                      # (bs, n_heads, qlen, dim_per_head)
        elif cache is None or self.layer_id not in cache:
            k = v = kv
            k = shape(self.k_lin(k))                                          # (bs, n_heads, qlen, dim_per_head)
            v = shape(self.v_lin(v))                                          # (bs, n_heads, qlen, dim_per_head)

        if cache is not None:
            if self.layer_id in cache:
                if kv is None:
                    k_, v_ = cache[self.layer_id]
                    k = torch.cat([k_, k], dim=2)                             # (bs, n_heads, klen, dim_per_head)
                    v = torch.cat([v_, v], dim=2)                             # (bs, n_heads, klen, dim_per_head)
                else:
                    k, v = cache[self.layer_id]
            cache[self.layer_id] = (k, v)

        q = q / self.temperature                                              # (bs, n_heads, qlen, dim_per_head)
        scores = torch.matmul(q, k.transpose(2, 3))                           # (bs, n_heads, qlen, klen)
        mask = (mask == 0).view(mask_reshape).expand_as(scores)               # (bs, n_heads, qlen, klen)
        # to line is to solve :
        # /pytorch/aten/src/ATen/native/cuda/LegacyDefinitions.cpp:19: UserWarning: masked_fill_ received a mask with dtype torch.uint8, this behavior is now deprecated,please use a mask with dtype torch.bool instead.
        mask = mask.bool()
        scores.masked_fill_(mask, -float('inf'))                              # (bs, n_heads, qlen, klen)

        weights = F.softmax(scores.float(), dim=-1).type_as(scores)           # (bs, n_heads, qlen, klen)
        weights = F.dropout(weights, p=self.dropout, training=self.training)  # (bs, n_heads, qlen, klen)
        context = torch.matmul(weights, v)                                    # (bs, n_heads, qlen, dim_per_head)
        context = unshape(context)                                            # (bs, qlen, dim)

        if need_weights :
            return self.out_lin(context), weights 
        else :
            return self.out_lin(context)
</code></pre>
<p>Test</p>
<pre><code class="lang-auto">self_attn = MultiHeadAttention(n_heads = num_heads, dim = hidden_size, dropout=0)
mask = torch.ones((bsz, seq_len))

mask = torch.ones((bsz, seq_len))
att, w = self_attn(input = x, mask = mask, kv=None, cache=None)
</code></pre>
<h2><a name="notebookhttpscolabresearchgooglecomdrive1-b9bqofwkpvk1yvybbu6e95gdu_u_lquauthuser4scrolltoefwz_sf7gqyruniqifier2-1" class="anchor" href="#notebookhttpscolabresearchgooglecomdrive1-b9bqofwkpvk1yvybbu6e95gdu_u_lquauthuser4scrolltoefwz_sf7gqyruniqifier2-1"></a><a href="https://colab.research.google.com/drive/1-B9BqofwKPVK1YvYBBU6e95gdU_U_LQU?authuser=4#scrollTo=Efwz_sF7GqYr&amp;uniqifier=2" rel="noopener nofollow ugc">Notebook</a></h2>

Vu_Nguyen: <p>Thank you for the detailed reply!</p>
<p>I’ll test it later.</p>

liyunhan1993: <p>Hey,I have the same problem as you, have you solved it?</p>