Source-linked AI summary

StablePrompt: Automatic Prompt Tuning using Reinforcement Learning for Large Language Models

Minchan Kwon, Gaeun Kim, Jongsuk Kim, Haeil Lee, Junmo Kim

arXiv:2410.07652v1cs.CL

TL;DR

Automatic prompt tuning must handle RL instability and restricted prompt search while producing effective natural-language prompts. StablePrompt formulates tuning as online RL and introduces APPO, which uses an anchor model to adapt policy updates. The method achieves strong performance across tasks and models, while its evaluation does not cover domains such as medicine or law.

  • Problem

    RL-based prompt tuning is unstable and environmentally dependent, while restricting the search space can reduce prompt performance.

  • Method

    StablePrompt formulates prompt tuning as online, on-policy RL and introduces APPO with an anchor model to stabilize updates while preserving flexible search.

  • Results

    StablePrompt achieves SoTA performance across various tasks and works with diverse agent and target LLMs, including models larger than 7B.

  • Takeaways & Limitations

    StablePrompt supports stable, flexible RL-based prompt tuning across multiple tasks, model sizes, and input-dependent prompting settings.

  • Takeaways & Limitations

    The study does not evaluate domains significantly beyond prior learning, such as medical and legal tasks.

Abstract

from arXiv · show

Finding appropriate prompts for the specific task has become an important issue as the usage of Large Language Models (LLM) has expanded. Reinforcement Learning (RL) is widely used for prompt tuning, but its inherent instability and environmental dependency make it difficult to use in practice. In this paper, we propose StablePrompt, which strikes a balance between training stability and search space, mitigating the instability of RL and producing high-performance prompts. We formulate prompt tuning as an online RL problem between the agent and target LLM and introduce Adaptive Proximal Policy Optimization (APPO). APPO introduces an LLM anchor model to adaptively adjust the rate of policy updates. This allows for flexible prompt search while preserving the linguistic ability of the pre-trained LLM. StablePrompt outperforms previous methods on various tasks including text classification, question answering, and text generation. Our code can be found in github.

1 Introduction

StablePrompt addresses instability and restricted search in RL-based prompt tuning by combining online RL with an anchor-based APPO update scheme. It reports strong performance across tasks, models, and prompt settings, including input-dependent prompting.

  • RL prompt tuning is vulnerable to overfitting and environmental dependence, while restricting prompt length or action space reduces search flexibility.
  • StablePrompt formulates prompt tuning as online, on-policy RL and uses APPO with an anchor model to stabilize updates while preserving flexible search.
  • StablePrompt uses one dataset-wide prompt, while TTE-StablePrompt generates an input-dependent prompt for each example.
  • The methods are evaluated across multiple tasks and LLMs ranging from 2B to 13B, including Llama, Mistral, Gemma, and Falcon.
  • StablePrompt achieves SoTA performance across text classification, question answering, and text generation tasks.

2 Related Work

Automatic prompt tuning searches for effective prompts without directly optimizing readable natural-language prompts, while discrete methods differ in how they generate or edit them. StablePrompt targets the stability, scalability, and search-space limitations of prior RL approaches.

  • 2.1 Automatic Prompt Tuning: Discrete prompt tuning searches for natural-language prompts, contrasting with soft prompting that uses gradients but often produces unreadable prompts and requires substantial data.
  • 2.2 Discrete Prompt Tuning: Generation-based methods such as APE, ProTeGi, and PromptAgent exploit LLM generation or search, but can struggle on tasks outside pre-training.
  • 2.2 Discrete Prompt Tuning: Prior RL methods depend on manual prompts or predefined action spaces, while RLprompt and TEMPERA face computational or scalability limits as model hidden sizes grow.
  • 2.2 Discrete Prompt Tuning: StablePrompt's training framework generates prompts from a task-agnostic meta-prompt and evaluates them using training-data rewards.
  • 2.2 Discrete Prompt Tuning: APPO maintains an anchor model so the agent remains robust to incorrect rewards instead of progressively deviating from the optimal prompt.

3 Method

StablePrompt formulates discrete prompt tuning as online, on-policy reinforcement learning and uses APPO with an adaptive anchor model to balance training stability against search-space flexibility. Its test-time variant generates input-specific prompts, while experiments cover multiple tasks, models, and prompt-tuning settings.

  • 3.1 RL Formulation: StablePrompt models prompt generation as online, on-policy RL in which an LLM agent autoregressively generates a discrete prompt and receives reward from the target model.The prompt is generated from sampled input-output pairs and a task-agnostic meta-prompt.
  • 3.2 StablePrompt: APPO combines an anchor-model KL penalty with adaptive updates to stabilize training while preserving a larger prompt search space than RLHF-style PPO.The anchor is updated after validated improvements and can roll back the agent when performance declines.
  • 3.2 StablePrompt: The method combines classification accuracy with softmax difference to distinguish prompts with tied accuracy, while text generation uses F1 directly as reward.The classification reward weights accuracy and softmax difference with coefficients c_a and c_s.
  • 3.3 Test Time Editing StablePrompt.: TTE-StablePrompt generates prompts conditioned on the current input and trains from instance rewards, targeting individual queries rather than the entire dataset.It is intended for tasks that are difficult to solve with a single prompt.
  • 3.3 Test Time Editing StablePrompt.: Table 1 reports results for six few-shot text-classification datasets, where StablePrompt outperforms other discrete prompt-tuning methods.
  • 3.3 Test Time Editing StablePrompt.: Figure 4 shows that StablePrompt works well with a variety of LLMs across few-shot text classification tasks on diverse target-agent pairs, averaging results over six datasets.

4 Experiment

Across few-shot classification, induction, and question answering, StablePrompt achieves strong performance with different target models and task formats. Its results include broad gains over prompt-tuning baselines, including black-box and larger-agent comparisons.

  • 4.1 Few Shot Text Classification: StablePrompt achieves state-of-the-art performance on all few-shot classification tasks except QNLI, where it remains best among discrete prompt-tuning methods.Its average score also outperforms APE.
  • 4.1 Few Shot Text Classification: StablePrompt outperforms manual prompts across nearly all tested agent-target pairs, including settings with a 2B agent and an 11B target.The only exception is the Gemma-2B agent with the Llama3-8B target.
  • 4.1 Few Shot Text Classification: With appropriate prompting, Mistral-7B outperforms Falcon-11B despite lower manual-prompt performance, indicating that prompt quality can narrow model-size differences.The comparison is reported for few-shot classification.
  • 4.2 Induction Task: StablePrompt achieves state-of-the-art results on both BBII and Instruction Induction, including improved performance on text-generation tasks.The authors attribute this result to the RL framework handling output format requirements.
  • 4.2 Induction Task: Against InstructGPT3.5 as a black-box target, StablePrompt outperforms baselines, and its 7B model produces better prompts than APE using agents larger than 100B.These results support use with inaccessible targets and smaller agent models.
  • 4.3 Question Answering: StablePrompt attains the highest performance among the question-answering baselines, with results comparable to STEM and better than the other methods.The evaluation uses MMLU and OpenbookQA with Gemma-7B.
  • 4.3 Question Answering: TTE-StablePrompt improves on a single StablePrompt by generating different instructions for questions within the same subject, and it outperforms Zero-Shot CoT.The generated prompts are input-dependent rather than shared across all questions in a subject.

StablePrompt

StablePrompt extends prompt tuning with adaptive PPO updates and input-dependent prompting. The ablation indicates that adaptive anchoring combines the strengths of original and RLHF-style PPO.

  • StablePrompt: TTE-StablePrompt generates prompts tailored to individual questions, whereas APE and StablePrompt produce broadly reusable prompts for a subject.The comparison is illustrated with MMLU machine-learning question-choice pairs.
  • 4.4 Ablation Study: APPO outperforms the PPO variants on average by adaptively updating its anchor model to combine original PPO and RLHF-style PPO behavior.It can use no updates, update every period, or adapt toward the better-performing variant.
  • 4.4 Ablation Study: When the PPO variants differ substantially, APPO adapts toward the better-performing model, while on SNLI it finds prompts better than either variant alone.The reported pattern appears on MRPC, QNLI, MNLI, and SNLI.

5 Conclusion

StablePrompt formulates prompt tuning as online, on-policy reinforcement learning and introduces APPO. Across target models and tasks, it outperforms other methods and demonstrates RL-based prompt tuning beyond 7B models.

  • 5 Conclusion: StablePrompt defines prompt tuning as an online, on-policy RL problem and introduces APPO for this setting.The paper presents this as its central methodological contribution.
  • 5 Conclusion: StablePrompt outperforms other methods across various target models and tasks, demonstrating the potential of integrating existing RL methodologies into prompt tuning.The authors identify it as the first RL-based prompt-tuning method for models larger than 7B.

Limitation

The study does not evaluate domains substantially beyond prior learning, and it notes that the method could be abused for specific purposes, particularly through commercial black-box LLM APIs.

  • Limitation: The evaluation excludes domains significantly beyond prior learning, including medical and legal applications.The authors expect the training-based method to scale to such domains in future work.
  • Limitation: StablePrompt could be used to abuse LLMs for specific purposes, posing a particular threat to commercial LLMs exposed through APIs and black-box optimization.This is identified as a security-related limitation of the study.

A.3 Dataset Details

The paper evaluates StablePrompt across classification, induction, question answering, and smaller-target settings, with datasets, baselines, and training configurations detailed across the experiments.

  • A.3 Dataset Details: The evaluation covers few-shot classification, BIG-Bench instruction induction, Instruction Induction, and MMLU question answering datasets.BIG-Bench includes 21 tasks, Instruction Induction includes 24 tasks, and MMLU contains 57 subject subsets with 14,079 test questions.
  • A.4 Baseline Details: Baseline comparisons include APE, ProTeGi, RLprompt, and PromptAgent with task-specific scaling, model, or usage settings.APE is given the same prompt-generation budget as StablePrompt, RLprompt uses GPT2-XL, and PromptAgent is evaluated only for text classification.
  • A.5 Training Details: Training uses one A100 GPU, with 100 epochs for classification and 30 epochs for question answering and induction tasks.Reported training time is 2–3 GPU hours per classification task and 1–2 GPU hours per question-answering or induction task, depending on input length.
  • B.1 Text Classification in Small Target Model: The small-target experiment fixes RoBERTa-Large at 330M parameters and uses GPT-2 for RLprompt versus Mistral-7B for StablePrompt.The MR dataset replaces MRPC and evaluates sentiment classification rather than natural-language inference.
  • B.1 Text Classification in Small Target Model: StablePrompt achieves the highest performance across all evaluated datasets except MR, where it performs comparably to TEMPERA.The comparison uses a RoBERTa-Large target model and reports mean accuracy and standard deviation across three random seeds.

B.2 Ablation Study

The ablations examine training dynamics and the reward design used by StablePrompt. Rewards improve steadily during training, while softmax difference helps distinguish prompts with tied classification accuracy.

  • Training curve: Mean reward increases steadily and value loss decreases steadily during SST2 few-shot classification training.The falling value loss indicates that the value head becomes aligned with the reward model.
  • Reward Function Ablation: Softmax difference ranks prompts with identical classification accuracy, avoiding confusion caused by batches containing many tied prompts.The authors report a performance penalty when softmax difference is removed.
  • Reward Function Ablation: The reward ablation uses continuous F1 for text generation but adds softmax difference for text classification.The ablation changes the agent and target models to Mistral-7B while retaining the section 4.1 setting.

C Generated Prompt

The compact appendix provides generated prompts and examples across classification, induction, and question-answering tasks. These prompts range from concise output instructions to multi-step reasoning and domain-specific guidance.

  • C.1 Few-Shot Text Classification: Generated classification prompts specify task labels through direct instructions and demonstrations, such as movie-entertainment judgments and sentence-paraphrase decisions.Examples cover SST2, MRPC, RTE, QNLI, MNLI, and SNLI.
  • C.1 Few-Shot Text Classification: The MNLI and SNLI prompts guide premise–hypothesis analysis before predicting entailment-related outputs.The MNLI prompt decomposes analysis into premise, hypothesis, and output-prediction steps, while SNLI includes an example pair.
  • C.2 BigBench-Hard Insstruction Induction: The BBH-II examples include causal judgment, navigation, and tense-change prompts with outputs presented as demonstrations or corrections.The appendix states that three BBH-II tasks were selected at random.
  • C.3 Instruction Induction: Instruction Induction examples cover negation, number-to-word conversion, and animal-taxonomy extraction.The appendix states that four tasks were selected at random from the Instruction Induction dataset.
  • C.4 Question Answering: MMLU question-answering prompts provide domain-specific guidance for moral scenarios, machine learning, biology, and chemistry.The examples emphasize identifying relevant concepts, relating inputs to outputs, and selecting among answer choices.
  • C.5 Full prompt of Figure 5: The full QA prompt contrasts APE, StablePrompt, and TTE-StablePrompt on a linear-regression question about a high negative coefficient.StablePrompt gives general input-output instructions, whereas TTE-StablePrompt explicitly invokes the bias–variance trade-off.

D Full Experiment Results

The full experiments report results for MMLU question answering and instruction-induction benchmarks, including BBH-II and Instruction Induction datasets. Tables provide full task-level result scopes and compare prompting methods across target and agent models.

  • D.1 Question Answering: The full question-answering performance is reported for the entire QA dataset in Table 13.Table 13 is identified as the full-results table for the MMLU QA datasets.
  • D.2 Instruction Induction: The instruction-induction evaluation reports full-dataset results for BBH-II and Instruction Induction in Tables 10 and 11.The tables cover experiments with Gemma-7B as the target model.
  • D.2 Instruction Induction: Additional instruction-induction results provide detailed accuracy for 24 tasks with InstructGPT3.5 as target, alongside APE, StablePrompt, and human-prompt results for target and agent models.These comparisons are reported in Tables 12 and 11.
Loading 2410.07652v1…