Finetuning llama for classification

I have some domain specific datasets, so I want to fine-tune llama or other LLMs.
Other than putting classification head on the top, I want to train the model in the generative way, letting models output text that can be thought of labels.

But the issues is that, how can I restrict model generated text to the labels that I want? I’m worried if the model would generate texts that are invalid (text that are not label)

It would be a great help!

1 Like

You could look at the probabilities on the lm_head logit associated with the first token in each of your answers to turn your lm head into a classifier. It works well if you can formulate your problem such that the answers have unique token ids try yes no answers or multiple choice a), b), c)

1 Like

Hi, could you please elaborate more. I am fine-tuning LLAMA on a multiple-choice question-answering (MCQA) dataset. During the training phase, would it be a good approach to trim the model’s output head to just four tokens corresponding to the answer options, so that during the generation phase, the model is constrained to generate only the labels? Are there any alternative strategies I could consider for achieving this?

1 Like