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

TL;DR

This paper introduces IAPO, a unified framework that jointly optimizes prompts and inference strategies for large language models, improving alignment by considering inference costs and user preferences.

Contribution

It presents a novel inference-aware prompt optimization framework and a fixed-budget training algorithm with theoretical guarantees, addressing a key gap in existing methods.

Findings

IAPO improves alignment across multiple tasks.

Incorporating inference-awareness enhances prompt optimization effectiveness.

PSST provides finite-budget guarantees on error probability.

Abstract

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.