How to implement Bidirectional Alibi with padding using flex attention?

[sphmel]

Hi, I want to use FlexAttention for alibi with padding(no bias)

If seq_len is 5 I want to make alibi tensor like below, which is alibi tensor with seq_len, and last item is not penalized

0 -1 -2 -3 0
-1 0 -1 -2 0
-2 -1 0 -1 0
-3 -2 -1 0 0
0 0 0 0 0

How can I implement score mod like this? seq_len can be different every forward. Such alibi is used in Voicebox paper. I'm new to BatchedTensor or maybe vmap API? I do not know how to implement it at all. Can you help me?

[sphmel]

q_idx - kv_idx wiil make tensor below, but i want to last row and column is not biased

0 -1 -2 -3 -4
-1 0 -1 -2 -3
-2 -1 0 -1 -2
-3 -2 -1 0 -1
-4 -3 -2 -1 0

[Chillee]

@sphmel I think this should work.

def score_mod(score, b, h, q_idx, kv_idx):
    bias = (q_idx - kv_idx)
    bias = torch.where((q_idx == seq_len - 1) || (kv_idx == seq_len - 1), bias, 0)
    return score + bias

[chinsengi]

q_idx - kv_idx wiil make tensor below, but i want to last row and column is not biased

0 -1 -2 -3 -4
-1 0 -1 -2 -3
-2 -1 0 -1 -2
-3 -2 -1 0 -1
-4 -3 -2 -1 0

I don't understand, shouldn't q_idx - kv_idx give you an skew-symmetric matrix rather than a symmetric matrix?

[Chillee]

@chinsengi I assumed it was just a typo and they meant for the lower-triangular part to be positive.