Raw Attention Flow Product

Problem

Compute the raw attention flow product across layers by multiplying the attention matrices in order: result = A_last @ ... @ A_1, starting from the identity matrix. Unlike attention rollout this does not mix in the identity at each layer.

The signature is attention_flow_product(attentions). Return the final (seq_len, seq_len) matrix.

Hint

Multiplying two row-stochastic matrices keeps rows stochastic, so a uniform matrix is a fixed point.

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.5,0.5],[0.5,0.5]]

case 3

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

expected: [[0.5,0.5],[0.5,0.5]]

case 4

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

expected: [[0.5,0.5],[0.5,0.5]]

Starter codePython

def attention_flow_product(attentions):
    # Your code here
    pass

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