BERT Representations › Layer Selection: Where's the Information?
Language AI HandbookAll chapters BERT Representations › Layer Selection: Where's the Information?One section · Sections in this chapter3 / 10 ↓
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 doesn't produce just one representation. BERT-base has 12 transformer layers, each outputting a full sequence of hidden states. Which layer should you use?
The answer depends on your task. Different layers capture different types of linguistic information. A rich body of research on "probing classifiers" has shed light on what each layer knows. The idea is simple: train a small linear classifier to predict a linguistic property, such as part-of-speech tags or syntactic dependencies, using only the representations from a single BERT layer. The layer where the linear classifier performs best is the layer that encodes that property most cleanly. Across many such probing studies, a consistent pattern has emerged:
- Lower layers (1-4) capture surface-level features like part-of-speech and simple syntax
- Middle layers (5-8) capture syntactic structures and dependencies
- Upper layers (9-12) capture task-specific and semantic information
Think of the layers as a hierarchy of abstraction, similar to what happens in convolutional neural networks applied to images. Early CNN layers detect edges and textures; later layers detect faces and objects. In BERT, early layers learn that "running" is a verb and "the" is a determiner, while later layers learn that "the bank" in a financial context means something fundamentally different from "the bank" in a geological context.
The key insight is that BERT's layers are not interchangeable. Using the last layer by default is a reasonable starting point, but it is not always the best choice. For tasks that depend on surface form or morphology, such as named entity recognition in a highly inflected language, representations from layers 3 to 5 may carry more signal than the final layer. For tasks that require deep world knowledge and semantic reasoning, layers 10 to 12 tend to perform better. The right layer is a hyperparameter worth tuning.
Let's visualize how representations evolve across layers.
# Extract all layer representations for a sentence
text = "The bank approved the loan application."
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs, output_hidden_states=True)
# hidden_states is a tuple of 13 tensors: embedding layer + 12 transformer layers
all_layers = outputs.hidden_states # tuple of [1, seq_len, 768]
num_layers = len(all_layers)
# Get [CLS] representation from each layer
cls_per_layer = [layer[0, 0, :].numpy() for layer in all_layers]Number of layers (including embedding): 13 Representation shape per layer: (768,)
For this sentence, the size of the change varies across depth rather than declining monotonically. Middle transitions are somewhat smaller, while the final transition is the largest. The example reinforces that layer choice changes the representation; it does not establish a universal ordering for every sentence or task.
The heatmap reveals a clear pattern. Adjacent layers are highly similar, shown by the bright diagonal band. But layer 1 and layer 12 are quite different, with similarity around 0.5. The information encoded at different depths varies substantially.
Task-Specific Layer Selection
Different NLP tasks benefit from different layers. Probing experiments have systematically tested which layers encode which linguistic properties.
POS tagging peaks at layers 3-4, while coreference resolution benefits from the deepest layers. This has practical implications. If you're building a part-of-speech tagger, using the last layer might hurt performance. Using layer 3 or 4 could give you a better starting point.
Layer Combination Strategies
Why limit yourself to one layer? Several strategies combine information across layers.
The motivation for combining layers comes directly from the probing results. If syntactic information peaks at layer 6 and semantic information peaks at layer 11, a task that requires both, such as semantic role labeling, which asks which arguments fill which roles in a sentence, might benefit from using both. Rather than picking one layer and discarding the others, you can fuse information across the stack. This idea predates BERT: ELMo (2018) combined representations from two LSTM layers using learned weights, achieving state-of-the-art results on a wide range of tasks by letting downstream models decide how much to weight each level of abstraction.
In practice, the last four layers are the most commonly combined because they carry the richest semantic signal and are most likely to help with downstream tasks. But for structured prediction tasks like dependency parsing, combining a broader range of layers can capture both low-level morphological features and high-level syntactic patterns in a single representation.
Concatenation stacks representations from multiple layers:
def concat_layers(hidden_states, layer_indices):
"""Concatenate [CLS] representations from specified layers."""
selected = [hidden_states[i][0, 0, :] for i in layer_indices]
return torch.cat(selected, dim=0)
# Combine last 4 layers
last_four = concat_layers(all_layers, [-4, -3, -2, -1])Single layer shape: (768,) Concatenated (last 4) shape: torch.Size([3072])
Concatenation produces a larger vector (3072 dimensions for last-4 concatenation) but preserves distinct information from each layer. This is useful when different layers capture different aspects relevant to your task.
Weighted sum learns to combine layers adaptively:
def weighted_sum_layers(hidden_states, weights):
"""Compute weighted sum of representations across layers."""
# Stack all layers: [num_layers, seq_len, hidden_dim]
stacked = torch.stack([h.squeeze(0) for h in hidden_states], dim=0)
# Normalize weights
weights = F.softmax(torch.tensor(weights), dim=0)
# Weighted sum: [seq_len, hidden_dim]
combined = torch.einsum("l,lsh->sh", weights, stacked)
return combined[0] # Return [CLS]
# Example: emphasize later layers
layer_weights = [0.5] * 6 + [1.0] * 7 # 13 weights (embedding + 12 layers)
weighted_cls = weighted_sum_layers(all_layers, layer_weights)Weighted sum shape: torch.Size([768]) Same as single layer: True
The weighted sum keeps dimensionality fixed while allowing the model to learn which layers matter most. Models like ELMo pioneered this approach, and it works well when you want to fine-tune the combination weights for a specific task.
Scalar mix (from AllenNLP) is a learnable variant:
class ScalarMix(nn.Module):
"""
Learnable weighted combination of layer representations.
Used in ELMo and adaptable for BERT.
"""
def __init__(self, num_layers):
super().__init__()
self.num_layers = num_layers
self.scalar_parameters = nn.Parameter(torch.zeros(num_layers))
self.gamma = nn.Parameter(torch.ones(1))
def forward(self, tensors):
# tensors: list of [batch, seq_len, hidden]
normed_weights = F.softmax(self.scalar_parameters, dim=0)
pieces = []
for weight, tensor in zip(normed_weights, tensors):
pieces.append(weight * tensor)
mixed = self.gamma * sum(pieces)
return mixed
# This would be trained alongside your task-specific head
scalar_mix = ScalarMix(num_layers=13)ScalarMix parameters: Layer weights (before softmax): torch.Size([13]) Gamma scaling: 1.0000 Total trainable params: 14
ScalarMix adds only 14 parameters (13 mixing weights plus gamma), making it efficient to learn the optimal layer combination during fine-tuning.