How to account for the no:of parameters in the Multihead self-Attention layer of BERT I have read the BERT paper for NLP (https://arxiv.org/abs/1810.04805), and trying to understand the keras implementation. Here's my code to load the BERT model:
import keras
from keras_bert import get_base_dict, get_model, compile_model, gen_batch_inputs

model = get_model(
    token_num=30000,
    head_num=12,
    transformer_num=12,
    embed_dim=768,
    feed_forward_dim=3072,
    seq_len=500,
    pos_num=512,
    dropout_rate=0.05
)
compile_model(model)
model.summary()

Note that all the parameters I used are the default of BERT Base implementation. In the keras model summary, I could see that there are 2,362,368 trainable parameters at each of the multi-head self attention layer. But I don't understand how to get this number.
There are total 12 attention heads, and in each head, there are Q, K, V vectors each with dimension 768 × 768.
So the dimension should have been (768 × 768 × 3 + 768 × 3) × 12 including biases = only 1,771,776 × 12.
But the actual number 2,362,368 seems to be equal to (768 × 768 × 4 + 768 × 4).
Can someone explain how I can account for this number?
 A: After doing the multi-head attention, you have 12 heads context vectors of dimension 768 and you need to project them back to the model dimension, this gives you another 12 × 768 × 768 + 768 parameters. In addition, there is layer normalization with 2 × 768 parameters.
A: I have found the answer after digging into a pytorch implementation and a few other blogs. Here's the explanation for the number of paramteres in the Transformer cell (only the mult-headed self-attention part):

We can see the inside of transformer cell in above picture. The input vector is transformed in multiple heads, then applied the self-attention operation, then all are concatenated, and then a fully connected dense forward layer is applied. In terms of dimensions, here's how it looks:
The input vector of dimension d_model (in  X) gets multiplied by three matrices WQ, WK, WV, 12 (=attention heads, or A) times to give (3A) pairs of vectors (Q, K, V). These vectors (Z0 to Z7 in the image) are each of length d_model/A. So dimension of each of these matrices is d_model * d_model/A  and we have 3 * A such matrices.
Including the bias for eah of Q, K, V matrices, total weights till now = d_model * d_model/A * 3A + d_model * 3. By this point, we have Z0 to Zi vectors from above image. These are then concatenated, and passed through the dense layer W0 which would have dimension d_model * d_model + d_model (with bias).
So total dimension of transformer cell:
A * (d_model * d_model/A) * 3 + 3*d_model  + (d_model * d_model + d_model). For BERT base, the values are A= 12, d_model = 786. So total parameters = 12 * ( 768 * 768/12) * 3 + 3*768 + 768*768 + 768  = 2,362,368
Edit: The output of this will be a vector of dim d_model. This then gets a residual connection to the input itself, which then is passed into another Dense layer where we get two matrices of dimensions (d_model * d_feed_forward). Those weights are not part of this calculation
Img Ref: http://jalammar.github.io/illustrated-transformer/
