Transformer Encoder Block Parameter Count

Problem

Count the parameters of one post-LN transformer encoder block, including biases and both LayerNorms. Attention: 4 * (d_model^2 + d_model). MLP: d_model * d_ff + d_ff + d_ff * d_model + d_model. LayerNorms: 2 * 2 * d_model.

The signature is transformer_block_params(d_model, d_ff). Return an integer.

Hint

Four attention projections plus a two-layer MLP plus two LayerNorms, all with biases.

Test cases

case 1

in: [512,2048]

expected: 3152384

case 2

in: [768,3072]

expected: 7087872

case 3

in: [64,256]

expected: 49984

case 4

in: [128,512]

expected: 198272

Starter codePython

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

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