Special Tokens › Token Type IDs and Segment Embeddings

Beyond the tokens themselves, BERT introduces another mechanism for handling multi-segment inputs: token type IDs (also called segment IDs). Token type IDs answer a question that the token sequence alone cannot: for each position in the input, which sentence does it belong to? This question matters for tasks like natural language inference, question answering, and textual entailment, where the model needs to reason about the relationship between two distinct texts.

Think of token type IDs as jersey numbers in a two-team game: players wearing jersey 0 belong to team A, and players wearing jersey 1 belong to team B. The number doesn't describe what the player does, but it tells the referee (the model) which team's goals the player is working toward. Without this signal, the model has no reliable way to know which part of a two-sentence input is the premise and which is the hypothesis in an inference task.

In BERT's embedding layer, the final embedding for each token is the sum of three components. Before we write the formula, it helps to understand what each component contributes. The token embedding encodes what the token is: its linguistic identity and semantic content. The positional embedding encodes where the token is: its position in the sequence. The segment embedding encodes which sentence the token belongs to: its structural role in the multi-sentence input.

where:

• $\mathbf{e}^{\text{token}}_i$: the token embedding for the $i$-th input token, looked up from a vocabulary embedding matrix of shape $[V \times d]$, where $V$ is the vocabulary size (30,522 for BERT-base) and $d$ is the hidden dimension (768 for BERT-base)
• $\mathbf{e}^{\text{position}}_i$: the positional embedding for position $i$, looked up from a position embedding matrix of shape $[L \times d]$, where $L$ is the maximum sequence length (512 for BERT-base)
• $\mathbf{e}^{\text{segment}}_i$: the segment embedding for token $i$, looked up from a segment embedding matrix of shape $[2 \times d]$, with row 0 for sentence A tokens and row 1 for sentence B tokens

Why does this formula make sense? Notice that all three embedding vectors have the same dimension $d$, so their element-wise sum produces a single embedding vector $\mathbf{e}_i$ of dimension $d$. This sum encodes three orthogonal kinds of information simultaneously: what the token is, where it is, and which segment it belongs to. The model can potentially disentangle these signals across different dimensions of the embedding space. The combined embedding is then passed into the first transformer layer, carrying all three kinds of structural information into every subsequent attention computation.

The segment embedding lookup table has only two rows: one for segment A (type ID 0) and one for segment B (type ID 1). The model learns these two vectors during pre-training as part of the next-sentence prediction objective. Unlike positional embeddings (which have one vector per position up to the maximum sequence length) or token embeddings (which have one vector per vocabulary entry), segment embeddings are an extremely compact table: just two learned vectors that modulate the entire input based on segment membership.

For single-sentence inputs, all tokens receive type ID 0. The segment B embedding is irrelevant but still part of the model. For sentence-pair inputs, the type IDs switch from 0 to 1 at the second segment:

Modern models like RoBERTa dropped segment embeddings entirely. Since RoBERTa discarded next-sentence prediction during pre-training, there was no objective that rewarded the model for learning meaningful segment embeddings. Removing segment embeddings simplifies the architecture without hurting performance. Interestingly, even models that handle sentence pairs for tasks like Natural Language Inference work fine without explicit segment embeddings, relying instead on the [SEP] tokens and positional context to infer segment boundaries. This empirical finding suggests that [SEP]'s boundary signal may already be sufficient, and segment embeddings are redundant rather than necessary.

Let's trace through the complete construction of a BERT input for a sentence-pair task, step by step, to make the three-component embedding concrete. Suppose the task is Natural Language Inference, and the input is:

• Premise: "The cat sat on the mat."
• Hypothesis: "A cat was resting."

Step 1: WordPiece tokenization. The tokenizer first splits each sentence into subword tokens. "Resting" might become ["rest", "##ing"] if "resting" is not a vocabulary entry but "rest" and the suffix "##ing" are.

Step 2: Special token insertion. The tokenizer prepends [CLS] to the entire sequence, inserts [SEP] after sentence A, appends sentence B's tokens, and adds a final [SEP]:

Step 3: Embedding summation. For each position $i$, three vectors are looked up and summed: the token embedding for the token at position $i$, the positional embedding for position $i$, and the segment embedding for the type ID at position $i$ (either row 0 or row 1 of the segment embedding table).

Step 4: Transformer processing. The summed embeddings enter the first transformer layer. All 16 positions attend to each other (with full bidirectional attention). After all 12 layers, each position has a contextualized representation that encodes information from the entire sequence.

Step 5: Classification. The [CLS] representation at position 0 is passed through a classification head with three outputs: entailment, contradiction, neutral. The class with the highest softmax probability is the predicted label.

The segment boundary created by the [SEP] at position 8 and the jump in type IDs from 0 to 1 gives the model two complementary signals about where sentence A ends and sentence B begins. The attention mechanism can learn to use either signal: some heads may track the [SEP] boundary while others track the type ID boundary. Having both signals makes the boundary information resilient to noise in either one.

The figure below visualizes this three-layer structure: token identity, positional index, and segment membership all encoded in parallel for the same sequence.