Resumen Técnico: POETS (Ensembles de Políticas para Muestreo Thompson)

Declaración del Problema

Equilibrar la exploración y la explotación sigue siendo un desafío central en la toma de decisiones secuencial y la optimización de caja negra, particularmente cuando se aplica a Modelos de Lenguaje Grandes (LLM). Una exploración efectiva requiere una cuantificación de la incertidumbre basada en principios. Aunque el ensamblaje de modelos entrenados independientemente es un método probado y agnóstico al modelo para estimar la incertidumbre epistémica, aplicar directamente un ensamblaje ingenuo a los LLM es computacionalmente prohibitivo debido a las restricciones de memoria y cómputo. Además, los métodos de exploración existentes para LLM a menudo dependen de modelos de recompensa explícitos (que introducen complejidades de optimización anidadas) o dependen en gran medida de las capacidades de seguimiento de instrucciones del modelo.

Metodología

El artículo introduce POETS (Ensembles de Políticas para Muestreo Thompson), un marco que une la cuantificación de la incertidumbre y la optimización de políticas sin requerir un modelo de recompensa probabilístico explícito.

Idea Central: Codificación Implícita de la Recompensa

POETS se basa en la dualidad entre funciones de recompensa y políticas óptimas bajo objetivos regularizados. Específicamente, las políticas entrenadas con regularización de Kullback-Leibler (KL) codifican implícitamente una función de recompensa subyacente. En lugar de aprender un modelo de recompensa y luego ajustar una política a este (un proceso de dos pasos), POETS ajusta directamente las funciones de recompensa implícitas codificadas por un ensemble de políticas a los datos en línea observados.

El Algoritmo

1. Ensemble de Políticas: POETS mantiene un ensemble de $n$ políticas independientes, $\Pi = \{\pi_1, \dots, \pi_n\}$.
2. Bootstrapping de Poisson: Para capturar la incertidumbre epistémica, los miembros del ensemble se diversifican mediante bootstrapping de Poisson. Para cada lote entrante de datos, cada miembro del ensemble se asigna un peso entero independiente extraído de una distribución Poisson(1). Esto permite que el ensemble aprenda una aproximación bayesiana posterior sobre las funciones de recompensa sin requerir un conjunto de datos estático.
3. Función de Pérdida: El método optimiza una función de pérdida que alinea la recompensa implícita de una sola política con la recompensa verdadera. La derivación teórica muestra que optimizar esta pérdida es matemáticamente equivalente a una actualización de gradiente de política (específicamente, una variante de GRPO) escalada por los pesos de Poisson.
4. Muestreo Thompson: Al ajustar un ensemble de recompensas implícitas a las recompensas verdaderas utilizando bootstraps independientes, POETS aprende una posterior sobre las funciones de recompensa. Ejecutar el muestreo Thompson regularizado por KL se logra simplemente muestreando uniformemente una política del ensemble $\Pi$ y actuando de acuerdo con ella.

Arquitectura Eficiente: "Tronco y Rama"

Para superar el costo prohibitivo de mantener $n$ copias completas de LLM, POETS utiliza una arquitectura eficiente en parámetros:

• Tronco Compartido: Una base preentrenada (la mayoría de los pesos) se comparte entre todos los miembros del ensemble y se evalúa solo una vez por paso hacia adelante/hacia atrás.
• Ramas Independientes: La diversidad del ensemble se encapsula completamente en módulos independientes de adaptación de bajo rango (LoRA) que actúan como las capas finales de lectura.
• Eficiencia: Este diseño asegura que generar respuestas y acumular gradientes para $n$ políticas tome virtualmente el mismo tiempo de reloj que una sola política, con una sobrecarga de memoria que escala solo con el tamaño de las ramas LoRA.

Contribuciones Clave

1. Marco Novel: Introducción de POETS, que captura la incertidumbre epistémica mediante ensembles de políticas y realiza implícitamente muestreo Thompson regularizado por KL sin un modelo de recompensa explícito.
2. Garantías Teóricas: Demostración de que POETS hereda los fuertes límites de arrepentimiento suave acumulativo de $O(\sqrt{T\gamma_T})$ asociados con el muestreo Thompson regularizado por KL, donde $\gamma_T$ es la ganancia de información máxima.
3. Eficiencia Computacional: Demostración de que una arquitectura "Tronco y Rama" LoRA permite un ensemble profundo con una sobrecarga mínima (por ejemplo, ~7.7% de aumento en el tiempo de ejecución para $n=16$ en una GPU H200).
4. Rendimiento Empírico: Eficiencia de muestra de vanguardia en diversos dominios de optimización de caja negra combinatoria (refinamiento de preguntas frecuentes, búsqueda de proteínas, diseño de circuitos cuánticos) y mayor robustez en entornos de aprendizaje por refuerzo, particularmente en escenarios fuera de política con replay de experiencia.

Resultados Experimentales

Los autores evaluaron POETS frente a líneas base que incluyen GRPO, TOSFIT (que utiliza un modelo de recompensa explícito), FIBO (muestreo Thompson en contexto) y EVOSEARCH (búsqueda evolutiva).

• Optimización de Caja Negra: POETS logró una eficiencia de muestra de vanguardia en la búsqueda de proteínas, el diseño de circuitos cuánticos y el refinamiento de preguntas frecuentes. Notablemente, superó a las líneas base incluso sin un búfer de replay, mientras que el GRPO estándar falló en mantener la diversidad y se sobreajustó a candidatos tempranos.
• Aprendizaje por Refuerzo (RL): En el escenario del conjunto de datos AIME (bandido contextual con recompensas verificables), POETS mejoró las trayectorias de optimización tanto para configuraciones en política como fuera de política (replay de experiencia). Mitigó con éxito los problemas de sobreajuste observados en GRPO cuando se expuso a contextos repetidos.
• Diversidad: El ensemble mantuvo la diversidad de generación (medida por la divergencia Jensen-Shannon) durante todo el entrenamiento, evitando el colapso de entropía común en los métodos estándar de gradiente de política.
• Ablaciones: Los experimentos confirmaron que el bootstrapping de Poisson es esencial para la diversidad, que ensembles pequeños ($n=8$) son suficientes para un alto rendimiento, y que la arquitectura "Tronco y Rama" captura efectivamente la incertidumbre con un costo computacional mínimo.

Significado y Afirmaciones

El artículo afirma que POETS representa un avance significativo al permitir una exploración basada en principios y consciente de la incertidumbre para LLM en el descubrimiento científico y el RL, sin la carga computacional del modelado explícito de recompensas ni la fragilidad del seguimiento de instrucciones en contexto.

• Teórico: Establece que los ensembles de políticas con bootstrapping son suficientes para implementar el muestreo Thompson regularizado por KL, heredando límites óptimos de arrepentimiento.
• Práctico: Proporciona una solución escalable que hace viable el ensemble profundo para LLM, permitiendo un uso efectivo del replay de experiencia y una exploración robusta en regímenes de alta incertidumbre.
• Alcance: El trabajo se posiciona como una contribución algorítmica fundamental para el post-entrenamiento de LLM y el descubrimiento científico, específicamente dirigido a problemas de optimización combinatoria donde las evaluaciones son costosas.

Los autores señalan que, aunque POETS mejora la eficiencia de muestra y la robustez, no aborda la mala especificación de recompensas, la equidad o la alineación, que siguen siendo responsabilidades de los diseñadores del sistema. Se sugiere trabajo futuro para extender el marco desde bandidos contextuales a configuraciones de múltiples pasos con asignación de crédito detallada.