Self Self-attention implementation results are 'a bit' suprising

dey: <p>I am studying attention so I was making some experiments on implementing. I was curious of each part’s (Residual connections, Layer normalization,<br>
Dropout, Multi-heads) importance on training so I tried to implement them myself and make some tests.</p>
<p>I saw that using attention without any residual connections was worse than using no attention at all. After adding the residual connections the performance got closer to base NoAttentionAttention version but still it was a bit worse off. Only after adding the Layer Normalization did I start to get better results.</p>
<p>My quetions is that is this behaviour normal? It seems using a simpler feed forward layer instead of attention with residual connections is a bit better and much faster. I have 3 explanations in my mind but I don’t know which one is correct (maybe couple of them are correct at the same time to a degree or all of them are wrong). 1 - Layer normalization is crucial part of the architecture and attention doesn’t quite work without it. 2 - My implementation of attention mechanism is wrong/I’ve overlooked some details. 3 - My setup is not quite suitable for showing the benefits of using transformers. I am training a character level model and my training data is Andrej Karpathy’s tiny_shakespeare dataset, actually I am not sure if he created the dataset or he just used it and made it popular, I am sorry if my source is wrong.</p>
<p>My results are as follows:</p>
<blockquote>
<p>NoAttentionAttention<br>
min_train_loss1.701506521344185<br>
min_test_loss1.5964061438009656<br>
max_train_accuracy49.58238525390625<br>
max_test_accuracy53.14371043238146<br>
duration174.00348663330078<br>
SimpleAttention<br>
min_train_loss2.2374376397132876<br>
min_test_loss2.024268022076837<br>
max_train_accuracy40.8646728515625<br>
max_test_accuracy46.13544332570043<br>
duration1405.215410232544<br>
ResidualAttention<br>
min_train_loss1.894147829771042<br>
min_test_loss1.66307471086239<br>
max_train_accuracy46.54041748046875<br>
max_test_accuracy52.319441170528016<br>
duration1434.7716104984283<br>
LayNormResAttention<br>
min_train_loss0.36205335560441015<br>
min_test_loss1.1429171343301905<br>
max_train_accuracy89.10374755859375<br>
max_test_accuracy67.51815269733298<br>
duration1451.627137184143</p>
</blockquote>
<p>These are my self-attention implementations:</p>
<pre><code class="lang-auto">class NoAttentionAttention(nn.Module):
    def __init__(self, emb_size, head_size):
        super(NoAttentionAttention, self).__init__()
        self.linear = nn.Linear(emb_size, head_size)
    def forward(self, x):
        x = nn.functional.relu(self.linear(x))
        return x

class SimpleAttention(nn.Module):
    def __init__(self, emb_size, head_size):
        super(SimpleAttention, self).__init__()
        self.emb_size = emb_size
        self.Q = nn.Linear(emb_size, head_size)
        self.K = nn.Linear(emb_size, head_size)
        self.V = nn.Linear(emb_size, head_size)
        self.linear = nn.Linear(head_size, head_size)
    def forward(self, x):
        q = self.Q(x)
        k = self.K(x)
        v = self.V(x)
        att_w = (q @ k.transpose(1, 2)) * self.emb_size ** (-1 / 2)
        seq_len = x.size(1)
        mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
        att_w = att_w.masked_fill(mask, float('-inf'))
        att_s = nn.functional.softmax(att_w, dim=-1)
        att_out = att_s @ v
        x = nn.functional.relu(self.linear(att_out))
        return x

class ResidualAttention(nn.Module):
    def __init__(self, emb_size, head_size):
        super(ResidualAttention, self).__init__()
        self.emb_size = emb_size
        self.Q = nn.Linear(emb_size, head_size)
        self.K = nn.Linear(emb_size, head_size)
        self.V = nn.Linear(emb_size, head_size)
        self.linear = nn.Linear(head_size, head_size)
    def forward(self, x):
        q = self.Q(x)
        k = self.K(x)
        v = self.V(x)
        att_w = (q @ k.transpose(1, 2)) * self.emb_size ** (-1 / 2)
        seq_len = x.size(1)
        mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
        att_w = att_w.masked_fill(mask, float('-inf'))
        att_s = nn.functional.softmax(att_w, dim=-1)
        att_out = att_s @ v
        x = x + att_out
        x = torch.nn.functional.relu(self.linear(x))
        return x

class LayNormResAttention(nn.Module):
    def __init__(self, emb_size, head_size):
        super(LayNormResAttention, self).__init__()
        self.emb_size = emb_size
        self.Q = nn.Linear(emb_size, head_size)
        self.K = nn.Linear(emb_size, head_size)
        self.V = nn.Linear(emb_size, head_size)
        self.linear = nn.Linear(head_size, head_size)
        self.layer_norm = nn.LayerNorm(emb_size)
    def forward(self, x):
        q = self.Q(x)
        k = self.K(x)
        v = self.V(x)
        att_w = (q @ k.transpose(1, 2)) * self.emb_size ** (-1 / 2)
        seq_len = x.size(1)
        mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
        att_w = att_w.masked_fill(mask, float('-inf'))
        att_s = nn.functional.softmax(att_w, dim=-1)
        att_out = att_s @ v
        x = self.layer_norm(x + att_out)
        x = torch.nn.functional.relu(self.linear(x))
        return x
</code></pre>
<p>Another suprising result I’ve come accross with the same setup is when I used multi heads.</p>
<p>With this implementation:</p>
<pre><code class="lang-auto">class TorchMhAttention(nn.Module):
    def __init__(self, emb_size, head_size, dropout=0.1):
        super(TorchMhAttention, self).__init__()
        self.emb_size = emb_size
        self.head_size = head_size
        self.multihead_attn = nn.MultiheadAttention(emb_size, num_heads=4, dropout=dropout)
        self.layer_norm = nn.LayerNorm(emb_size)
        self.dropout = nn.Dropout(dropout)
        self.linear = nn.Linear(emb_size, emb_size)
    def forward(self, x):
        x = x.transpose(0, 1) 
        seq_len = x.size(0)
        mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool().to(x.device)
        att_out, _ = self.multihead_attn(x, x, x, attn_mask=mask)
        x = self.layer_norm(x + att_out)
        x = torch.nn.functional.relu(self.linear(x))
        x = self.dropout(x)
        return x.transpose(0, 1)
</code></pre>
<p>when num_heads=1 I get these scores:</p>
<blockquote>
<p>min_train_loss1.2817998005747795<br>
min_test_loss1.1874091575371808<br>
max_train_accuracy61.50433349609376<br>
max_test_accuracy64.9322299299569<br>
duration1759.3112213611603</p>
</blockquote>
<p>but when I increased the num_heads to 4 I got this result:</p>
<blockquote>
<p>min_train_loss2.2550643470287324<br>
min_test_loss2.189865370043393<br>
max_train_accuracy41.768908691406246<br>
max_test_accuracy45.89278648639548<br>
duration2077.635309934616</p>
</blockquote>
<p>For all of these tests my emb_size == head_size == 256 and I use 6 layers of my attention implementations for each test, my context length is 128 characters.</p>