Lightweight reranking for language model generations (2024)

Abstract: Large Language Models (LLMs) can exhibit considerable variation in the quality of their sampled outputs. Reranking and selecting the best generation from the sampled set is a popular way of obtaining strong gains in generation quality. In this paper, we present a novel approach for reranking LLM generations. Unlike other techniques that might involve additional inferences or training a specialized reranker, our approach relies on easy to compute pairwise statistics between the generations that have minimal compute overhead. We show that our approach can be formalized as an extension of self-consistency and analyze its performance in that framework, theoretically as well as via simulations. We show strong improvements for selecting the best k generations for code generation tasks as well as robust improvements for the best generation for the tasks of autoformalization, summarization, and translation. While our approach only assumes black-box access to LLMs, we show that additional access to token probabilities can improve performance even further.

## introduction the rapid advancement and remarkable achievements of generative large-scale pre-trained language models (llms) have brought about a revolutionary transformation in the field of natural language processing (nlp). these models have demonstrated significant enhancements in various nlp applications, such as machine translation, summarization, and code generation. individual generations sampled from the models often yield high-quality results. however the quality of generated outputs can exhibit considerable variability. multiple output samplings for the same input can produce certain generations which are of substantially higher quality than the quality of the average generation from the model. several approaches have been proposed to exploit this phenomenon. one strategy involves improving the underlying models themselves to make the quality of the average generation consistently better. this can be achieved by taking existing model generations, ranking them based on a human feedback, automated evaluation metrics like bleu score, or execution feedback in case of code. the ranked generations can then be finetuned on directly or can be used to train a reward model that can be used in an rl loop @xcite @xcite @xcite . another common approach is best-of-n sampling or reranking. in this approach, the underlying model is not touchedwe instead take multiple samples from the model and select the best one post-facto using a reranking method @xcite @xcite @xcite @xcite . while this approach can often given strong improvements, most extant reranking techniques involve computationally intensive or cumbersome methods to compute the ranking criterion. these include methods like training an auxiliary model as a reranker, evaluating the probability of the query given the generated answer (query likelihood) but at the price of doubling the inference cost, etc. in case of code generation models, another alternative is executing the generated code on unit tests. while such an approach has been applied in various models such as alpha-code @xcite which is targeted towards contest coding problems, it becomes much less feasible as you move past the contest coding setting due to the complexity of setting up the build environment for arbitrary code as well as sandboxing it appropriately. recently, a simple approach, called selfconsistency was proposed for selecting the best answer from multiple generations @xcite for tasks where the set of possible answers is smallfor example multiple choice questions or math word problems where there is a unique answer consisting of a single or a very limited number of tokens. in that paper, the authors sample multiple chainof-thought generations from the llm, extract the predicted answer at end each generation and select the answer with the most number of votes. the motivation behind this is the observation that you can take different reasoning paths to get to the same answer. thus the method aims to marginalize over multiple different reasoning paths and rank the answers based on their marginal probability rather than their probability conditioned on a single reasoning path. while they achieve substantial improvements over existing baselines, it is not immediately clear how to apply this to open-ended generation tasks like code generation, summarization, or translation -where there is often no chainof-thought or reasoning path to marginalize over, nor is there necessarily a unique correct answer. we start off with two key observations -(1) we can have semantically equivalent or near-equivalent generations that are nevertheless not exact matches. these are one subset of generations we can marginalize over (2) for open-ended tasks, a generation can encompass multiple elements. for summarization, there might be multiple relevant facts in the text that a good summary should mention. for code, there might be multiple branch conditions that need to be present to generate a correct implementation. our generation set could be structured such that while different generations include a different subset of elements (different facts in case of summarization or different branch conditions in case of code), we have only a single generation that contains all of the relevant elements. in this case, simply marginalizing over semantically equivalent generations would not be sufficient as there is no semantically equivalent generation for the optimal generation. we develop these two observations in the next section into a minimal overhead reranking method for such open-ended tasks which does not require access to token probabilities. the rest of the paper is organized as follows. in section 2 we present our motivation. in section 3 we present our method and the similarity function. in section 4, we present and discuss our experimental results. in section 5, we describe the related work and we finally conclude in section 6. ## motivation consider the following coding problem from the mbpp datasetdef remove_dirty_chars(string, second_string): a solution to the above problem would be semantically equivalent to "iterate over the string skipping characters in second_string and then convert the result back to a string and return". two parts of the semantic meaning of this solution could then be (1) the return type should be a string (2) when iterating through the string, any character in second string has to be skipped over. these observations can be converted into predicates for the generations. specifically, for this prompt, we can define the predicates (1) p 1 = is the return value of the generated program a string? (2) p 2 = in the generated program, are all characters in second_string skipped over in the return string? these predicates capture properties of the semantic meaning of the generations. we sample three generations for the above prompt resulting in the following generations: now if we were able to evaluate the above predicates at inference time on the generations, we would be able to detect that generation 3 is the only one that satisfies both and is thus an optimal generation. however generating the relevant predicates, and then generating code to evaluate arbitrary predicates on code that confirms to the given natural language specification with high precision is an unsolved problem. is there a way transform the problem into something more tractable? let us look at the votes each predicate gets from the different generations (i.e. on how many generations the predicate evaluates to true). p 1 gets 2/3 votes (from the 2nd and 3rd generation) and thus the majority vote is that it should be true. p 2 gets 2/3 votes (from the 1st and 3rd generation) and thus the majority vote again says it should be true. generation 3 is the only one that agrees with the majority vote for p 1 , p 2 and is thus the consensus choice. in fact, we do not even have to do the step of first counting votes for p 1 , p 2 to figure out what their majority vote value is! we can just compute how much a generation agrees with the other 2 generations on the evaluation for p 1 , p 2 . to elaborate, generation 1 agrees with generation 3 on p 2 but not p 1 . it does not agree with generation 2 on anything. thus it has a total agreement score of 1. similarly generation 2 also has an agreement score of 1. generation 3 however agrees with generation 1 on p 2 and with generation 2 on p 1 resulting in an agreement score of 2. thus generation 3 has the highest agreement with all other generations and is the consensus choice. this transformation is depicted in figure 1 . there are a couple of points to note regarding the above. (a) the answer we want for both predicates ends up being also what the majority vote predicts. this is the connection to self-consistency that we will formalize shortly. (b) while the above procedure handles observation (2) in the mentioned in the introduction, does it also handle observation (1)? note that if we had a 4th generation as mentioned before, at inference time however, we have access neither to such predicates nor to what their evaluation on the different generations might be. however, as the previous example showed, we do not actually even need an evaluation of the predicates on the generations. we only actually need an understanding of how much a pair of generations agree on relevant predicates. in the next section we will describe simple similarity functions to approximate this agreement without any knowledge of the possible predicates and later show their empirical effectiveness. for now, we attempt to formalize the above intuition, connect it with self-consistency, and give it theoretical grounding. let v be a vector of length k where each element represents a predicate and the value represents the property's value that you want the generation to have. for example, for the example above, v would be of length 2 with a value of 1 in both. in general, we do not restrict v to be just a binary vector as some predicates can have multiple possibilities as an answer. let u i be the vector for the ith generation representing the predicates' values for that generation. we now make the self-consistency assumption which is that for each individual predicate, the most frequent response is assumed to be correct. formally if v l can take on m l values 1, . . . , m l and without loss of generality, v @xmath0, then @xmath1. now as mentioned, at inference time, we may not have access to either the predicates or their values for different generations. thus we only assume access to the agreement between the different generations on the predicates' values (later on we will approximate this with similarity functions between generations). in particular we assume we know what the pairwise fractional agreement on the predicates is between generations denoted as a(u i , u j ) = 1 k k @xmath2, @xmath3, n] where i indexes the generations and t the predicates. we then try to identify a generation i such that the average pairwise fractional agreement for that generation with all other generations is maximizedi.e. a(u i , v) is maximized. ## method as previously mentioned, we may not have the capability to compute predicates at inference time, thereby rendering the computation of the exact fractional agreement with v i.e. a(u, v), unattainable. however as we found out in the last section, choosing the generation that has the maximum average fractional similarity with all other generations can be a good approximation. however as we may not have predicates at inference time, we cannot always compute that either. intuitively however, if two generations are more similar to each other -for an appropriate definition of similarity -then they will tend to agree more on any possible predicates. surprisingly, we find that a very simple similarity function, which we will define shortly, is sufficient for our purposes. once we have our similarity function, we can define a generalized self-consistency score gsc sim (i) for each generation i, given by 1 m -1 m @xmath4,j̸ =i sim(i, j). here, sim denotes the similarity function, and m represents the number of generations. for generations with unique answers, if we have: one straightforward way to encode a generation is by using a binary vector that denotes the presence or absence of an n-gram. surprisingly, we find this simple encoding to be sufficient for defining a robust similarity function. for open-ended generation, we define our similarity function as follows. for each generation we define a vector v of size |v | where v is set of all possible n-grams for @xmath5. for the experiments in this paper, we simply use @xmath6. we show in section g, increasing k can be helpful though only up to a point. each element i of v is simply whether token i is present in the generation or not. we then take the inner product between two such vectors as similarity. we call this the ngram consistency score (ncs) and refer to the @xmath7. figure 2 shows a visualization of v for an example sentence. formally when token probabilities are available, we can leverage them to improve our approach. intuitively, if a generation has a low token probability for the generated token, then finding a match for that that token should count for less. in accordance with this intuition, we introduce two further variants. first we modify the definition of v as follows the mean log probability of a sequence is an oft-used ranking method. we can combine it with wucs by further weighting each generation by the per token probability as follows -for a generation i, finally, to rank the generations, we employ arg max i gsc sim (i) where sim can take the form of ucs, wucs, or consensus-ucs. the diagram illustrating the progression from u cs to consensus -w u cs is show in figure 3 . ## results we conducted experiments utilizing the codex family of models, specifically codex-davinci-001, codex-davinci-002, and codex-cushman as well as llama family of models. in addition we also evaluated gpt-j for xsum, minif2f, and wmt14. unfortunately due to the unexpected shutdown of the openai api, we were unable to obtain results for codex-001 and codex-cushman on the xsum, minif2f, and wmt14 datasets. we evaluated these models on a range of datasets for code generation tasks -in particular on the humaneval @xcite , mbpp, mbppsanitized @xcite datasets for code generation. for the autoformalization of minif2f to isabelle, we used the dataset provided by @xcite . for text summarization, we utilized the xsum dataset @xcite . for machine translation, we used the wmt14 french to english and german to english datasets @xcite . our primary evaluation metric for code generation is ranked pass@1 where we rerank a sample set of generations and assess whether the top-ranked generation successfully passes all unit tests. we also evaluate with ranked pass@k for @xmath8. for the minif2f autoformalization task, we measure the quality using the bleu score, following @xcite . for xsum we use the rouge-2 and rouge-l scores for evaluation. for all code generation datasets, we sample 125 generations from the models which serves as our dataset for the different experiments for minif2f and xsum, we sample 50 generations from the model. unless otherwise specified, for all experiments, we use the codex-davinci-002 model. following @xcite , we perform bootstrap sampling 50 times with a sample size of 25 to generate the results. our baselines are random selection, ranking by mean log probability, ranking using medoid in our confidence weighted unigram space, and for code generation -ranking using the coder reviewer ranker method @xcite . a full description of the datasets, experiments, and the baselines is in the supplement. we also have several additional results in the supplement. ## related work in @xcite , they use a perceptron based reranker to rerank model generated translations. summareranker @xcite use mixture of experts training to train their reranker to optimize for multiple automated evaluation metrics (like rouge or bleu score) at once. pairreranker @xcite uses automated evaluation metrics to rank model generations and then select the top few best and worse and train a model to classify the better summary between a pair of summaries. all of the previous reranking methods however require training an auxiliary model. more detailed related work is described in the supplement. ## conclusion we analyze the self-consistency method for problems that have fixed answers and develop a framework to extend it to open-ended generations. we establish connections between our framework and other code generation reranking functions and prove that if the optimal generation is present in our generation set, we can always recover it as well as prove bounds on how close we can get to the optimal generation under certain settings. our simulated tests reveal our ability to consistently recover the best or close to best possible generation in the set. we introduce several lightweight similarity functions and show that they give strong and consistent improvements over state of the art baselines. notably, our unigram consistency score (ucs) function, the most minimal of our similarity functions, requires only access to raw generations to effectively rerank. we show that the ucs variants uniformly enhance the performance of code and text generation and are competitive with strong baselines like coder reviewer reranker despite them needing a lot more compute resources as well as time. for code generation, we also leverage the fact that our reranking metric is based on pairwise similarity to improve performance for pass@k for @xmath9. additionally, we conduct multiple variations on our primary experiments to ascertain the robustness and reliability of our performance. we also find that similar to self-consistency, increasing the number of generations improves the performance of the model.