Prompt Optimization with Logged Bandit Data

Authors: Haruka Kiyohara, Daniel Yiming Cao, Yuta Saito, Thorsten Joachims
Primary area: causal reasoning
Keywords: off-policy evaluation, prompt tuning, large language models, contextual bandits

Abstract:
We study how to use naturally available user feedback, such as clicks, to optimize large language model (LLM) pipelines for generating personalized sentences using prompts. Naive approaches, which estimate the policy gradient in the prompt space, suffer either from variance caused by the large action space of prompts or bias caused by inaccurate reward predictions. To circumvent these challenges, we propose *Direct Sentence Off-policy gradient* (DSO), which estimates the policy gradient by leveraging similarity among generated sentences, substantially reducing variance while suppressing the bias. Empirical results on our newly established suite of benchmarks, called *OfflinePrompts*, demonstrate the effectiveness of the proposed approach in generating personalized descriptions for movie recommendations, particularly when the number of candidate prompts is large.

Decision: reject
Reviewer scores: 5, 3, 8, 8

Reviews and discussion:
<thread>
creation date:
1729173079595

thread_type:
official_review

<thread_object>
creator: reviewer_h6mu

<content>
summary:
The paper introduces a new method for offline prompt policy learning for LLMs. The main challenge in this setting is the distribution shift between the logged data and the target data. Importance sampling can correct the distribution shift but only at the cost of potentially very high variance. The key idea behind the new method is to exploit similarity relations between sentences to reduce the variance. The bias-variance trade-off of the new method is analyzed theoretically and the method is tested on synthetic data and a LLM movie description task.

strengths:
* The method is well-motivated and the theoretical analysis supports the desired variance reduction. Intuition for the analysis is provided. 
* Ablations w.r.t to differences in the setting (dataset size, number of actions, reward noise) and w.r.t to the hyperparameters (kernel type, kernel bandwidth) of the method are carried out.
* Plan to open-source a benchmark for offline prompt policy learning

weaknesses:
* Figure 6: there are 5 bars for each method. I was/am a bit confused about what the difference between these bars is. For now, I assume these are the results from the 5 random seeds, ordered by performance. But I think it would be good to have a label for this or mention it in the Figure caption. 
* Literature on contextual bandits/kernelized bandits is left out.
* The performance gain (in particular compared to regression) seems much stronger in the synthetic setting than in the full-LLM experiment.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730774514717

thread_type:
official_review

<thread_object>
creator: reviewer_mszw

<content>
summary:
The paper presents Direct Sentence Off-policy gradient (DSO) for optimizing large language model (LLM) pipelines using logged user feedback such as clicks. DSO addresses the challenges of high variance and bias in policy gradient estimation by leveraging the similarity among generated sentences. The paper provides theoretical analysis on the source of bias and variance reduction of DSO. Experiments on both synthetic environment and a proposed benchmark (OfflinePrompts based on MovieLens-10M) demonstrate the effectiveness of this method. OfflinePrompts is a new benchmark suite,to demonstrate DSO's effectiveness in generating personalized movie descriptions. This is an additional contribution of the paper by providing a practical solution for leveraging naturally logged feedback for prompt policy optimization in language generation tasks.

strengths:
- The algorithm DSO motivated by utilizing the information behind the sentence embedding is generally sound.
- The theoretically anslysis highlights the benefit of such algorithimic designs by indicating the source of bias and variance of such algorithms.
- The introduction of the OfflinePrompts benchmark suite is a valuable resource for the research community, facilitating further development and testing of off-policy learning methods for language generation

weaknesses:
The experiments for real-world validation is insufficient. (Indeed, we lack good benchmarks for this task.) How well does the real-world performance align with the score/reward in the simulated environment (OfflinePrompts)? I found Figure 11 in the appendix indicates the positive correlation between the simulated rewards and the click feedback from users. Is there other statistics (such as the accuracy)? I am curious on the click rate improvement using the policy trained by DSO in real-world settings.

questions:
How well the sythetic environments represent the real case? I note that there are some gaps between the sythetic environments and the target task. For example, reward is real-valued in synthetic case but it is binary in the real case (click or not); the policy is parameterized by an estimated reward function in the sythetic case.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730807851912

thread_type:
official_review

<thread_object>
creator: reviewer_qx8z

<content>
summary:
This paper addresses prompting policy optimization for large language model (LLM) pipelines by leveraging logged user feedback, such as clicks, to generate personalized sentences. The authors propose a novel method called Direct Sentence Off-policy gradient (DSO), which uses similarity between generated sentences to estimate the policy gradient. While this approach relies on importance sampling, it can reduce the variance of importance weights by treating them in a sentence space rather than the prompt space. Experiments on a synthetic task and an LLM-based task for personalized movie descriptions are shown to claim the effectiveness of the proposed DSO method.

strengths:
* Using similarity in the generated sentence space to control the bias-variance tradeoff through importance weights is an interesting approach.
* The paper evaluates the proposed method on two types of tasks, synthetic and LLM-based tasks, demonstrating applicability in varied settings.
* Theoretical analysis provides insights into the characteristics of DSO, although some detailed proofs could not be fully verified by the reviewer.

weaknesses:
**Lack of clarity in algorithmic steps**:
The specific steps for implementing the algorithm are unclear. It seems that gradient estimation would require sampling from both the prompt policy and the LLM. If this understanding is correct, how many samples would need to be generated per data point? Should this match the $m$  samples used to estimate $\pi_0(\phi(s_i)|x_i)$?

**Notation abuse and lack of clarity in definitions**:
This paper has some notation abuse, which leads to ambiguity. For example, the authors introduce $\pi_\theta(a|x, s)$ or $\pi_\theta(a|x, \phi(s))$ as a conditional distribution over prompts given the generated sentence $s$ in  Section 4 and Appendix D.1. However, this is problematic because $\pi_\theta$ is originally defined as a prompt selection policy and should not depend on $s$, which the LLM generates after selecting $a$. Additionally, while the expressions are somewhat interpretable, there is a lack of consistency in function arguments throughout the paper. For instance, $\pi_\theta(s|x)$ is used without explanation as $\sum_a \pi_\theta(a|x) p_{LLM}(s|x, a)$. To improve clarity, the authors should avoid redefining $\pi_\theta$ with different inputs and instead provide explicit auxiliary definitions where needed, along with a rationale for introducing these conditional probabilities.

**Unpractical setting  in Full-LLM Experiment with MovieLens**:
The LLM-based experiment in Section 7 lacks realistic user personalization. As shown in Figures 10 and 12, the prompt policy reduces user information to a single word (from a set of only 1000 words) before feeding it to the LLM. This simplistic representation raises concerns about whether the Full-LLM experiment setup can effectively capture real-world personalization. Without a richer prompt (e.g., short sentences) to convey nuanced user information, it is unclear if this approach offers any advantage over simply passing user attributes directly to the LLM. Consequently, this setup might be better categorized as a toy task rather than a realistic evaluation of the proposed method's applicability in real-world tasks.

**Concerns regarding the formulation of baseline approaches**:
The problem formulation in this work is novel; however, applying existing methods, particularly the regression approach, seems overly naive for this setup. Since the LLM that generates $s$ is available in this setup, it would be more appropriate for the reward predictor to take $(x, s)$ as input instead of $(x, a)$. Otherwise, the reward predictor would have to learn the LLM's inherent randomness (noise), which seems inefficient. Using $(x, s)$ would allow the reward predictor to avoid this redundancy and better capture the generated sentence features. A Nadaraya-Watson kernel regression (using the same kernel as in DSO) or a neural model like DistilBERT could be employed as the reward predictor to improve adaptability. In connection with the above, in the numerical experiments, using $(x, a)$ as the reward predictor's input in the regression approach may be unfair as a baseline comparison against DSO. DSO leverages (multiple) generated sentence(s) $s’$ for each context $x$ sampled from $\pi_\theta$ and the LLM. Thus, any observed performance gap between DSO and the regression approach may simply be due to this difference in formulation rather than any inherent advantage of DSO.

**Organization of the paper**:
The structure of the paper could be improved. For instance, details of the synthetic experiment setting and Section 4.2 (not cited in the main text) could be moved to the appendix, as these sections may be of lower priority for understanding the main contributions. Shifting these sections would allow more space for core elements like detailed algorithmic steps, problem setup, and full LLM experiment details in the main text.

questions:
**Figure 6 Interpretation**:
It seems that each bar in Figure 6 represents the results across 5 random seeds. Given the variation across seeds, can we still conclude that the proposed method (DSO) consistently outperforms the regression baseline? The performance between DSO and regression appears similar when accounting for this variability.

**Minor comments**
* Line 391: $\sigma_o$ should be $\sigma_s$?
* Line 989: MSE loss should be $\sum_{i=1}^{n} (r_i - \hat{q}(x_i, a_i))^2$ instead of $\sum_{i=1}^{n} (r_i - \hat{q}(x_i, a_i))$.
* Line 1075: $\nabla_{\theta} \pi_{\theta}$ should be $\nabla_{\theta} \log \pi_{\theta}$?
* In Section 3.1, the classification of "conventional approaches" into "regression-based methods" and "importance sampling (IS)" feels somewhat unclear. It may be more intuitive to categorize these as "reward predictor-based approaches" and "reward predictor-free approaches." This distinction clarifies that IS methods directly use observed rewards, whereas regression-based methods estimate rewards across all actions.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731254196617

thread_type:
official_review

<thread_object>
creator: reviewer_layf

<content>
summary:
This paper proposed a new policy gradient-based prompt optimization. The goal is to learn a policy that is able to generate prompts with good responses (as in good rewards). This paper proposed a new DSO that is better than the traditional policy gradient and IS-based method. Some experimental results provided by this paper show that the new method is able to outperform others.

strengths:
1. The idea of learning a policy to generate good prompts is new to me.
2. The proposed method clearly addressed the weakness of IS.

weaknesses:
1. The experimental session is the major weakness of this paper. This paper only contains a synthetic experiment and a single model experiment on a single dataset witha  simulated reward function. Experimental results on more datasets and models will make the paper more convincing.

2. The following work should be discussed in the related work since they study prompt optimization with human feedback by learning a reward function and hence related.

https://arxiv.org/abs/2402.00396
https://arxiv.org/abs/2405.17346

An similar line of work on prompt optimization should also be discussed:

https://arxiv.org/abs/2306.03082
https://arxiv.org/abs/2310.02905
https://arxiv.org/pdf/2402.09723

questions:
1. Can the author describe the main insight for the theorems in this paper? and how they are reflected in the performance of the new approach? There seems to be some disconnection between the theoretical section and empirical verification.

2. How does your method perform in a normal prompt optimization setting? like [1]?


[1] https://arxiv.org/abs/2306.03082
</content>
</thread_object>

</thread>
<thread>
creation date:
1732315506099

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revision

comment:
Dear reviewers,

Thank you once again for your valuable feedback on the paper. We have uploaded the revision based on your comments. We summarize the key updates below.

- Additional discussions of related works
  - Additional papers on online learning and RLHF (Section 2, reviewer LAYf)
  - Additional discussion on kernelized bandits (Appendix A, reviewer H6Mu)
  - Additional comparison with LLM-based recommendation (Appendix A, reviewer Qx8z)
- Additional clarifications on notations (Sections 3 and 4, reviewer Qx8z)
- Additional clarification on the full-LLM experiment (Section 7, Figure 6, and Appendix C.2, reviewers H6Mu and Qx8z)
- Additional discussion of future work regarding benchmarks (Section 8, reviewer LAYf)
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a method for learning a prompt policy using only offline data with bandit feedback. The proposed method is based on policy gradient and consists of specilaized techiniques to reduce both the bias and variance of the reward estimator. The paper also introduces a new benchmark for the studied problem setting of prompt optimization with offline bandit feedback.

The reviewers generally agree that the studied problem setting is interesting, and the theoretical results are insightful and useful since they provide support for practice.

The reviewers gave disparate scores for the paper even after the rebuttal. One common concern which is shared by 3 reviewers (Reviewers LAYf, Qx8z and MsZw) is that the real-LLM experiments may be insufficient, for example, the comparisons with previous methods are not enough or not entirely fair, the benchmark used in the full LLM experiments may not be representative of real scenarios. After reading the discussions between the authors and reviewers, I tend to agree with the points raised by Reviewer Qx8z. That is, the benchmark constructed using the MovieLens dataset may be too simplistic (see Reviewer Qx8z's initial review for details). I understand the authors' response stating that their method can be applied to more complex prompts, but unfortunately its performance in such experiments is not validated yet. I also agree with Reviewer Qx8z that in the regression baselines, the regression methods should take the generated sentence $s$ as input (in fact, I wonder can you let the regressors take both $a$ and $s$ as input?), and hence encourage the authors to follow this in both synthetic and full LLM experiments. The regression-based methods, including those taking $a$ and $s$ as inputs, have very close performances with the proposed DSP method in the full LLM experiments. This is another important concern since it puts into question whether the proposed method is indeed practically superior than the previous baselines. I also agree with Reviewer Qx8z's comment regarding using kernel-based regression as the regressor in the experiments for a fair comparison, since kernel similarity is also used in the proposed DSP.

I understand that the paper contains some theoretical contributions as well. However, in my opinion, for the topic studied in this paper (offline prompt optimization), having a practical algorithm (with appropriate and sufficient benchmarking) provides more value than theoretical contributions. Given the above, rejection is recommended. We encourage the authors to take into account the comments from the reviews, especially those related to the real-world LLM experiments, which will further strengthen the contributions of the paper.