Regression Aware Inference with LLM s (2024)

Abstract: Large language models (LLMs) have shown strong results on a range of applications, including regression and scoring tasks.Typically, one obtains outputs from an LLM via autoregressive sampling from the model’s output distribution. We show that this inference strategy can be sub-optimal for common regression and scoring evaluation metrics. As a remedy, we build on prior work on Minimum Bayes Risk decoding,and propose alternate inference strategies that estimate the Bayes-optimal solution for regression and scoring metrics in closed-form from sampled responses.We show that our proposal significantly improves over baselines across datasets and models.

## introduction large language models (llms) are currently the most capable models across many nlp tasks @xcite @xcite . owing to their remarkable few-and zero-shot abilities @xcite , pre-trained llms are often applied without any additional training on domain-specific datasets: instead, one may query the llm with a suitably crafted input prompt. more recently, llms have been successfuly applied to regression and scoring tasks. for example, @xcite explored zero-shot learning for time series prediction; @xcite showed how llms are remarkably strong at incontext learning for regression tasks; @xcite ; @xcite considered autoregressive fine-tuning over numerical targets applied to arithmetic tasks; and @xcite applied llms for listwise ranking. the quality of an llm is often assessed using an application-specific evaluation metric. one popular metric is the exact match (em), which penalises any response not exactly equal to the one in the dataset annotation. this is an analogue of the conventional classification accuracy. while em is an intuitive metric, there are many applications where it is not suitable. this includes tasks such relevance scoring @xcite and sentiment analysis @xcite , where the outputs are numerical or ordinal categories. in these cases, one instead prefers metrics such as the squared error, absolute error or ranking scores that take the outputs' ordinal nature into account. despite the wide variety of evaluation metrics, llm inference is typically performed in the same manner for every task: namely, one performs autoregressive sampling from the llm's underlying distribution (see §2). while intuitive, such inference does not explicitly consider the downstream evaluation metric of interest. this raises a natural question: is there value in adapting the inference procedure to the evaluation metric at hand for regression and scoring tasks? a prominent line of work takes a decisiontheoretic approach to the above problem. dubbed as minimum bayes risk (mbr) decoding, this approach seeks to optimize at inference time the metric of choice under the model's distribution @xcite @xcite . much of the work on mbr is focused on evaluation metrics for machine translation and text generation tasks, such as the bleu score @xcite . of particular interest in this literature are self-consistency based decoding strategies that take a (weighted) majority vote of sampled responses @xcite , which have shown to provide quality gains in arithmetic and reasoning problems. in this paper, we build on the existing literature on mbr to design metric-aware inference strategies for general regression and scoring tasks. we first observe that choosing the most likely target for an input corresponds to inherently optimizing for the em metric, and is consequently not opti-figure 1 : illustration of metric-aware llm inference for regression and scoring tasks. an input x is passed to the llm, and samples are drawn from the distribution over targets y conditioned on x. these are then used to find the target optimizing a metric m through a closedform decision rule φ (e.g., mean or median); ## when (naïve) llm inference fails on regression tasks we begin with the problem setting. for a finite vocabulary v of tokens (e.g., words in english), let d denote a distribution over inputs @xmath0, and targets @xmath1. let p(y | x) denote the conditional distribution over targets given an input. we consider a special case of this setting where y ⊂ r corresponds to numeric targets. here, we assume that each @xmath2; for example, the integer 1 has the string encoding "1". a language model (lm) takes a string x as in-put and predicts an output ŷ ∈ v * . typically, the lm first produces a distribution p(• | x) over targets. in a slight abuse of notation, we use p(y | x) a prediction from an lm is typically derived via a suitable inference (or decoding) procedure. perhaps the most common inference strategy is to choose the mode of p(• | x): in practice, one may approximate the mode via greedy decoding or beam search, or sampling multiple candidates and picking the among them the one with the highest likelihood score @xcite . in principle, the extracted target may not be numerical. in such cases, a possible strategy is to resort to predicting a default numerical value such as 0.0. in practice, we find the targets from high-quality llms tend to be numerical even under zero-shot settings, and so converting most likely targets from v * to y is usually possible. the quality of an lm's prediction is measured by some evaluation metric m(y, ŷ), where we assume that higher values are better. while the exact match (em), given by m(y, ŷ) = 1(@xmath3, is a commonly used evaluation metric, there are a range of other metrics popularly used to evaluate lms. these include the (negative) squared error m(y, ŷ) = -(yŷ) 2 or absolute error m(y, ŷ) = -|y -ŷ| for regression tasks. a natural goal is to then choose the inference strategy ŷ(x) to maximize the metric m of interest, i.e., to maximize the expected utility: for many choices of metric m(y, ŷ(x)), picking the mode of the predicted distribution (1) can be sub-optimal for (2). as an example, consider predicting the star rating (on the scale 1-5) associated with a review text. suppose m(y, ŷ) is the negative absolute error between the true and predicted ratings. given the review text "this keybord is suitable for fast typers", suppose the lm responses and the associated probabilities are {"1": 0.3, "2": 0.0, "3": 0.3, "4": 0.0, "5": 0.4}. the mode of the predicted probabilities is "5". in contrast, the maximizer of (2) is the median "3". algorithm 1 rail: regression-aware inference with llms 1: input: model m , #samples k, sampling temperature t , effective temperature t , decision rule φ 2: for @xmath4: ## experiments and discussion we experimentally evaluate our proposed approach on nlp tasks with different evaluation metrics. datasets. we use two datasets: (i) semantic samples with an effective temperature of @xmath5. in appendix d, we report results for variants of mali with no post-hoc scaling ( methods. we evaluate the following methods: (i) greedy decoding, (ii) a baseline inspired from the self-consistency decoding of sampling k candidates and picking the one with the maximum likelihood (argmax) @xcite , (iii) the pro-posed rail approach on the same k samples, and (iv) rail with temperature scaling ( §3.3). for (iv), we choose α yielding effective temperature 1 4 . metric-aware inference helps. in ## conclusions we have shown how regression and scoring-aware inference strategies can yield notable benefits for small and medium-sized llms. in the future, we wish to extend our approach to other less-explored evaluation metrics in the mbr literature.