Parameter count formula for RWKV

[Triang-jyed-driung]

According to the formula at the arXiv article, the parameter count of the RWKV-4-world 0.1B model ($L=12, D=768, V=65536$) would be:
$$2VD + 13D^2 L + D(11L+4)$$
which yields

>>> 2*65536*768 + 13*768*768*12 + 768*(11*12+4)
192780288

But the output is

emb.weight                        f32      cpu  65536   768 
blocks.0.ln1.weight               f32      cpu    768       
blocks.0.ln1.bias                 f32      cpu    768       
blocks.0.ln2.weight               f32      cpu    768       
blocks.0.ln2.bias                 f32      cpu    768       
blocks.0.att.time_decay           f32      cpu    768       
blocks.0.att.time_first           f32      cpu    768       
blocks.0.att.time_mix_k           f32      cpu    768       
blocks.0.att.time_mix_v           f32      cpu    768       
blocks.0.att.time_mix_r           f32      cpu    768       
blocks.0.att.key.weight           f32      cpu    768   768 
blocks.0.att.value.weight         f32      cpu    768   768 
blocks.0.att.receptance.weight    f32      cpu    768   768 
blocks.0.att.output.weight        f32      cpu    768   768 
blocks.0.ffn.time_mix_k           f32      cpu    768
blocks.0.ffn.time_mix_r           f32      cpu    768
blocks.0.ffn.key.weight           f32      cpu    768  3072 
blocks.0.ffn.receptance.weight    f32      cpu    768   768
blocks.0.ffn.value.weight         f32      cpu   3072   768 
....................................................................................................................................................................................
blocks.11.ln1.weight              f32      cpu    768
blocks.11.ln1.bias                f32      cpu    768
blocks.11.ln2.weight              f32      cpu    768
blocks.11.ln2.bias                f32      cpu    768
blocks.11.att.time_decay          f32      cpu    768
blocks.11.att.time_first          f32      cpu    768
blocks.11.att.time_mix_k          f32      cpu    768
blocks.11.att.time_mix_v          f32      cpu    768
blocks.11.att.time_mix_r          f32      cpu    768
blocks.11.att.key.weight          f32      cpu    768   768
blocks.11.att.value.weight        f32      cpu    768   768
blocks.11.att.receptance.weight   f32      cpu    768   768
blocks.11.att.output.weight       f32      cpu    768   768
blocks.11.ffn.time_mix_k          f32      cpu    768
blocks.11.ffn.time_mix_r          f32      cpu    768
blocks.11.ffn.key.weight          f32      cpu    768  3072
blocks.11.ffn.receptance.weight   f32      cpu    768   768
blocks.11.ffn.value.weight        f32      cpu   3072   768
ln_out.weight                     f32      cpu    768
ln_out.bias                       f32      cpu    768
head.weight                       f32      cpu    768 65536

So, $2VD$ comes from

emb.weight                        f32      cpu  65536   768
head.weight                       f32      cpu    768 65536

$13D^2L$ comes from

blocks.0.att.key.weight           f32      cpu    768   768 
blocks.0.att.value.weight         f32      cpu    768   768 
blocks.0.att.receptance.weight    f32      cpu    768   768 
blocks.0.att.output.weight        f32      cpu    768   768 
blocks.0.ffn.key.weight           f32      cpu    768  3072 
blocks.0.ffn.receptance.weight    f32      cpu    768   768
blocks.0.ffn.value.weight         f32      cpu   3072   768 

$11DL$ comes from

blocks.0.ln1.weight               f32      cpu    768       
blocks.0.ln1.bias                 f32      cpu    768       
blocks.0.ln2.weight               f32      cpu    768       
blocks.0.ln2.bias                 f32      cpu    768       
blocks.0.att.time_decay           f32      cpu    768       
blocks.0.att.time_first           f32      cpu    768       
blocks.0.att.time_mix_k           f32      cpu    768       
blocks.0.att.time_mix_v           f32      cpu    768       
blocks.0.att.time_mix_r           f32      cpu    768  
blocks.0.ffn.time_mix_k           f32      cpu    768
blocks.0.ffn.time_mix_r           f32      cpu    768

However! Where is $4D$? these are only $2D$:

ln_out.weight                     f32      cpu    768
ln_out.bias                       f32      cpu    768

Am I missing something?

[howard-hou]

I guess another 2D is from ln0
if self.layer_id == 0:
self.ln0 = nn.LayerNorm(args.n_embd)