Attention Head Dim And Scale

Problem

Given d_model and num_heads, compute the per-head dimension head_dim = d_model / num_heads and the attention scaling factor 1 / sqrt(head_dim).

The signature is attention_head_scale(d_model, num_heads). Return [head_dim, scale] as an integer and a float.

Hint

Keep more heads means a smaller head_dim and a larger softmax temperature scale.

case 1

in: [512,8]

expected: [64,0.125]

case 2

in: [768,12]

expected: [64,0.125]

case 3

in: [64,4]

expected: [16,0.25]

case 4

in: [512,16]

expected: [32,0.1767766953]

Starter codePython

import math
def attention_head_scale(d_model, num_heads):
    # Your code here
    pass

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