Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Learned Position Embeddings › BERT-Style Position Embeddings

Michael BrenndoerferUpdated June 1, 20253 min read

Language AI HandbookAll chapters

Part 14 of 67Positional Encoding

Chapter 98Learned Position Embeddings › BERT-Style Position Embeddings

Reading section 11 of 19 · BERT-Style Position Embeddings · 3 minSection 11 of 19 · 3 min

Section 11 of 19 · 3 min read

Learned Position Embeddings › BERT-Style Position EmbeddingsOne section · Sections in this chapter11 / 19 ↓

Sections in this chapter19

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

BERT also uses learned position embeddings, but with a few differences. BERT includes segment embeddings to distinguish between sentence pairs and uses a masked language modeling objective that trains the model to fill in masked tokens rather than predict the next token. These differences change how position embeddings are used but not how they are defined: BERT's position embedding table works identically to GPT's.

The segment embeddings in BERT represent a third type of input information beyond token identity and position. BERT processes pairs of sentences concatenated together (for tasks like question answering and natural language inference), and the segment embedding distinguishes which sentence each token belongs to. The three embeddings are summed together:

\mathbf{h}_i = \mathbf{e}_{w_i} + \mathbf{p}_i + \mathbf{s}_{g_i}

where:

$\mathbf{e}_{w_i} \in \mathbb{R}^d$ : the token embedding for word $w_i$
$\mathbf{p}_i \in \mathbb{R}^d$ : the position embedding for position $i$
$\mathbf{s}_{g_i} \in \mathbb{R}^d$ : the segment embedding for segment $g_i \in \{0, 1\}$
$\mathbf{h}_i \in \mathbb{R}^d$ : the combined input representation

In[24]:

Code

class BERTStyleEmbedding:
    """
    BERT-style embedding layer with token, position, and segment embeddings.
    """

    def __init__(
        self, vocab_size, max_seq_len, embed_dim, num_segments=2, seed=None
    ):
        if seed is not None:
            np.random.seed(seed)

        # Token embeddings
        self.token_embeddings = np.random.randn(vocab_size, embed_dim) * 0.02

        # Position embeddings
        self.position_embeddings = (
            np.random.randn(max_seq_len, embed_dim) * 0.02
        )

        # Segment embeddings (for distinguishing sentence A vs sentence B)
        self.segment_embeddings = (
            np.random.randn(num_segments, embed_dim) * 0.02
        )

        self.vocab_size = vocab_size
        self.max_seq_len = max_seq_len
        self.embed_dim = embed_dim

    def forward(self, token_ids, segment_ids=None):
        """
        Compute embeddings for a sequence.

        Args:
            token_ids: Array of token indices, shape (seq_len,)
            segment_ids: Array of segment indices (0 or 1), shape (seq_len,)

        Returns:
            Embeddings of shape (seq_len, embed_dim)
        """
        seq_len = len(token_ids)

        if segment_ids is None:
            segment_ids = np.zeros(seq_len, dtype=int)

        # Look up all three embedding types
        token_embeds = self.token_embeddings[token_ids]
        pos_embeds = self.position_embeddings[np.arange(seq_len)]
        seg_embeds = self.segment_embeddings[segment_ids]

        # Sum all three
        return token_embeds + pos_embeds + seg_embeds

The addition of segment embeddings allows BERT to understand sentence structure in tasks like next sentence prediction and question answering. The position embeddings work identically to GPT-style: a simple lookup and addition. The key difference is that BERT's position embeddings must encode position within the full 512-token input, which may span two distinct sentences, while GPT's position embeddings encode position within a single continuous text stream.

This structural difference means BERT's position embeddings potentially encode slightly different information than GPT's. BERT's embeddings must help the model recognize when a position is near the boundary between two sentences, since the [SEP] token that separates sentences provides explicit signal but position provides context about how far into each sentence each token sits.