BERT embedding layer

Question

I am trying to figure how the embedding layer works for the pretrained BERT-base model. I am using pytorch and trying to dissect the following model:

import torch
model = torch.hub.load('huggingface/pytorch-transformers', 'model', 'bert-base-uncased')
model.embeddings

This BERT model has 199 different named parameters, of which the first 5 belong to the embedding layer (the first layer)

==== Embedding Layer ====

embeddings.word_embeddings.weight                       (30522, 768)
embeddings.position_embeddings.weight                     (512, 768)
embeddings.token_type_embeddings.weight                     (2, 768)
embeddings.LayerNorm.weight                                   (768,)
embeddings.LayerNorm.bias                                     (768,)

As I understand, the model accepts input in the shape of [Batch, Indices] where Batch is of arbitrary size (usually 32, 64 or whatever) and Indices are the corresponding indices for each word in the tokenized input sentence. Indices has a max length of 512. One input sample might look like this:

[[101, 1996, 4248, 2829, 4419, 14523, 2058, 1996, 13971, 3899, 102]]

This contains only 1 batch and is the tokenized form of the sentence "The quick brown fox jumps over the lazy dog".

The first word_embeddings weight will translate each number in Indices to a vector spanned in 768 dimensions (the embedding dimension).

Now, the position_embeddings weight is used to encode the position of each word in the input sentence. Here I am confused about why this parameter is being learnt? Looking at an alternative implementation of the BERT model, the positional embedding is a static transformation. This also seems to be the conventional way of doing the positional encoding in a transformer model. Looking at the alternative implementation it uses the sine and cosine function to encode interleaved pairs in the input. I tried comparing model.embeddings.position_embeddings.weight and pe, but I cannot see any similarity. The in the last sentence under under A.2 Pre-training Procedure (page 13) the paper states

Then, we train the rest 10% of the steps of sequence of 512 to learn the positional embeddings.

Why is the positional embedding weight being learnt and not predefined?

The next layer after the positional embedding is the token_type_embeddings. Here I am confused about how the segment label is inferred by the model. If I understand this correctly each input sentence is delimited by the [SEP] token. In the example above there is only 1 [SEP] token and the segment label must be 0 for that sentence. But there could be a maximum of 2 segment labels. If so, will the 2 segments be handled separately or are they processed in parallel all the same as one "array"? How does the model handle multiple sentence segments?

finally the output from theese 3 embeddings are added togheter and passed through layernorm which I understand. But, are the weights in these embedding layers adjusted when fine-tuning the model to a downstream task?

Answer 1

Why are positional embeddings learned?

This was asked in the repo of the original implementation without an answer. It didn't get an answer either in the HuggingFace Transformers repo and in cross-validated, also without answer, or without much evidence.

Given that in the original Transformer paper the sinusoidal embedding were the default ones, I understand that during preliminary hyperparameter tuning, the authors of BERT decided to go with learned embeddings, deciding not to duplicate all experiments with both types of embeddings.

We can, nevertheless, see some comparisons between learned and sinusoidal positional embedding in the ICLR'21 article On Position Embeddings in BERT, where the authors observe that:

The fully-learnable absolute PE performs better in classification, while relative PEs perform better in span prediction.

How does the model handle multiple sentence segments?

This is best understood with the figure of the original BERT paper:

The two sentences are encoded into three sequences of the same length:

• Sequence of subword tokens: the sentence tokens are concatenated into a single sequence, separating them with a [SEP] token. This sequence is embedded with the subword token embedding table; you can see the tokens here.
• Sequence of positional embedding: sequentially increasing positions form the initial position of the [CLS] token to the position of the second [SEP] token. This sequence is embedded with the positional embedding table, which has 512 elements.
• Sequence of segment embeddings: as many EA tokens as the token length of the first sentence (with [CLS] and [SEP]) followed by as many EB tokens as the token length of the second sentence (with the [SEP]). This sequence is embedded with the segment embedding table, with has 2 elements.

After embedding the three sequences with their respective embedding tables, we have 3 vector sequences, which are added together and used as input to the self-attention layers.

Are the weights in these embedding layers adjusted when fine-tuning the model to a downstream task?

Yes, they are. Normally, all parameters are fine-tuned when fine-tunine a BERT-based model.

Nevertheless, it is also possible to simply use BERT's representations as input to a classification model, without fine-tuning BERT at all.

In this article you can see how these two approaches compare. In general, for BERT, you obtain better results by fine-tuning the whole model.