MultiheadAttention after LSTM returns the same output for all input, please watch me!

liyunhan1993: <aside class="quote quote-modified" data-post="1" data-topic="122035">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/vu_nguyen/48/38605_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/multiheadattention-after-lstm-returns-the-same-output-for-all-input/122035">MultiheadAttention after LSTM returns the same output for all input</a> 
  </div>
  <blockquote>
    Hi, 
I would like to use MultiheadAttention as self-attention after applying LSTM on a single sequence. 
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.…
  </blockquote>
</aside>

<p>Here is the question，I meet the same question too. I have no idea why this happen and how to solve it, After a few rounds of training, I also observed the attention socres and found that<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/1/614d067c416ea69b2a86bd404268435460f1931c.png" data-download-href="https://discuss.pytorch.org/uploads/default/614d067c416ea69b2a86bd404268435460f1931c" title="截屏2022-04-25 下午12.58.39"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/614d067c416ea69b2a86bd404268435460f1931c_2_690x411.png" alt="截屏2022-04-25 下午12.58.39" data-base62-sha1="dSLl7sj0eTMeSa8egunTXWUipfC" width="690" height="411" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/614d067c416ea69b2a86bd404268435460f1931c_2_690x411.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/614d067c416ea69b2a86bd404268435460f1931c_2_1035x616.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/614d067c416ea69b2a86bd404268435460f1931c_2_1380x822.png 2x" data-dominant-color="ECECEC"><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">截屏2022-04-25 下午12.58.39</span><span class="informations">1692×1008 48 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>

liyunhan1993: <p>And here is my attention out<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/6/1673c08cb784160ecb8516fca8255d2e14fb9210.png" data-download-href="https://discuss.pytorch.org/uploads/default/1673c08cb784160ecb8516fca8255d2e14fb9210" title="截屏2022-04-25 下午1.05.58"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/6/1673c08cb784160ecb8516fca8255d2e14fb9210_2_690x404.png" alt="截屏2022-04-25 下午1.05.58" data-base62-sha1="3cCv8nTBZzlWsyShNgGHEQZVgAg" width="690" height="404" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/6/1673c08cb784160ecb8516fca8255d2e14fb9210_2_690x404.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/6/1673c08cb784160ecb8516fca8255d2e14fb9210_2_1035x606.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/6/1673c08cb784160ecb8516fca8255d2e14fb9210_2_1380x808.png 2x" data-dominant-color="ECECEC"><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">截屏2022-04-25 下午1.05.58</span><span class="informations">1744×1022 55.8 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>