BERT Architecture › The Input Representation
Language AI HandbookAll chapters BERT Architecture › The Input RepresentationOne section · Sections in this chapter3 / 12 ↓
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's input representation is one of its most distinctive features. Unlike simpler models that use only token embeddings, BERT combines three embedding types to capture different aspects of the input. This design enables BERT to process sentence pairs for tasks like question answering and natural language inference.
Think of the three embedding types as answering three different questions about each token. The token embedding answers "what word is this?", storing the semantic and syntactic properties of the vocabulary item itself. The position embedding answers "where does this word appear?", encoding the sequential position within the input. The segment embedding answers "which sentence does this word belong to?", distinguishing sentence A from sentence B in paired inputs. All three answers are equally important, so BERT adds them together into a single dense vector rather than processing them separately.
The 30,522-token vocabulary is produced by WordPiece tokenization, the same subword tokenization scheme used in earlier Google systems. WordPiece splits rare or unknown words into smaller pieces. This ensures that even novel words can be represented as combinations of known subwords. A word like "embedding" might remain whole if it appears frequently in training, but "embeddings" might be split into "embedding" and "##s", where the "##" prefix indicates a continuation. This keeps the vocabulary manageable while still representing nearly any text.
Three Embedding Layers
Every input token receives three embeddings that are summed together:
- Token embeddings: Standard learned embeddings that map each vocabulary token to a dense vector
- Position embeddings: Learned embeddings for each position (0 to 511) encoding sequential order
- Segment embeddings: Two learned embeddings (A and B) indicating which sentence a token belongs to
The segment embeddings deserve special attention. BERT was designed for tasks involving sentence pairs: given two sentences, determine if the second follows the first (next sentence prediction), if the first entails the second (natural language inference), or find the answer span (question answering). The segment embeddings allow the model to distinguish which tokens belong to which sentence even after they're concatenated.
The concatenation of two sentences into a single sequence with special tokens might seem like a strange design choice. Why not process each sentence independently and combine their representations afterward? The answer is that cross-sentence attention patterns are precisely what many tasks require. In question answering, the model needs to attend from words in the question to matching words in the context passage. In natural language inference, it needs to identify contradictions between specific words in the premise and hypothesis. By putting both sentences in the same context window and letting attention operate freely between them, BERT can learn these cross-sentence relationships directly. The segment embeddings ensure that the model can always tell which sentence each token came from, even while attending freely across both.
Special Tokens
BERT uses several special tokens to structure its input:
- [CLS]: Prepended to every input. Its final representation is used for classification tasks
- [SEP]: Inserted between sentences and at the end of the input to mark boundaries
- [MASK]: Used during pretraining to indicate positions the model should predict
- [PAD]: Fills sequences shorter than the batch maximum length
- [UNK]: Represents tokens not in the vocabulary
Let's implement the embedding layer to see how these components combine:
class BertEmbeddings(nn.Module):
"""BERT embedding layer combining token, position, and segment embeddings."""
def __init__(
self,
vocab_size: int = 30522,
hidden_size: int = 768,
max_position: int = 512,
type_vocab_size: int = 2,
dropout: float = 0.1,
):
super().__init__()
# Three embedding tables
self.token_embeddings = nn.Embedding(vocab_size, hidden_size)
self.position_embeddings = nn.Embedding(max_position, hidden_size)
self.segment_embeddings = nn.Embedding(type_vocab_size, hidden_size)
# Layer normalization and dropout
self.layer_norm = nn.LayerNorm(hidden_size, eps=1e-12)
self.dropout = nn.Dropout(dropout)
# Register position ids as buffer (not a parameter)
self.register_buffer(
"position_ids", torch.arange(max_position).unsqueeze(0)
)
def forward(
self,
input_ids: torch.Tensor,
segment_ids: torch.Tensor | None = None,
) -> torch.Tensor:
"""
Combine embeddings for input tokens.
Args:
input_ids: Token indices, shape (batch_size, seq_len)
segment_ids: Segment indices (0 or 1), shape (batch_size, seq_len)
Returns:
Combined embeddings, shape (batch_size, seq_len, hidden_size)
"""
seq_len = input_ids.size(1)
# Get position ids for this sequence length
position_ids = self.position_ids[:, :seq_len]
# Default segment ids to all zeros (single sentence)
if segment_ids is None:
segment_ids = torch.zeros_like(input_ids)
# Look up embeddings
token_emb = self.token_embeddings(input_ids)
position_emb = self.position_embeddings(position_ids)
segment_emb = self.segment_embeddings(segment_ids)
# Sum all three
embeddings = token_emb + position_emb + segment_emb
# Normalize and apply dropout
embeddings = self.layer_norm(embeddings)
embeddings = self.dropout(embeddings)
return embeddingsInput shape: torch.Size([1, 8]) Output shape: torch.Size([1, 8, 768]) Embedding dimension: 768
The embedding layer maps our 8-token input into a tensor of shape (1, 8, 768), where each token now has a 768-dimensional representation that combines token identity with positional and segment information. This dense representation is what flows through the subsequent transformer layers.
The element-wise addition of three embedding types may seem unusual. Why not concatenate them? The answer lies in parameter efficiency and training dynamics. Addition keeps the hidden dimension fixed at 768, while concatenation would triple it. More subtly, addition forces the model to learn representations where token meaning, positional information, and segment identity can coexist in the same vector space. The layer normalization after addition rescales these combined representations to have consistent statistics.
In practice, this summation works because modern neural networks are remarkably good at learning linear decompositions within a shared vector space. During training, the model adjusts all three embedding tables simultaneously via gradient descent. The token embeddings learn to represent word meaning in a subspace that leaves room for positional and segment information to be added without collision. The layer normalization that follows stabilizes the combined signal by ensuring the activations entering the first transformer layer have zero mean and unit variance, which helps the subsequent attention computations behave consistently regardless of which specific tokens were summed.
Learned vs Sinusoidal Positions
BERT uses learned position embeddings rather than the sinusoidal encodings from the original transformer. Each of the 512 possible positions gets its own learned vector. This choice trades generalization for expressiveness: learned embeddings cannot extrapolate to positions beyond 512, but they can capture position-specific patterns that sinusoidal encodings cannot.
Recall from the transformer chapter that sinusoidal encodings use a fixed mathematical formula: position and dimension produce the value for even dimensions and for odd dimensions. This formula was designed so that the distance between any two position encodings depends only on their relative offset, not their absolute positions, which enables the model to generalize to sequence lengths it never saw during training. The key insight behind sinusoidal encodings is that any position can be expressed as a linear combination of nearby positions, which means a model trained on sequences up to length 512 can theoretically handle length 600.
Learned position embeddings abandon this mathematical guarantee. Each of BERT's 512 positions gets a 768-dimensional embedding vector that is randomly initialized and trained by gradient descent, exactly like word embeddings. This means position 37 and position 38 start with no particular relationship to each other; they only become related if the training data contains patterns that differ between those positions. In practice, this turns out to work well because NLP tasks rarely require generalizing to sequence lengths longer than those seen in training, and most downstream tasks use sequences well under 512 tokens. The tradeoff is that you cannot simply feed a 600-token sequence to BERT without modification, since positions 513-600 have no learned embeddings.
# Compare learned vs sinusoidal position embeddings
def sinusoidal_position_encoding(max_len: int, d_model: int) -> np.ndarray:
"""Generate sinusoidal position encodings."""
position = np.arange(max_len)[:, np.newaxis]
div_term = np.exp(np.arange(0, d_model, 2) * (-np.log(10000.0) / d_model))
pe = np.zeros((max_len, d_model))
pe[:, 0::2] = np.sin(position * div_term)
pe[:, 1::2] = np.cos(position * div_term)
return pe
# Simulate trained learned position embeddings:
# Real trained embeddings show irregular but structured patterns —
# a mix of low-frequency components (position clusters) and
# high-frequency variation (dimension-specific signals).
positions = np.arange(512)
dims = np.arange(768)
# Base: low-rank structure from 8 dominant components
n_components = 8
freqs = np.linspace(0.5, 6.0, n_components)
amplitudes = np.random.uniform(0.04, 0.12, n_components)
phases = np.random.uniform(0, 2 * np.pi, n_components)
learned_pe = np.zeros((512, 768))
for freq, amp, phase in zip(freqs, amplitudes, phases):
dim_pattern = np.sin(dims * freq / 768 * np.pi + phase)
pos_pattern = np.sin(positions * freq / 512 * np.pi)
learned_pe += amp * np.outer(pos_pattern, dim_pattern)
# Add moderate structured noise (dimension-specific offsets)
learned_pe += np.random.normal(0, 0.015, learned_pe.shape)
# Center and scale the simulation to the displayed color range so the
# irregular structure remains visible instead of saturating the heatmap.
learned_pe -= learned_pe.mean()
learned_pe *= 0.1 / np.max(np.abs(learned_pe))
sinusoidal_pe = sinusoidal_position_encoding(512, 768)
The learned embeddings appear less regular because they capture whatever positional patterns help with the pretraining objectives. The sinusoidal pattern's mathematical regularity means positions can be expressed as linear combinations of other positions, enabling some length generalization. BERT prioritized expressiveness over extrapolation since most downstream tasks don't require sequences longer than 512 tokens.