PRewrite: Prompt Rewriting with Reinforcement Learning (2024)

Abstract: Prompt engineering is critical for the development of LLM-based applications. However, it is usually done manually in a “trial and error” fashion that can be time consuming, ineffective, and sub-optimal. Even for the prompts which seemingly work well, there is always a lingering question: can the prompts be made better with further modifications? To address these problems, we investigate automated prompt engineering in this paper. Specifically, we propose PRewrite, an automated method to rewrite an under-optimized prompt to a more effective prompt. We instantiate the prompt rewriter using an LLM. The rewriter LLM is trained using reinforcement learning to optimize the performance on a given downstream task. We conduct experiments on diverse benchmark datasets, which demonstrates the effectiveness of PRewrite.

text: Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 594–601 August 11-16, 2024 ©2024 Association for Computational Linguistics
To address these problems, we investigate automated prompt engineering in this paper. Specifically, we propose PRewrite, an automated method to rewrite an under-optimized prompt to a more effective prompt. We instantiate the prompt rewriter using an LLM. The rewriter LLM is trained using reinforcement learning to optimize the performance on a given downstream task. We conduct experiments on diverse benchmark datasets, which demonstrates the effectiveness of PRewrite.

heading: 1 Introduction
text: With the right prompts, large language models (LLMs) can show impressive performance on various tasks in zero-shot or few-shot settings (Brown et al., 2020; Srivastava et al., 2022). However, manual prompt engineering is done on a trial-and-error ad-hoc basis and there are limited guiding principles on writing good prompts.
To address the problems, we investigate methods to automate the process of prompt engineering, often called “automated prompt engineering” or “prompt optimization”. Automated prompt engineering is important due to the wide and fast adoption of LLM applications. Moreover, LLMs themselves are evolving, and as a result, we also need effective automated methods to update existing prompts to adapt to new models.
Several previous works have explored automated prompt engineering. AutoPrompt (Shin et al., 2020) uses a gradient-based search method to iteratively edit prompts, but requires gradient access to
the language model. RLPrompt (Deng et al., 2022) optimizes prompts using reinforcement learning (RL), but often produces uninterpretable gibberish prompts. Also using RL, TEMPERA (Zhang et al., 2022) allows editing prompts based on task input, but its small action space might hinder exploration. Another common limitation is that they are based on relatively small-size language models like BERT (Devlin et al., 2019) and RoBERTa (Liu et al., 2019). It is not clear how well the proposed methods can generalize to larger models, especially with API-only model access.
More recent works like APE (Zhou et al., 2023), OPRO (Yang et al., 2023) and Promptbreeder (Fernando et al., 2023) use larger models from the PaLM 2 (Anil et al., 2023) and the GPT1 model families. These works leverage LLMs themselves to propose prompt candidates, and search for a better prompt from them via validating performance on a given training dataset. We follow a similar idea but aim to use RL instead of search to improve the optimization process.
In this work, we propose PRewrite, prompt rewriting with reinforcement learning, to address the limitations above. Our idea is to train a prompt rewriter to rewrite an initial under-optimized prompt to a more effective prompt. The prompt rewriter itself is a LLM, trained using RL to optimize for a downstream task. We give an overview in Figure 1. Specifically, given an initial prompt, the prompt rewriter LLM is instructed to generate a rewritten prompt, which in turn is used by the task LLM to generate the final output. Using a reward computed on the final output against the ground-
1https://platform.openai.com/docs/models
594
truth output, the rewriter LLM is finetuned with RL. As compared to previous RL-based methods, PRewrite produces interpretable prompts (cf. RLPrompt), allows unconstrained exploration without manually defined action space (cf. TEMPERA), and leverages larger models (PaLM 2).
Our contributions are summarized as follows: • We propose PRewrite, a novel automated prompt
engineering approach. It optimizes prompt via rewriting, in an end-to-end manner using reinforcement learning. • We develop two rewriting strategies, including one that searches for an optimal rewritten prompt from candidates generated by the RL-trained prompt rewriter (Section 2.4). This often further improves the prompt optimization performance. • We conduct experiments on diverse benchmark datasets, which testify the effectiveness of PRewrite and demonstrate its state-of-the-art performance.

heading: 2 PRewrite

heading: 2.1 Problem Formulation
text: We formulate our prompt rewriting problem more formally in this section. Given a text generation task, we denote the task input and output by x and y respectively. To solve the task, one can use a LLM for prediction, y = LLM(p), where p is the input to the LLM, also known as prompt. Prompts are usually constructed using a template that incorporates the input x and a task instruction t. For the example, t can be “Write a brief answer for the following question” for a question answering task. A prompt can be constructed using template p="{t}: {x}" (Python f-string), where x is the input question. Please refer to Table 4 in Appendix for a complete example.
Prompt rewriting aims to rewrite a given initial prompt to another prompt p† = δ(p), in order to optimize the task output. We call the rewritten prompt p†. Since prompts can be constructed using an instruction, we simplify the prompt rewriting problem to only rewriting the instruction t† = δ(t). In fact, most prior works (Fernando et al., 2023; Zhou et al., 2023; Yang et al., 2023) optimize prompts via optimizing instructions, and do not differentiate between prompt and instruction. In this paper, we use the two terms interchangeably wherever it is clear.
Prompt rewriting can be performed independent or dependent of the task input, i.e., δ(·) and δ(·|x).
We focus on input-independent prompt rewriting, following most prior works. In this case, the instruction is rewritten offline and prompts can be constructed cheaply online using the rewritten instruction.

heading: 2.2 Overview
text: Directly searching for an optimal rewritten prompt is challenging due to the large search space of natural language. So, we propose PRewrite to optimize prompt rewriting, as illustrated in Figure 1.
First, the prompt rewriter takes in an initial prompt p and rewrites it to another prompt p†. The initial prompt is usually crafted manually and can be sub-optimal. Observing the remarkable capability of LLMs, we instruct a LLM (e.g., PaLM 2-S) with a meta prompt m for rewriting as follows:
p† = LLMR("{m}\nInstruction: {p}"). (1)
We call LLMR, rewriter LLM, which is to be differentiated from the task LLM, used for the end task. We list our meta prompts in Appendix B.
Second, the rewritten prompt p† is then used by the task LLM to generate the task output. The task LLM is assumed to be a blackbox accessed via API and can be larger than the rewriter LLM.
Third, we compute rewards based on the task output in comparison with the ground-truth output and use reinforcement learning (RL) to finetune the rewriter LLM on a training set (Section 2.3). This is critical because our meta prompt is very generic. As a result, the rewriter LLM and the rewritten prompt are unlikely to perform well on the downstream task initially.
Lastly, we use the RL-trained prompt rewriter to rewrite the initial prompt according to Equation 1 based on two strategies outlined in Section 2.4.

heading: 2.3 Finetuning Rewriter LLM with RL
text: This section provides more details on RL finetuning for our rewriter LLM, which is very similar to other RL-based LLM alignment work (Ouyang et al., 2022). Action space consists of all tokens in the rewriter LLM’s vocabulary, allowing arbitrary text rewriting. State is defined as the concatenation of all the decoded tokens so far. Reward is the task LLM’s performance on the downstream task when using the rewritten prompt. We measure this using the end task metric, but also explore other rewards like perplexity and F1 in our experiments (see Appendix D). We use Proximal Policy Optimization
(PPO) (Schulman et al., 2017) with KL penalty as the RL algorithm for its robustness.
A key difference between our work and previous RL-based methods is that we use a capable LLM (PaLM 2-S) as our rewriter model. Because of this, our model is less likely to produce uninterpretable gibberish prompt as in RLPrompt (this can be also attributed to the KL penalty in PPO). We also don’t need to define a constrained action space manually as in TEMPERA.

heading: 2.4 Rewriting via Inference and Search
text: Once the rewriter LLM is trained, we use it for prompt rewriting following Equation 1. We design two specific rewriting strategies. For the inference strategy, denoted as PRewrite-I, we set temperature to zero, in which case the model greedily decodes and generates one single rewritten prompt. For the search strategy, denoted as PRewrite-S, we prompt the rewriter LLM K-times with temperature=1 to generate a set of prompts, {p†i}Ki=1. We then select the best p†i based on their end task performance on a dev dataset.

heading: 3 Experiments & Analysis

heading: 3.1 Experimental Setup
text: We evaluate PRewrite on diverse benchmark datasets, spanning from classification with AG News (Zhang et al., 2015) and SST-2 (Wang et al., 2018), question answering with Natural Questions (NQ) (Kwiatkowski et al., 2019) to arithmetic reasoning with GSM8K (Cobbe et al., 2021). We use the standard train/dev/test splits. As GSM8K doesn’t come with a dev split, so we randomly sample 10% examples from the train split as the dev split. Data statistics are reported in Appendix C.
Our initial prompts, prompt templates and meta prompts are listed in Appendix E, G and B respectively. We experiment with PaLM 2-S and PaLM 2-L (Anil et al., 2023) as the frozen task LLMs with zero temperature. We use PaLM 2-S as the rewriter LLM and set temperature to 1 for both the policy and value model during RL training. We use standard PPO algorithm for online policy optimization with GAE. The model is trained until convergence on the dev set. We test both the inference and search strategy for rewriting, denoted as PRewriteI and PRewrite-S respectively. For PRewrite-S, we search from K=10 rewritten prompts (Section 2.4).
For baselines, we cite evaluation results for AutoPrompt (Shin et al., 2020), RLPrompt (Deng
et al., 2022), and TEMPERA (Zhang et al., 2022), out of which the last two are RL-based methods; APE (Zhou et al., 2023), OPRO (Yang et al., 2023) and Promptbreeder (PB) (Fernando et al., 2023), which use LLMs of same size as ours. We report standard metrics on test: accuracy for AG News, SST-2, GSM8K; and Exact Match (EM) for NQ.

heading: 3.2 Results
text: We first present PRewrite results based on PaLM 2-S task model in Table 1.
First, PRewrite consistently improves over the initial prompts, demonstrating the effectiveness of the proposed method. We repeated the PRewrite experiments 5 times and the results were consistent. We list the rewritten prompts in Table 9 in Appendix. Second, we observe larger improvement for PRewrite when there is more headroom. For example, the performance gain on SST-2 is minimum, but we observe 80%, 22% and 10% relative improvement with PRewrite on GSM8K, NQ and AG News respectively. Third, PRewrite-S consistently shows improvement over PRewrite-I, suggesting that search strategy can be more helpful. We find the two strategies often produce prompts with small differences. For example, “sentiment classification” from PRewrite-I and “sentiment classification from text” for Prewrite-S on SST-2. Lastly, baseline models underperform PRewrite. However, this can be largely due to the smaller task model, RoBERTaLarge (Liu et al., 2019), being used for the baselines. That said, it is not straightforward to apply some of the baseline methods on larger models like PaLM 2, especially in case of API-only access.
Next, we compare PRewrite with baselines on GSM8K, all based on PaLM 2-L task model in Table 2 (PRewrite-S only due to space constraints). PRewrite-S not only dramatically improves the ini-
tial prompt, but also outperforms strong baselines like APE and OPRO, and is on par with Promptbreeder. This result is especially impressive in that the PRewrite setup is relatively simple with minimal customization, in comparison with the baselines. For example, APE proposes to use task input and output to induce instructions for most tasks but has a special treatment to GSM8K. It collects a customized dataset with questions and reasoning steps via prompting InstructGPT for instruction induction – this is more likely to induce chain-ofthought instructions. Promptbreeder uses 56 mutation prompts and 39 thinking style prompts including ones that contain phrases like steps required, taking a break or suggesting explanation. In comparison, we only use one generic meta prompt for GSM8K (Appendix B).
We also experiment with different rewards for PRewrite. Please refer to Appendix D for the results.

heading: 3.3 Case Studies
text: To showcase the capability of PRewrite, we present rewriting performed by it for two datasets in Table 3 (see Appendix E, F for more results).
For NQ, PRewrite not only learns the task needs a short answer, but also impressively adds an incontext example. For GSM8K, PRewrite rewrites the simple initial prompt to a creative chain-ofthought (CoT) prompt. This CoT prompt is different from previous human created ones (Kojima et al., 2022) – it does not instruct the LLM to think/write step by step, but instead assumes there already exists a solution with steps, that follows after the prompt.
Moreover, we find Prewrite always produces interpretable rewritten prompts, unlike RLPrompt, which often generates gibberish text. This is due to the LLM-based rewriter and KL-divergence penalty in PPO we have used (Section 2.3).

heading: 4 Related Work
text: We survey related work on automated prompt engineering for discrete prompts. Please refer Liu et al.
(2023) for a more comprehensive literature review. Some earlier works optimize prompts via paraphrasing (Jiang et al., 2020; Yuan et al., 2021; Haviv et al., 2021). In contrast, we adopt a powerful LLM to rewrite prompts, providing more capacity for prompt optimization. Shin et al. (2020); Wallace et al. (2021) propose gradient-based search approach which is challenging for larger API-access only models as it requires model gradient access.
Prior work have also explored RL based solutions. RLPrompt (Deng et al., 2022) optimizes prompts using RL, but often produces uninterpretable gibberish prompts. TEMPERA (Zhang et al., 2022) allows prompt editing at test time based on task input using RL, but it defines a small action space. By leveraging more capable LLMs, our method produces interpretable prompts and allows unconstrained exploration without a manually defined the action space.
More recent works use blackbox LLMs such as PaLM 2 and GPT models, similar to ours. These include APE (Zhou et al., 2023), Promptbreeder (Fernando et al., 2023) and OPRO (Yang et al., 2023). These works use LLMs in different ways to propose prompt candidates and search for the optimal one via validating performance on a given training dataset. We follow a similar idea but instead use RL for prompt optimization.

heading: 5 Conclusions
text: In this paper, we present PRewrite, a prompt rewriter trained with reinforcement learning (RL) for prompt optimization. We instantiate the rewriter with a LLM (PaLM 2-S) and finetune it using RL to optimize the end task performance. To further improve the performance, we develop a rewriting strategy that searches from the rewritten prompts generated by the trained rewriter. Our experiments testify the effectiveness of PRewrite and demonstrate its state-of-the-art performance.

heading: 6 Limitations
text: In this work, we only test with limited initial and meta prompts (see Table 8 and 5 in Appendix) on four benchmark datasets. It would be interesting to experiment with more initial-meta prompt combinations to understand their implications, and on more datasets to test the generality of PRewrite. Moreover, we do not investigate the use of multiple meta/initial prompts to diversify exploration in prompt rewriting, which may further improve PRewrite. We leave these ideas for future work.
Due to resource constraints, we only experiment with PaLM 2 models. However, we believe that our conclusions should generalize to other LLMs as well.

heading: A Problem Formulation Examples
text: In Table 4, we show an example of the task input (x), output (y), prompt (p), and instruction (t) for a question answering task.

heading: B Meta Prompts
text: In Table 5, we show the meta prompts used for prompting the rewriter LLM.

heading: C Dataset Statistics
text: Data statistics for train/dev/test splits are given in Table 6. GSM8K doesn’t come with a dev (or validation) split, so we randomly reserve 10% examples from the train split as the dev split.

heading: D Results based on Different Rewards
text: We test different rewards for all datasets and report the results based on PRewrite-I and PaLM 2-S task model in Table 7. Perplexity uses perplexity of the ground truth labels as the reward. F1 use wordlevel F1 measure as the reward. Perplexity+F1 sums perplexity and F1 as the reward.
First, we find that using the final task metric, accuracy or EM, performs well in general. In other
words, RL is able to directly optimize for these task metrics. Second, we find that F1 is in general be more stable than accuracy/EM. This is because F1 can provide more fine-grained feedback for RL – accuracy/EM is either 0 or 1 for a datapoint while F1 provides a fractional score. Third, we find perplexity can sometimes be harmful (see NQ and AG News results in the table), as it is not directly linked to the final task metrics. However, combining perplexity with F1 gives the best performance on NQ.
E Initial Prompts
Table 8 lists all initial prompts we used for prompt rewriting in our experiments.

heading: F Rewritten Prompts
text: Table 9 lists the produced rewritten prompts for experiments using PaLM 2-S as the task LLM. This includes rewritten prompts produced using both rewriting strategies (Section 2.4) for all datasets.
Table 10 lists the rewritten prompts produced by PRewrite-S for GSM8K experiments using PaLM 2-L as the task LLM.

heading: G Prompt Templates
text: Table 11 lists the prompt template we used for each dataset.

heading: Future Work
text: We leave these ideas for future work. Due to resource constraints, we only experiment with PaLM 2 models