Adaptation with Self-Evaluation to Improve Selective Prediction in LLM s (2023)

Abstract: Large language models (LLMs) have recently shown great advances in a variety of tasks, including natural language understanding and generation. However, their use in high-stakes decision-making scenarios is still limited due to the potential for errors. *Selective prediction* is a technique that can be used to improve the reliability of the LLMs by allowing them to abstain from making predictions when they are unsure of the answer. In this work, we propose a novel framework for adaptation with self-evaluation to improve the selective prediction performance of LLMs. Our framework is based on the idea of using parameter-efficient tuning to adapt the LLM to the specific task at hand while improving its ability to perform self-evaluation. We evaluate our method on a variety of question-answering (QA) datasets and show that it outperforms state-of-the-art selective prediction methods. For example, on the CoQA benchmark, our method improves the AUACC from 91.23% to 92.63% and improves the AUROC from 74.61% to 80.25%.

## introduction large language models (llms) have recently demonstrated impressive capabilities in many natural language understanding, reasoning and generation tasks, such as question answering @xcite , summarization @xcite , semantic classification, and code generation @xcite . as llms improve their remarkable performance, they are being increasingly considered to replace humans to perform high-stakes tasks. for example, llms can be used for medical qa to assist patients @xcite . however, llms are not guaranteed to be accurate for all queries, so it is important to understand which queries they are reliable for. this information can be used to direct human oversight to the queries with the lowest selection score. selective prediction @xcite , broadly refers to the deployment scenario for ai models where humans are involved to maintain overall accuracy by reviewing ai-generated, lowconfidence outputs. in this scenario, both human and ai performance are considered together to minimize human involvement cost. llms should be used in the real-world with enhanced selective prediction performance. they should be able to assess the accuracy of their predictions and refrain from making wrong predictions. if an llm detects that an answer might be wrong for a question, it should be able to generate an answer with the sentiment of "i don't know!" (as shown in fig. 1 ) or defer the answer to a human for manual inspection. this will help to ensure that llms are used in a reliably, especially for high-stakes applications. selective prediction for llms is challenging because llms are just trained to predict the next to-ken given a context but are not guaranteed to always predict the correct next token. also, since llms generate an output sequence in an auto-regressive way, they don't directly produce a confidence score for the output sequence. thus, obtaining selection scores from llms for their output sequences is not straightforward. although there is some research on selective prediction for llms, these studies have their own shortcomings. @xcite to use heuristic prompts (e.g., adding prompts like "is the proposed answer true or false?") to trigger self-evaluation of llms. however, those prompts may only work for the llm used in @xcite and may not generalize to other types of llms (e.g., opt and gpt2 models evaluated in our work). some approaches proposed using semantic entropy @xcite or selfconsistency @xcite as a measure of uncertainty for selection score. however, they usually require generating multiple output sequences to obtain the uncertainty measure for an input sequence, which introduces high computational cost and latency at test time. fine-tuning llms on training data from the target question answering task using the standard llm training loss can improve selective prediction performance. this is because fine-tuning can improve the accuracy of the predictions and maximize the likelihood of the ground-truth answer for a given question. however, maximizing the likelihood of the ground-truth answer is not the same as minimizing the likelihood of the wrong answers, since llms generate output sequences in an auto-regressive way. even after fine-tuning, some wrong answers may still have high likelihood and be generated by the llm at test time. therefore, distinguishing correct and incorrect answers based on likelihood scores alone is a challenging task. to address these challenges of self-evaluation and uncertainty estimation, we propose a novel framework -adaptation with self-evaluation to improve selective prediction in llms (aspire). unlike previous methods that rely on hand-crafted heuristics or multiple output sequences, our framework learns to self-evaluate from target-task data. we do this by training llms on a subset of the training data from the question-answering tasks. this allows the llms to learn to distinguish between correct and incorrect answers on their own. we then define a selection score that combines the likelihood of the generated answer with the learned self-eval score (see eq. ( 11 )) to make selective predictions. this makes our method much less computationally expensive than solutions that require generating multiple output sequences to obtain the uncertainty measure. thus, the proposed method is useful for practical applications where high selective prediction performance and low inference costs are desired, after deploying the llm. in such applications, practitioners prefer collecting some training data to fine-tune smaller llms to achieve high selective prediction performance rather than directly deploying very large pre-trained llms with limited selective prediction performance for specific tasks. we conduct extensive experiments to evaluate our proposed framework, aspire. we show that aspire achieves the state-of-the-art selective prediction performance on three question answering datasets: coqa, triviaqa and squad, using opt and gpt-2 models. we also provide empirical analysis to delve deeper into our proposed technique. ## related work selective prediction for llms. recently, llms (e.g., @xcite and palm (chowdhery et al., 2022)) have achieved great success in solving various kinds of natural language generation (nlg) tasks. however, llms are still not very reliable and may generate wrong outputs when solving nlg tasks. due to this, selective prediction (or sometimes called selective generation @xcite ) is critical for safely deploying llms in the real-world. different from selective prediction for classification tasks (e.g., natural language inference (nli) tasks) @xcite , selective prediction for llms in solving nlg tasks is fundamentally different since the prediction is done auto-regressively over many steps and the possible answer set has an infinite size. recently, several work propose some uncertainty measures for llms, which can be used for selective prediction @xcite @xcite @xcite . some recent work studies selective prediction for solving question answering tasks where questions are ambiguous @xcite . @xcite propose a selective prediction method that at inference time leverages an auxiliary model which is trained to distinguish the correct predictions of the qa model from the incorrect ones. different from previous work, our work proposes to improve selective prediction performance of llms in solving question answering tasks by learning self-evaluation during fine-tuning. parameter efficient fine-tuning. fine-tuning pretrained llms on downstream datasets can bring huge performance gains when compared to using the pretrained llms out-of-the-box (e.g., kshot inference). however, as llms get larger and larger, full fine-tuning becomes very expensive in terms of computational cost and memory requirements. in addition, massive models might not be data efficient and overfitting issues might be observed, yielding suboptimal generalization. to address these issues, parameter-efficient fine-tuning (peft) approaches have been proposed. peft approaches only fine-tune a small number of (extra) model parameters while freezing most parameters of the pretrained llms, thereby greatly decreasing the computational and storage costs. it has also been shown that peft approaches are better than fine-tuning in the low-data regimes and generalize better to out-of-domain scenarios. existing peft approaches include lora @xcite , prefix tuning @xcite , soft prompt tuning @xcite and p-tuning @xcite . in this work, we use soft prompt tuning to learn self-evaluation to improve selective prediction performance of llms. ## problem setup suppose we have a pre-trained llm f for an arbitrary generative modeling task such as question answering. the output can be represented as a sequence of tokens from the vocabulary v. let v * be the space of sequences of tokens. suppose the logits of f on @xmath0. the likelihood of the next token following x being v is defined as: we use f to generate the output sequence for the given input x by solving the following objective: it is impossible to solve this objective exactly since the output sequences can be arbitrarily long. however, we can employ some decoding strategy like greedy decoding or beam search to solve it. to evaluate if the generated output ŷ is correct or not, we need a set of reference outputs s and an evaluation metric m : v * × v * → [0, 1] that can evaluate the similarity of the generated output ŷ compared to the reference output y @xmath1. with a threshold γ, we can determine the correctness of the generated output -if max @xmath2, y r ) > γ, then the generated output is correct; otherwise, the generated output is wrong. we discuss the specific choices of m and γ in section 6. in selective prediction, we need a rejection option, which is denoted by ⊥. given a training dataset d @xmath3, y i )} ntr @xmath4, we aim to build a selective predictor f s : v * → v * ∪ {⊥} that can achieve strong selective prediction performance on the test dataset d @xmath5, s i )} nte @xmath6, where s i is the set of reference outputs for the input x i . the selective predictor f s is composed of a predictor f : v * → v * and a selection scoring function g : v * → r. with f and g, the selective predictor f s is proposed as: ## aspire framework we propose that llms should have the selfevaluation ability such that they should be able to distinguish whether their proposed answers for a given question are correct or not. although some previous work @xcite show that llms have good self-evaluation ability with specially designed prompts, those prompts may not transfer to different kinds of llms (as shown by our experiments and in @xcite al. (2023) ) and hand-crafting prompts for different kinds of llms can be expensive. a more effective approach is to collect some training data to employ selfevaluation. towards this end, we propose a novel framework -adaptation with self-evaluation to improve selective prediction in llms (aspire). fig. 2 illustrates the proposed framework and the details are explained next. given a training dataset for a generative task, we can fine-tune the pre-trained llm on the training data to improve its prediction performance. towards this end, parameter efficient tuning techniques (e.g., soft prompt tuning @xcite and @xcite ) might be employed to adapt the pre-trained llm on the task, given their effectiveness in obtaining strong generalization with small amount of target task data. specifically, the model parameters θ of the llm are frozen and adaptable parameters θ p are added for fine-tuning. only θ p are updated to solve the following training objective: to further improve selective prediction performance, we propose to fine-tune the llm to learn self-evaluation. we first use the llm with the learned θ p to generate different answers for each example (x, y) ∈ d tr . suppose the decoding algorithm used to generate output sequences for each input x is a. a would produce a list of generated output sequences: after sampling high-likelihood outputs for each query, we add adaptable parameters θ s and only tune θ s for learning self-evaluation. since the output sequence generation only depends on θ and θ p , freezing θ and the learned θ p can avoid changing the prediction behaviors of the llm when learning self-evaluation. let z c and z w be a pair of tokens that represent the words "correct" and "wrong" respectively. we can then optimize θ s using the following training objective: we use the beam search decoding method towards this. we define the likelihood of the output ŷ * being correct for the query x as: this score p (z c | x, ŷ * ) is referred as the learned self-eval score. overall, the selection scoring function is proposed as: ## implementation via soft prompt tuning in the proposed framework, θ p and θ s can be trained using parameter efficient tuning approaches. in our work, we focus on soft prompt tuning, as illustrated in fig. 3 . the driving force behind this approach lies in the recognition that if we can develop prompts that effectively stimulate self-evaluation, it should be possible to discover these prompts through soft prompt tuning in conjunction with targeted training objectives. we first briefly introduce the soft prompt tuning method proposed by @xcite . we consider llms based on the transformer architecture @xcite . given a query @xmath7, . . . , x mq ), transformers first embed the tokens, forming a matrix @xmath8, where d e is the dimension of the embedding space. the softprompts are represented as parameters θ ∈ r l×de , where l is the length of the prompt. the prompt is then concatenated to the embedded input forming a single matrix [ θ; x] ∈ r (mq+l)×de , which then flows through the transformer as normal. in the proposed framework, we need to train two portions of the prompts θ @xmath9. utilizing soft prompt tuning, the training objective ( 6 ) is proposed as: the learned self-eval score p (z c | x, ŷ * ) becomes: vitamin k likelihood: 0.7 q embed and a embed are input embeddings for the question and answer respectively. we first generate the answer and the likelihood of the answer, and then compute the learned self-eval score. we can cache the states when generating the answer and reuse those states when computing the learned self-eval score to save computational costs. to generate the output sequence and obtain the selection score for a given input sequence, we employ two stages: first, we obtain the generated output and the likelihood for the generated output and then, we obtain the learned self-eval score. since the query of the second stage is constructed by appending some additional tokens to the query of the first stage, the second stage can reuse the states in the first stage instead of recomputing them to save some computational cost (see fig. 3 ). lastly, we note that the computational complexity of the proposed method at test time is o(l max ) with l max being the maximum length of the generated output sequence. in appendix f, we provide a more detailed analysis of the computational complexity of different methods. the predictive entropy and semantic entropy methods have a complexity of o(m • l max ) where m is the number of output sequences sampled for uncertainty estimation, which is much larger than that of our method. ## experiments our experimental evaluation is focused on the following questions: (q1) could a learning-based system using selfevaluation improve selective prediction in llms compared to other post-hoc selective prediction alternatives? (a1) by learning self-evaluation, we can significantly improve selective prediction performance across different datasets and llms (see ## conclusion in this paper, we proposed a novel framework for adaptation with self-evaluation to improve selective prediction in llms. we implemented the framework via soft prompt tuning and demonstrated its superior performance over existing methods through extensive experiments. in future work, one could explore implementing our framework via other parameter efficient tuning approaches and applying our method to larger llms.