Transformer Segmentation Head

Problem

Apply a linear segmentation head to every patch token and return the argmax class per token, breaking ties by the smaller class index.

The signature is seg_head_forward(features, W, b). features is (num_tokens, d), W is (num_classes, d) and b has length num_classes.

Hint

A per-patch linear classifier turns transformer tokens into a segmentation map.

case 1

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

expected: [0,1]

case 2

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

expected: [0]

case 3

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

expected: [1]

case 4

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

expected: [1,0]

Starter codePython

def seg_head_forward(features, W, b):
    # Your code here
    pass

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