Extracting Embeddings from All Encoder Layers of BERT Using Hugging Face Transformers

This guide details how to extract contextual embeddings from all encoder layers of a pre-trained BERT model using the Hugging Face transformers library. While using only the final hidden state is a common practice, exploring embeddings from all layers can offer richer contextual information and potentially improve performance on various downstream Natural Language Processing (NLP) tasks.

Why Consider All Hidden Layers?

BERT's architecture comprises multiple encoder layers (12 in BERT-base). Each layer learns and encodes different levels of linguistic and semantic information.

• Layer-specific information: Earlier layers might capture more syntactic or word-level features, while deeper layers tend to capture more semantic and contextual relationships.
• Research insights: Studies, such as those presented in the original BERT paper, have demonstrated that concatenating embeddings from multiple layers can lead to performance improvements. For instance, on the Named Entity Recognition (NER) task, concatenating the last four encoder layers achieved a higher F1 score (96.1%) compared to using only the final encoder layer.

BERT Embedding Flow

BERT processes input text through a series of layers:

Input Embedding → Encoder 1 → Encoder 2 → ... → Encoder 12 → Final Output

Each encoder layer generates a "hidden state," which is a vector representation for each token in the input sequence. By enabling output_hidden_states=True, we can access the hidden states from all these layers.

Step-by-Step Guide to Extracting All Hidden States

Follow these steps to extract the hidden states from all encoder layers of a BERT model.

Step 1: Import Required Modules

First, import the necessary classes from the transformers library and torch.

from transformers import BertModel, BertTokenizer
import torch

Step 2: Load the Pre-trained Model and Tokenizer

When loading the BERT model, it's crucial to set output_hidden_states=True. This parameter instructs the model to return the hidden states from all encoder layers in addition to the final hidden states.

# Load a pre-trained BERT model and tokenizer
# 'bert-base-uncased' is a common choice, but others can be used.
# output_hidden_states=True is the key parameter.
model = BertModel.from_pretrained('bert-base-uncased', output_hidden_states=True)
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

Step 3: Preprocess the Input Text

Prepare your input sentence by tokenizing it, adding special tokens ([CLS] for classification and [SEP] for separation), converting tokens to their corresponding IDs, padding the sequence, and creating an attention mask.

# Example sentence
sentence = "I love Paris"

# Tokenize the sentence
tokens = tokenizer.tokenize(sentence)

# Add special tokens: [CLS] at the beginning, [SEP] at the end
tokens = ['[CLS]'] + tokens + ['[SEP]']

# Convert tokens to their corresponding IDs
token_ids = tokenizer.convert_tokens_to_ids(tokens)

# Define a target sequence length and pad if necessary
# For demonstration, let's pad to a fixed length of 7
max_seq_length = 7
padding_length = max_seq_length - len(token_ids)

# Pad with [PAD] token (ID 0)
if padding_length > 0:
    token_ids += [tokenizer.pad_token_id] * padding_length
    tokens += [tokenizer.pad_token_id] * padding_length # Update tokens for mask creation

# Create an attention mask: 1 for real tokens, 0 for padding tokens
attention_mask = [1 if token != tokenizer.pad_token else 0 for token in tokens]

# Convert lists to PyTorch tensors
# unsqueeze(0) adds a batch dimension, making the shape [1, sequence_length]
token_ids = torch.tensor(token_ids).unsqueeze(0)
attention_mask = torch.tensor(attention_mask).unsqueeze(0)

Step 4: Pass Inputs to the Model and Extract Hidden States

Feed the preprocessed token_ids and attention_mask to the BERT model. The outputs object will contain various attributes, including hidden_states.

# Pass the inputs to the model
# The model returns an object that contains the outputs
outputs = model(token_ids, attention_mask=attention_mask)

# Extract all hidden states
# 'hidden_states' is a tuple containing the output from each layer,
# including the initial embeddings and the 12 encoder layers.
hidden_states = outputs.hidden_states

# The number of hidden states will be number_of_encoder_layers + 1 (for initial embeddings)
print(f"Number of hidden states: {len(hidden_states)}")
# Expected output for BERT-base: Number of hidden states: 13

# Each hidden state tensor has the shape: [batch_size, sequence_length, hidden_size]
# For BERT-base, hidden_size is 768.
# In our example: [1, 7, 768]
print(f"Shape of the first hidden state (initial embeddings): {hidden_states[0].shape}")
print(f"Shape of the last hidden state (final encoder output): {hidden_states[-1].shape}")

Step 5: Concatenate the Last Four Layers (Optional for Better Results)

As per research findings, combining the last few layers can be beneficial. Here's how to stack and reshape the last four encoder layers to create a concatenated embedding for each token.

# Select the last four hidden states (which correspond to the last four encoder layers)
# hidden_states[-1] is the last encoder output, hidden_states[-2] is the second last, etc.
last_four_layers = hidden_states[-4:] # This gets the last 4 elements from the tuple

# Stack these layers. Shape will be [4, batch_size, sequence_length, hidden_size]
stacked_hidden_states = torch.stack(last_four_layers) # Shape: [4, 1, 7, 768]

# Permute dimensions to group token information together.
# We want to move the layer dimension to be before the hidden_size dimension.
# Original shape: [num_layers, batch_size, seq_length, hidden_dim]
# Target shape for reshaping: [batch_size, seq_length, num_layers, hidden_dim]
permuted_hidden_states = stacked_hidden_states.permute(1, 2, 0, 3) # Shape: [1, 7, 4, 768]

# Reshape to concatenate the last dimension (hidden states of the last four layers)
# This creates a single, wider embedding vector for each token.
# The new hidden size will be 4 * 768 = 3072.
token_embeddings_concat = permuted_hidden_states.reshape(
    permuted_hidden_states.size(0), # batch_size
    permuted_hidden_states.size(1), # sequence_length
    -1 # Automatically calculates the concatenated dimension
) # Shape: [1, 7, 3072]

print(f"Shape of concatenated embeddings from last 4 layers: {token_embeddings_concat.shape}")

This process yields a 3072-dimensional embedding vector for each token, incorporating richer contextual information from the final four encoder layers.

Conclusion

Leveraging embeddings from all BERT encoder layers, or a selection of them, can provide a more nuanced representation of text compared to using only the final layer. The Hugging Face transformers library simplifies this extraction process. These richer embeddings can be instrumental in improving performance on tasks like Named Entity Recognition, sentiment analysis, and more.

In subsequent guides, we'll explore how to practically apply these multi-layer embeddings to various NLP tasks.

SEO Keywords

• BERT embeddings from all layers
• Hugging Face BERT hidden states
• Extract BERT layer-wise embeddings
• Use all BERT encoder outputs
• BERT last four layers embeddings
• Contextual embeddings from BERT
• Transformers output hidden states
• BERT NER performance improvement

Interview Questions

• What are hidden states in the BERT architecture, and how are they different from attention weights?
• Why might it be beneficial to use embeddings from multiple BERT layers instead of just the final layer?
• Explain the role of output_hidden_states=True in Hugging Face Transformers.
• How can you extract the last four encoder layers’ embeddings in BERT using PyTorch?
• What is the dimensionality of token embeddings when concatenating the last four BERT layers (for BERT-base)?
• How do different BERT layers capture different types of linguistic information?
• What is the impact of using multiple encoder layers on Named Entity Recognition (NER) performance, according to research?
• Can you compare the typical results of using only the final BERT layer vs. using a combination of layers for downstream tasks?
• What considerations should be made when using all encoder layers for real-time applications (e.g., computational cost)?
• How would you visualize or interpret embeddings extracted from multiple BERT layers?