LLM Bandit: Cost-Efficient LLM Generation via Preference-Conditioned Dynamic Routing

Thank you for your detailed review and for raising several important points. We appreciate your feedback, as it will help us improve the presentation and ensure a clear understanding of our contributions. Please allow us to clarify the sections you highlighted:

Clarity of Writing in Section 2:

• In equation 2, the model identity vector $I_k$ and the the neural network $f$ are jointly trained.
• In equation 4, the distributions $q(I_k)$ and $p(I_k)$ represent the variational posterior and prior distributions, respectively, over the model identity vectors. This is similar to a typical VAE formulation, and here we parameterize $q(I_k)$ and $p(I_k)$ as Gaussian distributions. The KL divergence between posterior ad prior distribution acts as a regularization term on the learned representations so that it can generalize to unseen models.
• In equation 5, $I_{k^{\prime}}^T$ is the transpose of the model identity vector $I_{k^{\prime}}$. The policy $\pi(k^{\prime} \mid \cdot)$ outputs the probability of selecting model $M_{k^{\prime}}$.

Furthermore, Algorithm 1 describes the method in two stages: pretraining and training. However, these terms are not clearly defined in the main text, leaving their motivation unclear. It would be helpful if the author elaborated on the specifics of these stages and how they relate to the overall algorithm.

The pretraining stage is meant to pretrain the routing policy on the large collection of prompts using existing pairwise comparison datasets. Please see details in line 270-288. The training stage is RL training for the routing policy following equation 7.

Bandit Problem Formulation: The choice to frame the problem as a bandit problem seems misaligned with the algorithm itself. When I encounter the term "bandit," I expect the use of specific online bandit algorithms or techniques. However, Algorithm 1 appears more akin to a supervised-learning approach, as it seems to use prompt-reward pairs to train the IRT model (f) and the model identity vector (I_k). So what is the motivation for framing this task as a bandit problem? It would be helpful if the author could more explicitly justify the bandit analogy.

First, let us clarify that Algorithm 1 is for training the routing policy $\pi$, not for training the model identity vectors $I_k$ and IRT model $f$. Second, the motivation for framing the routing task as a multi-armed bandit is that the routing policy will choose one LLM to send the query. Third, the pretraining stage in Algorithm 1 is supervised training, but the training stage follow the RL optimization procedure.

In the second contribution, the author says they proposed a preference-conditioned mechanism. However there is no clear definition or explanation of what preference-conditioned means in the context of the paper.

The preference conditioned mechanism is explained in line199-205. The motivation to condition the routing policy on user preference is to allow users to specify preference at test time.

Similarly, the term "action-space aware" in the third contribution is introduced but not sufficiently explained. This lack of clarity hinders understanding and prevents a full appreciation of the contribution. A more thorough explanation of these terms and their relevance to the method would be beneficial.

The explanation for “action-space aware'” is around line182-198. Since user may specify an arbitrary set of LLMs to be selected at test time, the routing policy needs to adapt to arbitrary action spaces. We accomplish this by condition the routing policy on the set of model identity vectors, thus making the policy aware of the available actions.

In Contribution 4, the author claims that the method generalizes across prompts by "leveraging a wide range of data." However, generalization typically refers to a model's ability to perform well on unseen data—not merely data it has been trained on. Therefore the usage of "generalization" here might not be proper.

Thanks for pointing it out, we will rephrase it to make it clear. What we really mean is training on a wide range of data helps the routing policy gain generalizability to unseen prompts, which is enabled by the proposed pretraining procedure.

We sincerely appreciate your thorough review and the opportunity to clarify these points. Addressing these concerns will significantly improve the clarity and presentation of our work, ensuring a better understanding of our contributions. Please let us know if you have any further comments or suggestions.

Thank you for your insightful review and for acknowledging the strengths of our proposed preference-conditioned routing mechanism for LLMs, particularly its ability to generalize across diverse models, tasks, and user preferences. We appreciate you highlighting the significance of the model identity vector concept and its potential applicability to other tasks beyond routing.

Regarding the weakness pointed out, we thank you for bringing the concurrent EmbedLLM work to our attention. We were not previously aware of this work, which does indeed propose a similar concept of learning compact representations (embeddings) for LLMs. In the revised version of our paper, we will include a detailed discussion of the EmbedLLM approach and compare our method for generating model identity vectors to their systematic framework.

To your question about potentially marrying our identity vector approach with the EmbedLLM method, we find this to be an intriguing direction for future work. One possibility could be to use the EmbedLLM framework to learn a general-purpose embedding for each LLM, and then fine-tune these embeddings using our IRT-based approach to obtain task-specific identity vectors. Alternatively, we could explore incorporating the pairwise comparison data and KL regularization from our models into EmbedLLM training process to improve generalizability of the learned representations.

We thank you again for bringing this related work to our attention and for the suggestion to explore potential synergies between the two approaches. Integrating ideas from EmbedLLM could further enhance the effectiveness and generalization capabilities of our routing mechanism.

Please let us know if you have any other comments or suggestions. We appreciate your feedback and will incorporate it to improve our work.

Thank you for your thorough review and insightful comments on our paper. We appreciate the acknowledgment of the novelty and significance of our proposed preference-conditioned dynamic routing mechanism for LLMs. We will strive to clarify the concerns raised and address the weaknesses identified in the revised version of the paper.

Regarding the complexity of the training procedure, we agree that the method involves several components and data sources. To improve clarity, we will restructure the relevant sections and provide a more comprehensive overview of the training pipeline, explicitly detailing the data sources and their roles in pretraining and training each component. We will also consider adding a supplementary diagram or pseudocode to better illustrate the flow of data and components. Additionally, we plan to open-source our code to allow for transparent inspection of the data pipelines.

Concerning the handling of uncertainty in reward prediction, we acknowledge that our current approach does not explicitly model this uncertainty. However, the prediction scores $p\hat{p}_k = \text{sigmoid}(f(x, I_k))$ obtained from the model identity vectors are provided as auxiliary information to the routing policy during training. By conditioning the policy on these prediction scores, it can implicitly learn to account for the uncertainty in the predictions. When the prediction scores are less confident (closer to 0.5), the policy will place less weight on them during routing decisions. Conversely, for prompts where the predictions are more certain, the policy can rely more heavily on the prediction scores. In this way, uncertainty is implicitly communicated to the policy through the prediction score values themselves.

To address the question about joint training model identity vector and routing policy, we believe it can make training unstable. Since the model identity vectors effectively define the action space for the routing policy, allowing them to change during training makes the action space non-stationary, violating the Markov assumption and causing training to diverge. That said, we agree that end-to-end training could be an interesting future direction to explore more fully. Potential approaches could involve alternating optimization, where the identity vectors are kept fixed while training the policy, and then the vectors are updated while keeping the policy weights fixed. Techniques from the literature on non-stationary reinforcement learning may also provide insights.

Regarding the novelty of the training algorithm and data used, we acknowledge that some components, such as the IRT model and the use of existing evaluation leaderboards, are adapted from prior work. However, we believe the primary novelty lies in the formulation of the preference-conditioned routing policy, the action-space-aware architecture, and the integration of diverse data sources (evaluation scores, pairwise comparisons, and user preferences) into a unified training pipeline. We will endeavor to better highlight these novel aspects in the revised version.

Finally, concerning the exploration-exploitation balance, there are a couple of mechanisms at play. First, our routing policy outputs a distribution over the action space (selecting LLM $M_k$). During training, trajectories are sampled from this stochastic policy distribution, allowing for exploration of different actions. Second, the entropy regularization term in PPO explicitly encourages the policy to explore different actions. Even if the current value estimate $V_{\pi}$ for an action is low, there is a non-zero probability for the policy to still explore that action, which allows updating the value estimate if the actual reward contradicts the current estimate. Furthermore, by conditioning the policy on predicted scores $\hat{p}_k = f(x, I_k)$ from the model identity module, we provide an inexpensive signal about the expected reward for each action. This predictive information assists with guiding exploration towards more promising actions without requiring actual environment rollouts. Together, the stochastic policy, entropy regularization, and inexpensive predictive model scores enable the routing policy to inherently balance exploration of uncertain context-action pairs and exploitation of high-value estimates from past experience.

We appreciate the reviewer's insightful comments and constructive feedback, which will undoubtedly improve the clarity and quality of our paper. We look forward to addressing the raised concerns in the revised version.

Thank you for your insightful review and feedback. We appreciate you raising these important points, and we will incorporate your suggestions into the revised manuscript to strengthen our work. Here are our responses to your specific comments:

1. Regarding the presentation of experimental results, we agree that including tabular results alongside the figures would enhance clarity and interpretability. In the revised manuscript, we will incorporate tables presenting relevant metrics for the different benchmarks and LLM configurations, complementing the existing figures.
2. We acknowledge your concern about the potential complexity introduced by integrating the multi-armed bandit framework with dynamic routing based on user preferences. To address this, we will provide an open-source implementation of our code, allowing others to examine and potentially extend our work. Additionally, we will include a discussion on the modular design and extensibility of our framework, highlighting how it can be efficiently implemented and scaled.
3. Concerning the exploration of the routing policy components, you make a fair point. We did not extensively tune hyperparameters for certain components, such as model quizzing and incorporating pairwise comparison data. The exact hyperparameters we used are provided in appendix. In the revised manuscript, we will acknowledge this limitation and emphasize that further improvements are possible through more comprehensive hyperparameter tuning.
4. Regarding the predictor baseline's performance, you raise an excellent point. In theory, if the performance prediction is accurate, the predictor policy should be optimal. Our routing policy aims to improve upon the predictor when those predictions are inaccurate or uncertain. For settings where our routing policy does not significantly outperform the predictor, it is likely due to the inherent difficulty in accurately predicting certain types of scores, such as discontinuous metrics like accuracy. We will include this analysis in the revised manuscript, as it provides valuable insights into the strengths and limitations of our approach.
5. Concerning the efficiency analysis, we appreciate you highlighting this aspect. Both the IRT model and routing policy are relatively lightweight, consisting of only a few linear layers. Therefore, the computational overhead introduced by our framework, compared to running LLM inference, is relatively low. We will quantify this overhead in the revised manuscript.
6. Regarding the MT-Bench comparison, thank you for bringing this up. We do have results on this benchmark, but we left them out of the initial submission as it is a relatively small dataset. Our routing policy obtains similar performance to RouteLLM. We agree that including these results would provide a more comprehensive evaluation. In the revised manuscript, we will add results on this dataset.

We sincerely appreciate your thorough review and constructive feedback. Addressing these points will significantly strengthen our work and contribute to a better understanding of the proposed routing framework's capabilities, limitations, and potential areas for further improvement.

Please let us know if you have any additional comments or suggestions.

Thank you for the feedback. I decide to keep my score.