PAK-UCB Contextual Bandit: An Online Learning Approach to Prompt-Aware Selection of Generative Models and LLMs

We thank Reviewer 5QGu for the thoughtful feedback on our work. Below is our answer to the reviewer's comments and questions.

1. Regret and complexity of PAK-UCB and RFF-UCB

First, we would like to clarify that the numerical results in Section 6 are reported for PAK-UCB (Alg.2) and RFF-UCB (Alg.3) in the main text.

Also, as we stated in Theorem 1, the regret bound is shown for the variant of PAK-UCB in the Appendix, which we titled Sup-PAK-UCB (Alg.4). We note that our analysis of PAK-UCB and Sup-PAK-UCB parallels the analysis of LinUCB [1] and KernelUCB [2] references. Specifically, [1] introduces two versions of the LinUCB algorithm: "LinUCB" (Alg.1 on p.g. 3 in [1]) is recommended for usage in practice, whereas [1]'s' regret analysis is performed for "SupLinUCB" (Alg.3 on p.g. 4). Similarly, Reference [2] also introduces two versions of KernelUCB, "KernelUCB" (Alg.1 on p.g. 5) and "SupKernelUCB" (Alg.2 on p.g. 5), where the numerical application and theoretical analysis are aimed for the algorithms. Our analysis also follows a similar approach. We will include this clarification in the revised paper.

2. Regarding Lemma 2

We would like to clarify that Lemma 2 still holds for the variant of RFF-UCB analyzed in Theorem 2 (described in Appendix B.2). Note that Sup-PAK-UCB (Alg.4) computes the UCB values on the mutually exclusive subsets $\\{\Psi\_g^m\\}\_{m \in [M]}$, which satisfies $\sum_{m,g}|\Psi_g^m| \le t$ at the $(t+1)$-th iteration. Therefore, Sup-PAK-UCB using Compute_UCB_RFF (Alg.3) requires time at most $\Theta(\sum_{m,g}|\Psi_g^m| s^2) = O(ts^2)$ and space $\Theta(\sum_{m,g} |\Psi_g^m| s) = O(ts)$. We will add this clarification to the revised paper.

3. Performance of RFF-UCB

As pointed out by the reviewer, there could be a performance gap between PAK-UCB-poly3 and RFF-UCB (e.g. in Figure 2). Note that this is not in contradiction with the regret bound of Theorem 2, because PAK-UCB-poly3 and RFF-UCB use different kernel functions to predict the scores: RFF-UCB uses the Gaussian kernel, whereas PAK-UCB-poly3 uses a polynomial kernel with degree 3. Therefore, RFF-UCB is not supposed to match the result with PAK-UCB-poly3.

3. Discussing Remark 1 in the Related work

Thank you for the suggestion. We will update the related work with a summary of Remark 1.

4. Typo in lines 212-213

We thank the reviewer for pointing this out. We would like to clarify that the tilde in the notation $\widetilde{\Phi}$ is a typo, which we will correct in the revision. In the equation, $\widetilde{\Phi}$ should be replaced with $\Phi$.

5. Results on regret

We appreciate the reviewer's question about the numerical performance in terms of regret values. We note that the evaluations based on O2B and the (average) regret are equivalent. This is because the average regret is computed as $\text{Avg.Regret}(T) := \frac{1}{T} \sum_{t=1}^T (s_\star (y_t) - s_{g_t}(y_t))$, while O2B is computed as $\text{O2B}(T) := \frac{1}{T}\sum_{t=1}^T (s_{g_t}(y_t) - s_{g^\star}(y_t)) = -\text{Avg.Regret}(T) + C,$ where $g^\star$ is the best single model with the highest expected score. Therefore, the O2B scores of different policies are all shifted with the same constant $C := \frac{1}{T}\sum_{t=1}^T (s_\star (y_t) - s_{g^\star}(y_t))$ compared to their average regret. As a result, the O2B rankings of the policies in the plots are identical to the regret-based rankings.

6. Regret lower bound

First, we note that the regret (upper) bound derived in Theorem 1 matches the regret of the LinUCB [1] and KernelUCB [2] algorithms up to a factor of $\sqrt{G}$, where $G$ is the number of models. On the other hand, we anticipate that the regret lower bound could scale with $\Omega(\sqrt{dGT})$ for a kernel function with a finite dimension $d$ (e.g., by slight modification of [Theorem 2, 1] for linear bandits without arm-specification). Formally proving a regret lower bound for the arm-specific setting is an interesting future direction for our work. We will discuss this in the revised conclusion.

[1] Chu, et al. "Contextual bandits with linear payoff functions." AISTATS 2011.

[2] Valko, et al. "Finite-time analysis of kernelised contextual bandits", UAI 2013.

We thank Reviewer 3Dpj for the thoughtful feedback on our work. Below is our answer to the reviewer's comments and questions.

1. Details of the datasets

We note that the details of experiment settings are discussed in Appendix D. The following is a summary of the details asked by Reviewer 3Dpj, which we will include in the revised main text to improve the clarity:

• Setup 1: Prompts are uniformly randomly selected from the MS-COCO dataset under two categories: 'dog'/'car' (Fig.2), 'train'/'baseball-bat' (Fig.3a), 'elephant'/'fire-hydrant' (Fig.3b), and 'carrot'/'bowl' (Fig.3c).
• Adaptation to new models (Fig.4a): Prompts are uniformly randomly selected from the MS-COCO dataset under categories 'train' and 'baseball-bat',
• Adaptation to new prompt types (Fig.4b): In the first 1k iterations, the prompts are uniformly randomly selected from a pool that initially includes categories 'person' and 'bicycle' in the MS-COCO dataset. Then, categories 'airplane', 'bus', 'train', and 'truck' are added to the pool after each 1k iterations,
• Synthetic T2I and image-captioning experiments (Fig.19 and Fig.21): The prompts are uniformly randomly selected from the MS-COCO dataset under categories 'dog', 'car', 'carrot', 'cake', and 'bowl',
• Synthetic T2V task (Fig.22): The captions are uniformly randomly selected from the MSR-VTT dataset under categories 'sports/action', 'movie/comedy', 'vehicles/autos', 'music', and 'food/drink'.

2. Evaluation scores

In the numerical experiments of this paper, we primarily focus on text-to-image generation tasks and use CLIPScore as the evaluation score. We note that the online selection framework can be applied to other prompt-guided generation tasks as long as we know the score values assigned to the generated samples.

3. Features for kernel methods

The input of the PAK-UCB method and other baselines is the embedded prompt that is output by the pretrained CLIP-ViT-B-32-laion2B-e16 model from the open_clip repository (https://github.com/mlfoundations/open_clip/tree/main). Only for LinUCB and KernelizedUCB baselines, we also concatenate the one-hot encoded vector of the model index to the CLIP-embedded prompt.

4. Performance of the random selection strategy

We thank the reviewer for the suggestion of including the random selection strategy as a baseline. The random selection strategy would be expected to underperform for prompt-based model selection. Below, we provide the results in Setup 1 (Figure 3b) where we use PAK-UCB-poly3 as the competing strategy.

5. Expressivity of score estimation functions

The use of a polynomial kernel with degree 3 is inspired by the kernel Inception distance (KID) in the literature of generative models [1, 2]. Please note that a higher degree can lead to higher expressivity while increasing the risk of overfitting the data. Below, we conduct an ablation study to test the effect of degree on the performance of the PAK-UCB algorithm using a polynomial. We observe that a degree of 3 can achieve a better tradeoff between expressivity and generalization.

[1] Stein et al. "Exposing flaws of generative model evaluation metrics and their unfair treatment of diffusion models." NuerIPS 2023.

[2] Bińkowski et al. "Demystifying mmd gans." ICLR 2018.

6. Comparison to KernelUCB

In the following, we will compare the process in PAK-UCB and KernelUCB side-by-side to highlight their differences:

• Problem Setting of PAK-UCB (ours): We have $G$ arms where each arm represents a fixed generative model that remains unchanged across rounds: for example, Arm 1 represents the Stable Diffusion model in all rounds. At each round, the arms observe one shared context variable (i.e., the text prompt). We learn $N$ separate kernel-based models with different weights to predict the CLIPScore of a shared incoming prompt (context) for the $G$ fixed generative models.
• Problem Setting of KernelUCB [1]: At every round, we have $N$ arms where the expected reward of each arm is fully characterized by its context variable. The arms have different context variables at each round. We learn one shared set of weights to predict the expected reward for the $N$ observed contexts (i.e., arms) in the next round.

As explained above, in the KernelUCB [1] setting, there is not a fixed model corresponding to one arm across iterations, and the arms will perform independently across iterations depending on their context. However, in the setting of PAK-UCB (our method), each arm will represent one fixed generative model in all the learning rounds.

[1] Valko, et al. "Finite-time analysis of kernelised contextual bandits", UAI 2013.

We thank Reviewer MgCP for the thoughtful feedback on our work. Below is our answer to the reviewer's comments and questions.

1. Arm-specific reward model

We thank the reviewer for pointing out the bandit algorithm in [1] that utilizes arm-specific prediction functions. To the best of our knowledge, the arm-specific bandit algorithms in the literature (References [1-3]) consider a linear pay-off. However, in our experiments, we observed that the PAK-UCB algorithm performs better when non-linear kernel functions (e.g., polynomial kernel and RBF kernel) are applied. We think this phenomenon is due to the normalized output of the CLIP embedding that is on the surface of the unit sphere, where a linear classification could be sub-optimal compared to a non-linear rule provided by the kernel-based approach. In the revision, we will discuss the existing arm-specific implementations of LinUCB and the necessity of including non-linear kernel functions to obtain better results in the case of text-to-image model selection.

[1] Li et al. "A contextual-bandit approach to personalized news article recommendation." WWW 2010.

[2] Fang et al. "Networked bandits with disjoint linear payoffs." KDD 2014.

[3] Xu et al. "Contextual-bandit based personalized recommendation with time-varying user interests." AAAI 2020.

2. Performance on ImageReward DB dataset

Thank you for the suggestion. Due to the limited time, we report preliminary results on the ImageRewardDB ReFL dataset [1]. In the experiment, the algorithm selects among three T2I models considered in our paper: StableDiffusion v1.5, UniDiffuser, and PixArt-$\alpha$, which attain a CLIPScore of 35.54, 34.40, and 37.20 on (a subset of) the dataset, respectively. After 5k iterations, our proposed PAK-UCB-poly3 and RFF-UCB algorithms attain an OPR (ratio of picking the best model PixArt-$\alpha$) of 63.76% and 42.94%. On the other hand, the KernelUCB-poly3 baseline achieves only 34.18%. We will cite and discuss the dataset in the revised paper.

[1] Xu et al. "Imagereward: Learning and evaluating human preferences for text-to-image generation." NeurIPS 2023.

3. Comparison with the DiffusionGPT framework [1]

We thank the reviewer for pointing out the alternative DiffusionGPT method in [1], which leverages an LLM agent for model selection. Our proposed PAK-UCB approach can be viewed as complementary to the DifussionGPT method, which employs the UCB bandit approach to address this task. We will cite and discuss the potential combination of the PAK-UCB and DiffusionGPT approaches as a future direction in the revised text.

[1] Qin et al. "DiffusionGPT: LLM-driven text-to-image generation system."

4. Handling unseen categories

We agree with the reviewer that adaptation to unseen prompt categories would be challenging if the new prompts were fully orthogonal to previous prompts. On the other hand, in practice, it is often the case that the incoming prompt has some correlation with some previous prompts. For example, a user who has generated images of cats and dogs is likely to generate images of other pets or animals in the future. Assuming that the optimal model choice changes continuously with input prompts, the PAK-UCB online learning approach would be capable of predicting the optimal arm after observing previously correlated prompts.

To test this hypothesis, we conducted a numerical experiment to predict the CLIPScore of StableDiffusion v1.5. We train a poly3-based prediction model on $n=1,2,3$ categories and compute the prediction MSE on a new category after observing 10 samples in this new category. The results show that the prediction function can generalize effectively to unseen but correlated prompt categories.

5. Performance on a large number of arms

We appreciate the reviewer's comment on setups with a large number of arms. Please note that the primary goal of the online selection task is to minimize the regret. Our regret bound in Theorems 1 and 2 are on the order of $\widetilde{O}(\sqrt{GT})$ for $G$ arms and time horizon $T$. In practice, the effective number of arms will be lower if the performance scores change more smoothly across arms.

To numerically test the effect of a larger number of arms, we conducted an experiment to evaluate PAK-UCB-poly3 and RFF-UCB in a setting with ten arms: We included five more arms to the synthetic T2I task (Setup 3), which generate a clean image less frequently (only $10\%$ of the time). The results show that PAK-UCB-poly3 and RFF-UCB can outperform the best single arm and the KernelUCB baseline.