Flamingo Gated Cross-Attention

Problem

Compute a gated cross-attention output: standard scaled dot-product attention from query over keys and values, multiplied by tanh(gate). The tanh gate starts at zero so the new branch is an identity at initialization.

The signature is flamingo_gated_cross_attention(query, keys, values, gate).

Hint

Gated cross-attention layers are inserted between frozen language-model blocks and trained on interleaved data.

Test cases

case 1

in: [[1,0],[[1,0],[0,1]],[[1,2],[3,4]],0]

expected: [0,0]

case 2

in: [[1,0],[[1,0],[0,1]],[[1,2],[3,4]],1]

expected: [1.2646095042,2.0262036601]

case 3

in: [[1,1],[[1,0]],[[5,6]],2]

expected: [4.8201379004,5.7841654805]

case 4

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

expected: [-0.380797078,-0.380797078]

Starter codePython

import math
def flamingo_gated_cross_attention(query, keys, values, gate):
    # Your code here
    pass

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