BERT Architecture › Output Representations
Language AI HandbookAll chapters BERT Architecture › Output RepresentationsOne section · Sections in this chapter6 / 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 produces contextualized representations at every position. Different downstream tasks use these representations in different ways. Understanding the structure of these outputs is essential for knowing how to apply BERT effectively to new tasks.
After all 12 (or 24) encoder layers have processed the input, you have a matrix of shape (sequence length, 768). Each row of this matrix is a contextualized representation of the corresponding input token. The representation at position 5 has been updated by every attention head in every layer, incorporating information from every other token in the sequence. This is fundamentally different from a static word embedding: the representation of the word "bank" depends on which other words appeared in the input. The same word gets a different 768-dimensional vector in every different context.
This contextuality is the source of BERT's power, but it also creates a practical question: how do you extract task-relevant information from a sequence of 512 potentially useful vectors? The answer depends on what your task requires.
The [CLS] Token
The [CLS] token's final hidden state is designed to aggregate sequence-level information. During pretraining with the Next Sentence Prediction task, the [CLS] representation must contain enough information to determine whether two sentences are consecutive. This encourages [CLS] to capture global semantic content.
The [CLS] token is special in that it is prepended to every input regardless of the task. It has no inherent meaning as a word; its entire representation is built up through attention during the forward pass. Because [CLS] can attend to every other token in every attention head, it has the opportunity to aggregate information from the entire sequence. By the final layer, the [CLS] representation is a compressed summary of the input that has been shaped by the pretraining objective to capture sequence-level semantics.
In practice, the quality of [CLS] representations for downstream tasks depends heavily on whether the pretraining objective required sequence-level understanding. BERT's NSP pretraining task did require this, but later work by Lan et al. (2019) on ALBERT and by Liu et al. (2019) on RoBERTa questioned whether NSP helped. RoBERTa, which removed NSP from pretraining, sometimes matched or exceeded BERT's performance on sentence-pair tasks, suggesting that [CLS] can learn useful sequence-level representations even without an explicit sequence-level pretraining signal, as long as the MLM training corpus is large and diverse enough.
For classification tasks, you typically pass the [CLS] representation through a task-specific linear layer:
class BertClassificationHead(nn.Module):
"""Classification head using [CLS] token representation."""
def __init__(
self, hidden_size: int = 768, num_classes: int = 2, dropout: float = 0.1
):
super().__init__()
self.dense = nn.Linear(hidden_size, hidden_size)
self.dropout = nn.Dropout(dropout)
self.classifier = nn.Linear(hidden_size, num_classes)
def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
"""
Args:
sequence_output: BERT output, shape (batch, seq_len, hidden)
Returns:
Classification logits, shape (batch, num_classes)
"""
# Take [CLS] token (first position)
cls_output = sequence_output[:, 0, :]
# Project through dense layer
pooled = self.dense(cls_output)
pooled = torch.tanh(pooled)
pooled = self.dropout(pooled)
# Classify
logits = self.classifier(pooled)
return logitsToken-Level Representations
For sequence labeling tasks like named entity recognition or part-of-speech tagging, you use the representation at each token position:
class BertTokenClassificationHead(nn.Module):
"""Token classification head for sequence labeling."""
def __init__(
self, hidden_size: int = 768, num_labels: int = 9, dropout: float = 0.1
):
super().__init__()
self.dropout = nn.Dropout(dropout)
self.classifier = nn.Linear(hidden_size, num_labels)
def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
"""
Args:
sequence_output: BERT output, shape (batch, seq_len, hidden)
Returns:
Token logits, shape (batch, seq_len, num_labels)
"""
output = self.dropout(sequence_output)
logits = self.classifier(output)
return logitsSpan Representations
For extractive question answering, BERT predicts start and end positions of the answer span. Two linear layers project each token's representation to start and end logits:
class BertQuestionAnsweringHead(nn.Module):
"""QA head for extractive question answering."""
def __init__(self, hidden_size: int = 768):
super().__init__()
self.start_classifier = nn.Linear(hidden_size, 1)
self.end_classifier = nn.Linear(hidden_size, 1)
def forward(
self, sequence_output: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Args:
sequence_output: BERT output, shape (batch, seq_len, hidden)
Returns:
start_logits: shape (batch, seq_len)
end_logits: shape (batch, seq_len)
"""
start_logits = self.start_classifier(sequence_output).squeeze(-1)
end_logits = self.end_classifier(sequence_output).squeeze(-1)
return start_logits, end_logitsClassification output shape: torch.Size([2, 3]) Token classification output shape: torch.Size([2, 20, 9]) QA start/end logits shape: torch.Size([2, 20]) torch.Size([2, 20])
Each head produces output with the appropriate shape for its task. The classification head reduces the sequence to a single prediction per sample. The token classification head produces a label prediction for each position. The QA head generates start and end scores for every token, allowing span extraction by finding the highest-scoring start-end pair.
The flexibility of BERT's output representations is key to its success. The same pretrained model can power classification, sequence labeling, question answering, and many other tasks by simply changing the task-specific head.
Suppose you want to classify product reviews as positive or negative using BERT. The process is:
- Tokenize the review text using BERT's WordPiece tokenizer
- Prepend [CLS] and append [SEP]:
[CLS] This product is amazing [SEP] - Pass through BERT to get the sequence output of shape (seq_len, 768)
- Extract the [CLS] representation at position 0: shape (768,)
- Pass through a linear layer: 768 inputs, 2 outputs (positive/negative)
- Apply softmax to get class probabilities
During fine-tuning, you train the linear head and update BERT's weights with a small learning rate (2e-5 to 5e-5 is typical). The small learning rate is important: BERT already contains powerful language representations, and you want to gently adjust them for the task rather than overwrite them. Too large a learning rate destroys the pretrained representations through "catastrophic forgetting".
A common question is whether to fine-tune BERT's weights at all, or to freeze them and only train the task-specific head. The answer is almost always to fine-tune. Frozen BERT features, while powerful, do not adapt to the vocabulary or style of your specific domain. Fine-tuning allows BERT's internal representations to shift slightly toward representations that are most useful for your task. For small datasets (under a few thousand examples), fine-tuning can overfit, in which case freezing lower layers while fine-tuning upper layers provides a useful middle ground.