Multi-Head Attention Parameter Count with Bias

Problem

Count the parameters of a multi-head attention module: four projections Wq, Wk, Wv and Wo each with d_model x d_model weights, plus one bias vector of length d_model per projection when bias is enabled. The number of heads does not change the total.

The signature is mha_param_count(d_model, num_heads, bias=True). Return an integer.

Hint

Splitting into heads partitions each projection; it does not add parameters.

Test cases

case 1

in: [512,8]

expected: 1050624

case 2

in: [768,12,false]

expected: 2359296

case 3

in: [64,4]

expected: 16640

case 4

in: [128,8,true]

expected: 66048

Starter codePython

def mha_param_count(d_model, num_heads, bias=True):
    # Your code here
    pass

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