mask error about flash-pytorch

[marsggbo]

@lucidrains thanks for your quick reply. I double check the code and understand that the causal_mask will be created when causal==True and everything is well taken care of.

However, I'm curious about these parts below.

FLASH-pytorch/flash_pytorch/flash_pytorch.py

Lines 317 to 319 in 473ab62

	if exists(mask):
	lin_mask = rearrange(mask, '... -> ... 1')
	lin_k = lin_k.masked_fill(~lin_mask, 0.)

FLASH-pytorch/flash_pytorch/flash_pytorch.py

Lines 352 to 353 in 473ab62

	if exists(mask):
	attn = attn.masked_fill(~mask, 0.)

In summary, my two question is:
By default, mask=None, what is this variable for and what is the difference between mask and causal_mask?

[lucidrains]

@keyunluo oops! should be fixed here 8e0d2fd

[marsggbo]

In vanilla multi-head transformer, the mask is an upper triangle matrix and the shape of mask is [bs, n_head, seq, seq]. But here, the shape of mask is [bs, seq]. Is any explanation for this? How do we set the mask for mixed chunk attention, performing similarly to the upper triangle matrix?

[lucidrains]

@marsggbo just set causal = True and everything is taken care of

[lucidrains]

@marsggbo the mask is for non-causal variable lengthed sequences

[lucidrains]

mask is for bidirectional attention (like BERT)

you often have sequences of variable lengths, so you would use the mask to only attend to non-padding tokens