Weight Tying Parameter Count

Problem

Compare the parameter count of an embedding table and an output projection head when the weights are tied. The embedding costs vocab_size * d_model parameters; the output head costs the same amount unless tying removes it.

The signature is weight_tying_params(vocab_size, d_model, tied). Return [embedding_params, output_params, total_params].

Hint

Tying shares the embedding matrix with the output projection, saving the full table.

Test cases

case 1

in: [10000,512,true]

expected: [5120000,0,5120000]

case 2

in: [10000,512,false]

expected: [5120000,5120000,10240000]

case 3

in: [32000,768,true]

expected: [24576000,0,24576000]

case 4

in: [1,1,false]

expected: [1,1,2]

Starter codePython

def weight_tying_params(vocab_size, d_model, tied):
    # Your code here
    pass

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