MediumNLPnlp-069

Transformer Block Parameter Count for Language Models

Problem

Return the parameter count of one transformer block. Count four d_model x d_model attention projections with biases, a feed-forward network d_model -> d_ff -> d_model with biases, and two layer norms each with a weight and a bias of size d_model.

That is 4*d^2 + 4*d + (d*d_ff + d_ff) + (d_ff*d + d) + 4*d.

Hint

Add the FFN terms as two matrix products plus their bias vectors.

Test cases

case 1

in: [4,16]

expected: 244

case 2

in: [8,32]

expected: 872

case 3

in: [1,1]

expected: 16

case 4

in: [768,3072]

expected: 7087872

Starter codePython

def transformer_block_param_count(d_model, d_ff):
    # Your code here
    pass

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