Source-linked AI summary
Robust Prompt Optimization for Large Language Models Against Distribution Shifts
Moxin Li, Wenjie Wang, Fuli Feng, Yixin Cao, Jizhi Zhang, Tat-Seng Chua
TL;DR
Prompt optimization methods may not generalize across distribution shifts, even though such shifts occur in real-world settings. GPO uses labeled source data and unlabeled target inputs in prompt optimization. Experiments report better target-group performance with comparable source-group performance.
Problem
Current gradient-free prompt optimization methods ignore distribution shifts, while optimized prompts may face unlabeled target groups whose data differ from labeled source data.
Source: §1 · pp. 1–2Method
GPO incorporates unlabeled target inputs by using LLM zero-shot labeling and prompt ensembling alongside labeled source data.
Source: §1 · p. 2, §3.2 · p. 5Results
GPO improves target-group performance while preserving comparable source-group performance in experiments on three tasks.
Source: §1 · p. 2Takeaways & Limitations
The experiments support GPO's effectiveness for prompt robustness across source and target groups when target inputs are available.
Source: §6 · p. 9The work assumes target-group inputs are available and leaves generalization to completely unseen groups without those inputs underexplored.
Source: Limitations · p. 10
Abstract
from arXiv · showhide
Large Language Model (LLM) has demonstrated significant ability in various Natural Language Processing tasks. However, their effectiveness is highly dependent on the phrasing of the task prompt, leading to research on automatic prompt optimization using labeled task data. We reveal that these prompt optimization techniques are vulnerable to distribution shifts such as subpopulation shifts, which are common for LLMs in real-world scenarios such as customer reviews analysis. In this light, we propose a new problem of robust prompt optimization for LLMs against distribution shifts, which requires the prompt optimized over the labeled source group can simultaneously generalize to an unlabeled target group. To solve this problem, we propose Generalized Prompt Optimization framework, which incorporates the unlabeled data from the target group into prompt optimization. Extensive experimental results demonstrate the effectiveness of the proposed framework with significant performance improvement on the target group and comparable performance on the source group.
1 Introduction
Prompt optimization can fail under distribution shifts, and existing gradient-free methods do not account for them. GPO incorporates unlabeled target data into prompt optimization; experiments report improved target-group performance with comparable source-group performance.
1 Introduction
- Current gradient-free prompt optimization methods ignore distribution shifts between labeled optimization data and data served by an LLM.Gradient-free methods are the focus because LLMs are evolving into black-box API services.Source: §1 · p. 1
- Experiments on 30 paired data groups from six tasks found significant performance gaps under some distribution shifts.The experiments used prompts optimized by APE with gpt-3.5-turbo-0301.Source: §1 · p. 2
- The proposed problem is to optimize prompts using labeled source data and unlabeled target data so they generalize across both groups.The target group is kept unlabeled because shifts may arise over time, and labeling new data can add labor and latency.Source: §1 · p. 2
- GPO improves target-group performance while preserving comparable source-group performance across three tasks.The introduction reports experimental validation of the framework's effectiveness.Source: §1 · p. 2
- GPO uses LLM zero-shot labeling and prompt ensembling to incorporate target data into prompt optimization.The framework aims to improve target performance while keeping source performance comparable.Source: §1 · p. 2
2 Preliminary Experiments
The preliminary experiments test whether prompts optimized on source groups transfer to shifted target groups, using paired datasets and controlled evaluation. They find that generalization varies across shifts, with significant gaps in some groups but not others.
2.1 Evaluation Protocol
- The evaluation pairs source and target groups across 16 datasets from six NLP tasks, treating differences in dataset backgrounds and topics as distribution shifts.The groups include sentiment, NLI, textual entailment, commonsense QA, multi-turn dialog reasoning, and numerical QA.Source: §2.1 · pp. 2–3
- Experiments use APE in a zero-shot setting with gpt-3.5-turbo-0301; classification uses accuracy, DROP uses F1, and results average five runs.Prompt optimization follows APE's N-shot training and validation sampling setup.Source: §2.1 · p. 3
2.2 Experimental Results
- Generalization gaps may be absent for previously studied OOD and adversarial NLI and textual-entailment groups, yet remain significant for some sentiment, commonsense-QA, and DROP groups.The tables cover tasks with and without significant gaps.Source: Table 2 · p. 3, Table 1 · p. 3, §2.2 · p. 3
- Source-optimized prompts sometimes outperform target-optimized prompts, including OpenbookQA prompts on SocialIQA and DROP Spans prompts on Number.The reported comparisons are specific to those target groups.Source: §2.2 · p. 4
- Smaller label shifts and greater input similarity generally coincide with better target generalization, but these metrics do not explain every task outcome.The study measures label shifts with KL divergence and input similarity with n-gram similarity.Source: §2.2 · p. 4
- The authors conjecture that OpenbookQA's diverse input and coverage of SocialIQA and PIQA partly explain its stronger prompt generalization on those groups.Their analysis measures unique n-gram percentages and target n-grams covered by the source group.Source: §2.2 · p. 4
3 Robust Prompt Optimization
Section 3 formulates robust prompt optimization as using labeled source data and unlabeled target inputs to optimize one task-specific prompt across distribution-shifted groups. GPO labels target inputs with LLM-generated prompts, filters labels by agreement, and jointly optimizes on source and filtered target data.
3.1 Problem Definition
- GPO targets one task-specific prompt that performs across source and target groups with different distributions, while target labels remain unseen during optimization.The objective uses labeled source examples and unlabeled target inputs.Source: §3.1 · p. 5
- Unlabeled target inputs are practical to collect, while labeling them may be time-consuming and costly.In deployment, user queries can provide extensive target inputs without labels.Source: §3.1 · p. 5
- Separate prompts for each group add computation and require classifying each test example by its distribution, while source groups may not cover future target groups.The framework instead seeks one prompt that generalizes across groups.Source: §3.1 · p. 5
3.2 GPO Framework
- GPO uses LLM zero-shot labeling to incorporate target data into joint optimization, then applies two strategies intended to improve labeling accuracy.The motivation is that noisy target labels may damage optimized prompts.Source: §3.2 · p. 5
- A Meta Prompt uses labeled source examples to generate K candidate prompts for target labeling.Different source-data splits are used to generate the prompts.Source: §3.2 · p. 5
- Prompt ensemble labeling selects the most-consistent candidate label and accepts it only when agreement exceeds threshold T.This produces a filtered labeled target set.Source: §3.2 · p. 5
- GPO runs APE on the labeled source set mixed with filtered target examples, randomly upsampling the latter to match source-set size.The resulting joint optimization produces the final prompt.Source: §3.2 · p. 6
4 Experiments
Across three tasks, GPO improves target-group performance while keeping source-group performance comparable to competing generated prompts. The experiments also examine human-written prompts, ablations, labeling accuracy, and different backbone LLMs.
4.1 Setup
- Experiments cover sentiment analysis, commonsense QA, and DROP, comparing GPO with APE, APO, APE-ut, an Upper Bound, and human-written prompts.Top 1 uses the single prompt with best validation performance; Ensemble labels with all K prompts and selects the output with greatest agreement; results are averaged over five runs.Source: §4.1 · p. 6
4.2 Performance Comparison
- GPO achieves superior performance on all target groups under both Top 1 and Ensemble testing, with comparable source-group performance across tasks.The results retain room for improvement toward the target-label Upper Bound; on Spans, APE-ut underperforms GPO, highlighting the role of accurate target labels for difficult cases.Source: §4.2 · p. 6, Table 5 · p. 6
- GPO outperforms human-written prompts and PromptPerfect on sentiment analysis and commonsense QA, but underperforms human-written prompts on DROP.The authors suggest inaccurate Spans labels may hinder optimization on DROP.Source: §4.2 · pp. 6–7, Table 6 · p. 7
4.3 Ablation Study
- GPO generally beats the no-consistency-threshold variant on target groups, except Flipkart under Top 1 evaluation.Removing target training data helps source-group Top 1 results but harms target-group Ensemble results; target Top 1 is less affected because both variants retain target validation data.Source: §4.3 · p. 7, Table 7 · p. 7
4.4 In-depth Analysis
- The consistency threshold improves target-labeling accuracy across groups, and increasing the proportion of incorrect labels generally lowers target-group performance.The authors report larger performance gains from improved labeling accuracy for low-accuracy Spans than for groups already highly accurate without the threshold.Source: §4.4 · p. 7, Table 8 · p. 7, Figure 3 · p. 8
- Across tested backbones on Flipkart with Yelp as source, GPO consistently outperforms APE, and GPO with Vicuna 7B reaches Vicuna 13B's performance level.GPT-4 achieves the best GPO performance across the tested LLMs.Source: §4.4 · p. 8, Table 9 · p. 8
5 Related Work
Related work covers LLM generalization and robustness, then situates this paper within prompt optimization research. The paper focuses specifically on gradient-free prompt optimization under within-task distribution shifts.
Generalization Ability and Robustness of LLM
- Prior work reports LLM gains on OOD and adversarial tasks but also vulnerabilities to shortcuts, perturbations, prompt sensitivity, and in-context example selection.The paper narrows its focus to prompt-optimization generalization under distribution shifts within one task.Source: §5 · p. 8
Prompt Optimization
- Prompt optimization includes gradient-based tuning and gradient-free approaches such as reinforcement learning, search, evolutionary algorithms, and boosting.Recent black-box methods use LLMs for prompt generation and evaluation or prompt editing from critiques.Source: §5 · p. 8
- Unlike prior prompt-tuning work on cross-task, cross-model, or domain generalization, this paper focuses on gradient-free optimization under distribution shifts.The cited prior work includes cross-task and cross-model prompt tuning and domain adaptation.Source: §5 · p. 8
6 Conclusion
The conclusion identifies a generalization problem: prompts optimized on source data can lose performance on shifted target groups. GPO jointly considers source and target groups, and experiments support its robustness; future work will study targets without available inputs.
6 Conclusion
- Prompts optimized on source groups may show a performance drop on target groups under distribution shifts.The conclusion identifies this as a generalization issue for prompt optimization.Source: §6 · p. 8
- The analysis considered label-distribution shift and input-distribution similarity as factors correlated with variation in generalization across groups.Source: §6 · p. 8
- GPO jointly considers source and target groups, and experiments validate improved prompt robustness on both; future work will examine unseen targets without available inputs and in-context examples from different groups.Source: §6 · p. 9
Limitations
The study's scope is limited by omitted LLM inputs, assumptions about target-group data, and the models considered.
Limitations
- The analysis considers prompt generalization but ignores how in-context examples may affect LLM robustness.Future work could examine prompts together with in-context examples.Source: Limitations · p. 10
- The framework assumes target-group inputs are available; generalization to completely unseen groups without those inputs remains under-explored.The authors suggest extending the work to multiple heterogeneous groups.Source: Limitations · p. 10
- The study focuses on black-box LLMs capable of understanding instructions, rather than smaller models such as BERT and T5.The authors note that smaller models are generally not black-box and may be better suited to gradient-based prompt optimization.Source: Limitations · p. 10
A.1 Dataset Details
The experiments construct dataset splits and task-specific groupings for six NLP tasks.
A.1 Dataset Details
- The original training set is divided into 1,000 training samples and the remaining validation samples, with test subsets sampled from the original validation set.This is done because labels for the original test set are sometimes unavailable.Source: §A.1 · p. 13
- For MNLI, validation data are sampled to match the test set's matched and mismatched counts, while ANLI uses R2.Source: §A.1 · p. 13
- Review scores are mapped to negative (0–1), neutral (3), and positive (4–5) for Yelp and Flipkart.Source: §A.1 · p. 13
- Dialog-task instances are limited by turn count, and all three datasets use four answer choices.MuTual is limited to five turns; Ubuntu and DSTC7 to seven.Source: §A.1 · p. 13
A.2 Additional Implementation Details for Preliminary Experiments
The appendix describes APE's prompt search procedure and the settings used in preliminary experiments.
A.2 Additional Implementation Details for Preliminary Experiments
- APE generates prompts from input-output examples, selects them on validation performance, and optionally searches by iteratively paraphrasing effective prompts.Source: §A.2 · p. 13
- The sampled N-shot training data are split into K parts to generate candidate prompts; N is tested at 9, 16, 25, and 36, with K chosen among N's factors.Monte Carlo search is ablated because it is optional and was not significant for these tasks.Source: §A.2 · p. 13
- Table 10 reports each dataset's train, validation, and test split statistics alongside task-specific N and K values.The training and validation data are further divided into 1,000 training samples and the remainder for validation.Source: Table 10 · p. 13
- The backbone LLM uses temperature 0, top p 1.0, and a 100-token prompt-generation limit, with other parameters following the official APE implementation.Source: §A.2 · p. 13
A.3 Additional Details and Results for the Exploration on the Factors Affecting Prompt Robustness.
This appendix defines metrics for examining prompt robustness, reports how they relate to generalization, and documents baseline, case-study, and candidate-count details.
A.3 Additional Details and Results for the Exploration on the Factors Affecting Prompt Robustness.
- Label-distribution shift between source and target groups is quantified using KL divergence, while input similarity is measured by Spearman’s rank-order correlation between sampled bag-of-word vectors ranked over the target vocabulary.Source: §A.3 · pp. 13–14
- In the additional commonsense QA and dialog results, higher input similarity aligns with better generalization for commonsense QA, while smaller label shifts do not; dialog results show subtle differences.The metrics cannot be compared across target groups or tasks; examples include MuTual-to-Ubuntu performance of 74.7 versus PIQA-to-SocialIQA performance of 68.9.Source: §A.3 · p. 14
A.4 Details for Baseline Implementation
- The baseline comparison uses one optimization step for APO; GPO's consistency thresholds are 0.83 for sentiment and commonsense QA and 0.33 for DROP.APE and APO do not use unlabeled target data, whereas APE-ut and GPO use N-shot source and target data.Source: §A.4 · p. 15
- APE-ut's meta prompt adds unlabeled target examples to the source examples when requesting an instruction from the LLM.Source: §A.4 · p. 15
A.5 Case Study
- The case study reports that single-group APE prompts can contain group-specific background information, while GPO prompts are more general and perform well on both groups.For Spans, the optimized prompt is also described as general enough to transfer to Number.Source: §A.5 · p. 16, Table 14 · p. 16
A.6 Study on the Impact of the Number of Candidate Prompts
- On the 36-shot sentiment task's Flipkart target group, GPO performance stabilizes at K=6, with further candidate prompts unlikely to yield significant gains.Table 15 presents results for K values 3, 6, 9, 12, and 18.Source: §A.6 · p. 16, Table 15 · p. 16