AutoMix: Automatically Mixing Language Models
Authors: Pranjal Aggarwal
Aman Madaan
Ankit Anand
Srividya Pranavi Potharaju
Swaroop Mishra
Pei Zhou
Aditya Gupta
Dheeraj Rajagopal
Karthik Kappaganthu
Yiming Yang
Shyam Upadhyay
Manaal Faruqui
Mausam .
Venue: NeurIPS 2024 poster (2024)
Keywords: Few-shot learning
Zero-shot learning
Self-Verification
cost-quality optimization
Decision making
Prompting
LLMs
Primary area: natural_language_processing
TL;DR: AutoMix robustly routes queries among language models of varying sizes, efficiently balancing computational cost and solution accuracy

Abstract: Large language models (LLMs) are now available from cloud API providers in various sizes and configurations. While this diversity offers a broad spectrum of choices, effectively leveraging the options to optimize computational cost and performance remains challenging. In this work, we present AutoMix, an approach that strategically routes queries to larger LMs, based on the approximate correctness of outputs from a smaller LM. Central to AutoMix are two key technical contributions. First, it has a few-shot self-verification mechanism, which estimates the reliability of its own outputs without requiring extensive training. Second, given that self-verification can be noisy, it employs a POMDP based router that can effectively select an appropriately sized model, based on answer confidence. Experiments across five language models and five challenging datasets show that Automix consistently surpasses strong baselines, reducing computational cost by over 50\% for comparable performance.

AutoMix:
Automatically Mixing Language Models
Pranjal Aggarwal
Aman Madaan
Ankit Anand
Srividya Pranavi Potharaju
Swaroop Mishra* Pei Zhou
Aditya Gupta
Dheeraj Rajagopal
Karthik Kappaganthu
Yiming Yang*
Shyam Upadhyay?
Manaal Faruqui"
Mausam
Carnegie Mellon University
+ XAI
i Google f Google DeepMind
c UT Delhi A University of Southern California
auttom/x-moodeessggooglleo.
com
Abstract
Large language models (LL.Ms) now available from cloud API providers in various sizes and configurations. While this diversity offers a broad spectrum of choices, effectively leveraging the options optimize computational cost and performance remains challenging. In this work, we present AutoMis x, an approach that strategically routes queries larger LMs, based on the approximate correct- ness of outputs from a smaller LM. Central to AutoMix x are two key technical contributions. First, it has few-shot self-verification mechanism, which estimates the reliability of its own outputs without requiring extensive training. Second, given that self-verification can be noisy, employs a POMDP based router that can effectively select an appropriately sized model, based on answer confidence. Experiments across five language models and five challenging datasets show that AutoMix consistently surpasses strong baselines, reducing computational cost by over 50% for comparable performance.
Introduction
The landscape of Large Language Models (LLMs) is rapidly evolving, with wide array of models now available various sizes, capabilities, and computational requirements [Touvron et al., 2023, OpenAI, 2023, Jiang et al., 2023a]. While larger models generally exhibit superior performance, their substantial computational costs render them unaffordable for many simpler tasks. Moreover, the vast array of available options makes it challenging for end-users determine the optimal model configuration for their specific needs. This challenge is further compounded by the intrinsic complexity and variability of real-world tasks, ranging from simple (e.g.. binary classification on separable data) to complex (c.g., code generation) and potentially unsolvable tasks (c.g., certain forms of multi-step reasoning). To address these issues and ensure that end-users can obtain the best performance within their budget constraints, the development of model-switching techniques has become increasingly important. These techniques involve dispatching queries models of disparate sizes and capabilities, allowing for more efficient allocation of computational resources [Liu et al., 2020, Zhou al., 2020, Madaan and Yang. 2022, Geng et al.. 2021, Schuster al., 2022].
Contemporary model-switching strategies often rely on separate routing models trained for fixed of tasks [Chen et al., 2023, Ding et al., 2024). Moreover, modern LLMSs are frequently accessible only through black-box APIs, restricting direct model optimization due to the unavailability of fine-tuning capabilities and weight access. This constraint, coupled with the expectation of access to large amounts of task-specific data, creates a challenge that existing routing approaches fail to address adequately. In response, we introduce AutoMix x, method that enables users to mix models
**Equal Contribution. Work started and partly done during Aman's 's internship at Google.
Coode available at github. com/ con/automss automi x- 1lmmautomi
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

LLM
AutoMix
Generate
CO
8 LM, Verifier
Budget X
LM,ILLaMA
Self-Verify
Easy or Unsolvable
POMOP Router
Accuracy: X
Route
Hars?
Automix automotically
Budget:
optimizes cost-perto tracieoff for users by mising block-box LMss
Accuracy:
SLM
LM,/GPT-4
Figure I: Representative example for 2 model setup in Aut oMix. Instead of relying only on small model (SLM) with low performance large model (LLM) with high cost, Aut attcMix x automatically mixes multiple black-box language models, based on user desired cost-quality tradeoff. Aut oMix works in 3-step process: 1.) generation by small model (LM1), 2.) self-verification of the generated answer, 3.) using confidence assessments from self-verification to do appropriate routing larger model (L.M2). For N-model setup, the process is repeated till the final answer is reported.
of various sizes and capabilities, assuming only access to black-box LLM APIs. As illustrated in Figure I, Aut oMix x consists 3 steps designed within the constraints black-box access: solution generation (small model to generate initial answer), self-verification (same smaller model assesses difficulty), and selective routing (routing to larger models when suggested by self-verification)) At high level, this process mirrors human booliins-ooiig, which inherently follows multi-step process: generate a solution, verify its validity, and refine it further based on verification outcomes.
An ideal router in AutoMix x must incorporate several key characteristics. Firstly, it should be capable of identifying the difficulty of query based on confidence assessments of the smaller model's answers. Secondly, it should be able to route to the appropriately sized model, for instance, easy queries are routed to the small language model (SLM), hard queries to large language model (LLM), and unlike previous works [Chen et al., 2023, Ramírez et al., 2024], unsolvable queries that no model can solve should not routed to any LM, thereby saving costs. Thirdly, it should be able to learn from small amount of data, expected in real-world tasks. While ceff-verfiiattin seems a natural choice to provide confidence assessments and learn from limited data, prior works have shown self- verification is not particularly effective for reasoning tasks and is often noisy and ill-calibrated [Tyen et al., 2023, Huang et al., 2023b] challenge that needs to be addressed. Finally, the router should generalize different scenarios in terms of the number of models with varying costs and capabilities.
To address these requirements, we propose two novel contributions in AutoMix. First, we formulate self-verification an entailment problem, where the consistency of the generated answer with the context evaluated to estimate confidence levels [Poliak, 2020, Dagan et al., 2022]. For example, an answer discussing "apples" in context focused on "tea" would be flagged as highly inconsistent. Second, we introduce a Partially Observable MDP (POMDP) based router [Äström, 1965]. POMDPs extend standard MDPs to decision problems where observations velffieeiiiintionn probabilities) are unreliable and provide noisy estimate of states (question difficulty). Modeling the router POMDP agent ensures all our requirements are met, as the POMDP can implicitly model various difficulties and assess them based on noisy self-verification outputs. Furthermore, POMDPs provide a principled formulation to learn robust routing policies in different scenarios with varying numbers of models, different costs, and capabilities, while learning from as few as 50 examples.
We conduct extensive evaluations of AutoMix on five different dialogue and context- grounded reasoning tasks, with five different models. Our results demonstrate that AutoMi x consistently outperforms baselines while reducing cost by over two times while achieving the same performance, showcasing significant improvements and efficiency gains over existing model-switching strategies.
2 Background and Related Work
Self-Verification The concept of self-verification in reasoning problems has been explored in various works, such Weng et al. [2023], Jiang et al. [2023b], Pan et al. [2023a]. These approaches typically

---

Context: [contexL]
Queations (question]
Generated Answer: |generated_enswer
Instructions Your is evaluate the Generated Answer correct, based - provided context and question. Provide the judgement and reasoning for each case. Chouse between Correct oz Tncorrect.
Evaluation:"
Verification Prompt. The verification process is framed as a natural language entailment task, where the model determines the validity of the model-generater answer with respect to the context and question. We use a generic few-shot prompt for all tasks.
Listing I:
rely on the LLM's knowledge [Dhuliawala et al., 2023], a method that can pose challenges for reasoning problems [Madaan et al., 2023, Huang et al., 2023b]. In contrast, Aut oMix leverages context for verification and introduces a POMDP router to mitigate potential noise from the verifier. Another line of work collects a corpus of past mistakes made by models [Madaan et al., 2022], and external knowledge bases for verification [Peng et al., 2023, Gao et al., 2023, Pan al., 2023b]. In contrast, Aut oMi uses the question's 's context to verify the answer. While previous works find self-verification unreliable to repair the model's output [Huang al., 2023a], we demonstrate that self-verification provides a valuable signal used for routing to an appropriate model.
Mixing Models Several works have sought to optimize LLM inference cost through model switching. employing specifically trained verifiers [Chen et al., 2023, Zhu et al., 2023, vSakota et al., 2023, Ding al., 2024]. Aut oMix eliminates the need for expensive verifier training through few-shot SLM model prompting and does not require upfront access to all input queries. The router, trained with as few as 50 samples, outperforms specialized models. Our work thus aligned with recent work that aims at composing different models and external tools for inference time improvement of language models [Khattab et al., 2023, Press et al., 2022, Yao et al., 2022, Zhou al., 2022].
Adaptive Computation Adaptive computation and model routing methods often preempt computa- tion via intermediate representations [Liu al., 2020, Zhou et al., 2020, Geng al., 2021, Schuster et al., 2022, Madaan and Yang, 2022]. Unlike these methods, Aut oMix requires no architectural changes and assumes only black-box access to APIs. While some black-box methods like Adaptive- Consistency [Aggarwal et al., 2023] aim to optimize inference for reasoning tasks, they are limited to a single LLM model, whereas AutoMix flexibly optimizes between two or more models.
Background on Partially Observable Markov Decision Processes POMDPs extend Markoy Decision Processes (MDPs) to scenarios where agent's 's state observation is partial or noisy. Defined as tuple (5, A, R, S., O), S denotes states, A actions, and Ω possible observations. The transition function T gives transition probabilities between states given an action, while the observation function P connects actions and states to observations. The reward function R assigns rewards to actions in specific states. Agents in POMDPs maintain a belief (b), a probability distribution over states. This belief updates based on actions and observations. The objective in a POMDP is to find a policy : + A that maximizes the expected cumulative long-term reward. POMDPs have been used in various domains, including robotics, automated navigation, crowdsourcing workflows, and strategic planning [Kaelbling al., 1998, Schwarting al., 2018, Dai al., 2010, Meuleau et al., 2013]. However, Aut oMix x uses POMDPs in novel way to route queries between LMs. Our noisy self-verification outputs act as observations to the POMDP, which assesses question difficulty and routes to the appropriate LM to maximize a reward function of cost and performance.
3 Problem Formulation
We address the problem of selecting the appropriate language model (LM) to maximize user-defined cost-quality tradeoff. We assume access to N distinct language models: LM, LM2 .. LM numbered in increasing order of number of parameters. Each model has an associated cost of C; and (unknown) performance of P; for each input query. Our objective to maximize the total performance for any total cost while dynamically invoking any LM as appropriate for given test
3

---

Context
Which country grows the most tea? The answer India. grows three times much China. Which country drinks the most tea? It's neither China nor Japan. It's Great Britain. [..]] One customer put the bag into pot. Then he just poured hot water over it. And the tea bag was born. Shen Nong was the first to drink tea. (Shen was Chinese emperor.) This was about 2737 B.C. Shen had bad digestion. So he drank several cups of hot water daily[...] Tea became the drink of China.
Question
Generated Answer (by ALAMA2-13BL
When did Shen Nong drink tea?
He took in 1990.
Verifier Output (by LLAMA2-13B)
The context does not mention that he took in 1990. The Al-generated answer Incorrect.
Figure Goruext- Self- Verificciii using LLAMA2-13B in Action. The example showcases the verifier, utilizing the same model as the answer generator, identifying and rejecting an inaccurate answer He took 1990-by effectively leveraging the context.
point. In experiments, compare cost-quality curves of various methods and also evaluate them using a newly proposed IBC metric (See Section 5.1).
To test our approach, consider context-groundee reasoning tasks, such as comprehension-bice based QA and a variety of dialogue reasoning tasks, where given a context C (e.g., stories, newswire, or dialogue history) and a question q, the model must generate accurate and coherent answer consistent with the provided context. Our choice of tasks is motivated by two key considerations: (1) longer queries are more computationally demanding, underscoring the need for approach like AutoMix x to navigate the cost-accuracy tradeoff, and (2) the context allows for cross-checkine preliminary answers with available information using self-verification (descrihed shortly). We assume only black-box access is available to the LM APIs. For training any router models, we assume access small training/validation dataset Drrain consisting of < C, 9, U. ULM, triplets, where the ground truth answer and YLM, the answer generated by L.M..
4 AutoMix
At a high level, Aut oMix consists three steps: solution generation using small LM, say LM; (initially i = 1), to generate an answer A,; self- verificiiiic using the same LM; to assess A.; and selective routing - picking a larger LM, LM,(j > i), when suggested by self velffieiifi,,tion, else returning A, as final answer. Figure shows a representative example of the process for the two-LM case. Next, discuss our two technical contributions in detail.
4.1 Self- Verification
To assess the trustworthiness of Ag, AutoMix employs few-shot verifier, V. which validates LM,'s 's output. Unlike existing works that perform verification by creating a new question [Weng et al., 2022, Jiang al., 2023b], we frame verification Dagan et al., 2022], to determine whether the answer generated by L.M, aligns with the provided context. Specifically, the verifier gauges p(correct 1 A.,, C, q). with correct that A, is correct. To estimate the probability, we sample & > 1 times from the verifier (LM,) with high sampling temperature, and probability is then computed & 1 [correett= = 1}. We use the same 4-shot verification prompt for all tasks and not train a verifier. Figure shows the prompt and example of velf-verification in action. We refer readers to Appendix B for all prompts used.
entailment task [Dagan et al., 2005, Poliak, 2020,
indicating

---

4.2 Router
Routing follows solution generation and self velf-verification.. The router decides whether the output from L.M; should be accepted or the query be routed to some I.M, (j > i) for improved performance. The router can also be interpreted as meta-verifier, providing an extra layer of confidence assessment on the few-shot verifier's 's evaluations. Specifically, V ascertains whether LM; 's answer is entailed by the context, making its decision without considering the inherent difficulty of the problem. For example, when handling Unsolvable queries, invoking a larger LM will be resource-inefficiee and would not enhance performance. A good router can address this by not routing such a query further, and this decision needs to be made using the verification probability and trends from the training data.
Addressing the notable challenges of ceff-correction in large language models [Madaan et al., 2023, Huang et al., 2023b], Aut oMix employs a non-LLM setup for routing and avoids escalating issues such as hallucination and reasoning errors [Dziri et al., 2023]. The router can, in principle, adopt various learning strategies, including supervised learning, reinforcement learning, and symbolic reasoning. Subsequent sections provide details of two different routing strategies in AutoMix.
Thresholding In this simplistic routing approach for the two-model case (N = 22), the decision to route to LM2 is based on the probability of the L.V, verifier and a threshold t. If e 2 1, it returns LM1's 's answer, else routes the query to LM2. Intuitively, a high probability indicates the verifier is confident its decision and can be trusted. Varying can help explore cost-performance tradeoff'.
POMDP-based Router A router should direct a query to larger LMs only when the performance gap justifies the cost-quality tradeoff.. Given the inherent uncertainty in the true state of system performance, which remains unobserved, we formulate the router as a Partially Observable Markov Decision Process (POMDP) [Áström, 1965]. POMDPs are particularly well-suited for scenarios where observations, such as self-verification probabilities, may not be entirely reliable.
Recall that (Section 2) POMID is characterized by (S, A, T, R, S, O). In our application, the states S represent the current selected LM, and performance metrics (c.g.. accuracy or F-score) of various LMs on a data point, denoted as {i, Perr fL either retaining the current LM's (LM,) answer or routing to one of the larger LMs. Observations Q in the form of the verifier, outputs v from LM,, enabling the POMDP to ascertain its belief state b: probability distribution over S. The observation probabilities P(o s). which indicate the likelihood of observing o (verification output) given state s, are crucial for defining the POMDP model. E.g., high verifier confidence might suggest that the current LM's performance Per frAf is high enough, reducing the necessity to switch to more costly L.M The observation probabilities are estimated directly on the train set, by computing the expectation of verification probabilities for
fLMx) ). The actions include
each state, i.e P(ols) E,g,, E PPeaain 1[8;=s and vy=o)
However, since the states can be continuous,
22, CDerain 1(s,=a}
we use non-parametric Gaussian kernel density estimation to estimate the observation probability for new state. Instead of directly estimating P(o]s), we first learn a joint distribution P(S, O) and P(S), by drawing KDE at each training point. Then conditional probability P(o)s) is computed as: P(ols) P(e.o)
P(a)
The objective of our POMDP is to optimize reward function R = P X C, where P is overall performance, C is overall cost, user preferences. We utilize the AdaOps POMDP solver [Wu et al.. 2021]. which employs particle filtering for representing belief, where each particle corresponds to a particular state. At inference. the POMDP solver starts with an initial belief (uniform distribution). updates its belief state based on verifier observations, and, based on the updated belief state b', takes an action maximize the expected reward. We provide further details and the complete POMDP formulation in Appendix A.3.
tuning parameter that balances two criteria, according to
Experiments
Through our experiments, we wish to answer the following research questions. How does AutoMix compare with other model-switching strategies? How well does AutoMi perform when varying (1) cost ratios between models, and (2) the amount of available training data? Following recent work [Ding et al., 2024], we perform most experiments in the setting when N = 2, but also report
5

---

70
procedure ANSWERQUERY(C.
SLM
C. Context, q: Question, SLM/ LLM:
LLM
65
Small/large language model
Performance
A, f SOLVE(SLM, C,
V f- SELF-VERIFY As, C,
60
if META-VERIFY (v, A., C. then
return A.
55
else
A t (LLM,, return A
50
20
40
60
end if
end procedure
Cost
Figure 3: Left: Aut oMix algorithm. Right: Performance vs. Cost curve. The slope between SLM and L.LM provides way to the Incremental Benefit per Cost (TBC) for methods that mix models. Methods with steeper slope than this reference when plotted against SLM have positive IBC (green region), whereas those below the reference have a negative IBC (red region).
additional results in the three model cases (in Section 5.5). For N = 2, we use the terms SLM and LLM to denote small (LM,) and large (I.M2) language models, respectively.
5.1 A Metric for Cost-Performance Efficiency Analysis
Incremental Benefit Per Cost (IBC) In our approach to leveraging model performance, it essential to consider not only the raw accuracy of predictions but also the associated computational or monetary costs. To that end, we introduce a metric to understand the performance of the models in terms of cost. We introduce methods, denoted by M, to optimally integrate SLM and LLM. For each method M, we associate a cost CM and performance PM- To quantify the utility of M over SLM, we define the metric Incremental Benefit Per Cost (IBC) as IBCa (Equation (1)).
Pas PSLM
PLLM PSLM
IBCM IBC 100
IBCM
IBC
AIc((M)
Cas CSLM"
CLLM CSLM
IBC
(1)
The IBC metric captures the efficiency of performance enhancement relative to the additional cost. For comparative evaluation, LLM over SLM. Finally, we compare methods using AIRC, which compares the IBC of a specific method with IBCRase- A positive IBC lift suggests that M achieves performance increments more cost-effectively than a standalone L.M,, whereas a negative lift indicates reduced efficiency (Figure 3)
baseline IBC, IBCRASE, representing the benefit of always using
Geometric Interpretation On a Performance vs. Cost plot, consider the line segment joining the data points of the small language model (SLM) and the large language model (LLM). This segment's s slope represents the rate of performance increase per unit of cost. The Incremental Benefit per Cost (IBC) for any method M is the slope of the line from the SLM point the point representing M (Figure 3). A method M that lies above the SLM-LLM segment provides a steeper slope, indicating favorable
(and positive Amc). Conversely, if M lies below the segment, suggests an unfavorable or negative IBC. Our primary objective is to develop methods that yield consistently positive IBC, maximizing performance enhancements for each additional unit of cost.
5.2 Setup
Models and Cost Calculation We use GPT-3.5, LLAMA2-13R, and MISTRAL-7B (instruct v0.2 version) [Touvron et al., 2023, Jiang et al., 2023a] as the smaller language models (SLMs), and GPT-4 [OpenAl, 2023] as the larger language model (LLM). For each input query, we model fixed costs for the models and verification, denoted by OSLM for smaller models, CLAM for the larger model, and Cver the verification model. As the verification performed by CSLM and the major cost driver is the question context, which remains the same in both models, set C CSLM. At

---

DIPLOMAT
MUTUAL
QUALITY
0.8
0.96
0.8
0.78
0.70
0.75
0.74
0.65
0.65
0.72
0.6
50
100
150
100
150
200
100
150 200
COQA
QASPER
AutoMi +POMDP
0.55
0.35
AutoMix +Threshold
FrugalGPT [Chen et al., 2023]
HybridLLM [Ding et al., 2024]
LLAMA2-13B (SLM)
GPT-4 (LLM)
0.35
0.25
Random Mixing
50
100
150
200
50
100
150
Figure Main Results: performance (y-axis) vs. cost (x-axis) for different methods the small and large (LAMA2-13/GPTE POMDP based meta-verifier is consistenly above the linear interpolation (random mixing) SLM-LLM, signifying a higher incremental benefit per unit cost (IBC).
inference, the total cost is computed as the sum of individual costs, = CSLM + Over + w CLIM- where w indicates whether the larger model invoked.
While the pricing these APIs [Dehghani et al., 2021] is influenced by various complexities, our focus the black-box utilization of language models leads to represent cost based the actual API price disparity between these models." For each of GPT-3.5, LLAMA2- 13B. MISTRAL-7B: we assign relative cost of 1 unit for the SLM and 60, 100, 200 units for the LLM respectively. It's important to note that the cost ratio between models can vary significantly depending on specific deployment scenarios. For example, for 13B might incur virtually no cost, while utilizing GPT-4 could be prohibitively expensive. We simulate this scenario in Section 5.4. We refer readers Appendix B for additional details.
with access to single A6000 GPU, running I.LAMA2-
Datasets We experiment with a diverse set of datasets: QASPER [Dasigi al., 2021]: Question answering over research papers; ii) QUALITY [Pang et al., 2022]: Multiple-choice questions (MCQ) on long articles and stories; iii) COQA [Reddy et al., 2019]: Conversational comprehension requiring coreference and pragmatic reasoning: iv) MUTUAL [Cui et al., 2020]: Multi-turn dialogue reasoning (next response prediction); v) DIPLOMAT [Li al., 2023]: Pragmatic identification and reasoning questions on multi-turn dialogues. We use the F1 score for QASPER and COQA, and accuracy for the remaining datasets. We use the default validation splits and utilize prompts from Shaham al. [2023] for QASPER and QUALITY, and adapt the QUALITY prompt for other datasets. We use identical prompts for all models. We refer readers Appendix B for more details.
Baselines We compare against FrugalGPT [Chen et al., 2023] and HybridLLM [Ding al., 2024], two state-of-the models, as our baselines. FrugalGP7 uses a finetuned DistillBert model [Sanh et al., 2019] as router. If the router's confidence for given question, context, and SLM answer falls below a threshold, the query is routed to the L.M.. Hybri dLLM uses trained DeBERTa [He et al., 2021] as router which directly chooses between SLM and LLM without
Phttps::/<penaal wttp://c./p//aeeeeeeeeeeee https://together:

---

Model
Method
DIPLOMAT
MUTUAL
COQA
QASPER
QUALITY
FrugalGPT
16.8
20.3
-16.7
7.7
8.1
HybridLLM
67.1
3.4
10.5
0.4
2.7
MISTRAL-7B
Aut oMix + T
149.7
46.7
12.1
67.9
33.6
AutoMix +
156.8
46.7
12.4
69.3
51.6
FrugalGPT
-7.0
-8.7
2.6
9.4
447
HybridLLMM
3.8
2.2
-0.5
7.2
6.5
LLAMA2-13B
Auix +
50.1
11.8
86.5
-0.2
9.4
AutoMix
58.5
12.4
83.1
8.5
10.3
FrugalGPT
30.1
11.1
37.0
87.2
10.1
HybridLLM
8.3
1.8
20.1
21.7
11.8
GPT-3.5
Aut AMix + T
168.2
18.3
62.7
109.8
23.5
AutoMix
151.2
18.8
65.0
114.7
24.1
Table I: ABC values: AutoMix T and oMtoMix + P are variations of our proposed method with thresholding (T) and POMDP-based routers, respectively. Best numbers bold, and second-best are underlined. AutoMix + POMDP demonstrates robust and consistent ABc across all datasets, implying judicious utilization of computational resources. Despite domain-spectie training and O-cost verifier, F FrugalGPT and Hybr Hybriddi riLLLM underperform Aut oMix in almost all scenarios.
generating SLM response. Further, training labels for Hybr dLLM's router are generated based on the probability of SLM outperforming L.M by margin, computed by drawing multiple samples. Note, that training each of these baselines along with AutoMi requires running inference all models, however, it is only a training time cost. However, at inference, unlike Aut oMix we assign cost of the two baselines' routers due to their lower operational costs.
5.3 Main Results
Figure 4 illustrates the performance versus cost curves for various datasets and model-mixing methods using MISTRAI.-7B SL.M. Across all datasets, Threshold consistently outperform F FrugalGPP and HybridLIM, staying above the SLM-LLM curve and yielding better performance per unit cost. This is particularly notable both FrugalGPT and Hybricn iIMM dLLM leverage domain-specific trained routers and do not incur any verification costs. Comparable trends are observed with other SLMs LLAMA2- 13B and GPT-3.5 Figure 18, 19.
x-PMMDP and AutoMi
Table presents a comparison of the average ABE across five equally sized cost regions for MerrMMix-POMIE consistently outperforms
each method, dataset, and three different SLMs. across all datasets and models, with performance equalling
Frugal and HybridLin Fruga GPT in the qasper dataset for the I.AMM2--5 2-BBB model. The gains vary ent settings, with the best performance observed on the DIPLOOMA dataset, averaging 122% across all models. Importantly, Aut oMix P-MMDP is the only method that yields positive gains across all configurations and consistently matches or exceeds the performance of Aut oMi x Threshold. This is despite being tested variety of datasets and models representing different costs, task difficulty (accuracy ranging from 30% to 90%), performance gap between the models (8% to = 50%), different answer types (MCQ, Open-ended). The results demonstrate that Aut cMix, utilizing verfficriii and appropriate routing, can effectively mix different SLMs and GPT-4 on a wide range of tasks without access model weights or domain-specific routing data. These substantial gains translate directly into cost savings, underscoring the economic relevance of our approach.
differ-
5.4 Additional Experiments
Effect of Cost Ratio on AutoMix Our main experiments assumed different cost ratio ranging from 1:60 to 1:200 for different SLMss. Subsequently, we conduct an analysis to understand how alterations the cost ratio influence the ABc values across different cost ratios. In Figure (right), we vary cost-ratio and report A normalized by the maximum value obtained on the dataset (a higher value denotes higher gains). The results suggest that even for a cost ratio low as 1:10, AutoMi begins to deliver good gains for many datasets, which naturally improve cost-ratio gets further skewed.

---

Macro Averaged IBC Lift
20
Nomalzzdd LIR
10
DiPlomat
0.4
MuTual
Automix w/ POMDP
0.2
CoQA
Frugal
QASPER
HybridLLM
Quality
250
500
1000
Train Size
Figure 5: Left: Comparison of AutoMi with Fruga GPT and Hybri over varying training data sizes shows that despite O-cost verifier and domain-specifie training, baselines underperform AutoMix x. AutoMix x excels in limited data settings. Right: Normalized Anc across different cost ratios. Aut oMix show robust performance gains even when cost ratio is low.
We supplement the analysis with cost-performance curves for different cost ratios in Appendix D.2. We note that cost need not only be monetary cost, and could, in general, represent other considerations such as latency or energy usage. These results demonstrate that Aut oMix provides a novel method to balance the pesforerfom tradeoff across wide range of cost scenarios and is robust to changes without any modifications to the method.
AutoMix Effective in Low-Resource Scenarios Figure (left) demonstrates the performance dynamics of Aut oMix, FrugalGPT and HybridLLM with varying validation set size. Notably, our method significantly outperforms FrugalGPT and HybridLLM with limited data, despite the latter's 's domain-speciff training and zero verifier cost. E.g., with just 50 training examples, AutoMix x maintains 15% ABc and beats the baselines by more than absolute 20%. This underscores that AutoMi particularly advantageous in the real-world scenarios where training data is scarce.
When Does POMDP Routing Help? POMDP implicitly categorizes question difficulty into three types: (a) easy queries solvable by the small model, (b) complex queries only larger models can solve, and (c) potentially unsolvable queries for any model. It leverages sefffverification probabili- ties ties-which which provide a noisy estimate of difficulty to make non- non-linear decisions. For example, in the QASPER dataset using MISTRAL-78 correspond to such cases, and instead of routing to the LLM as other methods do, returns the SLM answer, saving cost. Furthermore, in setups with more than two models (as discussed in Section 5.5), POMDP makes more nuanced decisions, for instance by combining information from small and medium-sized verifiers, resulting in significantly superior performance.
SLM, POMDP identifies that lower confidence values
Other Experiments We perform various other experiments, whose details in appendix due to space limitations. We first evaluate few-shot self verificciin quantitatively and qualitatively and observe that the verffieciiination can effectively use context identify errors in answers generated by SLM in many cases (see Appendix A).
While AutoMix x is outperforms diverse set of datasets, to further demonstrate our router's generalizability, we evaluate out-of-domain generalization by training on one dataset and evaluat- ing on held-out datasets. Results in Appendix D.3 show that Aut consistently outperforms FrugalGPT I and HybridLLM by significant margins. Overall, the task-agnostic design POMDP and these results highlight our method's generalizability.
We also study the latency overhead of AutoMix. At inference time, POMDP takes less than Ims for every query in the two-model setup. Moreover, network latency (usually approximately 10ms) is much less than the average solution generation time by SLM and LLM (on the order seconds). Thus, AutoMi x adds no additional computational or latency overhead compared to language model inference. Finally, incorporating the POMDP router for end-users convenient, as our library can be incorporated in less than five lines of code.

---

5.5 Results of Automix with Three Models
We evaluate the performance of AutoMix ap- plied to a mhree-model scenario (N = 3). Specif- ically, LLAMA2-1 3B SLM, LLAMA2- 70B as MLM, and GPT-4 as LLM. Results presented in Figure 6. AutoMix consis- tently outperforms baselines across cost regions. We first compare AutoMix to FrugalGPT, showing significant improvements across all cost regions considered. We also compare AutoMix to baseline, Union Aut oMix, which selects between Aut MLMM and M cost requirements. For instance, if the de- sired average cost is less than that of the XSLMSSSSMMMM. MLM, AutoMi AutoMixsTMM-M MLM is employed; otherwise, AutoMixMINT Further, we consider a baseline, Chained AutoMix: by chaining SLM-MM The query first
Pedormante (73 Scweea
0/44
Ll 138
based on the user's
Comprstational Cost
Figure 6: AutoMix with 3 models: LLAMA2- 1 13B, GPT-4 and GPT-4. AutoMix method shows consistent IBC lifts for both SLM-MLM and MLM-LLM regions. Further, compared to baselines: FrugaIGPT, chaining two Aut oMix models or using the union of two Aut oex, AutoMixg x3 provide significant improvements.
is utilized.
with Aut XMXXMMM-LMM goes to the SLM, and an Aut oettoMiisssn MMixxLM-MLMM decides between reporting the SLM answer or routing the MLM. In the latter case, second AutoMi XMLM- LLM repeats the procedure using the MLM and LLM models. Chained AutoMix underperforms across the board, as it cannot directly route queries from the SLM to the LLM. Additionally, whenever Chained Aut OMix prompts the MLM, it invariably uses the costly verifier, even when it might not be necessary. Further, in Figure 9, 10, study two additional cases: 1.) If the MLM underperforms the SLM, and its verifier is uninformative, then unlike F Fruga GPT, AutoMi x learns the irrelevance of MI.M and routes directly from SI.M LLM needed; 2.) In scenarios where MLM underperforms the SLM, but its verifier provides useful information, Aut oMix x leverages information from MLM's verifier to outperform all baselines considered. The experiments show the efficacy of AutoMi x in providing optimum performance across diverse situations without additional intervention. We refer readers Appendix C for more details.
Conclusion
AutoMix x integrates multiple black-box LLM APIs into a multi-step problem-solving framework, optimizing the computational cost and performance trade-offs. Our work interweaves Good Old- Fashioned Artificial Intelligence (GOFAI) approaches with LLMs, demonstrating that the incorpora- tion of POMDP can lead to effective routing between LLMs. Our work provides a novel method to balance cost rest-performacce tradcoffs across a wide variety of models over different desired cost ranges, consistently outperforming the baselines by significant margins. Aut oMix opens avenues for several interesting research directions. While self-verification and correction are challenging for LLMs in general, we find promising results using context-gres grooudddd few-shot verification coupled with POMDPs, indicating that similar approaches may also help other scenarios. The work may be extended to generation tasks, where the notion performance subjective.
Limitations
While our empirical evidence across various models and datasets demonstrates the effectiveness of Aut outoMix, x, its broader applicability might vary depending on the specific models and datasets used. Furthermore, Aut oMix is designed with a dialogue- or context- groonddd reasoning setup in mind for effective self-verification. which does not include tasks like factual question-answerit and commonsense reasoning. In the future, as open-source models become more powerful and inference costs decrease, it might be feasible to serve strong model for all queries. However, there are likely to still be availability trade-offs that could be managed using Aut oMix.
10

---

7 Acknowledgements
This work is supported by the IBM AI Horizons Network grant, grants by Google, and Microsoft, an IBM SUR award, and the Jai Gupta chair fellowship by IIT Delhi. We thank the IIT Delhi HPC facility for its computational resources. We are grateful to Microsoft AFMR for supporting this work. We also thank Kalpesh Krishna, Prakhar Gupta, Rahul Gupta, Siddharth Gopal, and Yang Song for their valuable feedback.
References
Pranjal Aggarwal, Aman Madaan, Yiming Yang, and Mausam. Let's sample step by step: Adaptive-consistency for efficient reasoning and coding with LLMs. In Houda Bouamor, Juan Pino. and Kalika Bali, editors, Proceedings of the 2023 Conference in Natural Language Processing, pages 12375, Singapore, December 2023, Associa- tion for Computational Linguistics. doi: 10.18655/w1/2023.emm http2222..m..mmmaimmm..mm caaumlppmain.61.... // //aclanthology. org/ 2023. emmlp-main. 761.
Empirical Methods
URL https:
Lingjiao Chen, Matei
Zaharia, and James Y. Zou. Frugalgpt: How to use large language models while reducing cost and improving performance. ArXiv, abs/2305.05176, 2023. https:
//api. semant commanticscoolln iescholar ottp/oorps11:::555555552
Leyang Cui, Yu Wu, Shujie Liu, Yue Zhang, and Ming Zhou. Mutual: A dataset for multi-turn dialogue reasoning, ArXiv, abs/2004.04 2020. URL htt.t::t/////e.c. //pii sew.nticscconitt
org/Corpust 2D:155482111
Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entailment challenge. Machine learning challenges workshop, pages 177-190. Springer, 2005.
Ido Dagan, Dan Roth, Fabio Zanzotto, and Mark Sammons. Recognizing textual entailment: Models and applicatioos. Springer Nature, 2022.
Peng Dai, Mausam, and Daniel S. Weld. Proceedings of the AAAI Conference on Artificial Intelligence, 2010. URL https: /api.
-eeerornttcoreticcccco control of crowd-sourced workflows.
semanticschon
wrg/Corpussrrrraaoooooo ID: :D:073027 8730278.
Pradeep Dasigi. Kyle Lo, Iz Beltagy, Arman Cohan, Noah A Smith, and Matt Gardner. A dataset of information-sen -eeekiiiio--eekingg questions and answers anchored in research papers. In Proceedings of the 2021 Conference the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, pages 4599-4610, 2021.
Mostafa Dehghani, Yi Tay, Anurag Amnb,, Lucas Beyer, and Ashish Vaswani. The efficiency misnomer. In International Conference on Learning Representations, 2021.
Shchzaad Dhuliawala, Mojtaba Komeili, Jing Xu, Roberta Raileanu, Xian Li, Asli Celikyilmaz, and Jason Weston. Chain-of-verrficatti reduces hallucination in large language models. arXiv preprint arXiv:2309. /495, 2023.
Dujian Ding, Ankur Mallick, Chi Wang, Robert Sim, Subhabrata Mukherjee, Victor Rühle, Laks V. S. Lakshmanan, and Ahmed Hassan Awadallah. Hybrid LLM: Cost-efficient and quality-aware query routing. The Twelfth International Conference on Learning Representations, 2024. URL https:
// /tttps://pperreee http://ieeeeeeeeeeeeemee/ee./..e wwtp://eeeettttttttt...e.. faarrmmiidd-2233mttttmm
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jian, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D Hwang, al. Faith and fate: Limits of transformers on compositionality. arXiv preprint arXiv:2305. 1854, 2023.
Luyu Gao, Zhuyun Dai, Panupong Pasupat, Anthony Chen, Arun Tejasvi Chaganty, Yicheng Fan, Vincent Zhao, Ni Lao, Hongrae Lee, Da-Cheng Juan, et al. Rarr: Researching and revising what language models say, using language models. In Proceedings of the 61 Annual Meeting of the Association for Computational Linguisties (Volume 1: Long Papers), pages 16477-16508, 2023.
11

---

Shijie Geng, Peng Gao, Zuohui Fu, and Yongfeng Zhang. Romebert: Robust training of multi-exit bert. arXiv preprint arXiv:2 10l.09901 2021.
Pengcheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen. Deberta: Decoding-enhanee bert with disentangled attention, 2021.
Jie Huang, Xinyun Chen, Swaroop Mishra, Huaixiu Steven Zheng, Adams Wei Yu, Xinying Song, and Denny Zhou. Large language models cannot icasonnn yet. ArXiv, abs/2310.011988 2023a. URL https:/ /tt///appa appi...........///.. semant www.sseeeatatstaaeeeeeeeem www.icccchhc..c.coo www.co/aa....g.oooraaaaaaaa.aaaa 2
Jie Huang, Xinyun Chen, Swaroop Mishra, Huaixiu Steven Zheng, Adams Wei Yu, Xinying Song, and Denny Zhou, Large language models cannot self-correct reasoning yet. arXiv preprint arXiv:2310001798. 2023b.
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023a.
Weisen Jiang, Han Shi, Longhui Yu, Zhengying Liu, Yu Zhang, Zhenguo Li, and James T Kwok. Backward reasoning in large language models for verification. arXiv preprint 23iv:2330..0775: 2023b.
Leslie Pack Kaclbling, Michael L. Littman, and Anthony R. Cassandra. Plamning and acting in partially observable stochastic domains. Artificial Intelligence. 101(1):99-134 1998. ISSN 0004-3702. doi: http://oooo000010001000000000000 sciencedirect.com
URL https: /www.
Omar Khattab, Arnav Singhvi, Paridhi Maheshwari, Zhiyuan Zhang, Keshav Santhanam, Sri Vard- hamanan, Saiful Haq, Ashutosh Sharma, Thomas T Joshi, Hanna Moazam, et al. Dspy: Compiling declarative language model calls into self-improving pipelines. arXiv preprint arXiv:2310.03714 2023.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023.
Hengli Li, Songchun Zhu, and Zilong Zheng. Diplomat: A dialogue dataset for situated pragmatic reasoning, ArXiv, 0bs/2306.090300 2023. URL https: / //ppi/ppp////////////// semanticscho. org/ CoxpusID:2591644
Weijie Liu, Peng Zhou, Zhiruo Wang, Zhe Zhao, Haotang Deng, and Qi Ju. Fasthert: a self-distilling bert with adaptive inference time. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 6035-6044, 2020.
Aman Madaan and Yiming Yang. Flowgen: Fast and slow graph generation. arXiv preprint arXiv:2207.07656 2022.
Aman Madaan, Niket Tandon, Peter Clark, and Yiming Yang. Memory prompt editing to improve GPT-3 after deployment. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 2022. Association for Computational Linguistics. URL httpps:/accaatthlllee /rg/ 2022. commmppmmiiie 183.
Abu Dhabi, United Arab Emirates, December
Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri Alon, Nouha Dziri. Shrimai Prabhumoye, Yiming Yang. et al. Self-refine: Iterative refinement with self-feedback. arXiv preprint arXiv:2303.17653 2023.
Nicolas Meuleau, Leonid Peshkin, Kee-Fung Kim, and Leslie Pack Kaelbling. Learning finite-state controllers for partially observable environments, 2013.
12

---

OpenAI. Gpt-4 technical report, 2023.
Liangming Pan, Michael Saxon, Wenda Xu. Deepak Nathani, Xinyi Wang, and William Yang Wang. Automatically correcting large language models: Surveying the landscape of diverse self-correction strategies. arXiv preprint 1rXiv:2308.03/88. 2023a.
Liangming Pan, Xiaobao Wu, Xinyuan Lu, Anh Tuan Luu, William Yang Wang, Min- Yen Kan, and Preslav Nakov. Fact-checking complex claims with program-guie guided reasoning. arXiv preprint arXiv:2305. 12744, 2023b.
Richard Yuanzhe Pang, Alicia Parrish, Nitish Joshi, Nikita Nangia, Jason Phang, Angelica Chen, Vishakh Padmakumar, Johnny Ma, Jana Thompson, He He, et al. Quality: Question answering with long input texts, yes! In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 5336-5358, 2022.
Baolin Peng, Michel Galley, Pengeheng He, Hao Cheng, Yujia Xie, Yu Hu, Qiuyuan Huang, Lars Liden, Zhou Yu, Weizhu Chen, et al. Check your facts and try again: Improving large language models with external knowledge and automated feedback. arXiv preprint arXtv:2302.128133 2023.
Adam Poliak. A survey on recognizing textual entailment as an nip evaluation. arXiv preprint arXiv:2010.03061, 2020.
Ofr Press, Muru Zhang, Sewon Min, Ludwig Schmidt, Noah A Smith, and Mike Lewis. Measuring and narrowing the compositionality gap in language models. arXiv preprint arXwv:2210003350, 2022.
Guillem Ramírez, Alexandra Birch, and Ivan Titov. Optimising calls to large language models with uncertainty-based two-tier selection, 2024.
Siva Reddy, Danqi Chen, and Christopher D Manning. Coqa: A conversational question answering challenge. Transactions of the Association for Computational Linguistics, 7:249-2666 2019.
Nils Reimers and Iryna Gurevych. Sentence-bert Sentence embeddings using siamese bert-networks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 11 2019. URL https: /ttas:///arxiiii /awpivivviv../////.//a./////aava wtwp////////oooooee//////aaeeee/ 1908. 10084.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint 1xXiv:1100001188, 2019.
Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Q Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. arXiv preprint arXin:2207.07066 2022.
Wilko Schwarting, Javier Alonso-Mora, and Daniela Rus. Planning and decision-making for autonomous vehicles. Annu. Rev. Control. Robotics Auton. Syst., 1:187-210, 2018. URL https::/aapp www.nticsmoolle.eee..eeeeeceeeee org/ I 64853900.
Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. SCROLLS: Standardized CompaRison over long language sequences. Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 12007-12021, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. URL https: httpp:://acaatthaaaeeeae http://ccaaanttoooy.oo org/2022. emnlp-main. 823.
Uri Shaham, Maor Ivgi, Avia Efrat, Jonathan Berant, and Omer Levy. Zeroscrolls: A zero-shot benchmark for long text understanding, 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Erie Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint 13977 2023.
13

---

Gladys Tyen, Hassan Mansoor, Peter Chen, Tony Mak, and Victor Cärbune. Llms cannot find reasoning errors, but can correct them! arXiv preprint arXiv:231 1.085 16, 2023.
Marija vSakota, Maxime Peyrard, and Robert West. Fly-swat cannon? cost-effective lan- guage model choice via meta-modeling. ArXiv, abs/2308.06077, 2023. URL htw.s:/tp///ec /api. semantticccoll. wwwp:rorggoopgoree ID:
Yixuan Weng, Minjun Zhu, Shizhu He, Kang Liu, and Jun Zhao. Large language models are reasoners with self-verification, arXiv preprint arXiv:22/2.0956/ 2022.
Yixuan Weng, Minjun Zhu, Fei Xia, Bin Li, Shizhu He, Kang Liu, and Jun Zhao. Large language models are better reasoners with self-verification CoRR, abs/2212.09553 2023.
Chenyang Wu, Guoyu Yang, Zongzhang Zhang, Yang Yu, Dong Li, Wulong Liu, and Jianye HAO. Adaptive online grccingg-gided search for POMDPs. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL https::////o / httpt:::///peeeeieeeeee net forum weerummiii-iaee -iteenddd-iiiaai 21V2VTTB0VVVVV
Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. React: Synergizing reasoning and acting in language models. In The Eleventh International Conference on Learning Representations, 2022.
Shuyan Zhou, Uri Alon, Frank F Xu, Zhengbao Jiang, and Graham Neubig, Docprompting: Gen- crating code by retrieving the docs. In The Eleventh International Conference on Learning Representations, 2022.
Wangchunshu Zhou, Canwen Xu, Tao Ge, Julian McAuley, Xu, and Furu Wei. Bert loses patience: Fast and robust inference with early exit. Advances in Neural Information Processing Systems, 33: 18330-18341, 2020.
Banghua Zhu, Ying Sheng, Lianmin Zheng, Clark W. Barrett, Michael I. Jordan, and Jiantao Jiao. On optimal caching and model multiplexing for large model inference. ArXiv, abs/2306. 0220003 2023. URL https:
appi. semant icscholar
259075212.
K.J Åström. Optimal control of markov processes with incomplete state information. Journal of Mathematical Analysis and Applications, 10(1):174-005, 1965. ISSN 0022-247X. doi: https: 1oop.ooo/1..11100002222441116611 sttinee//rrrraee//11/22227777600
URL https: www. wttp://wwwwneeeeeeeetececo
14

---

A How effective is few-shot self-verification?
A.1 Quantitative Analysis
One notable contribution of this work is the concept of few-shot self-verification of outputs. Self- Verification, especially for reasoning problems, poses its own set of challenges; however, our setup has unique advantage: the capacity to utilize context to validate answers. For instance, model can identify factual inaccuracies in the answer or discern apparent contradictions that might not have been evident during the initial response. But does this advantage translate to effective self- verification in practice? As depicted in Figure 12, for all but once few-shot self verfiiciiiin succeeds in accurately identifying correct examples by assigning them higher probabilities across all other datasets. Specifically, the y-axis the graph denotes the precision of identifying the correct responses (defined as F1 > 0.1). For all but QUALITY and QASPER higher verification probability, always higher precision. Even for QASPER, barring 2/9 outliers, the trend generally holds.
A.2 Qualitative Analysis
Representative Examples from our qualitative analysis are shown in Tables 2, 3, and 4.
Attribute
Content
Context
Which country grows the most tea? Shen sipped it. He liked it. He drank it all. Shen was proud of his new drink. Truncated for brevity
Question
When did he take it?
Generated
An-
He took it in 1990.
swer
Gold Answer
about 2737 B.C
Verifier Output
The context does not mention that he took tea in 1990. Verification Decision: The Al generated answer is Incorrect.
Table 2: An example where the generated answer is verifiably incorrect, citing an unsupported date (1990), and the verifier successfully catches this discrepancy.
Attribute
Content
Context
Fort Lauderdale, Florida (CNN) - Just taking a sip of water or walking to the bathroom excruciatingly painful for 15-year-old Michael Brewer, ... He suffered second- and third-degree burns over about two-thirds of his body, according to the hospital's 's associate director, Dr. Carl Schulman. Trancated for brevity
Question
How long before he recovers from his injury?
Generated
An-
he will never recover from his injury.
swer
Gold Answer
He has lifelong recovery.
Verifier Output
The context clearly states that the teen is recovering from his injuries. Verifi- cation Decision: The AI generated answer is Incorrect.
Table 3: The example demonstrates a challenging case for the verifier. The generated answer is rephrasing of what's 's mentioned in the context, however the verifier fails to catch it.
FrugaIGPT vs. AutoMix at different levels of data availability Figure 7 demonstrates the performance dynamics of AutoMix x and FrugaIGPT with varying validation data sizes. Notably, our method significantly outperforms FrugalGPT with limited data, despite the latter's domain- specific training and zero verifier cost. However, training data increases, FrugalGPT narrows the
15

---

Attribute
Content
Context
Pruning keeps the plant only three or four feet tall. This is an easy height for tea picking. Green tea is steamed right after the leaves are picked. Green tea its healthy qualities. For example, it may prevent heart disease. How did we get tea bag? The answer: by accident. Truncated What did they do to green tea after picking it?
Question
SLM Answer
steamed.
Gold Answer
prune it
Verifier Output
The context mentions that green is steamed after being picked. Verification Decision: The AI generated answer is Correct.
Table An instance where the verifier deems the answer correct, and the gold label was incorrect. The verifier is able to correctly infer that the context mentions tea being steamed after picking.
performance gap by leveraging its domain- sppecifc training. This pattern indicates that Aut oMix provides a particularly advantageous solution in real-workh scenarios where data may be scarce.
30
Macro Averaged IBC Lift
20
10
Automix w/ POMDP
Frugal
-10
HybridLLM
50 100
250
500 Train Size
1000
Figure 7: Comparison of AutoMix with FrugalGPT over varying Training Dataset Size. Despite zero-cost verifier and domain-speciff training. FrugalGPT underperforms AutoMix Aut oMix is especially useful for limited data settings, with higher gains visible when dataset size is less than 1000.
A.3 POMDP-based Router in AutoMix
router should direct a query to larger Language Models (LMs) only when the performance gap justifies the cost-quality tradeoff. Given the inherent uncertainty in the true state of system performance, which remains unobserved, we formulate the router as Partially Observable Markov Decision Process (POMDP) [Åström, 1965]. POMDPs are particularly well-suited for scenarios where observations, such self-verification probabilities, may not be entirely reliable.
Recall that (Section POMDP characterized by the tuple (S, A,T, R, O). our application, we define these components follows:
16

---

States (S): The state space represents the current selected LM, (stage) and performance met- rics (e.g., accuracy or F-score) of various LMs on a data point. We denote this as = (i, LLM, Per fLM Per fLM),, where is the index of the current LM and PerfLM, the performance of the j-th LM. The terminal state is represented by 5 = (N,...), where N is the total number of LMs.
The action space includes either retaining the current LM's answer (a LM,) or routing to one of the larger LMs (denoted by Roules.j > i). For LMs except LMN, routing to an L., implies performing both answer generation and self-verification.
Actions (A):
Transition Function (I'): The transition probabilities are deterministic since the true performance of language models on a question is static, albeit unknown. However, the first dimension in the state vector, representing the current LMM, changes to if action Route, is taken or N if the current LM's answer is retained. As the stage state always increases and is bounded by N, our problem has finite horizon.
Reward Function (R): The objective of our POMDP is to optimize a reward function R = P- C.C, where P is overall performance, C' is overall cost, and is a tuning parameter that balances these two criteria according to user preferences.
Observations (?2): Observations come in the form of verifier outputs from LM,, enabling the POMDP to ascertain its belief state b, which is a probability distribution over S.
Observation Function (O): The observation probabilities P(o]s) indicate the likelihood of observing (verification output) given state s. These probabilities are crucial for defining the POMDP model. For instance, high verifier confidence might suggest that the current LM's 's performance Per fLM, is sufficiently high, reducing the necessity switch to a more costly LM.
To estimate the observation probabilities, we first consider the discrete case. In this scenario, we can directly compute P(ols) using:
P(ols) Eis 1s; and ej = o}
Σx 11s, = s}
where represents the set of state-observation pairs from the training data.
However, in our continuous state space, we need to adapt this approach. We employ Kernel Density Estimation (KDE), a non-parametric method for estimating probability density functions. Instead of directly estimating P(ols). we first learn joint distribution P(S, O) using KDE. For this, we draw KDE at each training point:
1 nh
X₁ h
fx(x)
K
where K is Gaussian kernel function, h is the bandwidth, and X, are the sample points. The conditional probability P(ols) is then computed as:
P(s,0) P(o|s) P(s)
where P(s) is computed similar P(s, o) These observation probabilities estimated directly the training set by drawing KDEs for each training point.
By leveraging this POMDP formulation and the KDE-based observation model, our router make informed decisions about when to route queries to larger language models, effectively balancing performance and cost considerations in the face of uncertainty.
A.3.1 Solving POMDP
Once all the different components of our POMDP are defined as shown above, finding an optimal policy requires suitable solver. An optimal policy in this context is strategy n² B A.
17

---

where B is the set of all possible belief states that maximizes the expected cumulative discounted reward. However, our formulation presents unique challenges and opportunities that traditional POMDP algorithms struggle address efficiendly. Our POMDP models continuous state space in an (N + +1)++++00000 -)-dimensional space, where N is the number of language models. This continuous nature poses difficulties for many traditional POMDP algorithms, such as point-based value iteration methods, which are typically designed for discrete state spaces and can become computationally intractable in continuous domains.
Nonetheless, our particular formulation has several desirable characteristics that distinguish it from traditional POMDPs: 1. Finite and small horizon: Our problem has limited number of decision stages (maximum = N), bounding the depth of any decision tree we need to explore. 2. Finite and small number of observations: The verifier outputs, which serve our observations, are discrete and limited in number (9), allowing enumeration of all possible observation sequences. 3. Directed Acyclic Graph (DAG) structure: In our formulation, particular state is encountered only once, and the problem progresses unidirectionally through the stages,
These characteristics allow start by presenting a simple yet effective algorithm to solve our POMDP, which leverages these properties to achieve decent speed and accuracy. The approach combines particle filtering for belief state representation with a tree-based search strategy, as detailed in the following section.
detail more tailored and efficient solution approach. We first
POMDP Solver We start with description of simple POMDP solver to understand the intuition behind how to solve POMDPs with our proposed formulation. Because of the multidimensional continuous state space, we cannot represent all possible belief states explicitly, A common method to approximate belief states such cases is the use of particle filters.
Specifically, each particle represents a possible state of the system, and the belief state is represented as a set of such particles. Mathematically, if we denote the set of particles by [M Ji=1 where M is the number of particles, the belief state b at any time is approximated
b(s) ~
8(s )(1)
where 8 is the Dirac delta function.
To solve our POMDP, we construct decision tree that represents the possible sequences of actions and observations. This tree consists of two types of nodes: action nodes, which represent the selection of language model, and observation nodes, which represent the possible verifier outputs following an action. After each observation node, we update the belief particles using the observation probabilities, refining our estimate of the system's state.
Assuming that we can explore and expand all nodes for small trees, we need to compute the value of each node, We can compute the reward directly from the reward model for terminal nodes. For any other node, the value (bb is defined as the expected reward of the best action, which can be recursively computed using:
(b max R(b,a) Σ | b.
where D' is the updated belief state after observing following action a, and 7 the discount factor set to I in our case.
One final step remains: how to update the belief states. This is done using particle filtering, as described next. In particle filtering. belief states are updated by sampling from the current belief state based on the observation probabilities.
Mathematically, unweighted particle filtering updates can be described as follows. Given set of (0).4
particles {st It representing the belief state at time f, and an observation o,, the updated set of
particles (5/=)
(4) M obtained by:
1. For each particle sf sample new state $t41 according to the transition model P(s (8)
(4)
8 .a).
18

---

2. Weight each new particle by the observation probability P(oe ((41)-
3. Resample M particles from the weighted set form unweighted particle set.
The algorithm can be summarized as follows:
Algorithm 1 Unweighted Particle Filtering Update
1: Input: Current particles {s, Ji=1 action az, observation o,
2: Output: Updated particles ($
AM
3: for each particle
(1) do
4: Sample $141
(4)
P(s++1 1% a))
5: Compute weight P(o, s(+1)
6: end for
7: Resample M particles from {0 we ) to obtain unweighted particles (0 (82+1) )
Given our small and finite horizon problem, we set the discount factor as 1. Finally, we can simulate all possible scenarios and learn deterministic policies offline because of finite number of observation sequences. Thus, POMDP is reduced to a simple lookup table at inference, and the inference cost is negligible.
Using better POMDP solvers
Note that the simple solver design introduced in the previous section is sufficient and fast enough for our purpose. However, to handle more complicated scenarios in future works, we use a more advanced and recently introduced AdaOps POMDP solver [Wu et al., 2021]. AdaOps, which stands for Adaptive Online trcking-guidee Search for POMDPs, offers several ad advantages, such as efficient handling of larger state spaces and improved scalability through adaptive online search techniques. Additionally, it incorporates packing-guie -uidde exploration strategies to focus computational resources on the most promising parts of the belief space, leading to more accurate policy estimations in complex environments or where search and state space is large. However, given our simpler formulation, both AdaOps and the simple algorithm would approximately converge the same solution as the number of particles used in belief state representation increases.
B Additional Details on the Experimental Setup
For evaluation, we utilize the validation sets from Shaham et al. [2022] for QASPER, and QUALITY, and use the prompts from Shaham et al. [2023]. For COQA, MUTUAL, and DIPLOMAT, we employ its validation split and adapt the QUALITY prompt. For consistency, 1000 instances are sampled from the validation set of each dataset. Regardless of the dataset, identical input prompts are dispatched to both SLM and potentially LLM, ensuring consistent input processing costs. The output length is fixed in multichoice datasets like CNLI and QUALITY, and the brevity of responses in other datasets allows us to assume uniform output processing costs. We use greedy decoding (temperature 0) and draw single sample for both the SLM and LLM. For verification, we generate eight samples per question (temperature = 1), which has negligible cost owing to the large context. In Figure 6, we normalize ABc by a scaling factor such that for all datasets, the maximum is set to 1.
For running our experiments, we use LLAMA2- 13B and GPT-4 models from huggingface We use vllm [Kwon et al., 2023] for hosting models for inference.
Cost Ratio:
We have considered cost ratio of 1:100 between GPT-4 and GPT-4, reflecting the
API price disparity between the models, which stands at S0.225 for LLAMA2- 1- vs $30 for GPT-4 per IM tokens at the time of writing. Additionally, for self-verification purposes, samples. is important to note, however, that the cost of generating 8 samples is negligible compared to the cost of single sample, primarily because the major cost driver is the length of the context
generate 8
MModels available at: https:://huggingtace co /tt////eae//// a Lmaa/ LamaaLLaa Llama-2 2-13b bf and ht hpsps: //hgggginggaceeeee co/meta//aa./eeeeeee
19

---

(e.g., generation is 60 times and 50 times smaller for QASPER and QUALITY, respectively than base context). Therefore, invoking the verifier 8 times is considered equivalent in cost to calling it once. Furthermore, in Section 6, we explore different cost ratios and observe that, even with a ratio low as 1:25, oMix begins to yield non-trivial gains across most datasets.
Datasets We experiment with a diverse set of datasets: i) QASPER [Dasigi et al., 2021]: Question answering over research papers; ii) QUALITY [Pang et al., 2022]: Multiple-choice questions (MCQ)
on long articles and stories; iii) COQA [Reddy et al., 2019]: Conversational comprehension requiring coreference and pragmatic reasoning: iv) MUTUAL [Cui et al., 2020]: Multi-turn dialogue reasoning (next response prediction); v) DIPLOMAT [Li et al., 2023]: Pragmatic identification and reasoning questions on multi-turn dialogues. QUALITY, QASPER and MUTUAL are licensed under CC BY, while DIPLOMAT is licensed under CC BY-NC-SA. COQA uses multiple licenses for several splits and are detailed in https: : st an github. io/ loqa/o We use the F1 score for QASPER and COQA, and accuracy for the remaining datasets. To manage imput complexity, retain a context subset (max 3500 tokens) retrieved using the question as a key. Retrieval is performed with all aal--inni -M---------------- L6622 sentence embedding model [Reimers and Gurevych, 2019].
We utilize the validation sets from Shaham et al. [2022] for QASPER and QUALITY, and use the prompts from Shaham et al. [2023]. For COQA, MUTUAL, and DIPLOMAT, we employ their validation splits and adapt the QUALITY prompt. Regardless of the dataset, we provide identical input prompts to both SLM and LLM to ensure consistent input processing costs. The output length is fixed in multi-choice datasets like QUALITY, and the brevity of responses in other datasets allows us to assume uniform output processing costs. We greedy decoding (temperature 0) and draw single sample for both the SLM and LLM. For serffierii,,,tion, we use temperature=0.7 and draw 8 samples.
C Expanding AutoMix to Three-Models
crucial use-case of AutoMix to route between multiple models of varying scales in terms of cost and performance. In previous sections, we considered a two-model scenario, where SLM was 2 orders of magnitude smaller than the LLM. this section, we evaluate how well AutoMix performs if we consider third medium-sized language model (MLM). Specifically, we aim to address several questions: 1.) Does AutoMi x perform better than baselines in the n-model scenario, 2.) Can Aut oMix automatically learn to skip models if they are non-performant 3.) How beneficial is POMDP router in Aut oMix, 4.) What are the failure cases of AutoMix in the 3-model scenario?
To answer these questions, we employ LLAMA2- ILAMA2-1BB/MISRRAN 3BB.///TAAA-77//GTTTT.T
as the SLMs, LLAMA2- 70B as the MLM, and GPT-4 as the LLM. We consider the following baselines:
FrugalGPT: Frugal GPT proposed using a cascade style routing for models, where the query is first sent to the SLM, and if the confidence is below a threshold, it is sent to the MLM, and so on. The thresholds are tuned on the validation set, so as to maximize the performance across different cost regions. Note that Fr TrugalGPT trains N - 1 separate verifiers for each pair in a sequence of models. Similar to previous sections, we consider the cost of verifier as O.
Union AutoMix:
Union Auto is simple baseline, where select between the two- model variants out -MMMMMMMMM- AutoMixsIas XSLM-LLLLLLLL and MuttoMiMMM----MMM
depending on the cost requirements specified by the end-user. For instance, if the desired average cost is less than that of the MLM, AutoMixsLM-MMAN may be employed, whereas AutOMiXMIM-IEA LLM or Aut MauxxxxxMM-L... is utilized for cost regions exceeding that of the MLM. Precisely, for each cost region in validation set, the best-performing variant (highest ABCC) is chosen: maZ region (Per faaatoiixs MMMM Perffuuuose Jauzoxixm -L4Mr Per TaatcmisssMMMLL. Note, the formula easily generalizes to larger N by considering all possible pairs.
Chained AutoMix: Chained AutoMix is baseline where we chain the two-model
The query first goes to the SLM, and an Aut oMi XSLM - ML. decides between reporting the SLM answer or routing to the MLM. the latter case, a second AutoMi XML M-LLM repeats the procedure using the MLM and LLM models.
variants AutoMi X MeM and AutoMi
20

---

C.1 Extending IBC N=3 models
In order to extend IBC metric to > 2 models, is important to consider the cost region being considered. For instance, for N 3 consider two separate cases: when the SLM-MLM-LLM curve is convex, and 2) when the curve is concave. In the convex case, choosing between the MLM and SLM in low-cost regions is advantageous, while it is beneficial choose between the MLM and LLM in high-cost regions. Accordingly, the suitable (more competitive) IBC curve is selected for evaluation accordingly. However, in the second case, when the curves are concave, it would be more favorable to choose between the SLM and LLM and completely ignore the MLM, as in terms of incremental performance per cost, it consistently presents disadvantage. Thus, the BBCSLM-LLA, is chosen for evaluation throughout. Note that for general N models, for each cost region, N2 combinations need be considered.
With the baselines, and evaluation metrics well defined, we next aim to answer each of the four answers mentioned above:
C.2 Results of Automix with Three Models
C.2.1 Does Aut oMix outperforms baseline methods?
Figure 8: AutoMix outperforms baselines in = 3-model scenario. Representative examples shows QASPER, DIPLOOMT,, COQA datasets, where AutoMi x consistently outperforms the baselines: FrugalGPT, Union AutoMi x, Chained AutoMi x across all cost regions considered.
Figure presents results on three representative cases all from different datasets and SLMs. Across, all datasets, Aut oMix consistently outperforms the basclines, F ruugalGPT, Union Aut oMix, and Chained Aut oMix, across cost regions considered. Since FrugalGPT and Chained Aut oMix are cascade style routing, they are unable to directly route from SLM to LLM, and always first have invoke the MLM, thus demonstrating poor performance. While Union Aut x is able to alleviate this issue, it still needs to always invoke the MLM model and its costly Verifier to make decision, while Aut oMix x is able to use information available from SLM verifier, to decide whether to route to MIM or LLM. Even in cases when routing to MLM takes place, POMDP in AutoMix leverages information from both SLM and MI.M verifiers by updating belief, providing it with more nuanced decision decision-making capability.
21

---

C.2.2 Can M skip distractor models?
0.80
GFL4
0.75
Performance (F1 Score)
0.65
Mistral-7B
0.60
25
50
75
100
125
150
175
200
L1AMA-70B
Computational Cost
Figure AutoMix automatically skips distractor MLM models that underperforms the SLM. However, F FrugalGPT and Chained oMix struggle to discern this, resulting in poor performance across the board. Results on MUTUAL dataset with MISTRAL- SLM.
As robustness, we now consider scenario, where MLM underperforms the SLM. Such a scenario common in practice since models come from different families, trained on different kinds of data, and it is possible that one model is better than the other despite the scale. Figure 9 considers two/three scenarios, where MLM underperforms the SLM. Note that in such cases, FrugalCPT performs significantly worse since it always routes to the MLM for higher cost regions, thus wasting resources. Thus FrugalGPT requires an extra consideration on validation set, weed out such distractor models. However, AutoMix x is able to learn the irrelevace of MLM, and learns to route directly from SL.M to L.M,, as needed. The performance of AutoMi similar to Union AutoMi x in such cases since the MLM is never used.
22

---

C.2.3 Should Low Performing MLMs be ignored?
Performance (F1 Score)
Methods
0.78
Automixs
Chained Automix
0.77
Union Automix
Frugal
GPT3.5
GPT-3.5 GPF4
LLaMA-70B
0.75
GPT4
IBCaPT
0.74
IBCr0а/GPa
0.73
50
150
Computational Cost
LLaMA-70B
Figure 10: AutoMix can leverage even worse performing MLM by identifying questions, where MLM better than SLM and making use of MLM veirififcation probabilities. All other baselines, are unable to leverage this information, and perform worse than AutoMix
While in the previous section, we established that low-performing MLMS should ignored, this may not always be the case. This is because there may be certain types of questions where MLM indeed better than SLM, and if, through self-verification probabilities, it is possible to exploit such information, it can further improve the cost-quality of Aut oMix. Figure 10 demonstrates one such case. Specifically, Aut oMix x identify that low verification probabilities on SLM implies that MLM is likely more better than SLM. Thus, it routes to MLM models appropriately. Further, the combined verification information from SLM and MLM, allows Aut oMix to make a more nuanced decision, and thus outperforms all other baselines considered. So much so, oMix is able to identify patterns that result in statistically overperforming even the LLM even at roughly half the cost. The experiment, along with the previous ones, demonstrates the efficacy and usefulness of the POMDP router Aut.Mi x, as automatically able to handle a wide variety of cases and consistently provides the best cost-quality tradeoff.
C.2.4 Does Aut oMix always provide significant gains?
Autonits
Tertemiico
Chvinw Artmnts
Uniss Autozic
Frogel
Natrrl 75
Compratilooal CARE
Competational Cool
Figure 11: Aut oMix perform similarly baselines without any statistically significant gains in the two representative examples, owing relatively poor self-verification.
23

---

In previous subsections, considered cases, where AutoMix significantly better than Fruga 1 GPT and often the other variants of AutoMix. x. However, this section, we consider a couple of cases, where Aut oMix do provide significant improvements. Figure 11 represents the two cases. Specifically, we note that in both cases, Aut oMix is similar to F FrugalGGa GPTT. Further, the advantage over lines negligible or even negative. We qualitatively analyze and find that the self-verification by MLM in such scenarios is not particularly useful, and thus, POMDP is not able to provide significant gains. We believe future work should address such challenges by considering contextual information or other forms of verification to improve the performance of our methods in such cases.
Percentage of Correct Responses by Verifier Probability
1.0
80-
Parcentage Cood Elamples 18)
30
20
diglomet
mutual
Percentage of correct responses across distinet verifier probability bins for LLAMA2-133 13B as SLM. Each bin represents a range of verifier probabilities and the corresponding accuracy of the responses within that probability range across various datasets. Notably, for all datasets, excluding QUALITY and QASPER, a higher verification score generally corresponds larger proportion of correct examples, indicating that the verifier is, to extent, capable of discerning the reliability of responses generated by itself. We use a router behaving like a meta-verite verifie to get around these noisy predictions.
Figure 12: Verifier Probability and Correctness:
Figure 13, 14, 15, 16, 17 shows results for N scenario for AutoMix along with baselines.
Poga
Figure 13: AutoMix in N = 3-model scenario. Dataset: COQA.

---

Lakene
Figure 14: AutoMix in
3-model scenario. Dataset: QASPER.
Figure 15: AutoMix in N 3-model scenario. Dataset: QUALITY.
Rraad
Depplcaben
Figure 16: AutoMix in N = 3-model scenario. Dataset: MUTUAL.
Figure 17: AutoMix N 3-model scenario. Dataset: DIPLOMAT.
D Additional Results
D.I Main Results
Figure 19, 18, 4 shows the results of AutoMix for different choices of SLM. Notice, for 13/15 configurations, Aut <Mix-POMDP is clearly better than all other baselines throughout different operating costs. Even in the remaining 2 cases, AutoMix is competitive and almost equivalent to the baselines. We notice similar trends from Section 5.3 for LLAMA2- 13B and GPT-3.5 as well.
D.2 AutoMix effective across different Cost-Ratios
In this section, we furhter compare how, AutoMix compare to baselines for different cost-ratios. Figure 20 demonstrates that at the API based cost ratio, AutoMix outperform all baselines. Even at 20:1 ratio (an unlikely, pessimistic scenario for AutoMi x), AutoMi x outperforms baselines, albeit with smaller margins. Conversely, 2000:1 ratio, AutoMi shows even higher improvements. The result demonstrate the robustness of Aut oMix to ever-changing costs of models and APIs.
25

---

DIPLOMAT
MUTEAL
QUALITY
00.8
0.74
0.75
0.5
0.7
0.68
0.60
0.6
0.64
0.62
0.8
50
100
50
100
50 100
COQA
QASPER
0.5
AutoMix +POMDP
Autokix +Threshold
0.58
FrugalGPT [Chen al., 2023]
0.56
0.5
HybridLLM [Ding et al., 2024]
LLAMA2- 13B (SLM)
0.54
GPT-4 (LLM)
112
Random Mixing
0.52
100
50
100
Figure 18: Main Results: performance (y-axis) vs, cost (x-axis) for different methods on the small and large ALAMA2-13/GPT-A POMDP based meta-verifier consistenly above the linear interpolation of SLM-LLM, signifying higher incremental benefit per unit cost (IBC).
D.3 Out-of-Domain Generalization of AutoMix
The POMDP router introduced in autooix x does not assume any specific task dataset characteris- tics to work, as it relies only on self-verificatio confidence as input. Thus, our POMDP router is generalizable to various datasets and tasks. To further demonstrate this, we evaluate out-of-domain others. Specifically, we train the router
generalization by training on one dataset and evaluating on one of the five datasets and evaluate it on the other four. We repeated the experiment for all five datasets and three SLMs. The results in different settings, as shown in the Table 5, indicate that our POMDP router consistently outperforms both FrugalGPT and HybridLLM. The design of the POMDP router, along with these results, highlights the strong generalizability of our proposed method.
Mistral-7b
LLama-13h
GPT-3.5
Automix
28.3
31.5
70.9
Frugal
12.5
0.0
14.3
Hybrid
2.4
-2.8
7.6
Table 5: Out-of-domain generalization across various SLMs for different methods. Scores are averaged across five datasets. AutoMix significantly outperforms other methods.
E Few-Shot Prompts
E.1
Verifier Prompts
26

---

DIPLOMAT
MUTUAL
QUALITY
0.66
0.75
0.78
0.77
0.6
0.78
0.75
20
30
40
00
10'0
20
30
40
50
50
100
COQA
QASPER
Autokix +PMMD
AutoMix +Threshold
0.55
AutoMix +SelfCoosisteney
0.55
FrugalGPT [Chen al., 2023]
HybridLLM [Ding et al., 2024]
LLAMA2-133 (SLM)
0.3
GPT-4 (LLM)
0.35
Random Mixing
20
30
40
10.0
30 40
50
performance (y-axis) vs. cost (x-axis) for different methods on the small and large 3//-3.5/ POMDP based router is consistenly above the linear interpolation of SLM-LLM, signifying higher incremental benefit per unit cost (IBC). AutoMix is the best performing method consistently across all datasets, often by large margins.
Figure 19: Main Results:
DIPLRRST
6.78
mm
Astcoltix &OMDD
Frige GITT [Chen al, 2003]
6.78
0/0
-Hyyisst LMM 12ing eral. 2004
11.0052- 13n (SLM))
0.74
074
Raaddon Misizg
OPn- (LLM)
180
1,000
1,000
2,000
(a) Cost Ratio: 20:1
Cost Ratio: 200:1
(c) Cost Ratio: 2000:1
Figure 20: We compare AutoMix to baselines various LLM SLM cost ratios. For all considered cost-ratios, AutoMi. outperform baselines with varying degree of margins.
Story:
(relevant parts of the story)
(instruction)
Question: (question)
Answer:
Listing Task Prompt. We experiment with long-context reasoning tasks, which require answering questions from stories, legal contracts, research papers, and novels.
27

---

Context: [context]
Question: (question)
AI Generated Answer:
avenerate__anseer
Instruction:
Your task is to evaluate the Generated Answer + correct, based on the provided context and question. Provide the judgement and reascning for each case. Choose between Correct or Incorrect.
"valuation::!
Listing 3: Verification Prompt. The verification process is framed as a natural language entailment task, where the model determines the validity of the genel-generated answer with respect to the context and question.
28

---

Context: The manuscript, discovered in 1980 in a dusty attic, turned out 4 to be a lost work of Shakespeare, \n Question: whose lest work was discovered in a dusty attic in 19 Al Generated Answer: Shakespeare'n task
Instruction:
evaluate if + correct, based on the provided context and question. Provide the judgement and reasoning for each case. Choose between Correct or Incorrect. An
Generated Answer is
Evaluation: The context specifically menticns that lost work of Shakespeare was discovered in 1980 in dusty attic.
Verification Decisien: The AI generated answer is Correct.
Context:
The celestial event, known as the Pink Moon, is unique to the month of April and has cuitural significance in many indigenous tribes. \n
Question: In which month does the celestial event, the Pink Moon, 4 on
AI Generated Answer: July'n
Instruction: Your task is to evaluate if the AI Generated Answer is correct, based the provided context and question. Provide the + judgement
reasoning for each case. Choose between Correct or + Incorrect. An
Evaluation: The context clearly states that the Pink Moon is unique to the month of April.
Veri Decision: The AT generated answer is Incorrect.
(truncated examples)
Context: (context } An Question: (question)\n An AI Generated Answer: ggenerated_anseer
Instruction:
Your task is to evaluate if the AI Generated Answer is correct, based en the provided context and question. Provide the judgement and reasoning for each case. Choose between Correct Incorrect.
Evaluation:
Listing 4: Few-Shot Verifier Prompts: 3-shot verifier prompt for evaluating the correctness of SLM's answer. The same prompt used for all datasets.
29

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All claims are correctly reflected, and demonstrated in abstract, introduction methodology, experimental setup, results, limitations appendix sections.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should refleet the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't It acknowledged judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
30

---

Justification:
Paper does not include any theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced
statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Complete details are mentioned in Section B,
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model limited in some way (e.g., registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
31

---

Answer: [Yes]
Justification: Yes code is released in supplementary material, all data sources and sufficient instructions are mentioned in Appendix and Main Paper. We use official data, from official data sources.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https:/ /niissp..... cc/
public/ /uwid/ggusssseeeeee Saliiees/CoeeSaaae CooeeSuoiei ssionPolicg for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to rum to reproduce the results. See the NeurIPS code and data submission guidelines (ht tpa //nips wttpp//ppppccccctpcccccccccccclt /eaiii//iiiiie des Soodee ProdeSubmat as onPolloooiiooooo ooloiycyyy. icy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: details are mentioned in Appendix B, 5
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Due to large computational costs, it is not feasible to run statistical tests. For sampling experiment, error bars (standard deviation) are reported in Figure
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
32

---

The assumptions made should be given (e.g., Normally distributed errors). should be clear whether the error bar is the standard deviation the standard error of the mean,
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: All details mentioned in Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /ttpuruurpp...e/ http:///wcccce/tteccceoeeee.e.e.
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: We propose optimizing inference cost and performance for deployed large language models. Thus there is direct societal impact of our work.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
33

---

The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification:
All models already deployed, and data is already publicly available.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All citations made, and license is mentioned in Appendix.
Guidelines:
existing assets. The authors should cite the original paper that produced the code package or dataset.
The answer NA means that the paper does
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, papers aw! thcode com da set. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
34

---

If this information is available online, the authors encouraged to reach out the asset's ss creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[NA]
Justification:
Guidelines:
The answer NA means that the paper does not release assets. Researchers should communicate the details the dataset/code/moodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
35


Reviews:
[{"review_id": "5mLauitknI", "reviewer": "Reviewer_zQB8", "summary": "The paper presents AutoMix, a method designed to optimize the use of large language models (LLMs) by strategically routing queries to different models based on performance and cost considerations. AutoMix relies on a few-shot self-verification mechanism which involves the estimation of correctness of a smaller model's outputs. Thereafter, a router model determines whether to query a more powerful, larger model to generate the output or accepts the smaller model’s outputs. This router is POMDP-based, trained with a reward function that couples both accuracy and computational efficacy (penalizing use of larger model). Experiments conducted on five language models across five challenging datasets show that AutoMix significantly reduces computational costs while maintaining comparable performance levels, outperforming other computational efficient baselines.", "questions": "Please discuss the weakness mentioned above.", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The proposed approach is sound, intuitive and works even in the black-box setting.\nEmpirical results from five different dialogue and context-grounded reasoning tasks.show that the proposed approach (AutoMix) outperform existing computationally efficient LLM approaches. \nThe paper is well-written and easy to follow.", "weaknesses": "There is a need for a RL-trained (POMDP) router agent. The generalizability of this router to other tasks/data is unclear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review the paper and provide feedback! We address your concern here:\n\n### Generalizability of POMDP router to other tasks/data\n\nThe POMDP router does not assume any specific task or dataset characteristics to work, as it relies only on self-verification confidence as input. Thus, our *POMDP router is generalizable to various datasets and tasks.*\n\nTo further demonstrate the generalizability of our router, we evaluate out-of-domain generalization by training on one dataset and evaluating on others. Specifically, we train the router on one of the five datasets and evaluate it on the other four. We repeated the experiment for all five datasets and three SLMs. The results in different settings, as shown in the table below, indicate that our POMDP router consistently outperforms both FrugalGPT and HybridLLM:\n\n|          | Mistral-7b | LLama-13b | GPT-3.5  |\n|----------|------------|-----------|----------|\n| Automix  | **28.3**   | **31.5**  | **70.9** |\n| Frugal   | 12.5      | 0.0       | 14.3    |\n| Hybrid   | 2.4       | -2.8     | 7.6    |\n\n\nTable 1: Out-of-domain generalization across various SLMs for different methods. Scores are averaged across five datasets.\n\nThe design of the POMDP router, along with these results, highlights the strong generalizability of our proposed router.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "THxJwt1atr", "reviewer": "Reviewer_E2SA", "summary": "AutoMix is an approach designed to optimize the performance and computational cost of LLMs by selecting the appropriate model based on the difficulty of the task. This is achieved through a few-shot self-verification mechanism and a Partially Observable Markov Decision Process (POMDP) based router. The few-shot self-verification estimates the correctness of outputs from a smaller LM, and the POMDP router decides whether to route the query to a larger LM. Experiments demonstrate that AutoMix can reduce computational cost by over 50% while maintaining performance across various datasets and models.", "questions": "I think it would be interesting to just randomly pick up a LLM for each question in dataset, and to see the total cost and the quality of results. Have anyone done this before? Because there might be a case where, even without any filtering strategy, just putting several models with different performance and different cost out there is itself effective. And I think may be it’s possible to derive it mathematically？", "limitations": "The author pointed out the limitations  in the discussion.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The author introduce a novel combination of few-shot self-verification and POMDP-based routing to optimize the use of large language models. It is very thoughtful to make SLM itself both the answerer and the judge, and to consider it as the first candidate when selecting a suitable model. And the author successfully uses POMDP to meet the need of selecting a suitable model.\n- They provide a reasonable metric for cost-efficiency analysis. When the cost grows or when the performance decreases, the results of the metric will become small.\n- The methodology is sound, and the results demonstrate the effectiveness of AutoMix in reducing computational costs by over 50% while maintaining same performance. The experiments are robust, and the comparisons with several baselines prove the advantages of AutoMix.", "weaknesses": "- While the paper demonstrates the effectiveness of AutoMix across various datasets, it primarily focuses on dialogue and context-grounded reasoning tasks. It is not clear how well the approach would generalize to other types of tasks, such as factual question answering or commonsense reasoning.\n- The paper mentions that the POMDP can learn from as few as 50 examples, but it does not elaborate on the conditions or types of data required for effective training.\n- Although the paper compares AutoMix with strong baselines like FrugalGPT and HybridLLM, the comparisons could be more detailed in terms of why AutoMix outperforms these baselines in specific cases.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review the paper and provide feedback! We believe all your questions can be addressed within this discussion period, but we would love to provide further clarification if needed.\n\n---\n\n### Q: I think it would be interesting to just randomly pick up a LLM for each question in dataset, and to see the total cost and the quality of results. Have anyone done this before? \n\nWe thank the reviewer for raising the point. As noted in our paper (Lines 199-201), the dotted IBC line in Figure 4, “signifies the cost-performance curve that one would obtain if each data point was routed randomly between SLM and LLM.” We denote this baseline as “Random Mixing” in main results (Figure 4). \nAll our comparisons are done with respect to this baseline, and we demonstrate that while FrugalGPT and HybridLLM often struggle to beat this simple baseline, Automix consistently outperforms across all datasets and models.\n\n---\n\n### The paper mentions that the POMDP can learn from as few as 50 examples, but it does not elaborate on the conditions or types of data required for effective training.\n\n**We make no assumptions on the condition or type of data required for learning POMDP.** Our method does not require specific conditions or types of data for POMDP training. We evaluated it across 5 datasets and 3 SLMs (totaling 15 settings), encompassing a wide variety of domains and task types (e.g., reasoning, next utterance prediction, QA), consistently demonstrating strong performance.\n\n---\n\n### The comparisons could be more detailed in terms of why AutoMix outperforms these baselines in specific cases.\n\nWe note that our Self-Verification mechanism and POMDP router collectively contribute to the superior performance of AutoMix. For instance, in the Diplomat dataset with Mistral-7B as the SLM, our method achieved significantly higher self-verification accuracy compared to FrugalGPT and Hybrid LLM. \n\nFurther, unlike previous methods, POMDP automatically identifies cases where the query is very complex such that both SLM and LLM would give an incorrect answer. In such cases, while the previous methods would still route to larger LLM, POMDP would report SLM’s answer, saving cost. Further, in more than >2 model setup, POMDP can make more nuanced decisions, for example by considering the combined information from both SLM and MLM verifier. We will make this discussion clearer in the revised paper.\n\n---\n\n### It is not clear how well the approach would generalize to other types of tasks\n\n1. We thank the reviewer for bringing up the interesting point. We evaluate our method on a variety of models and datasets related to comprehension QA, and dialogue reasoning, which themselves offer a great diversity, such as answering questions from research papers (testing model's *factuality*), *reasoning* questions in multi-turn dialogue, next utterance prediction, conversational QA. The datasets are of varying difficulty, with LLM performance ranging from 35% to 95%, and different output formats: MCQs and open-ended generation. Despite the diversity, automix consistently outperforms baselines, demonstrating its generalizability. Further, our POMDP formulation is domain-agnostic and automatically extensible to other task types. \n\n\n2. Moreover, as we have already note in our limitations: “AutoMix is designed with a dialogue-related or context-grounded reasoning setup in mind for effective self-verification” we note that the inability to do self-verification for reasoning tasks is a limitation of current LLMs [1, 2,3, 4] rather than our method. On the contrary, our work demonstrates how to successfully use any self-verification for context-grounded reasoning tasks using a suitable routing function. Therefore, we consider our contribution to context-grounded reasoning tasks to be a strength.\n\n3. Further, we evaluate our method on an *additional common-sense reasoning dataset*: CICERO [5]. We find automix outperforms baselines by more than absolute 30% points. The results demonstrate potential superiority of Automix on datasets not evaluated in our paper. This additional result is provided in Table 2 of the general response.\n\n|          | Mistral-7b |\n|----------|------------|\n| Automix  | **66.4**   |\n| Frugal   | 32.1       |\n| Hybrid   | 19.7       |\n\n\n---\n\n\n\n\n**References**\n\nHuang, J., Chen, X., Mishra, S., Zheng, H.S., Yu, A.W., Song, X., & Zhou, D. (2023). Large Language Models Cannot Self-Correct Reasoning Yet. ArXiv, abs/2310.01798.\n\nTyen, G., Mansoor, H., Chen, P., Mak, T. and Cărbune, V., 2023. LLMs cannot find reasoning errors, but can correct them!. arXiv preprint arXiv:2311.08516.\n\nStechly, K., Valmeekam, K., & Kambhampati, S. (2024). On the Self-Verification Limitations of Large Language Models on Reasoning and Planning Tasks. ArXiv, abs/2402.08115.\n\nMadaan, Aman, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri Alon et al. \"Self-refine: Iterative refinement with self-feedback.\" Advances in Neural Information Processing Systems 36 (2024).\nGhosal, D., Shen, S., Majumder, N., Mihalcea, R., & Poria, S. (2022). CICERO: A Dataset for Contextualized Commonsense Inference in Dialogues. Annual Meeting of the Association for Computational Linguistics.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uWh5W6pfCo", "reviewer": "Reviewer_NEfd", "summary": "This work introduces a novel solution called AutoMix to achieve an optimal balance between performance and cost when using various scales of large language models (LLMs). The paper presents two variants of AutoMix: AutoMix-T, which employs a thresholding strategy, and AutoMix-P, which uses a POMDP-based strategy. Additionally, a new metric, IBC, is proposed to measure cost-performance efficiency. Experimental results on five datasets and a two-model setting with three LLMs demonstrate the superiority of AutoMix compared to other baselines.", "questions": "1. Generally, the performance of small language models (SLMs) is inferior to that of large language models (LLMs). Adopting the outputs of SLMs as the final response may result in considerable performance decay. Could you provide a quantified analysis of this impact?\n\n2. Could you provide a detailed analysis of how the POMDP-based router influences the decision of when to call LLMs? Clearly, the POMDP-based method cannot provide perfect predictions.", "limitations": "No.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This work addresses an intriguing problem: how to leverage small language models as assistants to reduce costs while maintaining comparable performance when using large language models is expensive.\n2. The proposed solution is technically sound and well-described.\n3. The empirical analysis demonstrates that the proposed method outperforms the baselines.", "weaknesses": "1. There is a lack of analysis regarding absolute performance decay. In most cases, performance is more critical than cost. The authors should show the impact on absolute performance when applying AutoMix.\n2. The POMDP-based router appears to introduce additional learning costs; however, this aspect is not analyzed.\n3. As noted by the authors, the generalizability of AutoMix is not well demonstrated, as the experiments cover only limited settings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review the paper and provide feedback! We believe all your questions can be addressed within this discussion period, but we would love to provide further clarification if needed.\n\n---\n\n### There is a lack of analysis regarding absolute performance decay.\n\nWe would like to clarify that Figures 4, 13, and 14 exactly show how the absolute performance varies as the cost changes. Specifically, our task is a multi-variable optimization problem, where there is a trade-off between cost and performance. Thus, performance can be compared only at fixed costs. Across all dataset and model settings, the figures show that when evaluated at the same costs, our method has higher performance than any of the baselines. It should also be noted that while higher performance is desirable when possible, many applications aim to trade off due to cost/latency concerns. So, we believe optimizing for this trade-off is well motivated and an important concern from a practical perspective when deploying these models to many real-world applications.\n\n---\n\n\n### POMDP-based router appear to introduce additional learning costs; however, this aspect is not analyzed.\n\nThe **POMDP-based router does not introduce additional learning costs**. The POMDP router is computationally efficient, requiring less than 1 second to train on a single CPU thread. This is substantially less than the training time required by routers in FrugalGPT and Hybrid LLM, which necessitate GPU resources and longer training durations. We will clarify this point in the revised paper.\n\n---\n\n### As noted by the authors, the generalizability of AutoMix is not well demonstrated, as the experiments cover only limited settings.\n\nWe would like to clarify our limitations here. While our experiments were conducted on context-grounded reasoning datasets, they covered a wide range of settings (15 in total: 3 models x 5 datasets), including different domains, objectives, answer types, and SLM, LLM performances, and their performance gaps. AutoMix consistently outperformed all baselines across these varied settings, demonstrating its broad applicability. Furthermore, our POMDP router assumes access to only self-verification probabilities and is thus generalizable to any setting or domain of task.\n\n\n---\n\n### Generally, the performance of small language models (SLMs) is inferior to that of large language models (LLMs). Adopting the outputs of SLMs as the final response may result in considerable performance decay. Could you provide a quantified analysis of this impact?\n\nWe acknowledge that owing to scaling laws in LLMs, SLMs have inferior performance compared to LLMs but at the same time, the performance gains from LLMs comes at additional cost/latency. Using LLMs for every problem could be an overkill for easy problems(where SLM could solve) and problems which are very hard (which even LLMs can’t solve). Our approach aims to exploit these gaps by using a smart self-verification and POMDP routing where one can achieve considerably higher performance for a given fixed cost. We believe this is well motivated from a practical standpoint of deploying language models in the real world where billions of queries are passed to language models and one also needs to optimize for cost in addition to performance. **Our cost-performance trade off clearly demonstrates this in Figure 4, 13 and 15.**\n\n\n---\n\n### Could you provide a detailed analysis of how the POMDP-based router influences the decision of when to call LLMs? Clearly, the POMDP-based method cannot provide perfect predictions.\n\nWe understand that any ML mechanism like POMDP cannot provide totally perfect predictions. However, it is crucial to note that the POMDP router is able to take advantage of self-verification probabilities to avoid routing to LLMs in primarily two cases: a) where the query is easy such that both SLM and LLM would give a correct answer, and b) the hard queries where both SLM and LLM would be wrong. This could be much more difficult to determine than a simple thresholding baseline. For instance, in the Qasper dataset with Mistral-7B as SLM, POMDP understands that lower confidence values correspond to such cases, and while other methods would have routed to LLM, POMDP returns the SLM answer, saving cost. Further, in a >2 model setup, it can make more nuanced decisions, for instance using combined information from SLM and MLM verifier, as demonstrated in Figure 6 in the main paper, and Figure 1 in general response.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MuluzeFRBR", "reviewer": "Reviewer_RAHf", "summary": "This paper presents AutoMix, a method that routes query to language models (LMs) of various size and capabilities to optimize the performance within a cost budget. AutoMix has two key technical features, a few-shot self-verification mechanism which estimates the reliability of its own outputs, and a POMDP-based router that can select a model to optimize the reward. Experiments across five LMs and five datasets show that AutoMix outperform two baselines (i.e., FrugalGPT, HybridLLM).", "questions": "See my comments in \"Weaknesses\".", "limitations": "The paper has a \"Limitations\" section.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. Formulating the query routing optimization problem as a Partially Observable Markov Decision Process (POMDP) is interesting. Treating performance of different models as $S$ is quite smart since the most challenging part of the optimization problem is we want to make as less computation as possible while we cannot observe whether an LM's output is correct or not since the optimization happens in the testing stage. The property of the optimization problem matches very well with POMDP.\n2. The experimental results show that AutoMix outperforms FrugalGPT and HybridLLM which are two recent methods in this setup.", "weaknesses": "1. The experimental setup is not very rigorous. The major metric in this paper is $\\Delta_{IBC}$ which will be influenced by how you set the cost of different LMs in your experiment. Although the arguments provided in Line 216-219 is legitimate, setting the cost as 1, 60, 100, 200 respectively feels like setting \"hyperparameters\" to me. Given the results on different datasets and models in Table 1 are not consistent, I don't have confident with the conclusion from the current results. For example, will the trend change if I set the cost as 10, 60, 100, 200?\n2. While I praise the integration of POMDP as interesting and well-motived, I am not sure whether such integration adds much to the performance given the main results in Figure 4. It seems like \"AutoMix+POMDP\" and \"AutoMix+Threshold\" achieve similar results while the latter excels in simplicity. I am also wondering how much adding POMDP will increase the overhead since the system needs to run the AdaOps POMDP solver.\n3. The explanation of \"domain-specific training\" could be clearer. Since this phrase in LM literature usually refers to post-train LM on domain-specific data, it's necessary to clearly define what it means in your method. Although I believe I understand the training of router in your method, I have a clarification question: Does computing `P(o|s)` require running the inference on all the samples in the dataset and also running the self-verification step? If this is the case, should this part of the cost be considered?\n\n**Update:** Some of these points are addressed by the author response, so I updated the score from 4 (Borderline Reject) to 5 (Borderline Accept).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review the paper and provide feedback! We believe all your questions can be addressed within this discussion period, but we would love to provide further clarification if needed.\n\n---\n\n### Although the arguments provided in Line 216-219 is legitimate, setting the cost as 1, 60, 100, 200 respectively feels like setting \"hyperparameters\" to me.\n\n\nWe would like to address the concerns regarding cost values (ratios) below:\n1. *Cost values as hyperparameters:* We respectfully disagree with the reviewer that cost is a hyperparameter for our approach. Cost values signify the parameters of the problem. Moreover, these parameter values are not in any way optimized by our system but motivated by the cost of apis of popular and SOTA language models. Nevertheless, since these values could also change in future, we measure robustness and sensitivity of our approach if we vary cost ratio between different SLMs and LLMs in Section 5.4 (as described below).  To reiterate, we use conservative estimates to the pricing: 30 USD/M tokens for GPT-4, 0.1/M for Mistral7B, 0.225/M for Llama-13b, and 0.5/M for GPT3.5, giving cost-ratios of roughly 200, 100 and 60\n2. *Robustness of AutoMix to different cost scenarios:* To understand the effect of different cost values, in section 5.4, we further evaluated the performance of Automix with different cost values. As we describe in the paper, a huge cost ratio between LLM and SLM would be favorable to Automix, as relative cost of self-verification decreases, while lower cost would be detrimental. Nonetheless as stated in lines (267-269), “The results suggest that even for a ten times lower cost ratio, AutoMix begins to deliver good gains for many datasets” and outperforms all baselines at even half the cost ratio. Further, incase of lower SLM costs (a practical scenario as described in lines 221-223), Automix performs even better than the scores reported in main results. To further, supplement the analysis, we also include cost-performance curves for different cost-ratios in Figure 2 of general response. \n3. Thirdly, these costs values and methods should be seen beyond the pricing of LLMs. These could signify latency or energy costs of models. \n\nWe believe beyond just the exact cost values, Automix provides a novel method to balance cost performance trade off which works across a wide variety of cost ranges and is robust across a very wide range of cost ratios. In practical scenarios, it performs significantly better than other baselines. \n\n\n\n\n\n\n---\n\n### While I praise the integration of POMDP as interesting and well-motivated, I am not sure whether such integration adds much to the performance given the main results in Figure 4. \n\nPOMDP is a general setup, essential for generalizing to scenarios with more than two models, and sets a very robust baseline for future endeavors consisting of more complicated routing strategies. Further, as we discuss in Section 5.2 simpler methods perform much poorly than a POMDP model for settings with more than two models, as demonstrated in Figure 7. In addition to Figure 7, we include more such cases in Figure 1, general response, demonstrating the importance of POMDP.\n\nFinally, we note that implementing the POMDP model is straightforward and fast, requiring only a few lines of code (~5) for the end-user to incorporate our library. Additionally, as mentioned in line 285, it takes less than 1 ms to run the POMDP model for an input query. We will include this discussion in the revised paper to better justify our use of POMDPs.\n\n\n---\n\n### The explanation of \"domain-specific training\" could be clearer. Since this phrase in LM literature usually refers to post-train LM on domain-specific data, it's necessary to clearly define what it means in your method. \n\n“Domain-Specific training” refers to the usual meaning in LM literature of “post-training LM on domain-specific data” We note, that unlike FrugalGPT and HybridLLM, Automix does not require domain-specific training of **verifier** \n\n---\n\n### Although I believe I understand the training of router in your method, I have a clarification question: Does computing P(o|s) require running the inference on all the samples in the dataset and also running the self-verification step? If this is the case, should this part of the cost be considered?\n\nYes, computing P(o|s) requires running inference on all dataset samples and the self-verification step. However, note that this **cost is incurred only during training and not evaluation**. Further, even the **training cost is similar to baselines** since both FrugalGPT and HybridLLM require running the inference on all samples, and additionally requires training verifiers. While Automix requires running self-verification, verification is run only on smaller models, which are usually significantly cost-efficient than LLMs (eg, 60-200 times cheaper), thus adding negligible training costs. We will clarify this clearly in the revised version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pnXXip7RFb", "reviewer": "Reviewer_eJtm", "summary": "This paper investigates how to automatically and adaptively select the LLM with a smaller size but keep good performance so that it offers the potential to achieve a better cost-effectiveness trade-off. This work uses a smaller model to predict first and then uses the small model to do self-verification. If more computation is required, a POMDP-based router can effectively select an appropriately sized model. The proposed approach can reduce around 50% cost as reported in the paper.", "questions": "See Weakness.", "limitations": "NA", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) Adaptively selecting model sizes is smart and cool. Using a smaller model first makes sense because it can reduce the overhead when the task is difficult.\n2) The paper is very well-written and easy to understand. The Verifier Qualitative Analysis in the appendix is also helpful.", "weaknesses": "1) It is good that this work provides a detailed description of their cost modeling. However, it seems that there are many assumptions in it. How can authors justify that this is not biased?\n2) More importantly, this work requires more discussion about the real-world throughput  (response latency). Although the computation cost is reduced, the user experience would become worse if the latency is significant larger.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review the paper and provide feedback! We believe all your questions can be addressed within this discussion period, but we would love to provide further clarification if needed.\n\n---\n\n### It is good that this work provides a detailed description of their cost modeling. However, it seems that there are many assumptions in it. How can authors justify that this is not biased?\n\n**Our cost modeling directly reflects the real-world pricing of APIs, as detailed in Appendix D (Lines 649-657).** We did not introduce any subjective assumptions in the cost ratios. Furthermore, to handle evolving cost dynamics of LLM API providers, we conduct analysis with varying cost ratios in Section 5.4 (Figure 5, Right), confirming that AutoMix offers significant improvements even when LLM costs are reduced by a factor of 2-4. Our method is robust across a wide range of cost ratios of SLMs and LLMs. This robustness indicates that our results are not biased by the chosen cost model.\n\n\n---\n\n### More importantly, this work requires more discussion about the real-world throughput (response latency). Although the computation cost is reduced, the user experience would become worse if the latency is significantly larger.\n\nWe thank the reviewer for bringing up this point. As we note in Section 5.4, AutoMix is agnostic to the specific nature of the cost metric. Latency can be modeled similarly to how the cost is modeled, with the larger model incurring higher latency. As noted in our limitations, “AutoMix can automatically handle latency as another form of cost, but it might lead to less significant improvements” owing to lower latency differences between SLMs and LLMs. However, the exact latencies vary by different providers. As a real-world example, an ~8B model, through API providers such as Groq, provides generation speeds up to 1000 tokens/s, compared to roughly 30 tokens/s for GPT-4. In such scenarios, the latency of routing to LLMs would be far higher than a fixed number of more calls to SLMs for reaching a given performance. Therefore, AutoMix can improve latency for a given performance, especially by reducing the higher latency of calls to LLMs.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their valuable feedback! We are encouraged that they find our "approach sound and intuitive, our paper as easy to read and follow" (Reviewer zQB8), that "our approach to self-verification and POMDP-based routing is novel, our methodology is sound, and recognizing the importance of our proposed IBC metric in cost-efficiency analysis" (Reviewer E2SA), and that our paper "addresses an intriguing problem, and the proposed solution is technically sound and well-described" (Reviewer NEfd), and our "integration of POMDP is interesting and well-motivated" (Reviewer RAHf), and "our proposed approach is smart and the paper being very well-written and easy to understand" (Reviewer eJtm).

We have addressed reviewers' comments individually and look forward to further comments within the remainder of the discussion period.

As part of the rebuttal, we have also included certain additional analyses and experiments enumerated below:

1. We expand our analysis of the "Effect of Cost Ratio on AutoMix" (Section 5.4) to show how different cost ratios affect cost-performance curves in addition to the IBC metric.
2. To demonstrate the generalizability of our method to other domains, we evaluate our method on an additional commonsense reasoning dataset: CICERO, demonstrating significantly higher results than baselines.
3. We demonstrate the out-of-distribution generalization capability of our proposed POMDP router, while baselines performed very poorly.
4. We extend the analysis in Section 5.5 to show the effectiveness of POMDP over other routers and baselines on additional datasets and models.
5. There have been concerns on generalization ability of method. While the details are addressed in each reviewer comment, we reiterate that Automix was evaluated on 5 diverse datasets where performance of models range quite widely and with 3 different models. In addition, we have included another dataset in rebuttal period to further strengthen this. We would be happy to include any other dataset suggested by reviewer as applicable to our setting.


We enclose these additional experiments in the response PDF, along with individual reviewer responses.

Meta-review: This paper introduces a new technique (AutoMix) to optimize LLM inference costs, by training a router that can decide to query larger models only when required (while sticking to smaller models when not needed to properly answer a question).

After the discussion period, where authors clarified several points and included additional results, all reviewers agree this is a practically useful approach with clear demonstrated benefits and no major weakness. I agree this is a meaningful contribution, which anyone looking to decrease their API cost usage could be interested in. As a result, I am recommending acceptance.

I would like to ask the authors to be sure to include all additional results and important clarifications in the camera-ready version.
Decision: Accept (poster)