🤖 machine learning

Stabilizing Black-Box Prompt Optimization with Textual Regularization and Signal Aggregation

This paper introduces TRAS, a plug-and-play framework that stabilizes black-box prompt optimization by combining textual regularization from successful predictions with Monte Carlo signal aggregation to filter noise, thereby improving accuracy, convergence speed, and robustness against model migration compared to existing methods.

Original authors: MohammadReza Davari, Utkarsh Garg, Weixin Cai, Eugene Belilovsky

Published 2026-06-24

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: MohammadReza Davari, Utkarsh Garg, Weixin Cai, Eugene Belilovsky

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a very smart, but slightly unpredictable, robot how to solve puzzles. You can't see inside its brain (it's a "black box"), so you can only talk to it by writing instructions, called prompts.

For a long time, the best way to improve these instructions was a method called Automatic Prompt Optimization (APO). Think of this like a student taking a test, getting a question wrong, and then asking the teacher, "Why did I get this wrong?" The teacher gives a critique (a "textual gradient"), and the student changes their answer sheet to fix that specific mistake.

The paper argues that this "fix the mistakes" approach has a big flaw: it's too focused on what went wrong and ignores what went right.

The Problem: The "Over-Correction" Trap

Imagine you are tuning a radio to find a clear station.

The Old Way (Error-Only): Every time you hear static (an error), you twist the dial wildly to fix it. But because you only listen to the static, you might accidentally twist the dial past the clear station, or erase a setting that was actually working perfectly fine. You end up spinning in circles, making the signal worse before it gets better. In the paper's terms, this is called semantic drift—you fix a small error but accidentally delete a crucial instruction that was helping you overall.

The Solution: TRAS

The authors propose a new framework called TRAS (Textual Regularization with Aggregated Signals). It's like giving the student two types of feedback instead of just one.

1. The "Don't Touch That!" Signal (Textual Regularization)

Instead of just asking, "What's wrong?", TRAS also asks, "What's working?"

The Analogy: Imagine you are cooking a stew. The old method only tells you, "This carrot is too hard, chop it smaller!" TRAS also says, "Hey, the salt level is perfect, and the broth is rich—do not change those!"
How it works: When the model gets an answer right, TRAS analyzes why it was right and creates a "regularizer" (a rule) that says, "Keep these parts of the instruction exactly as they are." This prevents the robot from accidentally deleting the good stuff while trying to fix the bad stuff.

2. The "Committee of Experts" (Monte Carlo Signal Aggregation)

LLMs are a bit like humans; if you ask them the same question twice, they might give slightly different answers. Sometimes they are noisy or confused.

The Analogy: Imagine you ask one person for directions, and they say, "Turn left at the big oak tree." You ask a second person, and they say, "Turn left at the big oak tree." You ask a third, and they say, "Turn left at the big oak tree." You are confident. But if you ask a fourth person and they say, "Turn right at the gas station," you know that fourth person is an outlier or confused.
How it works: TRAS doesn't just ask the model for feedback once. It asks multiple times (sampling multiple "gradients" and "regularizers"). It then acts like a wise committee chair, listening to all the opinions and finding the common, consistent advice. It ignores the weird, noisy outliers and gives the robot a single, clear, reliable instruction.

The Bonus: "Moving House" (Automatic Prompt Migration)

The paper also tackles a practical problem called Automatic Prompt Migration (APM).

The Scenario: You have spent weeks perfecting a set of instructions for "Model A" (like GPT-3.5). Now, you want to use those same instructions on "Model B" (like GPT-4o).
The Risk: If you just take your perfect instructions and try to "re-optimize" them for the new model using the old "fix the mistakes" method, the new model might get confused and accidentally delete the parts of the instructions that were actually working. It's like moving your furniture to a new house and, in your excitement to rearrange, throwing away the pieces that fit perfectly.
TRAS's Fix: When moving to a new model, TRAS immediately activates the "Don't Touch That!" signal. It treats the old, working instructions as a precious foundation and only makes tiny, safe adjustments to fit the new model, ensuring you don't lose the "expert" knowledge you already built.

The Results

The paper tested this on various logic and reasoning tasks.

Better Accuracy: The robot solved more problems correctly.
Faster Learning: It didn't have to spin in circles trying to find the right settings; it got there faster.
Cheaper: Because it learned faster and made fewer mistakes, it used fewer "API calls" (which cost money).
Stability: The results were more consistent; it didn't have wild swings between good and bad performance.

In short, TRAS teaches us that to teach a black-box AI well, you shouldn't just scream at it when it fails. You should also whisper to it when it succeeds, telling it exactly what to keep, and you should ask for advice from a group to make sure you aren't listening to a confused voice.

Technical Summary: Stabilizing Black-Box Prompt Optimization with Textual Regularization and Signal Aggregation

1. Problem Statement

The paper addresses the challenges of Automatic Prompt Optimization (APO) for Large Language Models (LLMs) accessed via black-box APIs, where internal model states (gradients, hidden states) are inaccessible. While recent APO methods iteratively refine prompts using model-generated critiques (textual gradients) derived from failures, the authors identify two critical limitations:

Underutilization of Success: Existing methods predominantly optimize based on error correction, ignoring the information contained in successful predictions regarding which prompt components are effective.
Instability and Semantic Drift: Relying solely on error-driven updates can lead to "semantic drift," where edits intended to fix local failures inadvertently erase globally beneficial prompt components. This forces the optimizer into costly "remove-and-rediscover" cycles.
Automatic Prompt Migration (APM) Failure: In scenarios where an expert prompt optimized for one model (e.g., GPT-3.5-turbo) must be adapted to another (e.g., GPT-4o), naive re-optimization often results in instruction loss, where load-bearing instructions from the source prompt are overwritten during the adaptation process.

2. Methodology: TRAS Framework

The authors propose TRAS (Textual Regularization with Aggregated Signals), a feedback-centric framework designed to be plug-and-play with existing APO search backbones (e.g., MCTS-based strategies like PromptAgent). TRAS operates through three core mechanisms:

A. Dual-Signal Generation

Instead of relying solely on error signals, TRAS generates two complementary textual signals at each optimization step $t$ :

Textual Gradient (Corrective Signal): Derived from failures ( $B_t^-$ ). A critique policy $\pi_g$ analyzes incorrect predictions to generate natural language instructions on what to change (e.g., missing constraints, formatting issues).
Textual Regularization (Preservation Signal): Derived from successes ( $B_t^+$ $B_{t}^{+}$ ). An attribution policy $\pi_r$ $π_{r}$ analyzes correct predictions to identify which prompt components contributed to success. This generates a regularizer $r_t$ $r_{t}$ prescribing "do-not-remove," "keep-unchanged," or "strengthen" constraints to preserve beneficial structures.
- Warm-up Schedule: To prevent anchoring to poor initializations, the regularization signal is disabled for the first $\tau_{warmup}$ iterations, allowing the optimizer to explore freely before stabilizing.

B. Monte Carlo Signal Aggregation (MCSA)

Since both the textual gradient and regularizer are stochastic LLM outputs, they can be noisy. TRAS employs MCSA to reduce variance:

The system samples $K$ independent signals for both the gradient and the regularizer.
An aggregation operator (LLM summarizer) consolidates these $K$ samples into single actionable directives ( $\bar{g}_t, \bar{r}_t$ ).
This process emphasizes consistent, actionable advice while filtering out outliers. The authors note a trade-off: increasing $K$ reduces variance but risks "dilution" (generic summaries) if too many signals are compressed.

C. Prompt Update and Migration

The next prompt $p_{t+1}$ is generated by an update policy that applies the aggregated corrective changes while respecting the preservation constraints.

For APM: When migrating a prompt from a source model to a target model, TRAS activates textual regularization immediately ( $\tau_{warmup} = 0$ ) and uses a smaller sample count $K$ (due to the lower variance of warm-started expert prompts). This ensures that target-specific mismatches are corrected via $\bar{g}_t$ while the transferable structure of the expert prompt is preserved via $\bar{r}_t$ .

3. Key Contributions

The paper formalizes and contributes the following:

Formalization of Automatic Prompt Migration (APM): The authors define APM as the practical problem of adapting expert prompts across model versions or providers while explicitly avoiding "instruction loss."
Textual Regularization: The introduction of a success-conditioned signal that encourages the preservation of beneficial prompt components, stabilizing the optimization trajectory.
Monte Carlo Signal Aggregation (MCSA): A technique to aggregate multiple stochastic signals to reduce variance and manage the dilution of actionable advice.
TRAS Framework: A unified method that combines these elements to improve robustness in both standard APO and APM scenarios.

4. Experimental Results

The authors evaluated TRAS on five reasoning tasks (Causal Judgment, Geometric Shapes, Penguins, Biosses, CB) using GPT-3.5-turbo for standard APO and GPT-4o for APM (migrating from GPT-3.5-turbo).

Standard APO (GPT-3.5-turbo):
- TRAS achieved consistent and statistically significant accuracy improvements over the baseline (PromptAgent), ranging from 4.9% to 21.5% across tasks.
- It reduced standard deviation across random seeds, indicating more stable optimization.
- It reduced LLM API calls by 0.5% to 3.3% by avoiding inefficient remove-and-rediscover cycles.
Automatic Prompt Migration (GPT-3.5-turbo $\to$ GPT-4o):
- Naive re-optimization of expert prompts often eroded initial gains due to instruction loss.
- TRAS consistently improved final accuracy over the baseline, with gains ranging from 2.2% to 16.0%.
- Notable improvements were seen in Geometric Shapes (+12.5%) and Biosses (+16.0%), where preserving domain-specific structure was critical.
- TRAS reduced total API calls by 4.2% to 6.2% in the migration setting.
Ablation Studies:
- Warm-up: Performance peaked when regularization activated after 3–4 iterations; activating too early anchored weak prompts, while activating too late allowed drift.
- Aggregation Count ( $K$ ): Accuracy improved up to $K \approx 6$ (reducing variance) but degraded for larger $K$ due to signal dilution.

5. Significance and Claims

The paper claims that TRAS offers a robust solution to the instability inherent in black-box prompt optimization. By shifting from a purely error-driven paradigm to a feedback-centric approach that balances correction with preservation, TRAS:

Stabilizes Optimization: Reduces semantic drift and variance in the prompt search trajectory.
Enables Effective Migration: Solves the practical problem of adapting prompts across model versions without losing critical instructions, a failure mode common in naive re-optimization.
Improves Efficiency: Achieves higher accuracy with fewer API calls by converging more reliably and avoiding redundant cycles.

The authors position TRAS as orthogonal to search strategies, meaning it can enhance various existing APO backbones. They suggest future work could explore alternative aggregation mechanisms, adaptive regularization schedules, and extensions to multi-turn or continual deployment settings.