I'm trying to write a program that using Roberta to calculate word embeddings:

from transformers import RobertaModel, RobertaTokenizer
import torch

model = RobertaModel.from_pretrained('roberta-base')
tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
caption = "this bird is yellow has red wings"

encoded_caption = tokenizer(caption, return_tensors='pt')
input_ids = encoded_caption['input_ids']

outputs = model(input_ids)
word_embeddings = outputs.last_hidden_state

I extract the last hidden state after forwarding the input_ids to the RobertaModel class to calculate word embeddings, I don't know if this is the correct way to do this, can anyone help me confirm this ? Thanks


1 Answer 1


This was studied in the original BERT article, which concluded that the best approach was to concatenate the states of the last 4 layers:

enter image description here

Although BERT preceeded RoBERTa, we may understand this observation to be somewhat applicable to RoBERTa, which is very similar. You may, nonetheless, experiment with the precise number of layer states to concatenate to see what value gives the best results.

  • $\begingroup$ What does this table mean ? I don't understand $\endgroup$
    – user158782
    Jan 12 at 20:02
  • $\begingroup$ Ok, I have read the paper and understand, another question is how can I concatenate last 4 hidden layers in the code ? $\endgroup$
    – user158782
    Jan 12 at 20:06
  • $\begingroup$ First, you should invoke your model with model(input_ids, output_hidden_states=True) to get the hidden states. Then, you concatenate them with torch.cat, like torch.cat([outputs['hidden_states'][-i] for i in range(1,5)],dim=-1). $\endgroup$
    – noe
    Jan 12 at 20:14
  • $\begingroup$ So I did this and word embedding has shape of torch.Size([1, 9, 3072]), is this normal ? I thought the hidden size should be the same 768 why increase to 3072 ? $\endgroup$
    – user158782
    Jan 12 at 20:22
  • $\begingroup$ The hidden state size is 768, but you have concatenated the last 4 hidden states together in a single vector, so the resulting size is 4x. $\endgroup$
    – noe
    Jan 12 at 20:23

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.