Efficiently Deploying LLMs with Controlled Risk

## Abstract

Abstract Deploying large language models in production requires simultaneous attention to efficiency and risk control. Prior work has shown the possibility to cut costs while maintaining similar accuracy, but has neglected to focus on risk control. By contrast, here we present hierarchical chains with multi-level abstention (HCMA), which use model-intrinsic uncertainty to delegate queries along the LLM intelligence hierarchy, enabling training-free model switching based solely on black-box API calls. Our framework presents novel trade-offs between efficiency and risk. For example, deploying HCMA on MMLU cuts the error rate of Llama3 405B by 30% when the model is allowed to abstain on 20% of the queries. To calibrate HCMA for optimal performance, our approach uses data-efficient logistic regressions (based on a simple nonlinear feature transformation), which require only 50 or 100 labeled examples to achieve excellent calibration error (ECE), cutting ECE by 50% compared to naive Platt scaling. On free-form generation tasks, we find that chain-of-thought is ineffectual for selective prediction, whereas zero-shot prompting drives error to 0% on TruthfulQA at high abstention rates. As LLMs are increasingly deployed across computing environments with different capabilities (such as mobile, laptop, and cloud), our framework paves the way towards maintaining deployment efficiency while putting in place sharp risk controls.

## 1 Introduction

Given the recent excitement around large language models, researchers have been intensely focused on evaluating potentially transformative applications across domains. However, as the technology matures and its scale increases, efficiency starts to matter. In addition, real-world consequences begin to emerge, calling for greater risk control.

The NLP community has addressed the need for greater efficiency by optimizing the internal mechanics of the transformer-based neural networks (Dao et al., ()), developing new sampling schemes (Leviathan et al., ()), performing model compression and distillation (Hinton et al., (), Xu et al., ()), caching model responses (Bang, ()), and routing queries between different models (, Jiang et al., (), Hu et al., (), Aggarwal et al., (), Ding et al., ()). On risk control, researchers have considered probabilistic calibration and various approaches to selective prediction, in which a model is allowed to abstain from queries based on a confidence signal such as single-token or sequence-level probabilities (Xin et al., ()), entropy-based scores and consistency measures (Manakul et al., (), Kuhn et al., (), Chen et al., ()), lightweight probes on hidden layer embeddings (Azaria and Mitchell, (), Kossen et al., ()), or outputs of separate neural networks trained to predict correctness (Xin et al., (), Yoshikawa and Okazaki, (), Varshney and Baral, (), Gupta et al., (), Cobbe et al., (), Kadavath et al., ()).

However, previous work has not simultaneously addressed efficiency and risk control. Following previous routing approaches in which small models pass queries to larger models based on a confidence threshold (Kag et al., (), , Aggarwal et al., (), Wang et al., (), Ding et al., (), Sakota et al., ()), we incorporate selective prediction by adding “abstention thresholds” that let each model reject queries on behalf of the whole LLM system. In this paper, we explore the use of LLM token probabilities to instantiate this system, which we call “hierarchical chains with multi-level abstention.” Although the most performant LLM confidence signals – when considered in isolation – are based on hidden layer embeddings, repeated sampling, and neural-network correctness predictors, these approaches are not suitable for meeting our twin goals of 1) improving efficiency while simultaneously imposing risk control, and 2) creating a framework based on black-box API calls that is easily adaptable to new tasks without requiring much data. By contrast, model-intrinsic probabilities have the advantage that they do not require white-box access to model internals, are widely available as part of API calls to third-party LLM inference providers, do not require large amounts of training data, and avoid compute overhead from repeated sampling.

## 2 Related Work

- •
Uncertainty Quantification: our framework of hierarchical LLM chains relies on uncertainty quantification to distinguish easy queries from difficult ones. Previous work has shown that on text classification tasks (such as MMLU), the maximum softmax probability effectively predicts mistakes, even if it is not calibrated (Hendrycks and Gimpel, (), Plaut et al., ()). On natural language generation, the community has made progress with several different approaches including prompting an LLM to verify whether a response to a query is correct (Lin et al., (), Kadavath et al., (), Xiong et al., ()), assessing the consistency between multiple sampled LLM outputs (Lin et al., (), Manakul et al., (), Kuhn et al., (), Aichberger et al., (), Nikitin et al., ()), and training lightweight probes on models’ internal states (Ren et al., (), Azaria and Mitchell, (), Chen et al., (), Kossen et al., ()).
- •
Probability Calibration: several methods have been proposed to make probabilistic predictions more calibrated (such that a forecast with probability 0.8 comes true about 80% of the time), including Platt scaling (Platt, ()), temperature scaling (Guo et al., ()), isotonic regression (Zadrozny and Elkan, ()), Bayesian binning (Naeini et al., ()), and more. Recent work has shown that state-of-the-art large language models tend to have poor calibration, especially after reinforcement learning from human feedback (RLHF) (Ouyang et al., (), OpenAI et al., (), Plaut et al., ()). Due to its simplicity, temperature scaling has emerged as the most popular post-processing method for calibration in deep learning, though for the purposes of rigorous risk control it suffers from a lack of statistical grounding.
- •
Selective Prediction: this field studies the performance of machine learning models when they are allowed to abstain on difficult queries. With the machine learning community, foundational work of El-Yaniv and Wiener, (), Geifman and El-Yaniv, (), and others was centered on image classification. In NLP, selective prediction has originally attracted limited interest (Xin et al., (), Varshney et al., (), Varshney and Baral, (), Ren et al., (), Yoshikawa and Okazaki, (), ) but today has become largely subsumed in the emerging field of uncertainty quantification for LLMs, which uses similar performance metrics. We wish to highlight the SGR algorithm (Geifman and El-Yaniv, ()) as a compelling technique for endowing selective prediction methods (including ours) with provable risk guarantees.
- •
Routing between Language Models: recent literature has reported cost savings (and, sometimes, performance gains) from routing queries between different language models. In this area, many approaches focus on horizontal routing between a collection of small language models in the hope of leveraging complementary strengths (Lu et al., (), Hari and Thomson, (), Lee et al., (), Sakota et al., (), Wang et al., ()). Other works consider vertical routing in which queries are explicitly first sent to a smaller model, which forwards difficult queries to a larger model (Wang et al., (), Kag et al., (), , Yue et al., (), Aggarwal et al., (), Ding et al., ()). However, these approaches have focused on lowering the cost of LLM inference without simultaneously tackling selective prediction.

## 3 Our Contributions

First, we provide a novel formula explaining why uncertainty-based delegation between language models works. Second, we introduce a nonlinear feature transformation that makes Platt scaling a highly effective calibration technique for LLM token probabilities, providing an alternative to temperature scaling grounded in a rigorous statistical model (logistic regression). Third, we present hierarchical chains with multi-level abstention (HCMA), the first (to our knowledge) routing LLM model that incorporates selective prediction alongside cost efficiency. By computing the full two-dimensional Pareto frontier of efficient HCMA configurations on MMLU, using formulas we derive, we show that an LLM routing approach based on model-intrinsic probabilities improves over the selective prediction performance of single LLMs, in terms of strengthening risk control without paying more, or achieving similar risk control at the same (or lower) price. Finally, we note that chain-of-thought prompting yields poor selective prediction performance on TruthfulQA compared to zero-shot prompting, calling for caution when applying established prompting techniques in LLM uncertainty quantification.

## 4 Mathematical Framework

We present an LLM system called “hierarchical chains with multi-level abstention,” which is based on delegating queries from small to large models, but allowing each model to abstain on behalf of the entire chain. First, we examine the act of delegation and explain why it works.

### 4.1 Why Does Delegation Work?

By delegation, we mean passing queries from a small model to a larger model based on estimated difficulty. Prior work has shown that this strategy outperforms random query assignment, but has not provided an explanation (Xin et al., ()). Although the result is intuitively plausible, the situation is nuanced because differently sized models seem to share a common sense of difficulty. Figure 1 shows this trend by exhibiting logistic regressions that predict correctness of Llama3 8B, 70B, and 405B on MMLU solely based on the transformed probability of the 8B model.

Figure: Figure 1: Model-intrinsic uncertainty aligns across Llama3 8B, 70B, and 405B models on MMLU, seemingly reflecting a shared sense of difficulty, but with the larger models more resistant to it. The figure shows logistic regressions fit to the transformed probabilities of the 8B model.
Refer to caption: extracted/5897593/images/difficulty_scale_new.png

Given that differently sized models share a common notion of difficulty, the effectiveness of delegation depends on the net result of two opposing forces: the small model makes fewer mistakes because it only answers the easy queries, but the large model makes more mistakes since it answers only difficult queries. Thus, the effectiveness of delegation rests on the large model being less sensitive to incremental difficulty than the small model.

The proposition below frames our reasoning in mathematical terms:

###### Proposition 1 .

Compared to randomly assigning queries to a small language model $\mathcal{M}_{\text{sm}}$ and $\mathcal{M}_{\text{lg}}$, the reduction in error from forwarding queries based on a delegation decision $D$ is

$$ $\Delta E=\text{Cov}(\mathbf{1}_{D},\mathbf{1}_{\mathcal{M}_{\text{bg}}\text{~{ }makes an error}})-\text{Cov}(\mathbf{1}_{D},\mathbf{1}_{\mathcal{M}_{\text{sm }}\text{~{}makes an error}}).$ (1) $$

In this formula, the covariances $\text{Cov}(\mathbf{1}_{D},\mathbf{1}_{\mathcal{M}_{\text{sm}}\text{~{}makes an
error}})$ and $\text{Cov}(\mathbf{1}_{D},\mathbf{1}_{\mathcal{M}_{\text{bg}}\text{~{}makes an
error}})$ measure the impact of the delegation decision on each model’s propensity to make a mistake. Empirically, both covariances are positive if $D$ is based on query difficulty, since both models’ performance degrades with incremental difficulty. However, as long as $\text{Cov}(\mathbf{1}_{D},\mathbf{1}_{\mathcal{M}_{\text{sm}}\text{~{}makes an
error}})>\text{Cov}(\mathbf{1}_{D},\mathbf{1}_{\mathcal{M}_{\text{bg}}\text{~
{}makes an error}})$, meaning that the smaller model is more sensitive to difficulty, the change in error in (1) is negative, and delegation outperforms random query assignment.

###### Proposition 1 .

### 4.2 Hierarchical Chains with Multi-Level Abstention

A HCMA consists of a chain $\mathcal{M}_{1}\rightarrow...\rightarrow\mathcal{M}_{k}$ of LLMs. A query is first sent to $\mathcal{M}_{1}$, which makes a decision whether to provide an answer to the query (ACCEPT), delegate it to a larger model (DELEGATE), or reject the query outright (REJECT). This decision is based on an estimated probability $\hat{p}_{\mathcal{M}_{1}}(x)$ of correctness on query $x$. Models $\mathcal{M}_{2},...,\mathcal{M}_{k}$ behave analogously. Importantly, when a model rejects a query, the rejection applies to the whole HCMA.

Figure: Figure 2: Illustration of a hierarchical chain with multi-level abstention (left), which consists of three models M1-M3, and the action space of an individual model in the chain (right), as a function of the estimated correctness probability.
Refer to caption: extracted/5897593/images/hcma_diagram.png

Each model $\mathcal{M}_{j}$ ($j<k$) makes its decision based on the estimated correctness probability and two tunable thresholds $r_{j}$ and $a_{j}$. The model rejects the query if estimated correctness probability is below $r_{j}$, delegates the query if it is between $r_{j}$ and $a_{j}$, and accepts the query if the estimated correctness probability exceeds $a_{j}$. This strategy results in the model-specific policies

$$ $\pi_{\mathcal{M}_{j}}(\hat{p}_{\mathcal{M}_{j}})=\begin{cases}\texttt{REJECT}& \text{if }\hat{p}_{\mathcal{M}_{j}}<r_{j},\\ \texttt{DELEGATE}&\text{if }\hat{p}_{\mathcal{M}_{j}}\geq r_{j}\text{~{}~{}but ~{}~{}}\hat{p}_{\mathcal{M}_{j}}<a_{j},\\ \texttt{ACCEPT}&\text{if }\hat{p}_{j}\geq a_{j},\end{cases}$ (2) $$

for models $j<k$. The last model $\mathcal{M}_{k}$ cannot delegate to a larger model, so it only rejects or accepts queries based on a single threshold $r_{k}$. All in all, an HCMA with $k$ models has $2k-1$ tunable configuration parameters. Figure 2 illustrates the structure of an HCMA and the action space of an individual model.

### 4.3 Analyzing System Performance: Efficiency and Risk

To measure performance of an HCMA in terms of efficiency and risk control, we propose the following metrics: error rate, cost, and abstention rate (the fraction of queries on which the chain refuses to give an answer, in order to control risk). We measure cost in terms of either money ($) or latency (ms). Since a HCMA delegates queries sequentially along the chain, costs add up as a query propagates more deeply into the chain. For example, suppose the HCMA $\mathcal{M}_{1}\rightarrow\mathcal{M}_{2}\rightarrow\mathcal{M}_{3}$ has model-specific costs $c_{1}$, $c_{2}$, and $c_{3}$. Then, whenever the query ends at model $i$, the effective cost is $C_{i}=\sum_{j=0}^{i}c_{j}$.

The following proposition provides formulas for the HCMA performance metrics, based on the graph structure in Figure 2.

###### Proposition 2 .

Consider an HCMA $\mathcal{M}_{1}\rightarrow\mathcal{M}_{2}\rightarrow...\rightarrow\mathcal{M}_
{k}$ with parameters and $\{r_{j}\}_{j=1}^{k}$, $\{a_{j}\}_{j=1}^{k-1}$ and model-specific costs $c_{1},...,c_{k}$. Based on the true joint distribution $p(x,y)$ of queries $x$ and ground truth responses $y$, the population performance metrics are

$$ $\displaystyle\mathbb{P}(\text{Error})$ $\displaystyle=\sum_{j=1}^{k}\mathbb{P}_{(x,y)\sim p(x,y)}(\pi_{1}=\texttt{ DELEGATE},~{}...,~{}\pi_{j-1}=\texttt{DELEGATE},~{}\pi_{j}=\texttt{ACCEPT},~{} Y_{j}\neq y),$ (3) $\displaystyle\mathbb{P}(\text{Abstain})$ $\displaystyle=\sum_{j=1}^{k}\mathbb{P}_{(x,y)\sim p(x,y)}(\pi_{1}=\texttt{ DELEGATE},~{}...,~{}\pi_{j-1}=\texttt{DELEGATE},~{}\pi_{j}=\texttt{REJECT}),$ (4) $\displaystyle\mathbb{E}[\text{Cost}]$ $\displaystyle=\sum_{j=1}^{k}\mathbb{P}_{(x,y)\sim p(x,y)}(\pi_{1}=\texttt{ DELEGATE},~{}...,~{}\pi_{j-1}=\texttt{DELEGATE},~{}\pi_{j}\neq\texttt{DELEGATE })~{}C_{j},$ (5) $$

where $Y_{j}$ is the response of the $j$-th model and $y$ the assumed ground truth on query $x$. The constants $C_{j}=\sum_{\xi=1}^{j}c_{\xi}$ are the effective costs when model $\mathcal{M}_{j}$ resolves a query.

To evaluate HCMA performance in practice, we estimate the chain’s performance metrics using Monte Carlo approximations as

$$ $\displaystyle\widehat{\mathbb{P}(\text{Error})}$ $\displaystyle=\sum_{j=1}^{k}\sum_{i=1}^{n}\frac{1}{n}~{}\mathbb{I}[\hat{p}_{1} (x^{(i)})\in[r_{1},a_{1}),...,\hat{p}_{j-1}(x^{(i)})\in[r_{j-1},a_{j-1}),\hat{ p}_{j}(x^{(i)})\geq a_{j}]~{}(1-\hat{p}_{j}(x^{(i)})),$ (6) $\displaystyle\widehat{\mathbb{P}(\text{Abstain})}$ $\displaystyle=\sum_{j=1}^{k}\sum_{i=1}^{n}\frac{1}{n}~{}\mathbb{I}[\hat{p}_{1} (x^{(i)})\in[r_{1},a_{1}),...,\hat{p}_{j-1}(x^{(i)})\in[r_{j-1},a_{j-1}),\hat{ p}_{j}(x^{(i)})<r_{j}],$ (7) $\displaystyle\widehat{\mathbb{E}[\text{Cost}]}$ $\displaystyle=\sum_{j=1}^{k}\bigg{(}\sum_{i=1}^{n}\frac{1}{n}~{}\mathbb{I}[ \hat{p}_{1}(x^{(i)})\in[r_{1},a_{1}),...,\hat{p}_{j-1}(x^{(i)})\in[r_{j-1},a_{ j-1}),\hat{p}_{j}(x^{(i)})\notin[r_{j},a_{j})]\bigg{)}~{}C_{j},$ (8) $$

where the $\mathbb{I}[\cdot]$ are indicator variables and the $\hat{p}_{j}$ are the fitted correctness predictors (based on a small training set, independent from the data for which the performance estimates are computed). Note that in the formulas above, we write $a_{k}=r_{k}$ interchangeably for the rejection threshold of the last model, to maintain consistency and avoid a special case.

Importantly, the joint probabilities in (4)-(6) and (7)-(9) do not approximately factorize into independent terms, since the estimated correctness probabilities across models are highly correlated (see Figure 1).

###### Proposition 2 .

### 4.4 Making Model-Intrinsic Uncertainty Work For HCMAs

The policy (2) of each model $\mathcal{M}_{j}$ in a HCMA $\mathcal{M}_{1}\rightarrow...\rightarrow\mathcal{M}_{k}$ requires an estimated probability of correctness $\hat{p}_{\mathcal{M}_{j}}(\cdot)$ for $j=1,2,...,k$. As we rely on each model’s intrinsic token probabilities $p_{\text{raw}}$ for uncertainty quantification, obtaining $\hat{p}$ boils down to calibrating $p_{\text{raw}}$.

Different calibration techniques have been employed for post-processing the conditional probabilities of large language models. Originally proposed for computer vision networks, temperature scaling (Guo et al., ()) has emerged as a favored method because of its simplicity and high performance (Jiang et al., ()). However, Platt scaling (Platt, ()) is preferable for rigorous risk control, since it is based on a statistical model – logistic regression – for which statisticians hafve developed an established catalog of confidence intervals and model diagnostics. Unfortunately, standard Platt scaling does not work well for LLMs because its conditional probabilities tend to form a tight cluster near 1.0, reflecting overconfidence (OpenAI et al., (), Plaut et al., ()).

However, it is possible to make Platt scaling work very well by spreading out the clusters of overconfident probabilities. We propose simple nonlinear transformations that accomplish this by introducing asymptotes near $p_{\text{raw}}=0$ and $p_{\text{raw}}=1$. Specifically, on multiple-choice QA, we transform the maximum softmax probability by

$$ $p_{\text{tr}}(p_{\text{raw}})=\log(\frac{1}{1-p_{\text{raw}}}),$ (9) $$

to obtain transformed probabilities $p_{\text{tr}}$.

On open-ended QA, we reduce the correctness prediction problem to binary classification by adopting the P(True) strategy proposed by Kadavath et al., (). In this approach, after calling a model $\mathcal{M}$ to answer the query, we call $\mathcal{M}$ again with a verification prompt to output “Y” or “N” based on whether its originally generated answer is correct. Since this is a binary classification with the constraint $p_{\text{raw}}(``Y")=1-p_{\text{raw}}(``N")$, we consider a different transformation that spreads overconfident probabilities for “N” and “Y” into the negative and positive reals. In terms of $p=p_{\text{raw}}(``Y")$, we apply the transformation

$$ $p_{\text{tr}}(p)=\begin{cases}\log(\frac{1}{1-p})&\text{if }p\geq 0.5\\ \log(2)-\log(\frac{1}{p})&\text{if }p<0.5\end{cases},$ (10) $$

which is a symmetric function around $p=0.5$.

## 5 Results

Our exploration of hierarchical chains with multi-level abstention (HCMA), based on model-intrinsic uncertainty, yields several findings. First, our nonlinear transformations make Platt scaling much more effective in calibrating LLM output probabilities, yielding a statistically grounded way of performing LLM calibration. Second, mapping out the Pareto frontier of efficient HCMA configurations shows novel risk and efficiency trade-offs that outperform single-model strategies for selective prediction with model-intrinsic uncertainty. Finally, we apply our calibrated correctness prediction methodology to the open-ended QA benchmark TruthfulQA. We discover that chain-of-thought prompting reduces the utility of the abstention signal, highlighting the need for caution in applying established prompting techniques in uncertainty quantification.

### 5.1 Transforming Raw Probabilities Significantly Enhances Platt Scaling

Table 1 shows that our modified Platt scaling approach for risk control based on logistic regression strongly outperforms naive Platt scaling. The data shows the results for $n=50$ training examples on MMLU, showing that our approach is highly data-efficient. Each number is the result of randomly sampling 50 training examples from the MMLU validation set and evaluating on the remaining 1480, repeated 100 times (except in the case of Llama3 405B, where we repeated 500 times).

**Table 1: Evaluating logistic regression for correctness prediction using raw vs transformed token probabilities, on MMLU. Here, precision is the probability that the model actually gives the correct answer given that we predict so. In the context of our hierarchical chains with multi-level abstention, high precision is critical for risk control.**
| Model | Precision $\uparrow$ | F1 Score $\uparrow$ |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Raw | Transformed | %Change | Raw | Transformed | %Change |  |
| Llama3 1B | $0.671$ | 0.798 | $19.00$% | $0.651$ | 0.701 | $7.62$% |
| Llama3 3B | $0.669$ | 0.792 | $18.28$% | $0.653$ | 0.702 | $7.44$% |
| Llama3 8B | $0.657$ | 0.796 | $21.22$% | $0.654$ | 0.702 | $7.38$% |
| Llama3 70B | $0.663$ | 0.793 | $19.61$% | $0.656$ | 0.704 | $7.29$% |
| Llama3 405B | $0.665$ | 0.798 | $19.93$% | $0.651$ | 0.697 | $7.04$% |
| Model | Accuracy $\uparrow$ | Expected Calibration Error $\downarrow$ |  |  |  |  |
| Raw | Transformed | %Change | Raw | Transformed | %Change |  |
| Llama3 1B | $0.656$ | 0.705 | $7.53$% | $0.111$ | 0.066 | $-40.23$% |
| Llama3 3B | $0.651$ | 0.700 | $7.53$% | $0.117$ | 0.072 | $-38.58$% |
| Llama3 8B | $0.648$ | 0.699 | $7.92$% | $0.143$ | 0.064 | $-55.07$% |
| Llama3 70B | $0.645$ | 0.704 | $9.12$% | $0.113$ | 0.056 | $-50.78$% |
| Llama3 405B | $0.653$ | 0.702 | $7.55$% | $0.064$ | 0.053 | $-17.56$% |

### 5.2 HCMA Provides Novel Risk & Efficiency Trade-Offs

We study the risk and efficiency trade-offs attainable on MMLU using a HCMA consisting of the Llama3 8B, 70B, and 405B models. To do this, we first tabulate achievable performance metrics – error, cost, and abstention rate – for each 5-dimensional HCMA configuration (involving acceptance thresholds $a_{\text{8B}}$, $a_{\text{70B}}$ and rejection thresholds $r_{\text{8B}}$, $r_{\text{70B}}$, and $r_{\text{405B}}$) by performing a grid search along the quantiles of the estimated correctness probabilities with $2.5\%$ resolution, resulting in $>$50 million distinct configurations. For the simulation, we use cost values of 0.3, 0.8, and 5.0 dollars per million tokens, which are close to the current pricing of LLM inference providers (Artificial Analysis, ()). From the achievable performance profiles, we compute the efficient Pareto frontier using the Skyline operator (Börzsönyi et al., ()).

Figure 3 maps out the achievable error rates vs their cost, and colored by the abstention rate. Globally, we observe that selective prediction based on model-intrinsic uncertainty effectively lowers the achievable error rates. At an inference cost of around $1 per million tokens (the cost of Llama3 70B in the model), the graph shows a kink in the slope of the relationship between error and cost, suggesting that increasing usage of Llama3 405B yields diminishing returns in terms of error vs cost.

To evaluate whether the HCMA is useful, it is important to compare performance against baselines. As we are not aware of other approaches incorporating selective prediction into cost-efficient routing, we evaluate against the single-model strategies of running selective prediction with the off-the-shelf Llama3 models (8B, 70B, or 405B) using the same method for achieving calibrated correctness prediction. Figure 4 compares the selective prediction performance achieved by the Pareto frontier of the HCMA (dashed lines corresponding to various cost buckets) with those attained by the single-model strategies (solid blue lines).

The plot shows that the HCMA provides two benefits over the single-model strategies. First, the HCMA makes available new error-abstention curves whose price lies between the error-abstention curves attainable with the single-model selective prediction strategies. This means that practitioners have the option of paying “a little bit more” in exchange for better selective prediction performance, without having to jump to the price point of the next biggest model. Second, the HCMA’s error-abstention curves in some cases dominate the closest single-model strategy. For example, the error-abstention curve for the $\$0.6-\$0.9$ cost bucket strictly outperforms the Llama3 70B at a similar price. Most strikingly, the HCMA replicates the performance of Llama3 405B, costing $\$5.0$ per million tokens, at an average cost of below $\$3.0$ per million tokens. Alternatively, spending the same amount using the HCMA ($\$4.7-\$5.0$ cost bucket) considerably lowers error for any abstention rate.

Figure: Figure 3: Achievable performance profiles of a HCMA consisting of the Llama3 8B, 70B, and 405B models on MMLU. The graph shows the Pareto frontier of the most efficient HCMA configurations.
Refer to caption: extracted/5897593/images/pareto_on_mmlu.png

Figure: Figure 4: Cost-based view of the Pareto frontier of achievable HCMA performance profiles on MMLU. Each dashed line is the average performance of efficient HCMA configurations in a cost bucket, showing that HCMA outperforms the selective prediction performance of single LLMs (solid blue lines).
Refer to caption: extracted/5897593/images/hcma_cost_view.png

### 5.3 Early Abstention Makes Lowest Risk Cheaper

To ablate our finding that the HCMA outperforms the selective prediction performances of the off-the-shelf Llama3 models, we consider whether multi-level abstention yields a benefit. Specifically, we compare an HCMA to a constrained version of an HCMA in which only the last model in the chain may abstain, disallowing early abstentions from smaller models.

In the case of a two-model HCMA in which Llama3 8B delegates to 70B (“8B$\rightarrow$70B”), we report two findings. First, allowing early abstention can yield a cost advantage in both dollar cost and latency, yielding respective improvements of 7% and 6% on MMLU (where we measured latency by querying the Fireworks AI API from our location in California). Second, when imposing a constraint that the average cost must be less than or equal to some value, allowing early abstention strictly dominates the error-abstention curve without multi-level abstention at higher abstention rates between 20% and 50%.

### 5.4 Chain-of-Thought Prompting May Impede Selective Prediction

On TruthfulQA, we observe that token probabilities derived from chain-of-thought based verification perform poorly with our approach, since the transformed probabilities form tight clusters around 0 and 1. By contrast, using a zero-shot prompt results in a smooth unimodal distribution, which is well-suited as a confidence signal. See Figure 5.

To ablate this finding, we also tried using a few-shot (but not chain-of-thought) based prompt for verification. The results are intermediate: the distribution of transformed probabilities is unimodal (just as in the zero-shot case) but not quite as symmetric, and the precision/recall curve for predicting incorrect answers is worse. Notably, these results hold despite the fact that chain-of-thought gives the highest overall accuracy for correctness prediction. On TruthfulQA, we observed accuracies for correctness prediction of 0.79 for few-shot with CoT, 0.74 for few-shot without CoT, and 0.73 for zero-shot.

These results urge caution in applying established prompting techniques to the task of verifying LLM correctness.

Figure: (a)
Refer to caption: extracted/5897593/images/truthfulqa_cot_hist.png

## 6 Conclusion

Motivated by the observation that differently sized language models directionally agree on the difficulty of queries, we have presented hierarchical chains with multi-level abstention (HCMA), a routed LLM system that incorporates selective prediction alongside cost efficiency. Given the increased usage of LLMs via black-box API calls, and the efficiency-driven need to avoid repeated LLM sampling, we base our routing and abstention decisions on LLM token probabilities. With our improved Platt scaling approach to calibrating these probabilities, our framework is easily adaptable to new tasks without requiring much labeled data. We show that HCMAs efficiently outperform the selective prediction performance of single LLMs, and that multi-level abstention (as opposed to only abstaining at the end of the chain) is a beneficial part of our system’s architecture. Finally, we find that chain-of-thought prompting hinders correctness prediction on TruthfulQA, showing the need for caution in applying established prompting techniques in correctness prediction.

## Peer Reviews

confidence: 2
contribution: 1
initial_rating: 3
presentation: 1
questions: see the questions mentioned above
soundness: 1
strengths: The proposed HCMA method operates independently of model weights, which allows it to function within API-based LLM query setups.
summary: The paper introduces Hierarchical Chains with Multi-Level Abstention (HCMA), a framework aimed at improving both efficiency and risk control in deploying LLMs.
weaknesses: ### Unclear Motivation
- The motivation for a method that addresses both efficiency and risk control in LLM deployment simultaneously is not clearly explained. It is unclear why existing methods addressing efficiency or risk control separately are insufficient.
- The rationale behind the HCMA approach requires clarification.
- The paper would benefit from a stronger scientific argument that demonstrates a common challenge in efficiency and risk control in LLM deployment, justifying the simultaneous attention to efficiency and risk control of the proposed method.

### Questionable Evaluation
- The paper’s evaluation of “risk control” is primarily based on performance metrics from tasks like MMLU. This choice raises questions about how HCMA’s risk control distinguishes itself from other methods that optimize efficiency through similar performance-cost tradeoffs.
- No baselines from related works are included, limiting the ability to benchmark HCMA’s effectiveness against existing approaches.

### Confusing Presentation
- In Figure 1, the y-axis appears to change despite a fixed x-axis value of 1.0 on the right. The basis for this plot needs further explanation: Is it an extrapolation based on several sample points?
- The text references numerous terms without adequate explanation or citation, such as “abstention rate” (L21), “based on hidden layer embeddings, repeated sampling, and neural-network correctness predictors” (L49), “uncertainty-based delegation” (L92), and "... Platt scaling ... calibration technique" (L93). A clearer introduction to these terms would enhance readability.
- Table 1 lacks a comprehensive analysis of the presented results, leaving the interpretation of findings ambiguous.
- In Figure 3, the x-axis for $/Mtok is presented as a variable, though it is typically a fixed cost for each specific model. This discrepancy requires clarification.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. In general, if more than one family of LLMs exist of various sizes,  such as Mistral 7B, llama 2 7B, Gemma 2b, Opt models listed above as well, how would this work generalize to that setting?
2. As this kind of delegation is more expensive during runtime than training a meta-model to route the query to right model, can you elaborate what are the scenarios where this method could be preferred over that? Perhaps this generalizes better on the domains that are out of the domain of the meta-model’s training data? It would also be nice to bolster your argument with relevant empirical evidence.
soundness: 3
strengths: 1. The related work is covered in great detail. 
2. This paper tries to reduce the cost by utilizing smaller LLMs if they answer the given query correctly rather than always using larger LLMs for each query. They delegate the more difficult queries to larger LLMs or abstain from answering the query altogether if they are not confident enough.
summary: This paper proposes a new algorithm to route a given query to the model chain comprised of models of the same family but varying sizes such as 1B, 7B, 13B etc. These models are arranged in the increasing order of model size. As inference on larger models is expensive, the query is first routed it to the smallest model in the chain. This model then decides if it should abstain from answering the query altogether on behalf of all the models in the chain, or if it should delegate the query to the next larger model or if it is should answer the query. This routing is determined by the probability of the correctness of that model. If the probability is less than the rejection threshold, then it abstains from answering on behalf of all the entire chain, if it is between the rejection and the acceptance threshold, it delegates the query to the next model in the chain and if it is greater than the acceptance threshold, it answers the query. Thus by using smaller LLMs whenever possible in lieu of using the largest model all the time, they reduce the inference time and the cost for each query. 
In order to make sure that the probability computation is calibrated, they employ Platt scaling and adapt it to avoid clustering around probability  value of 1.0.
weaknesses: While the problem they are tackling is quite relevant, the paper lacks sufficient experiments and baselines to demonstrate the efficacy of the proposed method. I have listed a few of my concerns below. 
1. How does the modified Platt scaling work in comparison to other uncertainty quantification and probability calibration techniques such as semantic entropy (Kuhn et al.), P_true (Kadavath et al., 2022), Eigen values, Degree, Eccentricity (Lin et al. 2024) and other works listed in the uncertainty quantification part of the related work section. While the authors are performing probability calibration, it is also comparable to uncertainty estimation as the query can be rejected when the uncertainty is high.
2. While you showed the results for various threshold values on one dataset, in a real world scenario, how would one go about in setting the acceptance, rejection threshold values so that it works well for various kinds of queries?
3. In order to demonstrate the generalization, please evaluate the figure 3 and table 1 on more datasets such as helaswag (https://huggingface.co/datasets/Rowan/hellaswag), SQUAD (https://huggingface.co/datasets/rajpurkar/squad) etc and other families of models such as Mistral, Flan-t5, Gemma etc. It is also important to verify if this algorithm works when the differences between model sizes are large or if it would work when the model sizes are opt 350m, opt 1.3b, opt 2.7 b etc.
4. Could you also tabulate the regret wrt accuracy and cost? That is given a model chain comprised of llama 8B, 70B and 405B, we need the ground truth of the smallest possible model that could answer the question and if the model chain should reject the query. Based on this ground truth, you could compute the error and the cost introduced by routing it the wrong model. If the cost is lower than the ground truth, then u can use 0. Only routing it to more expensive models would be penalized. While the plot in Figure 3 demonstrates if the delegated model is correct or not, we are interested in understanding if the algorithm is indeed routing it to the right model.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: See weaknesses.
soundness: 2
strengths: * The topic of using model cascades to cut costs is of practical relevance.
* The paper uses cross-validation across 100-500 seeds.
summary: The paper proposes a cascade of LLMs where each stage can either reject an input, defer it to the next larger LLM, or answer it, depending on their answer token/P(True) probability. The probability is recalibrated using logistic regression and nonlinear transformations. The cascade shows a more favorable error-per-cost tradeoff than using single LLMs on MMLU.
weaknesses: * The font size is reduced and the margins are made smaller. This is a potential breach of the ICLR guidelines. 
* Conceptually, I cannot follow why it is required to recalibrate the LLM token / P(True) probabilities via a logistic regression with a nonlinear transformation of probabilities. All that the method uses in the end are the two threshold values. Since all transformations are monotonic, the thresholds could have also been computed for original probabilities. 
* There is no comparison against the routing / abstained prediction SOTA that the paper cites.
* The method works only on one dataset (MMLU) and one cascade of models (Llama 3 models). It fails on TruthfulQA.
* The method requires searching 39^5 (= 90M) hyperparameters. This could be greatly reduced by excluding impossible combinations and using Bayesian optimization. 
* Figure 1 has an odd choice of the x-axis values (0, 0.86, 0.982, 0.998, 1.0, 1.0, 1.0, 1.0, 1.0) and does not show any data, just the estimate logistic curves. This makes it hard to tell if the data actually supports the interpretation made from the figure, namely that "differently sized models share a common notion of difficulty"
* Some statements are heavily marketed and oversold. E.g., 
	* "we introduce a nonlinear feature transformation that makes Platt scaling a highly effective calibration technique for LLM token probabilities, providing an alternative to temperature scaling grounded in a rigorous statistical model"
	* "which require only 50 or 100 labeled examples to achieve excellent calibration error (ECE), cutting ECE by 50% compared to naive Platt scaling", I'd advice to remove "excellent" and "only" to make this more objective, since a remaining ECE of 0.05-0.07 is far from excellent. I'd suggest to change 50% to "between 17% and 55%". 
* I cannot follow why an arbitrary nonlinear transformation is labeled as statistically grounded ("our nonlinear transformations make Platt scaling much more effective in calibrating LLM output probabilities, yielding a statistically grounded way of performing LLM calibration". Because of the follow-up logistic regression? A logistic regression has no guarantees to produce calibrated values, it just minimizes its loss)
* There is no code released that would allow to replicate the experiments.
* It would be beneficial to report standard deviations in Table 1, since you already used multiple seeds (which is nice!)
* Proposition 1 lacks notation. Is $1_D$ a vector across all decision, with each entry being 1 or 0? 

Small notes that did not influence my score and don't need to be rebuttled, I just note them to make the camera-ready better:
* Besides font size and margins, I would suggest to reformat the figures. If you could give them a uniform text size, correct aspect ratio, and potentially use tikz where applicable, removing titles from figures (and put them into the captions). That would improve the presentation a lot. 
* Reformat equation 9
* Reformat page 8
* The reference section is misformatted ("Can llms express their un-
certainty? an empirical evaluation of confidence elicitation in llms"). Consider adding double brackets to the .bib entries.
* Typo in line 21: yields drive
* Typo in line 249: hafve
* It's more common in ML literature to have the contributions (§3) be part of the introduction (§1), that would make it easier to digest for a quick reader