6
$\begingroup$

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?

$\endgroup$

3 Answers 3

1
$\begingroup$

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.

$\endgroup$
1
$\begingroup$

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):

Here's a popular reference of the inside of Transformer cell

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/

$\endgroup$
2
  • $\begingroup$ This is not correct, as it assumes the dense layer has dim d_model and only one projection, but in the question feed_forward_dim=3072 != d_model=768 (which is typical). I think this description you found and included here is wrong or at least atypical. Typically the dense layer involves two projections: one d_model x feed_forward_dim (project to larger dim), and one feed_forward_dim x d_model (project that back down). $\endgroup$ Commented Mar 10, 2021 at 15:46
  • $\begingroup$ @EmmaStrubell This answer was only about the weights within the multi-headed self-attention layer, and doesn't include the feed-forward dense layer that comes after the multi-headed attention layer. Within the Multi-headed self attention layer, the last operation is a linear FC layer that combines the output of individual heads into one single vector. The final layer here (W0) is what converts vector of dim d_model to another vector of dim d_model itself. My answer is not explicitly calling this out. I am editing it to mention this. $\endgroup$ Commented Jul 5, 2021 at 6:56
0
$\begingroup$

I fount both answer to this questions are not clear and coherrent. So here I would like to add another solution.

In multi-head attention, the parameters are from the linear layers rather than from the scaled dot-product operation, where only multplication of the inputs are conducted. The number of heads doesn't add up to the number of attentions, because it is just reshape the input rather than doing addition calculations, as show in the code of minGPT.

In the first linear layer, there are 3 linear operations for 3 inputs, V, K and Q, both the input and output has a dimenson of d_model, so the parameters are: 3x (d_model*d_model + d_model).

In the last linear layer, only one linear operation with both input and output of length d_model, so the number of parameters is d_model * d_model + d_model.

So the totoal number of parameters is 4x (d_model*d_model + d_model) which matches what @Dileep Kumar Patchigolla has asked.

Multi-head attention

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.