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

1. El Problema: El Desacoplamiento entre Prompts e Inferencia

La mayoría de los modelos de lenguaje de gran tamaño (LLMs) actuales se acceden mediante APIs de "caja negra", lo que impide el ajuste fino (fine-tuning) de sus pesos. Para alinearlos con los objetivos del usuario, se utilizan dos estrategias principales:

1. Optimización de Prompts: Modificar las instrucciones de entrada.
2. Escalado de Inferencia: Aumentar el cómputo mediante técnicas como Best-of-N (seleccionar la mejor respuesta entre $N$ muestras) o Majority Voting (seleccionar la respuesta más frecuente).

La brecha metodológica: Los métodos actuales de optimización de prompts son agnósticos a la estrategia de inferencia. Es decir, optimizan un prompt asumiendo una ejecución única (single-shot), sin considerar cómo se comportará ese prompt cuando se apliquen técnicas de escalado o cuando el usuario tenga restricciones de presupuesto computacional. El estudio demuestra que un prompt óptimo para una sola respuesta puede ser subóptimo (o incluso erróneo) cuando se utiliza Majority Voting o Best-of-N.

2. Metodología: El Marco IAPO y el Algoritmo PSST

Para resolver este problema, los autores proponen un nuevo paradigma:

A. Marco IAPO (Inference-Aware Prompt Optimization):
A diferencia de los métodos tradicionales, IAPO no busca solo el "mejor prompt", sino la mejor política $\pi(c)$. Esta política recibe un contexto $c$ (que incluye las preferencias del usuario sobre múltiples objetivos, como utilidad vs. seguridad, y su presupuesto de cómputo) y selecciona simultáneamente:

• El prompt óptimo.
• La escala de inferencia ($N$ muestras) adecuada.

El problema se formaliza como un problema de Identificación del Mejor Brazo Contextual (Contextual Best-Arm Identification), donde cada "brazo" es una combinación de (prompt, escala de inferencia).

B. Algoritmo PSST (Prompt Scaling via Sequential Trimming):
Dado que el entrenamiento es costoso, proponen PSST, un algoritmo de eliminación de brazos con presupuesto fijo. Sus innovaciones técnicas son:

• Reutilización de muestras (Cross-scale reuse): Si se evalúa un prompt con $N=16$ muestras, los resultados pueden reutilizarse para estimar el rendimiento con $N=8, 4, 2$, etc., optimizando el presupuesto.
• Asignación consciente de la estructura: El algoritmo asigna el presupuesto de manera que los prompts con escalas de inferencia más altas (que consumen más tokens) no penalicen la exploración de otros prompts.
• Heurística Top-K Screening: Una técnica de pre-filtrado para eliminar prompts claramente mediocres antes de iniciar el proceso completo de PSST, acelerando la convergencia en presupuestos bajos.

3. Contribuciones Clave

1. Formalización Teórica: Demuestran que la optimización de prompts y la de inferencia están intrínsecamente acopladas y que tratarlas por separado conduce a la desalineación.
2. Garantías de Error: Establecen límites teóricos de error para el algoritmo PSST, demostrando que puede encontrar el brazo óptimo con una probabilidad alta bajo un presupuesto finito.
3. Nuevo Algoritmo de Entrenamiento: PSST permite optimizar políticas complejas en entornos de caja negra de manera eficiente y sin necesidad de ajustar hiperparámetros.

4. Resultados Experimentales

El marco se evaluó en seis tareas diversas (razonamiento matemático, razonamiento de sentido común, resumen de textos y alineación de utilidad/seguridad) utilizando modelos como Llama-3.3-70B.

• Superioridad de IAPO: Los métodos que consideran la inferencia (IAPO) superaron significativamente a los métodos agnósticos. En algunos casos, la optimización conjunta mejoró la recompensa en un 25% respecto a la optimización desacoplada y hasta un 50% respecto a la optimización de solo prompts.
• Efectividad de PSST: El algoritmo PSST y su variante con Top-K screening superaron consistentemente a estrategias clásicas de exploración como UCB, $\epsilon$-greedy y Softmax.
• Sensibilidad al Presupuesto: Se demostró que el prompt óptimo cambia según el presupuesto: un prompt "sencillo" puede ser mejor para presupuestos bajos, mientras que un prompt "complejo" o "estructurado" brilla cuando se permite un mayor escalado de inferencia.

5. Significado y Conclusión

Este trabajo cambia la perspectiva de la alineación de LLMs de caja negra. Sugiere que la calidad de un prompt no es una propiedad intrínseca del texto, sino una función de cómo se ejecuta ese texto. Al integrar el diseño del prompt con la estrategia de despliegue (inferencia), se logra una alineación mucho más robusta, económica y adaptada a las necesidades específicas de cada usuario.