Autoregressive sampling with Multi-head attention

username: malbergo
cooked: <p>I am trying to make sure I understand the implementation of torch.nn.MultiheadAttention(), as I want to use it for autoregressive sampling in a “decoder-only” image transformer.</p>
<p>In such case, one tries to predict the next pixel by attending on all previous pixels. Say my pixels are just binary (0,1).</p>
<p>In the Multihead attention forward function, the input is supposed to be:<br>
query: [target length, batch size, embed dim]<br>
key: [sequence length, batch size, embed dim]<br>
value: [sequence length, batch size, embed dim]</p>
<p>What is “target length” referring to in this context? If I am just trying to predict the next pixel at each pass, is this just length 1?</p>
<p>Also, is there any way to do batch_first with this module?</p>
<p>Thanks for your help!</p>

username: vainaijr
cooked: <p>maybe this is wrong, but for example if we have image,</p>
<pre><code class="lang-auto">input = torch.randn(1, 1, 28, 28)
</code></pre>
<p>we pass it through convolutional layer,</p>
<pre><code class="lang-auto">cnv = nn.Conv2d(1, 10, 3, padding=1)
cnv(input).shape
</code></pre>
<p>torch.Size([1, 10, 28, 28])</p>
<pre><code class="lang-auto">q = torch.randn(1, 1, 10) # target pixel we want
k = cnv(input).reshape(28*28, 1, 10) # input pixels
v = cnv(input).reshape(28*28, 1, 10) # input pixels
</code></pre>
<pre><code class="lang-auto">attn = nn.MultiheadAttention(10, 2)
attn(q, k, v)[0].shape, attn(q, k, v)[1].shape
</code></pre>
<p>(torch.Size([1, 1, 10]), torch.Size([1, 1, 784]))</p>
<p>first one is representation of pixel that we want, second one is how much this pixel representation depended on each of 28*28 pixels.</p>
<pre><code class="lang-auto">lin1 = nn.Linear(10, 1)
lin1(attn(q, k, v)[0])
</code></pre>
<p>tensor([[[0.2838]]], grad_fn=)</p>
<p>this is value of next pixel, let us call it p.</p>
<p>next time we run this, we want to use this value also, so,</p>
<pre><code class="lang-auto">lin2 = nn.Linear(1, 10)
lin2(p)
</code></pre>
<p>would give something like [1, 1, 10], our next k, v would be</p>
<pre><code class="lang-auto">k = torch.cat((k, lin2(p)), dim=-3)
v = torch.cat((v, lin2(p)), dim=-3)
</code></pre>
<p>here is one implementation,</p>
<pre><code class="lang-auto">class Net(nn.Module):
  def __init__(self):
    super().__init__()
    self.attn = nn.MultiheadAttention(10, 2)
    self.conv = nn.Conv2d(1, 10, 3, padding=1)
    self.lin1 = nn.Linear(10, 1)
    self.lin2 = nn.Linear(1, 10)
  
  def forward(self, input, nxt_pxl):
    if(input.ndim == 4):
      conv_out = self.conv(input) # out.shape -&gt; (1, 10, 28, 28)
      conv_out = conv_out.reshape(conv_out.size(2)*conv_out.size(3), conv_out.size(0), conv_out.size(1))
    else:
      conv_out = input
    out = self.attn(nxt_pxl, conv_out, conv_out)
    out = self.lin1(out[0])
    out1 = self.lin2(out) # next pixel
    out2 = torch.cat((conv_out, out1), dim=-3) # next input
    return out, out1.detach(), out2.detach()
</code></pre>
<pre><code class="lang-auto">model = Net()
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
input = torch.randn(1, 1, 28, 28)
target = torch.tensor([1.])
nxt_pxl = torch.randn(1, 1, 10) # first time we specify random next pixel
</code></pre>
<pre><code class="lang-auto">for i in range(10):
  optimizer.zero_grad()  
  print(input.shape)
  pred, nxt_pxl, input = model(input, nxt_pxl)
  loss = loss_fn(pred, target)
  print(pred.shape)
  print(loss)
  loss.backward()
  optimizer.step()
</code></pre>