ALiBi Bias Matrix

Problem

Build the ALiBi attention bias. Head h uses slope 1 / 2^(h+1); entry [h][i][j] equals -slope * (i - j) for j <= i and 0.0 for future positions j > i.

The signature is alibi_bias(seq_len, num_heads). Return a num_heads x seq_len x seq_len tensor.

Hint

ALiBi replaces positional embeddings entirely with a distance-proportional penalty.

case 1

in: [3,2]

expected: [[[0,0,0],[-0.5,0,0],[-1,-0.5,0]],[[0,0,0],[-0.25,0,0],[-0.5,-0.25,0]]]

case 2

in: [2,1]

expected: [[[0,0],[-0.5,0]]]

case 3

in: [1,3]

expected: [[[0]],[[0]],[[0]]]

case 4

in: [3,1]

expected: [[[0,0,0],[-0.5,0,0],[-1,-0.5,0]]]

Starter codePython

def alibi_bias(seq_len, num_heads):
    # Your code here
    pass

4 test cases · run your code against them in the editor