The Embedding Layer Learning for BERT

The Implementation of Embedding layer for BERT in Pytorch

import torch
from torch import nn

class Embedding(nn.Module):
  def __init__(self, vocab_size=10, d_model=2, maxlen=512, n_segments=2):
    super(Embedding, self).__init__()
    self.tok_embed = nn.Embedding(vocab_size, d_model)
    self.pos_embed = nn.Embedding(maxlen, d_model)
    self.seg_embed = nn.Embedding(n_segments, d_model)
    self.norm = nn.LayerNorm(d_model)
  def forward(self, x, seg):
    seq_len = x.size(1)
    pos = torch.arange(seq_len, dtype=torch.long)
    pos = pos.unsqueeze(0).expand_as(x)
    embedding = self.tok_embed(x) + self.tok_embed(pos) + self.seg_embed(seg)
    return self.norm(embedding)

Test the Embedding class

batch_size = 2
seq_len = 3
x = torch.randint(0, 10, (batch_size, seq_len))
seg = torch.randint(0, 2, (batch_size, seq_len))

embed_layer = Embedding()
print(f"tok_embed: {embed_layer.tok_embed.weight.size()}")
output_embedding = embed_layer(x, seg) 
print(f"output_embedding: {output_embedding}")

Corresponding Outputs:

tok_embed: torch.Size([10, 2]) output_embedding: tensor([[[-0.9996, 
0.9996],
         [-1.0000,  1.0000],
         [ 1.0000, -1.0000]],
        [[-1.0000,  1.0000],
         [-1.0000,  1.0000],
         [-1.0000,  1.0000]]], grad_fn=<NativeLayerNormBackward0>)

The illustration for above codes

This code defines a PyTorch module for creating embeddings for a transformer model. The model takes in three parameters: vocab_size, d_model, and maxlen. The vocab_size is the size of the vocabulary, d_model is the dimensionality of the embedding, and maxlen is the maximum length of the input sequence. The module also takes in a n_segments parameter, which is the number of segments in the input sequence.

---

The module contains three embedding layers: tok_embed, pos_embed, and seg_embed. tok_embed is an embedding layer for the input tokens, pos_embed is an embedding layer for the position of each token in the sequence, and seg_embed is an embedding layer for the segment of each token in the sequence.

---

The forward method of the module takes in two inputs: x and seg. x is a tensor of shape (batch_size, seq_size) containing the input sequence, and seg is a tensor of shape (batch_size, seq_len) containing the segment of each token in the sequence. The method first generates a position embedding for each token in the sequence using the pos_embed layer, and then combines this with the token embedding and the segment embedding using element-wise addition. The resulting embeddings are then normalized using layer normalization and returned.

The illustration for ‘pos’

In the given code, pos is a tensor of shape (seq_len,) that contains the positions of tokens in the input sequence. It is generated using the torch.arange function, which creates a 1-dimensional tensor containing a sequence of integers from 0 to seq_len-1. The dtype argument specifies the data type of the tensor, which in this case is torch.long.

---

Next, the unsqueeze(0) method is called on pos, which adds a new dimension of size 1 at the beginning of the tensor. This is done to make the shape of pos compatible with the shape of the input tensor x, which has shape (batch_size, seq_len).

---

Finally, the expand_as(x) method is called on pos, which creates a new tensor by repeating the elements of pos along the first dimension (batch_size) to match the shape of x. This allows the position embedding to be added element-wise to the token embedding using broadcasting.