Implementing 1D self attention issues

username: Federico_Ottomano
cooked: <p>I’m trying to implement the 1D self-attention illustrated in this <a href="https://easychair.org/publications/preprint/WWf8" rel="noopener nofollow ugc">paper</a></p>
<p>Specifically, I’m focusing on the following picture:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/8/08e6108bff50578a3ac9f4886e2afd5432ab0b29.png" data-download-href="https://discuss.pytorch.org/uploads/default/08e6108bff50578a3ac9f4886e2afd5432ab0b29" title="Attn"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/8/08e6108bff50578a3ac9f4886e2afd5432ab0b29_2_690x274.png" alt="Attn" data-base62-sha1="1gIJjRpM1lnGqhVJFQJAx7mAdCx" width="690" height="274" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/8/08e6108bff50578a3ac9f4886e2afd5432ab0b29_2_690x274.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/8/08e6108bff50578a3ac9f4886e2afd5432ab0b29_2_1035x411.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/0/8/08e6108bff50578a3ac9f4886e2afd5432ab0b29.png 2x" data-dominant-color="F0F1F0"><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">Attn</span><span class="informations">1358×541 83.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>Below you can find my attempt:</p>
<pre><code class="lang-auto">import torch.nn as nn
import torch

#INPUT shape ((B), n_channels, height, width)

class Self_Attention1D(nn.Module):

    def __init__(self, in_channels=1, out_channels=3):

        super().__init__()

        self.pointwise_conv1 = nn.Conv1d(in_channels=in_channels, 
        out_channels=out_channels, kernel_size=(1,1))
        self.pointwise_conv2 = nn.Conv1d(in_channels=out_channels, 
        out_channels=in_channels, kernel_size=(1,1))
        self.phi = MLP(in_size = out_channels, out_size=32)
        self.psi = MLP(in_size = out_channels, out_size=32)
        self.gamma = MLP(in_size=32, out_size=out_channels)

    def forward(self, x):
        x = self.pointwise_conv1(x)
        phi = self.phi(x.transpose(1,3))
        psi = self.psi(x.transpose(1,3))
        delta = phi-psi
        gamma = self.gamma(delta).transpose(3,1)
        out = self.pointwise_conv2(torch.mul(gamma,x))
        return out

class MLP(nn.Module):

    def __init__(self, in_size, out_size):

        super().__init__()

        self.in_size = in_size
        self.out_size = out_size
        self.layers = nn.Sequential(
            nn.Linear(in_size, 64),
            nn.ReLU(),
            nn.Linear(64,128),
            nn.ReLU(),
            nn.Linear(128,64),
            nn.ReLU(),
            nn.Linear(64,out_size))
    def forward(self, x):
        out = self.layers(x)
        return out

</code></pre>
<p>What it’s not clear to me is, <strong>how can we implement</strong> the <em>boxed</em> operation in such a way that the delta function takes each single entry <strong>and</strong> its neighbour <em>one at a time</em> as it is supposed to be by looking at the picture.</p>
<p>I think that my implementation is definitely ignoring <em>locality</em> of the operation.</p>