Confident or Seek Stronger: Exploring Uncertainty-Based Small LM Routing From Benchmarking to Generalization

Thank you for your tremendous efforts in reviewing our paper. We would like to further answer your questions as follows.

[Q-1] Results on Code-Related Tasks

[A-1] We ran additional experiments on MBPP (Python coding problems). Below we report ROC AUC (uncertainty–correctness alignment) for two SLMs across 8 UQ methods. Trained Probe and OOD Probe perform best for Mistral‑7B, while Jaccard Degree and OOD Probe perform best for Qwen3‑0.6B.

Table 1: Uncertainty–Correctness Alignment on MBPP

Next, we evaluate our Proxy Routing Data on MBPP. Since we are not able to provide figures like Figure 5 and Figure 6 here, we compute the average root mean squared (RMS) distance between the routing curve derived from the full downstream dataset and the curve obtained using our method. Lower RMS distances indicate better performance. Below we show results when routing from Mistral‑7B and Qwen3‑0.6B to GPT‑4.1 mini, using two UQ methods (OOD Probe and Perplexity). We observe that our method predicts accurate routing thresholds on MBPP without accessing it at all, leading to small RMS distances.

Table 2: Results of Proxy Routing Data on MBPP (RMS distance; lower is better)

[Q-2] No Evaluation of Multi-LLM Routing, such as RouterBench: Cost and Generalization to New Domains

[A-2] Unlike the settings in RouterBench, which focus on multi-LLM routing using an external router to select the most appropriate LLM on large-scale servers for a given query, our work addresses a different stage of the routing pipeline: deciding whether to process a query locally on the edge device with an SLM or to offload it to a more powerful LLM in the cloud servers. Multi-LLM routing is outside the scope of our current on-device routing scenario, but it represents a valuable follow-up step once the offloading decision has been made. In other words, after a query is determined to route to the cloud from edge device, multi-LLM routing strategies, such as the baselines in RouterBench, can be applied to further optimize inference efficiency and budget.

Thank you for your tremendous efforts in reviewing our paper. We would like to further answer your questions as follows.

[Q-1] Need More Analysis of Why Certain UQ Methods Fail on Certain Datasets/Tasks

[A-1] The effectiveness of uncertainty-based routing relies heavily on the alignment between uncertainty scores and prediction correctness, which is a point we emphasize and validate in Section 3 (especially shown in Figure 1 for alignment assessment). A well-calibrated uncertainty score does not guarantee strong routing performance, regardless of how powerful the underlying LLMs are or how simple the tasks may seem. Misalignment between uncertainty scores and actual prediction correctness is one of the key factors leading to routing failures across different datasets/tasks. Furthermore, the distribution of correct predictions varies significantly depending on the properties of the LLMs, such as pretraining data, decoding strategies, etc. In addition, different uncertainty estimation methods can yield varying levels of alignment, further influencing routing performance across different LLMs and datasets.

[Q-2] Despite the generalization claims, the proxy routing approach is still indirectly tuned using benchmarked datasets, which may bias the proxy data toward known distributions

[A-2] We would like to respond from two perspectives:

• First, our proxy routing approach is not parameterized, which is an untrainable method by natural. The proxy datasets are used for routing curve prediction.
• For known distribution settings, we have considered the "partially in-domain" settings to assess the performance of proxy dataset. The domain of the proxy dataset in this settings may partially overlap to targeted testing data. The results are shown in Figure 5 and Appendix, indicating good performance in terms of using proxy data for routing curve prediction.

[Q-3] The pipeline for generating proxy routing data is not deeply explained in terms of how it adapts to drastically different downstream domains (e.g., multimodal tasks, non-English text), which limits its universality. Can the authors propose some guidelines from the observations

[A-3] In the main text of our paper (Figure 6), we evaluate Proxy Routing Data in a highly OOD setting: we construct proxy data only from commonsense‑reasoning and conversational/contextual datasets and evaluate on math (AQuA), routing from four SLMs to DeepSeek‑R1. In this case, our method still achieves competitive performance, demonstrating superior generalization.

To further address the concern, we also evaluate on translation using WMT24++ (English → Mexican Spanish), where Spanish does not appear in any other datasets used to build the proxy. We compute the average root mean squared (RMS) distance between the routing curve derived from the full downstream dataset and the curve obtained using our method. Lower RMS distances indicate better performance. Below we show results when routing from Mistral‑7B and Qwen3‑0.6B to GPT‑4.1 mini, using two UQ methods (OOD Probe and Perplexity). We observe that our method predicts accurate routing thresholds without accessing the dataset at all, obtaining small RMS distances.

Table 2: Results of Proxy Routing Data on WMT24++ (RMS distance; lower is better)

The author-reviewer discussion is coming to a close. Would you please go through the authors' rebuttal ASAP, highlight any remaining concerns (if any) so that the authors have sufficient time to respond to them, and adjust your review/rating accordingly?

Thanks!

Your AC

Thank you for your tremendous efforts in reviewing our paper. We would like to further answer your questions as follows.

[Q-1] Insufficient Theoretical Explanation of Empirical Findings

[A-1] We would like to clarify that:

• The effectiveness of uncertainty-based routing depends on the alignment between uncertainty scores and prediction correctness from SLM; however, this alignment cannot be theoretically guaranteed by nature [1, 2]. Thus, we provide a comprehensive empirical study to discover the pattern in our two main directions with our explanations across Section 3, 4, and 5.
• We provide two toy examples to explain this misalignment intuitively: Assume we have a binary classification task and obtain results from two binary classification models with different predicted probabilities (see below). Example 1 achieves lower accuracy but has a lower ECE score, whereas Example 2 achieves higher accuracy but a higher ECE score. This example demonstrates that calibration metrics are not generally correlated with the prediction's correctness.

Example 1: 
Bin = 2
predicted prob. of "1" = [0.5, 0.5, 0.5, 0.5]
ground truth = [0, 0, 1, 1]
-> ECE(↓)= 0%
-> accuracy (↑)= 50%

Example 2: 
Bin = 2
predicted prob. of "1" = [0.4, 0.6, 0.9, 0.9]
ground truth = [0, 0, 1, 1]
-> ECE(↓) = 20%
-> accuracy (↑) = 75%

[1] Huang, et al. "Look before you leap: An exploratory study of uncertainty measurement for large language models." arXiv, 2023.

[2] Detommaso, et al. "Multicalibration for confidence scoring in llms." arXiv 2024.

[Q-2] Imbalanced Training Sizes Between Trained Probe and OOD Probe

[A-2] Thank you for raising this concern. Actually, in all experiments relevant to OOD Probe, we explicitly controlled for training volume: we down‑sampled the pool of remaining datasets so that the OOD Probe was trained on a similar number of examples as the Trained Probe. Specifically, rather than using all available examples from other datasets, we randomly sampled a small, fixed percentage from each to match the Trained Probe’s data volume. Under this controlled setting, it's safe to infer that the performance comes from data diversity instead of volume. We will make this important setup more explicit in the final version to avoid confusion.

[Q-3] Results on Translation Tasks

[A-3] We ran additional experiments on WMT24++ (English to Mexican Spanish translation). Below we report ROC AUC (uncertainty–correctness alignment) for two SLMs across 8 UQ methods. Perplexity and Average Token Probability perform best for Mistral‑7B, while Jaccard Degree and Trained Probe perform best for Qwen3‑0.6B.

Table 1: Uncertainty–Correctness Alignment on WMT24++

Next, we evaluate our Proxy Routing Data on WMT24++. Since we are not able to provide figures like Figure 5 and Figure 6 here, we compute the average root mean squared (RMS) distance between the routing curve derived from the full downstream dataset and the curve obtained using our method. Lower RMS distances indicate better performance. Below we show results when routing from Mistral‑7B and Qwen3‑0.6B to GPT‑4.1 mini, using two UQ methods (OOD Probe and Perplexity). We observe that our method predicts accurate routing thresholds on WMT24++ without accessing it at all, leading to small RMS distances.

Table 2: Results of Proxy Routing Data on WMT24++ (RMS distance; lower is better)

[Q-4] Entropy as an Uncertainty Quantification Method

[A-4] Thank you for the suggestion. We have implemented Entropy as a new UQ method. Below we report ROC AUC (uncertainty–correctness alignment) on OpenBookQA across multiple SLMs using uncertainty derived by Entropy. Compared with Figure 1 in our paper, Entropy performs similarly to token/sequence‑probability methods (Average Token Probability and Perplexity): all three work well on the Llama series and are less effective on the Qwen series, obtaining similar numbers across most SLMs. We will include these results in the final version.

Table 1: Uncertainty–Correctness Alignment (ROC AUC) Using Entropy on OpenBookQA

[Q-5] Sensitivity Analysis of Hyperparameters

[A-5] We conducted a sensitivity analysis on both the sampling ratio and the number of bins to assess our method’s generalizability and robustness. The first table below fixes the number of bins at 30 while varying the sampling ratio; the second fixes the sampling ratio at 0.1 while varying the number of bins. These experiments were performed on HellaSwag, routing from three SLMs to GPT-4o-mini using the OOD Probe. We computed the average RMS distance between the Routing Ratio vs. Accuracy curves obtained using thresholds from the full downstream dataset and those derived from our method—lower RMS distances indicate better performance.

Table 1: Varying Sampling Ratio (Fixed #Bins = 30)

Table 2: Varying Number of Bins (Fixed Sampling Ratio = 0.1)

From the tables above, we can see that our method is extremely stable and robust against various selections of sampling ratios and number of bins. We will ensure that this robustness characteristic is highlighted in the final version of our paper.

The author-reviewer discussion is coming to a close. Would you please go through the authors' rebuttal ASAP, highlight any remaining concerns (if any) so that the authors have sufficient time to respond to them, and adjust your review/rating accordingly?

Thanks!

Your AC

Thank you for your tremendous efforts in reviewing our paper. We would like to further answer your questions as follows.

[Q-1] Lack of Discussion of Distinct Contributions of Our Work

[A-1] We would like to clarify that there are two non-trivial contributions in our paper stated as follows:

• First, benchmarking routing performance based on uncertainty scores. This is a non-trivial result (we are glad that the Reviewer acknowledged our contribution in S1), as well-calibrated uncertainty scores do not always theoretically correlate with correctness [1, 2]. Therefore, a comprehensive benchmark and analysis are necessary, as previous work has often overlooked the importance of establishing fair and valid result comparisons.
• Second, we propose a proxy dataset construction method to predict the routing curve without the need for input queries, which enables efficient analysis of routing performance–cost trade-offs. In contrast, existing work [3, 4] typically relies on the (1) full set of test queries to construct the routing curve and (2) they observe these trade-offs between large-scale LLMs. Our work generates a proxy dataset and tests its ability for assisting SLM routing without any user queries. The constructed data can accurately predict the routing curve, allowing users to anticipate which UQ method and threshold offer the best trade-off even on SLM settings.

[1] Huang, et al. "Look before you leap: ...." arXiv 2023.

[2] Yona, et al. "Can Large Language Models Faithfully ...?." arXiv 2024.

[3] Ong, et al. "RouteLLM: Learning to Route LLMs from Preference Data." ICLR. 2024.

[4] Hu, et al. "Routerbench: A benchmark for multi-llm routing system." arXiv 2024.

[Q-2] Why Prior Work Faces Challenges on New Downstream Datasets

[A-2] We would like to clarify from two points:

• Classifier-style routers trained on historical prompt–performance data often struggle on new downstream tasks because the input distribution and the relative strengths of candidate LLMs shift. RouterBench [1] explicitly runs "out‑domain" experiments to test cross‑task transfer, observing that some routers fail to generalize to new tasks or new models. In contrast, our benchmark focuses on routing methods based on uncertainty quantification (UQ), which do not rely on training on historical prompt–performance data. However, even for training‑free UQ methods (e.g., average token probability, verbalization), they also face a challenge: how to set routing thresholds for unseen datasets to acheive effecitve routing performance. Our proposed Proxy Routing Data approach addresses this by effectively estimating routing curve for new downstream data.
• In our experiments, we test the Proxy Routing Data under the setting of "fully out-of-domain" to showcase the generalization ability on the challenges faced by the prior work. The results show that Proxy Routing Data can help predict the routing curve accurately without needing any out-of-domain or new downstream data.

[1] Hu, et al. "Routerbench: A benchmark for multi-llm routing system." arXiv 2024.

[Q-3] Unclear Notations in Algorithm 1

[A-3] We are sorry for the confusion. $X_i$ denotes the subset of proxy data $X$ followed the confidence score distribution $\\{F_{\mathbb{D}}\\}_{i=1}^{M}$ of $i$ datasest $D_i$. The element $x_j \in X_i$ denotes the $j$-th element in $X_i$. Thus, we can say, for every element $x_j \in X_i \subset X$. We will include this in our next version of paper.

[Q-4] Justification of Sampling in Proxy Routing Data Construction

[A-4] Our goal focuses on the routing configuration of on-device SLM. Unlike prior work conducted on large-scale server environments, on-device settings typically involve limited storage and computational resources, making it infeasible to store large-scale datasets on edge devices. Therefore, we evaluate routing curve prediction using a small proxy dataset designed to simulate real-world on-device scenarios. While it is possible to use 100% of the data for routing prediction under unconstrained resource conditions, such an approach contradicts the resource-limited assumptions of on-device SLMs explored in this work, which is the reason that our paper leverage 10% of data as an example. To better approximate the distribution of the original data, we adopt binned sampling, which ensures that the sampled proxy data more closely reflects the characteristics of the full population.

To further address the reviewer’s concern, we varied the sampling ratio used to construct proxy routing data while fixing the number of bins at 30. On HellaSwag, we routed from three SLMs to GPT‑4o‑mini using the OOD Probe. We report the average RMS distance between the Routing Ratio vs. Accuracy curves obtained with thresholds from the full downstream dataset and those derived from our method (lower is better). Results (in Table 1) show that our proxy data construction pipeline maintains its ability to predict routing curves even with a very small sample size (only sampling 1%), supporting its suitability for on‑device scenarios.

Table 1: Varying Sampling Ratio (Fixed #Bins = 30)

[Q-5] Sections 2 and 5 Can Be More Condensed

[A-5] We will further condense Sections 2 and 5 (~ 1.5 pages now) and include the discussions we had during the rebuttal in our next version. In Section 2, we include the introduction of method in our benchmarking settings, while in Section 5, we include the discussion of using proxy data to on-device scenarios, which are essential to highlight our contributions.

[Q-6] Sufficient Tweaking of SLMs; Insufficient Tweaking of LLMs

[A-6] We focus on deploying SLMs on edge devices, which motivates our emphasis on routing from SLMs to LLMs. Routing between LLMs falls outside the scope of this paper. We appreciate the reviewer’s insightful comment and will consider this as a promising direction for future work.

[Q-7] The cost of producing the routing dataset is at least one inference forward pass + UQ forward pass in the full routing set

[A-7] All routing data can be precomputed before deploying for edge device routing curve prediction, meaning that the cost of generating the routing dataset, consisting of one or two forward pass for inference, is incurred only once offline, and does not affect the efficiency during on-device inference serving.

The author-reviewer discussion is coming to a close. Would you please go through the authors' rebuttal ASAP, highlight any remaining concerns (if any) so that the authors have sufficient time to respond to them, and adjust your review/rating accordingly?

Thanks!

Your AC