masked self-attention

[LuisMoralesAlonso]

Again a question about implementation vs paper. I haven't found in the code any implementation of what they call "masked self-attention".

Hope your comments,

[khirotaka]

Oops, completely forgot to apply the mask to the MHA.
I certainly need to fix the code in EncoderBlock.

I'm currently focused on other tasks so I can't deal with it right away, but if you need it, you can use the following code to improve the forward method of EncoderBlock.

import torch
import numpy as np


def subsequent_mask(size: int) -> torch.Tensor:
    attn_shape = (size, size)
    mask = np.triu(np.ones(attn_shape, dtype=np.float32), k=1)
    mask = torch.from_numpy(mask) == 0
    return mask.float()


class EncoderModule:
    ...
    def forward(self, x):
        # x.shape == [N, L, E]
        mask = subsequent_mask(x.shape[1])    # mask.shape == [L, L]
        x = self.attention(x, attn_mask=mask)
        ...

It should work fine. Maybe.

[LuisMoralesAlonso]

i'll try it, thanks (maybe)!!!

;-)

[pedromingues]

Hi,

I tried it but it gave the error:

TypeError: forward() got an unexpected keyword argument 'attn_mask'

I'm a novice in pytorch and also in terms of dealing/programming complex models such as this, so maybe i'm making some silly mistake or not seen the easy solution. This is the code I used in the EncoderBlock:

class EncoderBlock(nn.Module):
    def __init__(self, embed_dim: int, num_head: int, dropout_rate=0.1) -> None:
        super(EncoderBlock, self).__init__()
        self.attention = ResidualBlock(
            nn.MultiheadAttention(embed_dim, num_head), embed_dim, p=dropout_rate
        )
        self.ffn = ResidualBlock(PositionWiseFeedForward(embed_dim), embed_dim, p=dropout_rate)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mask = subsequent_mask(x.shape[1])    # mask.shape == [L, L]
        x = self.attention(x, attn_mask=mask)
        x = self.ffn(x)
        return x

Did it work for you @LuisMoralesAlonso ? If so, what have you done differently?

[khirotaka]

@pedromingues It's because you wrapped MultiheadAttention with ResidualBlock.

[pedromingues]

I used the EncoderBlock from your code, but since you said that I tried:

class EncoderBlock(nn.Module):
    def __init__(self, embed_dim: int, num_head: int, dropout_rate=0.1) -> None:
        super(EncoderBlock, self).__init__()
        self.attention = nn.MultiheadAttention(embed_dim, num_head)
        self.ffn = ResidualBlock(PositionWiseFeedForward(embed_dim), embed_dim, p=dropout_rate)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mask = subsequent_mask(x.shape[1])    # mask.shape == [L, L]
        x = self.attention(x, attn_mask=mask)
        x = self.ffn(x)
        return x

But this also gave an error, this time:

TypeError: forward() missing 2 required positional arguments: 'key' and 'value

What should I specifically change for it to work? Can you provide the code of what you think is the solution?

Sorry for your time, and thank you very much for the patience.

[khirotaka]

I'm sorry. I was wrong, if you want to add an attention mask while using ResidualBlock, fix it here.

SAnD/core/modules.py

Line 45 in b579717

output, self.attn_weights = self.layer(src, src, src)

[pedromingues]

Sorry for the delay in writing back.

So instead of using the subsequent_mask function in the EncoderBlock, it should only be applied to the self.attn_weights in here right?

SAnD/core/modules.py

Line 45 in b579717

output, self.attn_weights = self.layer(src, src, src)

[khirotaka]

yes. like this.

attn_mask = subsequent_mask(x.shape[1])
output, self.attn_weights = self.layer(src, src, src, attn_mask=attn_mask)