jekbradbury
commented
7 years ago Yes, it's quite easy to introduce a bug in the masking process and it would lead to misleadingly high results. Some possibilities (from my experience writing a Transformer implementation) include softmax over the wrong dimension and the use of the wrong kind of triangular matrix. Here's an implementation of optionally masked attention I know to be correct that works for both training and inference (but depends on broadcasting/PyTorch master):
# F.softmax has strange default behavior, normalizing over dim 0 for 3D inputs
def softmax(x):
if x.dim() == 3:
return F.softmax(x.transpose(0, 2)).transpose(0, 2)
return F.softmax(x)
# torch.matmul can't do (4, 3, 2) @ (4, 2) -> (4, 3)
def matmul(x, y):
if x.dim() == y.dim():
return x @ y
if x.dim() == y.dim() - 1:
return (x.unsqueeze(-2) @ y).squeeze(-2)
return (x @ y.unsqueeze(-2)).squeeze(-2)
class Attention(nn.Module):
def __init__(self, d_key, drop_ratio, causal):
super().__init__()
self.scale = math.sqrt(d_key)
self.dropout = nn.Dropout(drop_ratio)
self.causal = causal
def forward(self, query, key, value):
dot_products = matmul(query, key.transpose(1, 2))
if query.dim() == 3 and self.causal:
tri = key.data.new(key.size(1), key.size(1)).fill_(1).triu(1) * INF
dot_products.data.sub_(tri.unsqueeze(0))
return matmul(self.dropout(softmax(dot_products / self.scale)), value)
kendricktan
larspars
jadore801120
ZiJianZhao
I get 98% accuracy after 10 epochs on the multi30k validation set using this 1-layer model:
This is a very small model (note -d_inner_hid 1), which should not get good results at all (98% accuracy is way too high in any case). Generating translations with translate.py produces non-sense. This makes me suspect that there is a problem with the masking code that allows the model to 'cheat' by looking at the target sequence.
I haven't been able to figure out where the problem is, but something seems wrong.