Add a classification head to a fine-tuned language model

Beginners
1

We I have fine-tuned a GPT-2 model with a language model head on medical triage text, and would like to use this model as a classifier. However, as far as I can tell, the Automodel Huggingface library allows me to have either a LM or a classifier etc. head, but I don’t see a way to add a classifier on top of a fine-tuned LM.

Am I mistaken in my understanding of the possibilities? We want to preserve the fine-tuning we have done for the LM task, so don’t want to swap this head for a classifier and lose all of those weights. I suppose we could unlock the top layers of the GPT2 model and try and push the training back into the model, but I suspect that would soon become untenable for our somewhat limited resources. I see this question has been asked here, and my fellow worker has posted a follow-up question in that thread, but we didn’t see any answer we could relate to, so I am asking again in this new topic.

2

You need to use the AutoModelForSequenceClassification class to add a classification head on top of your pretrained model.

3

Thanks for your reply @sgugger. I have a fine-tuned pretrained model, that is - a AutoModelForCausalLM.from_pretrained("gpt2-large"), which I have fine-tuned on my specific language (which is triage text and quite unlike standard English).

I thought that if I called an AutoModelForSequenceClassification("gpt2-large") then it would be just like reverting to the standard pretrained model, without any of the fine-tuning I have just done, since my fine-tuning resides in the head of the AutoModelForCausalLM.

My question was if it was possible to sit a classifier head on top of the fine tuned language generator model, so as to inherit the weights that I have obtained through the fine-tuning. If I load the weights of the AutoModelForCausalLM into the AutoModelForSequenceClassification then the fine-tuned weights are discarded, but would it be beneficial to add a linear layer to allow those weights to be preserved for use by a manually added classifier head?

4

Couple of thoughts, though I am relatively new to LLMs.

  1. Typically, if you are trying to have the model learn a new language or writing style (medical data), you may want to fine tune the tokenizer, not the model itself. Fine tuning the model should be done for the specific task - classification, summarization etc - not general language understanding.

  2. You could add whatever layers you like to the model.

Full example on my GitHub:

classifier_layer_names += ['extra_class' + str(layer_ind)]
extra_classifiers.add_module(name=classifier_layer_names[-1],
module=nn.Linear(input_channels,
                                       self.extra_class_layers[layer_ind]))
classifier_dropout_layer_names += ['extra_class_dropout' + str(layer_ind)]
extra_classifiers.add_module(name=classifier_dropout_layer_names[-1],
                                             module=nn.Dropout(p=self.fine_tune_dropout_rate))
# Replacing, rather than adding a layer makes it so that forward() does not need to be extended
self.model.classifier = extra_classifiers

https://github.com/bdzyubak/utils/blob/f0367f3c76cf38c790af43c25abc059d00557ba9/LLM_pytorch_lighting_wrapper.py->replace_layers_with_fc()