What is mask in Masked Multi Head Attention? - machine-learning

Question

So I am trying to build a transform , but i am stuck at masking in decoder. So i am using Multi head attention so my input shape of (64,64,128) is converted to (64,8,64,16) for head = 8. For attention i used scaled multiplicative attention. So my attention layer gets input shape (64,8,64,16). I am having trouble calculating mask for the decoder multihead attention. As in scaled multiplicative attention the shape changes after matrix multiplication from (64,8,64,16) to (64,8,64,64) so how should i create a mask for this matrix.
My scaled multiplicative attention layer:
class ScaledMultiplicativeAttentionLayer(Layer):
def __init__(self):
super(ScaledMultiplicativeAttentionLayer,self).__init__()
def call(self,query,key,value,dim_key,mask=None):
print(query.shape)
mul = tf.matmul(query,key,transpose_b=True) / tf.sqrt(tf.cast(dim_key,tf.float32))
print(mul.shape)
if mask is not None:
mul += -1e9 * tf.cast(mask,tf.float32)
sfm = K.softmax(mul)
return tf.matmul(sfm,value)
My multihead attention :
class MultiHeadAttentionLayer(Layer):
def __init__(self,heads,dim_key,dim_val,dim_model,**kwargs):
super(MultiHeadAttentionLayer,self).__init__(**kwargs)
self.attention = ScaledMultiplicativeAttentionLayer()
self.heads = heads
self.dim_key = dim_key
self.dim_val = dim_val
self.dim_model = dim_model
self.Wq = tf.keras.layers.Dense(dim_key)
self.Wk = tf.keras.layers.Dense(dim_key)
self.Wv = tf.keras.layers.Dense(dim_val)
self.Wo = tf.keras.layers.Dense(dim_model)
def reshape_tensor(self,x,heads,flag):
if flag:
x = tf.reshape(x,shape=(x.shape[0],x.shape[1],heads,-1)) # (batch_size,seq_len,dim_key) -> (batch_size,seq_len,head,-1)
x = tf.transpose(x,perm=[0,2,1,3]) # (batch_size,seq_len,head,-1) -> (batch_size,head,seq_len,-1)
else:
x = tf.transpose(x,perm=[0,2,1,3]) # (batch_size,head,seq_len,-1) -> (batch_size,seq_len,head,-1)
x = tf.reshape(x,shape=(x.shape[0],x.shape[1],-1)) # (batch_size,head,seq_len,-1) -> (batch_size,seq_len,dim_model)
return x
def call(self,query,key,val,return_score=False,mask=None):
q_reshaped = self.reshape_tensor(self.Wq(query),self.heads,True)
k_reshaped = self.reshape_tensor(self.Wk(key),self.heads,True)
v_reshaped = self.reshape_tensor(self.Wv(val),self.heads,True)
if mask is not None:
mask = self.reshape_tensor(mask,self.heads,True)
o_reshaped = self.attention(q_reshaped,k_reshaped,v_reshaped,self.dim_key,mask)
output = self.reshape_tensor(o_reshaped,self.heads,False)
if return_score:
return self.Wo(output),K.sum(output,axis=1)
return self.Wo(output)
For now i am creating mask using this function:
def create_mask(shape):
mask = np.triu(np.ones((shape[-2],shape[-1])),1)
mask = np.expand_dims(mask,0)
mask = tf.tile(mask,[shape[0],1,1])
return mask
this create mask of shape same as input shape
I get error shape mismatch error at line
mul += -1e9 * tf.cast(mask,tf.float32)
because mul shape and mask shape is different. I cant seem to get the jist of what actual mask should look like? Help me understand mask and how to calculate mask.
Tried reshaping but no use. reproduce error here Reproduce Error Collab