DRO-InstructZero: Robust Prompt Optimization

• DRO-InstructZero is a framework that uses distributionally robust optimization (DRO) to ensure LLM prompt performance under worst-case distribution shifts.
• It employs a robust Bayesian latent space search with adversarial weighting and Gaussian Process modeling to optimize soft prompts efficiently.
• Experimental results demonstrate significant accuracy gains, with improvements up to 30 percentage points in tasks like formality rewriting and code debugging.

DRO-InstructZero is a prompt optimization framework for LLMs that explicitly targets robustness under distribution shift and adversarial evaluation. Unlike existing methods such as InstructZero, which optimize prompt performance for a nominal evaluation distribution, DRO-InstructZero employs distributionally robust optimization (DRO) to maximize worst-case utility over an ambiguity set—typically defined by an $f$-divergence ball around the evaluation distribution. The methodology is plug-and-play and general; it yields substantial and consistent performance gains in zero-shot settings (e.g., formality rewriting, code debugging, translation) without sacrificing efficiency or in-distribution accuracy (Li, 17 Oct 2025).

1. Underlying Motivation and Robustness Principles

Prompt discovery in LLMs is highly sensitive to distributional variation. Methods such as InstructZero, which rely on Bayesian optimization over latent soft prompts, often produce instructions that fail under domain shift, adversarial querying, or minor lexical changes in input. DRO-InstructZero addresses this by reframing prompt optimization in the DRO paradigm. Here, the foundational principle is to optimize the prompt for worst-case expected performance across all distributions $Q$ within a divergence ball (such as KL divergence) centered on the nominal evaluation distribution $D^t$:

$$\max_{v \in \mathcal{V}} \inf_{Q \in \mathcal{U}(D^t)} \mathbb{E}_{(X, Y) \sim Q}\big[ h(f([v; X]), Y) \big]$$

where $\mathcal{U}(D^t)$ denotes the ambiguity set (e.g., those $Q$ such that $\mathrm{KL}(Q \,||\, D^t) \leq \epsilon$), $f$ is the black-box LLM, $v$ is the candidate instruction, $h$ is the performance metric (typically execution accuracy), and $(X, Y)$ are task inputs and targets. This robust optimization ensures the selected prompt retains high utility in every plausible shift scenario.

2. Methodology: DRO-Driven Bayesian Latent Space Search

DRO-InstructZero extends the InstructZero pipeline but modifies the Bayesian optimization step to be robust. The process is as follows:

• Latent Prompt Representation: Optimization is conducted over a continuous soft prompt $p \in \mathbb{R}^d$, projected via a random matrix $A$ into a token embedding space for prompt generation.
• Instruction Decoding: An open-source LLM (e.g., Vicuna) maps the projected soft prompt plus few-shot exemplars into a natural language instruction: $v = g([A p; \mathrm{exemplars}])$.
• Zero-Shot Evaluation: The decoded instruction is sent to the target black-box LLM, and its performance is evaluated on task data $\mathcal{D}_t$ using metric $h$.
• Robust Objective Modeling: A Gaussian Process (GP) is fit over the robust objective

$$H(p) = \inf_{Q \in \mathcal{U}(D^t)} \mathbb{E}_{(X,Y) \sim Q}\big[h(f([g([A p;\text{exemplars}]); X]), Y)\big]$$

capturing epistemic uncertainty about performance under worst-case distributional shifts.

• Robust Acquisition Function: Standard acquisition metrics (e.g., Expected Improvement, UCB) are adapted: for a batch of candidate prompts $p_m$, robust acquisition selects adversarial weights $w^*_m$ within the divergence ball (by solving $$\min_{w':\,\|w'-w_\text{ref}\|_{\mathcal{M}} \leq \epsilon(m)}\, \langle \text{ucb}_m, w' \rangle$$), and chooses the next query as the maximizer of robust expected utility $\langle \text{ucb}_m, w^*_m \rangle$.

This procedure incorporates worst-case risk directly into the BO loop, maintaining query efficiency while explicitly seeking reliability under natural or adversarial variation.

3. Experimental Protocol and Setting

DRO-InstructZero is experimentally validated on standard zero-shot prompt optimization tasks:

• Instruction-Induction Regime: Tasks adhere to the instruction-induction protocol of prior prompt search research, providing a fixed set of few-shot exemplars with matched query budgets during optimization.
• Evaluated Tasks: Performance assessments encompass (1) formality rewriting (e.g. BIG-Bench “informative-to-formal”), (2) code debugging under induced domain shift, and (3) multiple translation pairs.
• Query Efficiency: The search matches InstructZero in query cost (e.g., 25 prompt candidates per iteration), enabling direct comparative analysis.
• Robustness Controls: The methodology examines sensitivity of robust gains to the divergence radius $\epsilon$, divergence type (KL, total variation), and decoding temperature.

4. Performance Gains and Robustness Across Conditions

DRO-InstructZero yields pronounced improvements in both average and worst-case accuracy under various shift scenarios:

Task	Baseline Accuracy (%)	DRO-InstructZero Accuracy (%)	Absolute Gain (%)
Form-to-Formal Rewriting	61.3 ± 0.7	85–90	+25 to +30
Auto-Debugging (Shifted)	Baseline +0	Baseline +25	+25
Translation EN→DE	84	95	+11
Cause-and-Effect (Stable)	>96	>96	≈0

• Distribution Shift: DRO-InstructZero closes 25–30-point gaps in adversarial rewriting and debugging tasks.
• In-Distribution Stability: No degradation is observed for tasks already robust to shift, e.g., cause-and-effect remains >96%.
• Generalization: Robust performance is consistent across divergence metrics and choices of temperature, affirming the method's regularization effect and transferability over real-world conditions.

5. Ambiguity Set Specification and Theoretical Justification

The ambiguity set $\mathcal{U}(D^t)$ is typically taken as an $f$-divergence ball: for instance,

$$\mathcal{U}(D^t) = \{ Q : D_f(Q, D^t) \leq \epsilon \}$$

where $D_f$ is a chosen $f$-divergence (KL, TV, $\chi^2$). The choice of $\epsilon$ mediates between conservatism and adaptability: larger balls yield more robust but potentially less specific prompts. Empirical ablation demonstrates stability of gains over plausible $\epsilon$ ranges.

The robust optimization problem mirrors classical DRO regimes applied to prompt learning. At both the acquisition and evaluation stages, adversarial distributions within the ambiguity set are considered, directly targeting minimax prompt reliability.

6. Practical Implications and Plug-and-Play Applicability

DRO-InstructZero provides a direct, general extension to existing prompt search procedures (not limited to InstructZero) by replacing average-case Bayesian optimization steps with robustified objectives and acquisitions. This modification requires no fundamental alteration to standard pipeline components; it is compatible with any BO-driven soft prompt system.

Practitioners seeking reliable prompt discovery—particularly where input distributional stability is not guaranteed—can deploy DRO-InstructZero to mitigate variance under domain, adversarial, or temporal drift. Its query efficiency and consistency under multiple objective and decoding conditions obviate the need for complex uncertainty modeling at the pre-processing or post-processing level.

7. Broader Connections: Robust Prompt Learning in LLMs

DRO-InstructZero's approach extends recent trends toward risk-aware optimization in black-box LLM interaction. By formalizing zero-shot prompt induction as a minimax problem over ambiguity sets parameterized by $f$-divergences and leveraging robust Bayesian search, the method advances the push for reliable, transferable instruction alignment. Its empirical demonstration of robust improvement across diverse tasks and evaluation settings sets a new technical standard for prompt robustness, with direct applicability to translation, debuggers, rewriting, and more.

In sum, DRO-InstructZero integrates distributionally robust optimization principles into automatic prompt engineering, solving a previously open problem in reliable, zero-shot instruction induction for LLMs under real-world uncertainty (Li, 17 Oct 2025).