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

arXiv: https://arxiv.org/abs/2508.10030 · HTML: https://mlsanity.com/paper/2508.10030/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 16 of its 28 references that are in the map.
- First posted on arXiv 8 Aug 2025.

## Builds on (most cited first, 16 of 16)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2204.05862/) (2022)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [CommonsenseQA: A Question Answering Challenge Targeting Commonsense Knowledge](https://mlsanity.com/paper/1811.00937/) (2018)
- [Does Reinforcement Learning Really Incentivize Reasoning Capacity in LLMs Beyond the Base Model?](https://mlsanity.com/paper/2504.13837/) (2025)
- [Rewards-in-Context: Multi-objective Alignment of Foundation Models with Dynamic Preference Adjustment](https://mlsanity.com/paper/2402.10207/) (2024)
- [Black-Box Prompt Optimization: Aligning Large Language Models without Model Training](https://mlsanity.com/paper/2311.04155/) (2023)
- [BoNBoN Alignment for Large Language Models and the Sweetness of Best-of-n Sampling](https://mlsanity.com/paper/2406.00832/) (2024)
- [BOND: Aligning LLMs with Best-of-N Distillation](https://mlsanity.com/paper/2407.14622/) (2024)
- [RankGen: Improving Text Generation with Large Ranking Models](https://mlsanity.com/paper/2205.09726/) (2022)
- [GenARM: Reward Guided Generation with Autoregressive Reward Model for Test-time Alignment](https://mlsanity.com/paper/2410.08193/) (2024)
- [Efficient Prompt Optimization Through the Lens of Best Arm Identification](https://mlsanity.com/paper/2402.09723/) (2024)
- [Align-Pro: A Principled Approach to Prompt Optimization for LLM Alignment](https://mlsanity.com/paper/2501.03486/) (2025)
- [MORL-Prompt: An Empirical Analysis of Multi-Objective Reinforcement Learning for Discrete Prompt Optimization](https://mlsanity.com/paper/2402.11711/) (2024)

## Abstract (arXiv)

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.