PIAST: Rapid Prompting with In-context Augmentation for Scarce Training data

1. Problem Statement

Large Language Models (LLMs) are highly sensitive to prompt design. While manual prompt engineering is effective, it is labor-intensive and often requires intricate crafting of few-shot examples. Existing automatic prompt engineering (APE) methods face two major limitations:

• Computational Cost: Many methods require hours of compute time to explore large search spaces (e.g., evolutionary algorithms or extensive beam search).
• Data Dependency: Most methods rely on large training sets to score candidates or are restricted to selecting examples only from the existing training data.
• Focus Misalignment: Prior work often optimizes only the instruction string (via rephrasing) while ignoring the most impactful component of In-Context Learning (ICL): the few-shot examples.

The paper addresses the need for a fast, data-efficient method that can automatically synthesize high-quality few-shot examples not found in the training set, even under limited computational budgets.

2. Methodology: PIAST

PIAST (Prompting with In-context Augmentation for Scarce Training data) is an iterative automatic prompt construction algorithm. It consists of three main components, all instantiated as frozen LLMs:

A. Core Components

1. Example Proposer: Generates an initial diverse set of candidate few-shot examples based on a human-written instruction. It ensures diversity in topic and length.
2. Prompt Evaluator: Assesses the quality of the current prompt (instruction + examples) on a small, randomized subset of the data.
3. Example Improver: Iteratively refines the set of examples. It identifies the least useful example and decides whether to Replace, Drop, or Keep it based on utility estimation.

B. Optimization Loop & Shapley Values

The core innovation lies in how the method selects which example to remove or replace:

• Shapley Value Estimation: Instead of simple heuristics, PIAST uses Monte Carlo Shapley estimation to calculate the marginal contribution (utility) of each example to the overall performance.
  • The Shapley value $\phi_i$ for an example $i$ is approximated by sampling $K$ random permutations of the example set and measuring the performance difference when $i$ is added to a subset.
  • The example with the lowest Shapley value ($i^*$) is identified as the "least helpful."
• Replace/Drop/Keep Decision: Once the worst example $i^*$ is identified, the Example Improver proposes $m$ new candidate examples. The system computes three scores:
  • Replace: Max performance with a new candidate added.
  • Drop: Performance with $i^*$ removed.
  • Keep: Current performance.
  • The action with the highest score is selected (with a preference hierarchy: Replace > Drop > Keep).

C. Efficiency Mechanisms

To ensure the method is "Rapid" and suitable for scarce data:

• Aggressive Subsampling: Evaluations are performed on small, randomized subsets of the training data.
• Replay Buffer: To prevent overfitting to the current small subset, a small portion of data from previous iterations is stored in a replay buffer and merged with the current subset. This acts as a regularizer.
• System Optimizations:
  • KV-Cache Reuse: Caches keys/values for the shared instruction prefix to avoid recomputation.
  • PagedAttention: Manages memory efficiently to reduce fragmentation.
  • Continuous Token-Level Batching: Maintains high GPU utilization by dynamically forming batches.

3. Key Contributions

1. Conceptual: PIAST is the first method to rapidly synthesize new few-shot examples (not present in the training set) and optimize them using Shapley values, rather than just rephrasing instructions or selecting from existing data.
2. Implementation: The authors introduce a highly optimized pipeline using KV-cache reuse, PagedAttention, and continuous batching, enabling the method to run efficiently on a single GPU.
3. Anytime Performance: The algorithm exhibits strong "anytime" behavior. It yields competitive results even with a small computational budget and limited data access, outperforming prior methods that require significantly more resources.

4. Experimental Results

The authors evaluated PIAST on four task families: Text Classification, Summarization, Text Simplification, and Mathematical Reasoning (GSM8K, DeepMath, Math500), as well as domain-specific QA (MedQA).

• Performance vs. Baselines:
  • Limited Budget: PIAST outperformed previous state-of-the-art (SoTA) automatic prompting methods on Text Simplification and Mathematical Reasoning (GSM8K, DeepMath, Math500). It achieved second-best results on Classification and Summarization.
  • Extended Budget: With a modestly extended budget, PIAST set a new SoTA among automatic prompting methods on Classification, Simplification, and Mathematical Reasoning.
  • Efficiency: PIAST is significantly faster than competitors like PRL (which takes ~48 hours) and GRACE. For example, on classification tasks, PIAST achieved high accuracy in ~8 minutes, whereas PRL required 48 hours.
• Data Efficiency: PIAST achieves these results using only a small fraction of the training dataset (e.g., 20% or less), demonstrating robustness in data-scarce scenarios.
• Ablation Studies:
  • Removing the optimization loop (using only initial examples) still yielded competitive results, proving the strength of the synthesis step.
  • The Shapley-based selection was shown to be more robust than simple Leave-One-Out (LOO) heuristics, particularly on complex tasks like subjectivity classification.
  • The method showed strong cross-model robustness, transferring well from Qwen2.5-7B to Mistral-7B.

5. Significance

• Shift in Focus: The paper argues that optimizing in-context examples is a more dominant lever for prompt engineering than exhaustively searching for instruction rewrites.
• Practicality: By drastically reducing the compute time (from hours to minutes) and data requirements, PIAST makes automatic prompt engineering accessible for real-world applications where training data is scarce or expensive.
• Scalability: The "anytime" nature of the algorithm allows users to trade off computation time for performance, making it adaptable to various resource constraints.

In conclusion, PIAST represents a significant step forward in efficient prompt engineering, demonstrating that synthetic few-shot examples optimized via game-theoretic utility estimation can outperform manual and existing automated methods with a fraction of the computational cost.