Implementation of the Dense Synthesizer

username: alvations
cooked: <p>I’m trying to understand the Synthesizer paper (<a href="https://arxiv.org/pdf/2005.00743.pdf" rel="noopener nofollow ugc">https://arxiv.org/pdf/2005.00743.pdf</a>) and there’s a description of the dense synthesizer mechanism that should replace the traditional attention model as described in the Transformer architecture.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/a/0abbe98b8ffef3240e0294a6f635b18285e7facb.png" data-download-href="https://discuss.pytorch.org/uploads/default/0abbe98b8ffef3240e0294a6f635b18285e7facb" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/a/0abbe98b8ffef3240e0294a6f635b18285e7facb_2_690x399.png" alt="image" data-base62-sha1="1wXnfbWHdH5TPu5M1t4SZCTqk1t" width="690" height="399" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/a/0abbe98b8ffef3240e0294a6f635b18285e7facb_2_690x399.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/a/0abbe98b8ffef3240e0294a6f635b18285e7facb_2_1035x598.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/a/0abbe98b8ffef3240e0294a6f635b18285e7facb_2_1380x798.png 2x" data-dominant-color="F5F4F4"><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">image</span><span class="informations">1629×942 68.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>The <strong>Dense Synthesizer</strong> is described as such:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/5/250b4e2f4ca6b233510396fe91b2f2de225a6d14.png" data-download-href="https://discuss.pytorch.org/uploads/default/250b4e2f4ca6b233510396fe91b2f2de225a6d14" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/5/250b4e2f4ca6b233510396fe91b2f2de225a6d14_2_690x415.png" alt="image" data-base62-sha1="5hHT3Zf7U8RVVOvzU8oqucI149u" width="690" height="415" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/5/250b4e2f4ca6b233510396fe91b2f2de225a6d14_2_690x415.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/5/250b4e2f4ca6b233510396fe91b2f2de225a6d14_2_1035x622.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/5/250b4e2f4ca6b233510396fe91b2f2de225a6d14.png 2x" data-dominant-color="E9E9E9"><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">image</span><span class="informations">1122×675 72.9 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>So I tried to implement the layer and it looks like this but I’m not sure whether I’m getting it right:</p>
<pre><code class="lang-auto">class DenseSynthesizer(nn.Module):
    def __init__(self, l, d):
        super(DenseSynthesizer, self).__init__()
        self.linear1 = nn.Linear(d, l)
        self.linear2 = nn.Linear(l, l)
    
    def forward(self, x, v):
        # Equation (1) and (2)
        # Shape: l x l
        b = self.linear2(F.relu(self.linear1(x)))   
        # Equation (3)
        # [l x l] x [l x d] -&gt; [l x d]
        return torch.matmul(F.softmax(b), v) 
</code></pre>
<p>Usage:</p>
<pre><code class="lang-auto">l, d = 4, 5

x, v =  torch.rand(l, d), torch.rand(l, d)

synthesis = DenseSynthesizer(l, d)
synthesis(x, v) 
</code></pre>
<p>Example:</p>
<p><code>x</code> and <code>v</code> are tensors:</p>
<pre><code class="lang-auto">x = tensor([[0.0844, 0.2683, 0.4299, 0.1827, 0.1188],
         [0.2793, 0.0389, 0.3834, 0.9897, 0.4197],
         [0.1420, 0.8051, 0.1601, 0.3299, 0.3340],
         [0.8908, 0.1066, 0.1140, 0.7145, 0.3619]])

v = tensor([[0.3806, 0.1775, 0.5457, 0.6746, 0.4505],
         [0.6309, 0.2790, 0.7215, 0.4283, 0.5853],
         [0.7548, 0.6887, 0.0426, 0.1057, 0.7895],
         [0.1881, 0.5334, 0.6834, 0.4845, 0.1960]])
</code></pre>
<p>And passing through a forward pass through the dense synthesis, it returns:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; synthesis = DenseSynthesizer(l, d)
&gt;&gt;&gt; synthesis(x, v) 

tensor([[0.5371, 0.4528, 0.4560, 0.3735, 0.5492],
        [0.5426, 0.4434, 0.4625, 0.3770, 0.5536],
        [0.5362, 0.4477, 0.4658, 0.3769, 0.5468],
        [0.5430, 0.4461, 0.4559, 0.3755, 0.5551]], grad_fn=&lt;MmBackward&gt;)
</code></pre>
<p><strong>Is the implementation and understanding of the dense synthesizer correct?</strong></p>
<p>Theoretically, <strong>how is that different from a multi-layered perceptron that takes in two different inputs and makes uses of it at different point in the forward propagation?</strong></p>

username: alvations
cooked: <p>Comparatively, the original transformer attention mechanism is as such (from <a href="https://nlp.seas.harvard.edu/2018/04/03/attention.html" rel="nofollow noopener">https://nlp.seas.harvard.edu/2018/04/03/attention.html</a>):</p>
<pre><code class="lang-auto">def attention(query, key, value, mask=None, dropout=None):
    """Compute 'Scaled Dot Product Attention"""
    d_k = query.size(-1) # no. of hidden dimension of query tensor.
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn
</code></pre>
<p>Usage:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; l, d = 4, 5
&gt;&gt;&gt; q, k, v = torch.rand(l, d), torch.rand(l, d), torch.rand(l, d)
&gt;&gt;&gt; attention(q, k, v)
(tensor([[0.5511, 0.6426, 0.6486, 0.4216, 0.3083],
         [0.5560, 0.6441, 0.6665, 0.4169, 0.3312],
         [0.5590, 0.6426, 0.6647, 0.4173, 0.3283],
         [0.5678, 0.6450, 0.6291, 0.4012, 0.2998]]),
 tensor([[0.2579, 0.2811, 0.2288, 0.2322],
         [0.2724, 0.2851, 0.2585, 0.1840],
         [0.2993, 0.2884, 0.2382, 0.1740],
         [0.2769, 0.2465, 0.2166, 0.2600]]))
</code></pre>