Fine-tuning BERT with multiple classification heads
I need to train a model that has the same backbone such as BERT as a feature extractor and use multiple classification heads. This scenario is similar to multi-task learning but all the tasks are classification tasks so I will need multiple classification heads. does anyone have any similar notebook code that I can start with?
2022-10-24T18:35:27Z
HI SaraAmd, I’m looking at doing something very similar.  I currently have ~15 classification models that all use the same language model.  I have a feeling that I will have to write my own forward() function in the end.  I’m curious if you made any progress, and if so, I was wondering if you could provide me some wisdom on the subject.  I’m wondering if I need to re-train all of these models at the same time, or if I can extract the classification heads from my existing models and pack them together into a new model with a custom forward() function.  If you found any existing literature that could put me on the right path, that’d be awesome too
2022-10-24T19:01:26Z
Hi, Unfortunately, I haven’t made any progress and I haven’t found any literature. But I also assume that the forward function needs to be implemented as well as the loss function for each classification head (all will have cross entropy but how I am not sure). We can brainstorm together to move forward. My goal is to have one BERT model as the feature extractor and then add n-number of classification heads on top of it to train the model is this your goal as well? I feel like you want to train different models separately.
2022-10-24T19:49:34Z
Yes that is exactly my goal.  Currently I am just barely able to fit most of my models onto my GPU at inference time, but the goal will be to have a smaller footprint on my GPU so that I can process data in batches instead.  I don’t mind writing the forward function, I will likely begin tinkering with it next week.  I just need to figure out how to extract the classification heads and pickle them so that I can import them into one large model.As I mentioned, I have all my models trained separately alreadyBut if someone has pre-written code to do multi-head training all at once, I don’t mind re-training if it means re-using code.  I’m happy either way.  I’ll share code with you as soon as I start putting pen to paper.
2024-01-10T16:15:27Z
Did you folks managed to find a solution to this in the end? I’ve tried overwriting several classes to handle multihead classification for coarse and fine grained labels.
2024-01-16T17:34:13Z
No, I unfortunately did not get around to tackling this.  I ended up brute forcing it with a bigger GPU with more memory.  I still want to solve this in the long term though, but I think the only answer is to write a custom forward() function that creates multiple parallel output linear layers that each generate one set of prediction/classifications.
2024-01-17T09:01:24Z
I think you are right. In fact, I am able to get the model training by overwriting the following:compute_lossforwardthe model itself, to add the layersdatacollatorget_train_dataloader/ eval_dataloderI get it to train but I cannot get it to compute any metrics, for some reason the HF Trainer is skipping over compute_metrics at eval time.
2024-01-17T18:17:57Z
Have you considered Google? It’s a pretty good search engineMedium – 28 Sep 21Multi-Task Learning with Transformers: Transformers with Multiple Prediction...Since the inception of transformer architecture and the birth of hugging face transformers package of Natural Language Processing as a…Reading time: 3 min read
2024-01-19T12:17:32Z
I do appreciate the article, it may be useful for the other posters but it doesn’t give me any clues on getting the compute metrics function to work during the evaluation step. So far I have narrowed it down to the all_labels variable being set to None prior to this function being called in the eval_loop.On the whole, as a senior community member, I think its worth reflecting on how a response like this creates a culture in which newer community members lose the confidence to even ask questions.Thank you for sharing the article.
2024-01-19T12:40:14Z
Hi,Sorry for the earlier reply! One should be able to ask any question.One can add acompute_metricsfunction to the blog post above. Thecompute_metricsfunction takes in a named tuple, which one can split into labels and logits, e.g.:def compute_metrics(eval_preds):
    metric = evaluate.load("glue", "mrpc")
    logits, labels = eval_preds
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)Hence, in case of a multi-task model as shown in the blog above, the logits depend on the task head being used. To circumvent this, one could add an attribute to the model’s init called “task_name”, which you can change depending on which task you want to evaluate for. Then you can calltrainer.evaluate()depending on which task is set.
2024-01-19T13:22:13Z
Thanks again, this is helpful for me. Am I right in thinking that the task_name variable effectively chooses the right head task dependently? If so, what if a single task requires two heads? Such as in my own case I need to train the model to give a coarse and a fine-grained output.For some context, here is my loss function:def compute_loss(self, model, inputs, return_outputs=False):

    labels = inputs.pop('labels')

    labels_coarse = labels[0][0]
    labels_fine = labels[0][1]

    outputs = model(**inputs)
    logits_coarse, logits_fine = outputs

    loss_fct = torch.nn.CrossEntropyLoss()
    loss_coarse = loss_fct(logits_coarse.view(-1, num_labels_coarse), labels_coarse.view(-1))
    loss_fine = loss_fct(logits_fine.view(-1, num_labels_fine), labels_fine.view(-1))
    
    loss = (loss_coarse + loss_fine) / 2

    return (loss, outputs) if return_outputs else loss