PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2310.12963v5cs.CLcs.AI

TL;DR

Choosing among available language models to balance performance and computational cost is challenging, especially when APIs are black boxes and task-specific training data are limited. AutoMix combines few-shot, context-grounded self-verification with a POMDP router to select whether to use a larger model. Across five tasks and five models, it outperforms baselines while reducing cost by over two times at the same performance.

  • Problem

    The range of model capabilities and costs makes it challenging to choose configurations that meet user needs, while existing routing approaches face limits with black-box APIs and scarce task-specific data.

    Source: §1 · p. 1

  • Method

    AutoMix uses few-shot context-grounded self-verification and a POMDP router to decide whether to accept a small model’s answer or route to a larger model.

    Source: §1 · p. 2

  • Results

    Across five tasks and five models, AutoMix consistently outperforms baselines while reducing cost by over two times at the same performance.

    Source: §1 · p. 2

  • Takeaways & Limitations

    AutoMix balances cost and performance across a wide variety of models and desired cost ranges, consistently outperforming baselines.

    Source: §6 · p. 10

    AutoMix is designed for dialogue-related or context-grounded reasoning and does not include factual question-answering or commonsense reasoning tasks.

    Source: Limitations · p. 10

Abstract

from arXiv · show

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.

1 Introduction

The growing range of language models makes it difficult to balance performance and computational cost, while existing routing approaches face limits when models are black-box APIs and task-specific data are scarce. AutoMix combines few-shot self-verification with POMDP routing and outperforms baselines while reducing cost by over two times at comparable performance.

1 Introduction

  • The expanding range of model sizes and capabilities complicates choosing a configuration that fits users’ performance needs and budgets.Larger models generally perform better but can be unaffordable for simpler tasks.Source: §1 · p. 1
  • Existing routing approaches often use task-specific trained routers, but black-box APIs and limited task-specific data constrain their applicability.AutoMix is introduced to mix models using black-box API access.Source: §1 · p. 1
  • AutoMix generates an answer with a small model, verifies it, then routes to a larger model according to the verifier’s confidence.For setups with more models, the process repeats until a final answer is reported.Source: §1 · p. 2, Figure 1 · p. 2
  • AutoMix frames self-verification as context entailment and uses a POMDP router to handle noisy confidence estimates and learn policies from as few as 50 examples.The POMDP is intended to support routing across differing model counts, costs, and capabilities.Source: §1 · p. 2
  • Across five tasks and five models, AutoMix outperforms baselines while reducing cost by over two times at the same performance.The evaluation covers dialogue and context-grounded reasoning tasks.Source: §1 · p. 2

2 Background and Related Work

Prior work explores self-verification, model switching, adaptive computation, and POMDPs, but AutoMix distinguishes itself by using context-based verification and black-box model routing. Its POMDP treats verifier outputs as noisy observations for choosing which language model to use.

Self-Verification

Mixing Models

  • Model-switching methods often train specialized verifiers, while AutoMix uses few-shot SLM prompting and a router trained with as few as 50 samples [Chen et al., 2023, Zhu et al., 2023, vSakota et al., 2023, Ding et al., 2024].The passage reports that the router outperforms specialized models.Source: §2 · p. 3

Adaptive Computation

Background on Partially Observable Markov Decision Processes

  • AutoMix uses noisy self-verification outputs as POMDP observations to estimate question difficulty and route queries among language models.The POMDP’s reward balances cost and performance.Source: §2 · p. 3

3 Problem Formulation

The paper formulates model selection as choosing among language models with different costs and unknown performance to maximize performance under a total-cost constraint. Its evaluation focuses on context-grounded reasoning and assumes black-box API access.

3 Problem Formulation

  • The objective is to maximize total performance for a given total cost by dynamically invoking an appropriate model for each test point.The formulation assumes N models ordered by parameter count, each with an associated cost and unknown performance for each input.Source: §3 · p. 3
  • The experiments compare cost-quality curves and evaluate methods using the proposed IBC metric.IBC is introduced as a cost-aware evaluation measure.Source: §3 · p. 4
  • The evaluated tasks are context-grounded reasoning, including comprehension QA and dialogue reasoning, where answers should be accurate and consistent with supplied context.The context may consist of stories, newswire, or dialogue history.Source: §3 · p. 4

4 AutoMix

AutoMix generates an answer with a small model, verifies it with that model, and selectively routes the query to a larger model. Its verifier estimates answer trustworthiness from context, while its router uses potentially noisy verification outputs to decide whether escalation is worthwhile.

4 AutoMix

  • AutoMix’s core pipeline generates an answer with a small LM, verifies it using the same LM, and routes to a larger LM when the verifier suggests it.If no escalation is selected, the initial answer is returned; the process can repeat across multiple models.Source: §4 · p. 4

4.1 Self-Verification

  • The few-shot verifier treats checking an answer as an entailment task, estimating its correctness probability from the answer, context, and question.It samples the verifier multiple times at high temperature to estimate the probability and uses the same four-shot prompt across tasks without training a verifier.Source: §4.1 · p. 4

4.2 Router

  • The router decides whether to accept the current answer or escalate, and can avoid costly escalation for unsolvable queries.It uses verification probabilities and trends in training data; routing is handled outside the LLM.Source: §4.2 · p. 5
  • Threshold routing accepts the small model’s answer when verifier probability v meets threshold t; otherwise it routes to the larger model.Varying t explores different cost-performance tradeoffs.Source: §4.2 · p. 5
  • The POMDP router treats verifier outputs as uncertain observations of model performance and chooses among retaining the current answer or routing to a larger model.Its states encode the selected model and model performance; beliefs are updated from verifier outputs, and the reward is R = P − λ · C, balancing performance and cost according to user preferences.Source: §4.2 · p. 5

5 Experiments

The experiments define a cost-efficiency metric, compare AutoMix with model-switching baselines across varied models and tasks, and test its behavior under different costs, data amounts, and model configurations.

5.1 A Metric for Cost-Performance Efficiency Analysis

  • IBC measures performance improvement relative to added cost; ∆IBC compares a method’s IBC with the always-LLM-over-SLM baseline.A positive lift means the method improves performance more cost-effectively than using the LLM alone.Source: §5.1 · p. 6

5.2 Setup

  • Experiments use five datasets, three SLMs, GPT-4 as the LLM, and FrugalGPT and HybridLLM as baselines.The datasets cover dialogue and context-grounded reasoning; F1 is used for QASPER and COQA, and accuracy for the other datasets.Source: §5.2 · pp. 6–7

5.3 Main Results

  • AutoMix-POMDP outperforms baselines across datasets and models, yields positive gains across configurations, and averages 122% on DIPLOMAT across models.It consistently matches or exceeds AutoMix-Threshold, with one exception: it equals FrugalGPT on QASPER with LLAMA2-13B.Source: Table 1 · p. 8, §5.3 · p. 8

5.4 Additional Experiments

  • AutoMix begins to deliver good gains for many datasets at a 1:10 cost ratio, with gains improving as the ratio becomes more skewed.The paper also frames cost as potentially including latency or energy, not only monetary expense.Source: §5.4 · pp. 8–9
  • With 50 training examples, AutoMix maintains 15% ∆IBC and beats FrugalGPT and HybridLLM by more than 20 percentage points.The baselines use domain-specific training and have zero verifier cost.Source: §5.4 · p. 9
  • POMDP uses noisy self-verification probabilities to distinguish easy, complex, and potentially unsolvable queries, sometimes avoiding an unnecessary LLM call.For QASPER with MISTRAL-7B as SLM, low-confidence cases can lead it to return the SLM answer rather than route to the LLM.Source: §5.4 · p. 9
  • AutoMix outperforms FrugalGPT and HybridLLM in held-out datasets, while its two-model POMDP router takes less than 1 ms per query.The authors report no additional computational or latency overhead compared with language-model inference.Source: §5.4 · p. 9

5.5 Results of Automix with Three Models

  • In three-model tests, AutoMix outperforms baselines across cost regions; chained routing underperforms because it cannot route directly from SLM to LLM.When the medium model underperforms the SLM, AutoMix can either bypass it or use its verifier when that information is useful.Source: §5.5 · p. 10

6 Conclusion

AutoMix combines black-box LLM APIs in a multi-step framework to balance computational cost and performance. In the three-model comparison, it improves IBC across both model regions and significantly outperforms the listed baselines.

6 Conclusion

  • AutoMix improves IBC in both SLM-MLM and MLM-LLM regions and significantly outperforms FrugalGPT, chained AutoMix models, and the union of two AutoMixes.The figure describes a three-model comparison.Source: Figure 6 · p. 10
  • AutoMix combines black-box LLM APIs in a multi-step problem-solving framework to optimize computational cost and performance trade-offs.The conclusion describes this approach as combining GOFAI methods with LLMs and using a POMDP for routing.Source: §6 · p. 10
  • The authors identify generation tasks, where performance is subjective, as a possible direction for extending the approach.They also report promising results from context-grounded few-shot verification coupled with POMDPs.Source: §6 · p. 10

Limitations

The authors qualify AutoMix’s broader applicability by the models and datasets used, and by its focus on dialogue-related or context-grounded reasoning.

Limitations

  • AutoMix is designed for dialogue-related or context-grounded reasoning, not factual question-answering or commonsense reasoning.The authors state that this setup supports effective self-verification.Source: Limitations · p. 10
  • The authors caution that applicability beyond the evaluated settings might vary with the specific models and datasets used.Their empirical evidence spans various models and datasets.Source: Limitations · p. 10
  • If open-source models grow more powerful and inference costs fall, serving a strong model for every query might become feasible.The authors note that availability trade-offs could still be managed using AutoMix.Source: Limitations · p. 10

A.1 Quantitative Analysis

Few-shot self-verification generally assigns higher verification probabilities to correct examples, aligning those probabilities with precision across most tested datasets. QUALITY and QASPER are exceptions to the broad pattern, with QASPER showing a few outliers.

A.1 Quantitative Analysis

  • Across datasets, higher verification probability generally corresponds to higher precision for identifying correct responses, except on QUALITY and QASPER.Correct responses are defined as F1 > 0.1; QASPER has two outliers among nine cases, while its general trend holds.Source: §A.1 · p. 15
  • Few-shot self-verification uses context to check answers for factual inaccuracies or apparent contradictions that may have been missed during generation.The authors identify context-based validation as an advantage of their setup for reasoning problems.Source: §A.1 · p. 15
  • The reported precision pattern supports higher verification probabilities as an indicator of correct examples, with stated exceptions for QUALITY and QASPER.The figure's precision measure counts responses with F1 > 0.1 as correct.Source: §A.1 · p. 15

A.2 Qualitative Analysis

AutoMix performs especially well against FrugalGPT when validation data are limited, while the router uses verifier observations to choose among language models under a cost-performance objective. Its POMDP solver represents beliefs with particles and searches possible action-observation sequences.

A.2 Qualitative Analysis

  • Table 2 shows the verifier catching an answer that cites an unsupported date, while Table 3 shows it failing on a context rephrasing.Table 4 shows the verifier deeming an answer correct even though the gold label was incorrect.Source: Table 2 · p. 15, Table 3 · p. 15, Table 4 · p. 16
  • With limited validation data, AutoMix significantly outperforms FrugalGPT despite FrugalGPT’s domain-specific training and zero verifier cost.As training data increases, FrugalGPT narrows the performance gap; Figure 7 indicates larger AutoMix gains below 1,000 examples.Source: §A.2 · pp. 15–16, Figure 7 · p. 16

A.3 POMDP-based Router in AutoMix

  • The router models verifier outputs as potentially unreliable observations and weighs performance against cost using a user-tuned parameter.Its states track the current model stage and model-performance metrics; actions retain an answer or route to a larger model.Source: §A.3 · pp. 16–17
  • For continuous states, the router estimates the joint state-observation distribution with KDE and derives observation probabilities for its POMDP.The discrete case estimates observation probabilities directly from training data; the continuous case uses KDE at training points.Source: §A.3 · p. 17

A.3.1 Solving POMDP

  • The continuous state space is (N + 1)-dimensional, making traditional methods designed for discrete states difficult or computationally intractable to use.The formulation also has a finite, small horizon and a limited set of nine verifier observations.Source: §A.3.1 · p. 18
  • The POMDP solver searches action and observation sequences, updates belief particles after observations, and values nonterminal choices by expected reward.The reward balances performance and cost, and the discount factor is set to 1.Source: §A.3.1 · p. 18
  • Because the decision horizon and observation sequences are finite, the authors can learn deterministic policies offline and use a lookup table at inference.They describe inference cost as negligible.Source: §A.3.1 · pp. 18–19

Using better POMDP solvers

  • AdaOps is used for more complicated future scenarios, while the authors expect it and the simple solver to converge approximately to the same solution as particle counts increase.AdaOps is described as supporting larger state spaces through adaptive online search.Source: Using better POMDP solvers · p. 19

B Additional Details on the Experimental Setup

The appendix specifies the evaluation protocol and extends AutoMix to three-model settings, defining baselines and cost-region rules for comparing models.

B Additional Details on the Experimental Setup

  • Evaluation samples 1,000 validation instances per dataset across five tasks, using F1 for QASPER and COQA, accuracy elsewhere, and question-keyed retrieval capped at 3,500 tokens.The tasks cover research-paper QA, long-context multiple choice, conversational comprehension, multi-turn dialogue reasoning, and pragmatic identification.Source: B Additional Details on the… · pp. 19–20
  • Both models receive identical prompts; answer generation uses one greedy sample, while self-verification uses eight samples.The setup assumes uniform output-processing costs and treats verifier cost as negligible relative to context processing.Source: B Additional Details on the… · pp. 19–20

C Expanding AutoMix to Three-Models

  • The three-model evaluation uses LLAMA2-13B, MISTRAL-7B, or GPT-3.5 as SLM, LLAMA2-70B as MLM, and GPT-4 as LLM.It compares AutoMix with cascade-style FrugalGPT, pair-selection Union AutoMix, and sequential Chained AutoMix.Source: C Expanding AutoMix to Three-Models · p. 20

C.1 Extending IBC to N=3 models

  • For three models, IBC evaluation selects the more competitive model pair by cost region when the performance-cost curve is convex, but ignores MLM when the curve is concave.For general N, the method considers model pairs for each cost region.Source: §C.1 · p. 21

C.2.1 Does AutoMix outperforms baseline methods?

In three-model tests, AutoMix generally beats the baselines by using verifier information to choose among models, including when an intermediate model is weak. Its gains are not universal: poor self-verification can leave it no better than baselines.

C.2.1 Does AutoMix outperforms baseline methods?

  • AutoMix consistently outperforms FrugalGPT, Union AutoMix, and Chained AutoMix across the considered cost regions in three representative datasets.The authors attribute the advantage to using SLM verification to choose whether to route to MLM or LLM, with POMDP belief updates incorporating both verifiers when MLM is used.Source: Figure 8 · p. 21, §C.2.1 · p. 21

C.2.2 Can AutoMix skip distractor models?

  • When MLM underperforms SLM, AutoMix learns to skip it and route directly to LLM as needed, while FrugalGPT can waste resources by routing to MLM.AutoMix performs similarly to Union AutoMix in these cases because it never uses MLM.Source: §C.2.2 · p. 22, Figure 9 · p. 22

C.2.3 Should Low Performing MLMs be ignored?

  • AutoMix can also exploit questions where MLM is better than SLM, using low SLM verification probabilities and combined verifier information to improve cost-quality tradeoffs.The authors report statistically outperforming even LLM at roughly half the cost in this example.Source: §C.2.3 · p. 23, Figure 10 · p. 23

C.2.4 Does AutoMix always provide significant gains?

  • In two cases with relatively poor self-verification, AutoMix shows no statistically significant gains; verifier scores generally track correctness except on QUALITY and QASPER.The authors report results similar to FrugalGPT and negligible or negative advantage over IBC lines, and suggest contextual information or other verification as future directions.Source: Figure 11 · p. 23, §C.2.4 · p. 24, Figure 12 · p. 24

D.1 Main Results

Across model configurations, cost ratios, and dataset transfers, AutoMix generally performs better than the compared methods. The appendix also describes its verifier prompt as a natural-language entailment task.

D.1 Main Results

  • AutoMix-POMDP is clearly better than all baselines in 13 of 15 configurations and is competitive in the remaining two.The reported comparisons span three SLM choices and different operating costs.Source: §D.1 · p. 25, Figure 18 · p. 26, Figure 19 · p. 27

D.2 AutoMix is effective across different Cost-Ratios

  • AutoMix outperforms baselines across considered model cost ratios, including 20:1 and 2000:1, with smaller margins at 20:1 and larger improvements at 2000:1.The authors describe 20:1 as an unlikely, pessimistic scenario for AutoMix.Source: §D.2 · p. 25, Figure 20 · p. 27

D.3 Out-of-Domain Generalization of AutoMix

  • When trained on one dataset and evaluated on the other four, the POMDP router consistently outperforms FrugalGPT and HybridLLM across five datasets and three SLMs.Table 5 reports scores averaged across the five datasets.Source: §D.3 · p. 26, Table 5 · p. 26

E.1 Verifier Prompts

  • The verification prompt frames judging an SLM answer against its question and context as a natural-language entailment task, using the same three-shot prompt across datasets.The prompt asks the model to choose Correct or Incorrect.Source: §E.1 · pp. 28–29

Related papers

Loading 2310.12963v5…