Inference-Aware Prompt Optimization for Aligning Black-Box Large Language Models

Prompt optimization methods have demonstrated significant effectiveness in aligning black-box large language models (LLMs). In parallel, inference scaling strategies such as Best-of-N Sampling and Majority Voting have likewise been shown to improve alignment and performance by trading additional computation for better output. However, existing prompt optimization approaches are inference strategy agnostic; that is, they optimize prompts without accounting for the inference strategy. This constitutes a significant methodological gap, as our empirical and theoretical analysis reveals a strong interdependence between these two paradigms. Moreover, we find that user preferences regarding trade-offs among multiple objectives and inference budgets substantially influence the choice of prompt and inference configuration. To address this gap, we introduce a novel unified framework named IAPO (Inference-Aware Prompt Optimization) that jointly optimizes the prompt and inference scale, while being aware of the inference budget and different task objectives. We then develop a fixed-budget training algorithm for IAPO, called PSST (Prompt Scaling via Sequential Trimming), and establish finite-budget guarantees on the error probability. Finally, we evaluate the effectiveness of PSST on six tasks, including multi-objective text generation and reasoning, and demonstrate the critical role of incorporating inference-awareness in aligning black-box LLMs using prompt optimization.

💡 Research Summary

The paper addresses a critical gap in the alignment of black‑box large language models (LLMs): existing prompt‑optimization methods ignore the inference‑time strategies (e.g., Best‑of‑N sampling, Majority Voting) that are often employed to improve output quality at the cost of additional computation. The authors demonstrate both empirically and theoretically that prompt quality and inference scaling are tightly coupled; a prompt that performs best under single‑shot decoding may become sub‑optimal when multiple samples are aggregated, and vice‑versa. Moreover, user preferences over multiple objectives (helpfulness, harmlessness, exactness, etc.) and computational budgets further complicate the selection of an optimal prompt‑inference configuration.

To fill this gap, the authors propose IAPO (Inference‑Aware Prompt Optimization), a unified framework that jointly selects a prompt and an inference scale (the number of sampled completions N) conditioned on a context vector c that encodes user‑specified weights for each objective and a total inference budget. Formally, each arm a ∈ A is a tuple (p, N) where p is a prompt and N ∈