Attention Transfer Loss

Problem

Compute the attention transfer loss between student and teacher maps.

At each position the attention value is the sum of squared activations over channels; normalize each map so all positions sum to 1, then return the MSE between the two normalized maps rounded to 4 decimal places.

Test cases

case 1

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

expected: 0

case 2

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

expected: 0.0625

case 3

in: [[[[1,1],[1,1]],[[1,1],[1,1]]],[[[2,2],[2,2]],[[2,2],[2,2]]]]

expected: 0

Starter codePython

def attention_transfer_loss(student, teacher):
    # Your code here
    pass

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