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

Resumo Técnico: Otimização de Prompts Consciente de Inferência (IAPO)

1. O Problema: O Gap da Agnosticidade de Inferência

Atualmente, o alinhamento de Grandes Modelos de Linguagem (LLMs) de "caixa-preta" (acessíveis apenas via API) é feito principalmente por meio de otimização de prompts. Paralelamente, existem estratégias de escalonamento de inferência (como Best-of-N Sampling e Majority Voting) que melhoram o desempenho ao gerar múltiplas respostas e selecionar a melhor.

O problema central identificado pelos autores é que os métodos atuais de otimização de prompts são agnósticos à estratégia de inferência. Ou seja, eles otimizam um prompt sem considerar como as respostas serão agregadas ou qual o orçamento computacional disponível. A pesquisa demonstra que:

• A performance de um prompt é altamente sensível à estratégia de inferência escolhida.
• Otimizar o prompt e a estratégia de inferência de forma separada pode levar ao desalinhamento (um prompt excelente para uma única resposta pode ser medíocre quando se usa votação majoritária).
• Existe uma interdependência complexa entre o design do prompt, o número de amostras ($N$) e as preferências do usuário (ex: equilíbrio entre utilidade e custo).

2. Metodologia: O Framework IAPO e o Algoritmo PSST

Para resolver essa lacuna, os autores introduzem o IAPO (Inference-Aware Prompt Optimization), um framework unificado que otimiza conjuntamente o prompt e a escala de inferência, levando em conta o orçamento e os objetivos da tarefa.

Principais componentes metodológicos:

• Formulação de Problema: O problema é modelado como um problema de Identificação do Melhor Braço Contextual (Contextual Best-Arm Identification - BAI). Cada "braço" é uma combinação de um prompt ($p$), um hiperparâmetro de decodificação ($g$) e um número de amostras ($N$).
• Algoritmo PSST (Prompt Scaling via Sequential Trimming): Como o orçamento de treinamento é limitado, os autores propõem o PSST, um algoritmo de eliminação de braços baseado em orçamento fixo. O PSST é eficiente porque aproveita três propriedades estruturais:
  1. Reuso entre contextos: Uma única consulta fornece dados para múltiplos contextos de preferência.
  2. Reuso entre escalas (Nested Sample Reuse): Se você amostra 10 respostas para um prompt, pode usar essas mesmas respostas para simular uma estratégia de $N=2$ ou $N=5$ sem gastar mais tokens.
  3. Alocação Consciente de Custo: O algoritmo prioriza a exploração de prompts de forma que o custo assimétrico das amostras não desequilibre o aprendizado.
• Heurística Top-K Screening: Uma técnica prática para reduzir o custo de treinamento, realizando uma triagem rápida de prompts em escala unitária ($N=1$) antes de aplicar o PSST completo.

3. Principais Contribuições

1. Framework Teórico: A formalização do IAPO como um problema de bandit contextual, integrando custo, múltiplos objetivos e estratégias de agregação (BON e MV).
2. Algoritmo PSST: Um método de treinamento robusto, livre de hiperparâmetros, que utiliza o reuso de amostras para maximizar a eficiência do orçamento.
3. Garantias Teóricas: Demonstração de limites de erro para o PSST, provando que o algoritmo converge para a política ótima dentro de um orçamento finito.
4. Descoberta da Interdependência: Prova matemática e empírica de que a otimização de prompts não pode ser desacoplada da estratégia de inferência.

4. Resultados Experimentais

Os autores avaliaram o método em seis tarefas distintas (incluindo raciocínio matemático, compreensão de senso comum e geração de texto multi-objetivo) usando o modelo Llama-3.3-70B-Instruct.

• Superioridade do IAPO: O framework IAPO superou significativamente os métodos agnósticos. A otimização conjunta superou a otimização de prompt pura em até 50% e a otimização desalinhada (disjoint) em até 25% em termos de recompensa ajustada ao custo (ACR).
• Eficácia do PSST: O algoritmo PSST e sua variante com Top-K screening mostraram-se consistentes e superiores a estratégias clássicas de exploração como $\epsilon$-greedy, UCB e Softmax.
• Evidência de Desalinhamento: Os experimentos mostraram que prompts otimizados para $N=1$ (single-shot) frequentemente falham ao serem escalados para $N=16$ em tarefas de votação majoritária, confirmando a necessidade de consciência de inferência.

5. Significância

Este trabalho é fundamental para o desenvolvimento de sistemas de IA mais confiáveis e econômicos. Ao permitir que o prompt e a computação de inferência sejam ajustados dinamicamente às necessidades do usuário (ex: "quero o máximo de precisão possível dentro de um orçamento de 10 centavos"), o IAPO oferece um caminho para o alinhamento eficiente de modelos de caixa-preta, onde o ajuste fino (fine-tuning) é impossível ou proibitivo.