RoBERTa › Extracting Representations
Language AI HandbookAll chapters RoBERTa › Extracting RepresentationsOne section · Sections in this chapter11 / 17 ↓
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
For downstream tasks, you often want the hidden state representations rather than MLM predictions. The representations are what make RoBERTa useful: they encode rich contextual information about each token and about the sequence as a whole. Understanding how to extract and interpret them is the foundation of applying RoBERTa to real tasks.
RoBERTa produces two kinds of representations that are commonly used. The first is the token-level hidden states from the final transformer layer, which give you a contextual embedding for every token in the sequence. These are useful when your task requires token-level predictions, such as named entity recognition or question answering span extraction. The second is the <s> (CLS) token's hidden state, which is a summary representation of the entire sequence. This is the standard choice for classification tasks like sentiment analysis or document categorization.
Extract them as follows:
from transformers import RobertaModel
# Load the base model (without MLM head) for embeddings
encoder = RobertaModel.from_pretrained("roberta-base")
encoder.eval()
sentences = [
"The cat sat on the mat.",
"The dog lay on the rug.",
"Machine learning is fascinating.",
]
embeddings = []
for sent in sentences:
inputs = tokenizer(sent, return_tensors="pt", padding=True, truncation=True)
with torch.no_grad():
outputs = encoder(**inputs)
# Use [CLS] token representation as sentence embedding
cls_embedding = outputs.last_hidden_state[:, 0, :]
embeddings.append(cls_embedding)
embeddings = torch.cat(embeddings, dim=0)Embedding shape: torch.Size([3, 768]) Each sentence represented as a 768-dimensional vector Cosine similarities: 'The cat sat on the mat....' <-> 'The dog lay on the rug....': 1.000 'The cat sat on the mat....' <-> 'Machine learning is fascinatin...': 0.998 'The dog lay on the rug....' <-> 'Machine learning is fascinatin...': 0.998
The similarities make sense: the two sentences about animals on surfaces are more similar to each other than to the machine learning sentence. RoBERTa's representations capture semantic relationships effectively.
A few caveats apply here. The CLS token representation from a pretrained (but not fine-tuned) RoBERTa model is not always optimal for measuring semantic similarity. RoBERTa is trained with MLM, which does not explicitly optimize for similarity comparisons. The model learns to predict masked tokens, not to place semantically similar sentences near each other in embedding space. As a result, raw CLS embeddings from RoBERTa can sometimes perform surprisingly poorly on semantic similarity benchmarks compared to models like sentence-transformers, which are fine-tuned specifically to produce similarity-preserving embeddings.
For production sentence similarity tasks, you are better served by fine-tuning RoBERTa on a similarity dataset using a contrastive or triplet loss, or using a model like roberta-base fine-tuned as part of the sentence-transformers library. The embeddings produced by that pipeline will be far more reliable for nearest-neighbor search or semantic similarity scoring than raw pretrained RoBERTa representations.
