DETR Decoder Layer

Problem

Forward pass of a simplified DETR decoder layer: cross-attend from queries x to encoder memory, add the residual and the ReLU MLP residual. With Q = xWq, K = memoryWk, V = memoryWv: a = x + softmax(QK^T / sqrt(d)) V Wo, and out = a + relu(a Wf + bf).

The signature is detr_decoder_layer(x, memory, Wq, Wk, Wv, Wo, Wf, bf).

Hint

Object queries attend to the encoder memory, so each query can specialize in a different object.

Test cases

case 1

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

expected: [[3.3395230987,0.6604769013],[0.6604769013,3.3395230987]]

case 2

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

expected: [[2.6604769013,5.3395230987]]

case 3

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

expected: [[1,0]]

case 4

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

expected: [[4.1,2.1]]

Starter codePython

import math
def detr_decoder_layer(x, memory, Wq, Wk, Wv, Wo, Wf, bf):
    # Your code here
    pass

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