Teacher-Student Training for Debiasing: General Permutation Debiasing for Large Language Models (2024)

Abstract: Large Language Models (LLMs) have demonstrated impressive zero-shot capabilities and versatility in NLP tasks, however they sometimes fail to maintain crucial invariances for specific tasks. One example is permutation sensitivity, where LLMs’ outputs may significantly vary depending on the order of the input options. While debiasing techniques can mitigate these issues, and yield better performance and reliability, they often come with a high computational cost at inference. This paper addresses this inefficiency at inference time. The aim is to distill the capabilities of a computationally intensive, debiased, teacher model into a more compact student model. We explore two variants of student models: one based on pure distillation, and the other on an error-correction approach for more complex tasks, where the student corrects a single biased decision from the teacher to achieve a debiased output. Our approach is general and can be applied to both black-box and white-box LLMs. Furthermore, we demonstrate that our compact, encoder-only student models can outperform their larger, biased teacher counterparts, achieving better results with significantly fewer parameters.

## introduction recent advancements in large language models (llms) have led to dramatic shifts within natural language processing (nlp). unlike prior "pretrain and fine-tune" @xcite approaches, instruction-tuned llms combined with effective good prompting techniques has enabled llms to excel at unseen tasks without task-specific training @xcite . this has led to the current capabilities of llms, where they demonstrate great versatility, while also being powerful and displaying state-ofthe-art performance on many standard nlp benchmark leaderboards @xcite . despite their impressive general abilities, llms suffer from particular limitations. they are prone to hallucinating information @xcite , can have large sensitivity to the form of prompts @xcite and also demonstrate systematic biases such as gender bias @xcite . furthermore, due to the general nature of their pre-training and instruction-tuning @xcite , for certain applications, they may be unaware of particular important task invariances. one such invariance that llms may fail to maintain is permutation-invariance. ongoing work has demonstrated that llms can be sensitive to the input order of options, which has been observed for both question answering @xcite and pairwise assessment @xcite @xcite . for these tasks, varying the ordering of the input options may lead to different decisions by the llm, which can impact downstream performance and reliability. although debiasing approaches can be applied to enforce invariances, such methods can be computationally expensive or inapplicable to black-box settings @xcite . to address these challenges, this work introduces a general framework that can be used to adapt both black-box and white-box systems to follow a particular invariance, while being inference efficient. for a given invariance and debiasing scheme, our framework trains a compact student to emulate the debiased teacher, which during inference can be efficiently deployed. we investigate two variants of students, a simple knowledge-distilled student, as well as an error-correction student that takes in a single biased teacher sample and corrects it to learn the debiased teacher decision, applicable for more complicated tasks. we demonstrate the effectiveness of our framework on permutation invariance, and illustrate that small 330m parameter student models can outperform their larger biased teacher counterparts, while also maintaining particular embedded invariances. the contributions of this work are: 1) we provide metrics for assessing the sensitivity of models to the input ordering of options. 2) we showcase that llms can demonstrate large permutation sensitivity and that biases seem correlated to task performance. 3) we study several different debiased approaches that yield significant performance gains. 4) experiments on race++ and summeval demonstrate that the teacher-student training for debiasing framework yields effective students that perform better than their biased teacher while being inference-efficient and not expensive to train. ## multiple choice prompting prompting has been shown effective in leveraging the diverse zero-shot abilities of instruction-tuned large language models (llms). for a particular task, inputs can be rephrased into natural language queries which condition the llm to generate useful responses @xcite . for example in multiple choice classification tasks, there may be input information x (e.g. context and question), and a set of k possible answers, @xmath0, ..., a k }. given an ordered realisation of the possible answers a σ (e.g. (a 2 , a 4 , a 1 , a 3 )), a prompt t(x, a σ ) can be designed to convert the input information into a textual representation. note that a different ordering of answers a σ j , would lead to a different textual prompt representation. the work of @xcite conducted a systematic study into multiple choice prompting (mcp) in which the position of each ordered answer in a σ is bound to a symbol or option label w k within the prompt. instead of tasking the llm with generating the full correct answer a * @xcite , it only needs to predict the label w * of the correct position (often a single token such as "a", "b", "c" or "d"). this converts the significant problems that arise when comparing the probabilities of variable-length answers @xmath1. the predictive probability of the k-th option under the particular permutation of answers then becomes: the above approach assumes full access to the output probabilities, which may not be available. for black-box llms that are served through apis @xcite , one may only have access to the autoregressively generated output text. in such settings, one can instead randomly sample from the underlying distribution to get an approximate system decision ã: for well-designed prompts, the majority of the probability mass should be associated with the option labels w 1:k . one can therefore directly sample from w(i) ∼ p llm (w k |t(x, a σ )) and reject samples that do not belong to the options labels. ## teacher-student training for debiasing to address the computational inefficiencies linked with permutation debiasing, this section proposes using teacher-student training to investigate if a smaller inference efficient proxy system p θ could emulate the characteristics of the debiased teacher distribution. instead of performing k! calls to obtain a permutation debiased prediction, a proxy student could potentially achieve it in a single call. our approach is general and is applicable to both white and black-box systems, without the need for labelled data. although we focus on correcting permutation sensitivity, the framework can be applied for any task invariance and debiasing strategy. ## conclusions this paper explores the sensitivity of llms to the order of input options for multiple-choice questionanswering and comparative assessment. we illustrate the effectiveness of various debiasing approaches for mitigating these biases and the associated performance improvement. while these debiasing methods often entail high computational costs, we show that teacher-student training can yield inference-efficient student models capable of emulating a debiased teacher distribution. our approach is practical in both white-box and black-box settings, requiring a manageable number of training data points and black-box samples.