Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DRO-InstructZero: Distributionally Robust Prompt Optimization for Large Language Models

Yangyang Li

TL;DR

DRO-InstructZero raises the bar for prompt optimization by addressing distributional shift through distributionally robust optimization embedded in Bayesian search. By replacing average-case acquisitions with worst-case robust objectives, it produces instructions that maintain high performance under adversarial and domain-shift conditions without increasing query costs. Empirical results across 32 BIG-Bench tasks show substantial robustness gains (e.g., translation and debugging improvements) while preserving in-distribution accuracy, illustrating the practical value of robustness in prompt learning. The work bridges distributionally robust optimization with prompt design, offering a plug-and-play approach to reliable, transferable instruction alignment for LLMs in real-world uncertainty.

Abstract

Large language models are highly sensitive to prompt wording. However, popular automatic prompt search methods, including InstructZero, often degrade under distribution shift and adversarial evaluation because they optimize expected performance under a single evaluation distribution. Consequently, prompts that work in one setting frequently fail to transfer. To address this, DRO-InstructZero formulates zero-shot prompt optimization as robust Bayesian optimization. Specifically, an f-divergence ball defines an ambiguity set around the evaluation distribution, and a robust acquisition rule maximizes worst-case expected utility while retaining the query efficiency of Bayesian search. Therefore, the search explicitly targets reliability under distribution shift rather than average behavior alone. Experiments follow the instruction-induction protocol with matched query budgets across formality rewriting, code debugging, and translation. For example, on BIG-Bench informative-to-formal rewriting, accuracy improves from 61.3 +/- 0.7% to approximately 85-90%, yielding an absolute gain of about 25-30 points. Moreover, auto-debugging shows about +25-point gains under domain shift. Meanwhile, stable tasks such as cause-and-effect remain above 96%, indicating no loss on in-distribution cases. Furthermore, improvements are consistent across divergence choices and decoding temperatures. Overall, DRO-InstructZero connects distributionally robust optimization with prompt learning, offering a plug-and-play and general approach for reliable, transferable prompt alignment under real-world uncertainty.

DRO-InstructZero: Distributionally Robust Prompt Optimization for Large Language Models

TL;DR

DRO-InstructZero raises the bar for prompt optimization by addressing distributional shift through distributionally robust optimization embedded in Bayesian search. By replacing average-case acquisitions with worst-case robust objectives, it produces instructions that maintain high performance under adversarial and domain-shift conditions without increasing query costs. Empirical results across 32 BIG-Bench tasks show substantial robustness gains (e.g., translation and debugging improvements) while preserving in-distribution accuracy, illustrating the practical value of robustness in prompt learning. The work bridges distributionally robust optimization with prompt design, offering a plug-and-play approach to reliable, transferable instruction alignment for LLMs in real-world uncertainty.

Abstract

Large language models are highly sensitive to prompt wording. However, popular automatic prompt search methods, including InstructZero, often degrade under distribution shift and adversarial evaluation because they optimize expected performance under a single evaluation distribution. Consequently, prompts that work in one setting frequently fail to transfer. To address this, DRO-InstructZero formulates zero-shot prompt optimization as robust Bayesian optimization. Specifically, an f-divergence ball defines an ambiguity set around the evaluation distribution, and a robust acquisition rule maximizes worst-case expected utility while retaining the query efficiency of Bayesian search. Therefore, the search explicitly targets reliability under distribution shift rather than average behavior alone. Experiments follow the instruction-induction protocol with matched query budgets across formality rewriting, code debugging, and translation. For example, on BIG-Bench informative-to-formal rewriting, accuracy improves from 61.3 +/- 0.7% to approximately 85-90%, yielding an absolute gain of about 25-30 points. Moreover, auto-debugging shows about +25-point gains under domain shift. Meanwhile, stable tasks such as cause-and-effect remain above 96%, indicating no loss on in-distribution cases. Furthermore, improvements are consistent across divergence choices and decoding temperatures. Overall, DRO-InstructZero connects distributionally robust optimization with prompt learning, offering a plug-and-play and general approach for reliable, transferable prompt alignment under real-world uncertainty.
Paper Structure (32 sections, 12 equations, 2 figures, 3 tables, 2 algorithms)

This paper contains 32 sections, 12 equations, 2 figures, 3 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Instruction Optimization with InstructZero
  3. Pipeline Overview
  4. Bayesian Optimization of Soft Prompts
  5. Instruction-Coupled Kernel
  6. Algorithm
  7. DRO-InstructZero: Robust Instruction Optimization
  8. Problem Formulation
  9. From Structured Combinatorial Search to Continuous Robust Optimization
  10. Distributionally Robust Bayesian Optimization
  11. Bayesian Optimization of Soft Prompts
  12. Robust Acquisition under Ambiguity Sets
  13. Instruction-Coupled Kernel under DRO
  14. Algorithm
  15. Experiments
  16. ...and 17 more sections

Figures (2)

  • Figure 1: Per-task accuracy on 32 BIG-Bench tasks comparing InstructZero (blue) vs. DRO-InstructZero (orange).
  • Figure 2: Per-task accuracy on 32 BIG-Bench tasks comparing InstructZero (blue) vs. DRO-InstructZero (orange).