AutoMix: Automatically Mixing Language Models

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.

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. 2) 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? 2) 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.

Questions

See Weakness.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

NA

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).

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. 2. 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? 2. 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. 3. 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? **Update:** Some of these points are addressed by the author response, so I updated the score from 4 (Borderline Reject) to 5 (Borderline Accept).

Questions

See my comments in "Weaknesses".

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

The paper has a "Limitations" section.

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.

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. 2. The proposed solution is technically sound and well-described. 3. 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. 2. The POMDP-based router appears to introduce additional learning costs; however, this aspect is not analyzed. 3. As noted by the authors, the generalizability of AutoMix is not well demonstrated, as the experiments cover only limited settings.

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? 2. 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.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

No.

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.

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. - 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. - 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. - 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. - 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.

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？

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The author pointed out the limitations in the discussion.

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.

Strengths

The proposed approach is sound, intuitive and works even in the black-box setting. Empirical results from five different dialogue and context-grounded reasoning tasks.show that the proposed approach (AutoMix) outperform existing computationally efficient LLM approaches. The 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.

Questions

Please discuss the weakness mentioned above.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Thank you

I thank the authors for their response. I acknowledge the comments and decided to keep my positive score with more confidence.