How to resize RobertaLMHead with pretrained weights?

[yeounyi]

Hi, I'm trying to train my model with a new token 'name', but it keeps throwing size mismatch error.
I don't know how to resize RobertaLMHead while loading pretrained weights from 'roberta-base'

Setting Tokenizer

tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
tokenizer.add_tokens('<name>')

Setting Model

model = MaskedLM.from_pretrained('roberta-base')

Model Structure

class MaskedLM(RobertaPreTrainedModel):
    def __init__(self, config):
        super().__init__(config=config)
        self.roberta = RobertaModel(config)
        self.lm_head = RobertaLMHead(config)
        self.refinement_num = 5
        self.mask_id = 50264
        self.init_weights()

    def forward(...):
         self.roberta.resize_token_embeddings(50266)
         ## HOW TO RESIZE LM HEAD?! ## 
         # self.lm_head.resize_token_embeddings(50266)
         outputs = self.roberta(input_ids, attention_mask)
         prediction_scores = self.lm_head(outputs[0])
         ...

I tried _get_resized_lm_head from here
But it doesn't work as RobertaLMHead has no weight attribute.

    def _get_resized_lm_head(
        self, old_lm_head: torch.nn.Linear, new_num_tokens: Optional[int] = None, transposed: Optional[bool] = False
    ) -> torch.nn.Linear:
        """
        Build a resized Linear Module from a provided old Linear Module. Increasing the size will add newly initialized
        vectors at the end. Reducing the size will remove vectors from the end

        Args:
            old_lm_head (:obj:`torch.nn.Linear`):
                Old lm head liner layer to be resized.
            new_num_tokens (:obj:`int`, `optional`):
                New number of tokens in the linear matrix.

                Increasing the size will add newly initialized vectors at the end. Reducing the size will remove
                vectors from the end. If not provided or :obj:`None`, just returns a pointer to the input tokens
                :obj:`torch.nn.Linear`` module of the model without doing anything.
            transposed (:obj:`bool`, `optional`, defaults to :obj:`False`):
                Whether ``old_lm_head`` is transposed or not. If True ``old_lm_head.size()`` is ``lm_head_dim,
                vocab_size`` else ``vocab_size, lm_head_dim``.

        Return:
            :obj:`torch.nn.Linear`: Pointer to the resized Linear Module or the old Linear Module if
            :obj:`new_num_tokens` is :obj:`None`
        """
        if new_num_tokens is None:
            return old_lm_head

        old_num_tokens, old_lm_head_dim = (
            old_lm_head.weight.size() if not transposed else old_lm_head.weight.t().size()
        )

        if old_num_tokens == new_num_tokens:
            return old_lm_head

        if not isinstance(old_lm_head, nn.Linear):
            raise TypeError(
                f"Old language model head is of type {type(old_lm_head)}, which is not an instance of {nn.Linear}."
                f"You should either use a different resize function or make sure that `old_embeddings` are an instance of {nn.Linear}."
            )

        # Build new lm head
        new_lm_head_shape = (old_lm_head_dim, new_num_tokens) if not transposed else (new_num_tokens, old_lm_head_dim)
        has_new_lm_head_bias = old_lm_head.bias is not None
        new_lm_head = nn.Linear(*new_lm_head_shape, bias=has_new_lm_head_bias).to(self.device)

        # initialize new lm head (in particular added tokens)
        self._init_weights(new_lm_head)

        num_tokens_to_copy = min(old_num_tokens, new_num_tokens)

        # Copy old lm head weights to new lm head
        if not transposed:
            new_lm_head.weight.data[:num_tokens_to_copy, :] = old_lm_head.weight.data[:num_tokens_to_copy, :]
        else:
            new_lm_head.weight.data[:, :num_tokens_to_copy] = old_lm_head.weight.data[:, :num_tokens_to_copy]

        # Copy bias weights to new lm head
        if has_new_lm_head_bias:
            new_lm_head.bias.data[:num_tokens_to_copy] = old_lm_head.bias.data[:num_tokens_to_copy]

        return new_lm_head

[LysandreJik]

You should do model.resize_token_embeddings(50266).

Here is the documentation of that method.

[yeounyi]

When I do model.resize_token_embeddings(50266), embedding size changes from 50265 to 50266.

MaskedLM(
  (roberta): RobertaModel(
    (embeddings): RobertaEmbeddings(
      (word_embeddings): Embedding(50266, 768)
      (position_embeddings): Embedding(514, 768, padding_idx=1)
      (token_type_embeddings): Embedding(1, 768)
      (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
      (dropout): Dropout(p=0.1, inplace=False)

But out_features size in lm_head remains the same (50265) and throws error.
prediction_scores size is not [batch size, sequence length, 50266], it remains still [batch size, sequence length, 50265]

  (lm_head): RobertaLMHead(
    (dense): Linear(in_features=768, out_features=768, bias=True)
    (layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
    (decoder): Linear(in_features=768, out_features=50265, bias=True)
  )

[LysandreJik]

Are you sure? I just ran the following on master:

>>> from transformers import RobertaForCausalLM

>>> model = RobertaForCausalLM.from_pretrained("roberta-base")
>>> model.lm_head.decoder
Linear(in_features=768, out_features=50265, bias=True)

>>> model.resize_token_embeddings(50266)
Embedding(50266, 768)

>>> model.lm_head.decoder
Linear(in_features=768, out_features=50266, bias=True)

Please observe how the decoder is resized.

[patil-suraj]

@yeounyi
In your example lm_head is not resized because there are no get_output_embeddings and set_output_embeddings methods in your MaskedLM class. The resize_token_embeddings method needs these methods to get the lm_head.

You should add those methods and then call resize_token_embeddings on the instance MaskedLM class. See the implementation of RobertaForMaskedLM

transformers/src/transformers/models/roberta/modeling_roberta.py

Lines 984 to 1006 in d55e10b

	class RobertaForMaskedLM(RobertaPreTrainedModel):
	_keys_to_ignore_on_load_missing = [r"position_ids", r"lm_head.decoder.bias"]
	_keys_to_ignore_on_load_unexpected = [r"pooler"]
	def __init__(self, config):
	super().__init__(config)
	if config.is_decoder:
	logger.warning(
	"If you want to use `RobertaForMaskedLM` make sure `config.is_decoder=False` for "
	"bi-directional self-attention."
	)
	self.roberta = RobertaModel(config, add_pooling_layer=False)
	self.lm_head = RobertaLMHead(config)
	self.init_weights()
	def get_output_embeddings(self):
	return self.lm_head.decoder
	def set_output_embeddings(self, new_embeddings):
	self.lm_head.decoder = new_embeddings

[LysandreJik]

Ah, I indeed missed that this was a custom MaskedLM implementation, my bad.

[yeounyi]

Thanks all! After adding get_output_embeddings and set_output_embeddings methods, it works perfectly

[bestpredicts]

M