Attention Rollout

Problem

Compute attention rollout. Start from the identity matrix, then for each attention matrix A form aug = row-normalized(0.5 * A + 0.5 * I) and update result = aug @ result, so later layers multiply on the left.

The signature is attention_rollout(attentions). Return the final rollout matrix.

Hint

Adding the identity accounts for the residual connection that carries information around each attention layer.

Test cases

case 1

in: [[[[1,0],[0,1]]]]

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

case 2

in: [[[[0.5,0.5],[0.5,0.5]]]]

expected: [[0.75,0.25],[0.25,0.75]]

case 3

in: [[[[1,0],[0,1]],[[0.5,0.5],[0.5,0.5]]]]

expected: [[0.75,0.25],[0.25,0.75]]

case 4

in: [[[[0.5,0.5],[0.5,0.5]],[[0.5,0.5],[0.5,0.5]]]]

expected: [[0.625,0.375],[0.375,0.625]]

Starter codePython

def attention_rollout(attentions):
    # Your code here
    pass

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