Thompson Sampling via Fine-Tuning of LLMs

El artículo presenta ToSFiT, un método escalable de muestreo Thompson que utiliza el ajuste fino de modelos de lenguaje grandes para optimizar espacios discretos sin necesidad de maximizar funciones de adquisición, logrando una eficiencia superior tanto en muestras como en cómputo en tareas diversas como el diseño de circuitos cuánticos y la búsqueda de proteínas.

Lo esencial

¡Claro que sí! Imagina que eres un chef experto que quiere crear el plato más delicioso del mundo, pero tienes un problema: no tienes una receta fija y el mundo de los ingredientes es tan inmenso que probar cada combinación tomaría más tiempo que la vida del universo.

Este es el desafío que resuelve el paper "Muestreo de Thompson mediante el Ajuste Fino de Modelos de Lenguaje" (TOSFIT).

Aquí te lo explico como si fuera una historia:

1. El Problema: El Laberinto Infinito

Imagina que tienes que encontrar la combinación perfecta de ingredientes (o de código para una computadora, o de secuencias de proteínas para un medicamento). El espacio de posibilidades es gigantesco y desordenado.

• El método antiguo (Optimización Bayesiana clásica): Imagina que tienes un mapa del tesoro (un modelo matemático) que te dice dónde podría estar el tesoro. Pero para decidir a dónde ir, el método antiguo te obliga a hacer un cálculo matemático extremadamente difícil y lento para encontrar el "mejor" punto en el mapa. Es como intentar encontrar la salida de un laberinto mirando un plano gigante y calculando cada paso con una calculadora. En espacios desordenados (como escribir código o crear proteínas), este cálculo es imposible de hacer rápido.

2. La Solución: El Chef con Instinto (TOSFIT)

Los autores proponen una idea brillante: ¿Y si en lugar de calcular el mejor punto, le preguntamos a un chef experto (un Modelo de Lenguaje o LLM) qué cocinaría él?

• El Chef (LLM): Ya tiene "instinto". Ha leído millones de recetas (datos de entrenamiento) y sabe qué suena bien.
• La Estrategia (Muestreo de Thompson): En lugar de buscar el punto perfecto matemáticamente, el método le dice al chef: "Basado en lo que hemos probado hasta ahora, imagina un plato que podría ser el mejor. ¡Cocínalo!".
• El Ajuste Fino (Fine-Tuning): Aquí está la magia. Si el plato que cocinó el chef es delicioso (tiene una buena recompensa), le damos un aplauso y le decimos: "¡Hazlo un poco más así la próxima vez!". Si es malo, le decimos: "¡Oye, eso no estaba tan bien, intenta algo diferente!".

En lugar de calcular matemáticas complejas para encontrar el mejor camino, enseñamos al chef a mejorar su instinto directamente. El modelo de lenguaje se convierte en el "mapa" que nos guía hacia las mejores soluciones.

3. ¿Por qué es tan genial? (Las Analogías)

• El "Muestreo de Thompson" como un Sueño Lúcido:
  Imagina que sueñas que encuentras el tesoro. En lugar de despertar y tratar de calcular dónde estaba en la vida real, el método te dice: "Cree que ese sueño es real y actúa como si fuera verdad". El modelo de lenguaje genera una solución basada en una "versión optimista" de la realidad. Si esa solución funciona, el modelo aprende. Si no, ajusta su sueño.

• El "Ajuste Fino" como un Entrenador Deportivo:
  Piensa en un entrenador (el algoritmo) y un atleta (el modelo de lenguaje).

  • El atleta ya es bueno porque ha practicado mucho antes (pre-entrenamiento).
  • El entrenador no le dice al atleta cómo correr cada paso matemáticamente. Solo le dice: "Corre, mira qué tal te fue, y ajusta tu técnica un poquito para la próxima".
  • Esto es mucho más rápido y eficiente que intentar calcular la física perfecta de cada movimiento.

4. Los Resultados: Ganando en Tres Frentes

El paper probó esta idea en tres áreas muy diferentes, como si fuera un chef probando su técnica en tres cocinas distintas:

1. Mejorar respuestas de preguntas frecuentes (FAQ): Como un asistente que aprende a responder mejor a los clientes.
2. Buscar proteínas estables: Como un biólogo que intenta diseñar una proteína que no se rompa con el calor (útil para medicamentos). ¡El espacio de posibilidades aquí es más grande que el número de átomos en el universo!
3. Diseñar circuitos cuánticos: Como un ingeniero que escribe código para computadoras cuánticas, donde un error pequeño lo arruina todo.

El resultado: TOSFIT encontró las mejores soluciones usando menos intentos (más eficiente en muestras) y menos tiempo de computadora (más eficiente computacionalmente) que cualquier otro método, incluidos los que usan inteligencia artificial avanzada o evolución artificial.

En Resumen

El papel nos dice: "Deja de intentar calcular la salida perfecta en un laberinto gigante. En su lugar, usa una inteligencia artificial que ya sabe mucho, pídele que imagine la salida, prueba esa idea, y enséñale a mejorar su imaginación poco a poco."

Es una forma de combinar la sabiduría previa de una IA (lo que ya sabe) con la curiosidad científica (probar cosas nuevas) para resolver problemas que antes parecían imposibles de optimizar. ¡Y lo hace sin perderse en cálculos matemáticos infinitos!

Resumen técnico

Resumen Técnico: TOSFIT (Thompson Sampling via Fine-Tuning of LLMs)

1. El Problema: Optimización Bayesiana en Espacios Discretos No Estructurados

La Optimización Bayesiana (BO) es un marco fundamental para la optimización de funciones de recompensa "caja negra" costosas, donde se busca maximizar una función desconocida $R(x)$ sobre un dominio $\mathcal{X}$. Tradicionalmente, la BO utiliza un proceso de selección de candidatos que equilibra la exploración y la explotación mediante una función de adquisición (como la probabilidad de mejora o el límite de confianza superior).

Sin embargo, la BO enfrenta un desafío crítico en espacios discretos grandes y no estructurados (como secuencias de aminoácidos para diseño de proteínas, código para circuitos cuánticos o texto natural):

• Falta de Gradientes: A diferencia de los espacios continuos, no existen gradientes para realizar una búsqueda eficiente.
• Intratabilidad de la Maximización: La estrategia estándar de Muestreo de Thompson (Thompson Sampling - TS) requiere maximizar una función de adquisición (o muestrear de la distribución posterior y encontrar su máximo). En espacios combinatorios masivos (ej. $20^{100}$ para proteínas), iterar sobre todos los puntos o maximizar la función de adquisición es computacionalmente imposible.
• Costo Computacional: Los métodos existentes que intentan resolver esto a menudo requieren estructuras de dominio específicas o son ineficientes.

2. Metodología: TOSFIT

Los autores proponen TOSFIT (Thompson Sampling via Fine-Tuning), un enfoque escalable que evita la maximización explícita de funciones de adquisición. En su lugar, parametriza directamente la Probabilidad de Maximalidad (PoM) utilizando Modelos de Lenguaje Grandes (LLM).

Conceptos Clave:

• Muestreo de Thompson como Fine-Tuning: En lugar de buscar el punto máximo en una función de adquisición, TOSFIT trata la generación de candidatos por parte del LLM como muestras directas de la distribución de PoM ($P[R_x = R^* | \text{datos}]$).
• Inicialización con Priors Fuertes: El algoritmo no comienza desde una política uniforme. Utiliza un LLM pre-entrenado y condicionado por prompts (contexto) como punto de partida. Esto aprovecha el conocimiento previo del modelo sobre la estructura del dominio (ej. gramática de código, biología de proteínas).
• Adaptación Variacional (VBOS): El modelo se ajusta (fine-tuning) incrementalmente hacia la posterior utilizando el objetivo de Muestreo Optimista Bayesiano Variacional (VBOS).
  • El objetivo VBOS maximiza una función que combina la recompensa esperada ($\mu_x$) y un término de exploración basado en la entropía y la incertidumbre ($\sigma_x$).
  • La actualización de los parámetros del modelo se realiza mediante descenso de gradiente estocástico sobre este objetivo variacional.
• Estabilización de Gradientes: Para manejar la alta varianza en la estimación de gradientes, TOSFIT utiliza una variante del método RLOO (Reinforce Leave-One-Out) con normalización, similar a GRPO (Group Relative Policy Optimization), lo que permite un entrenamiento estable sin necesidad de un modelo de valor (critic) separado.
• Inferencia Escalable: Utiliza Procesos Gaussianos (GP) con kernels lineales sobre representaciones de embeddings (deep features) para mantener la complejidad computacional constante respecto al número de observaciones, independientemente del tamaño del espacio de búsqueda.

Algoritmo (Resumen):

1. Burn-in: Generar candidatos iniciales con el LLM pre-entrenado para ajustar los hiperparámetros del GP.
2. Bucle de Optimización:
   • Generar un lote de candidatos ($x_1, \dots, x_B$) usando la política actual $\pi_\theta$.
   • Evaluar recompensas.
   • Actualizar el GP (posterior).
   • Calcular el gradiente del objetivo VBOS respecto a los parámetros del LLM.
   • Realizar pasos de fine-tuning (gradiente ascendente) para adaptar el LLM hacia la PoM posterior.

3. Contribuciones Clave

1. Nuevo Límite de Regret Teórico: Los autores derivan un límite de regret acumulado para una formulación variacional de Thompson Sampling que mejora el estado del arte.
   • Mejoran el límite de $\tilde{O}(\sqrt{T|\mathcal{X}|})$ (que es vacuo en espacios grandes) a $\tilde{O}(\sqrt{T\gamma_T})$, donde $\gamma_T$ es la ganancia de información máxima. Este nuevo límite es independiente del tamaño del espacio $\mathcal{X}$ y depende de la estructura del kernel, coincidiendo con las garantías fuertes de TS estándar.
   • Demuestran que la divergencia Bregman entre la política aproximada (LLM) y la óptima (VBOS exacto) es el factor crítico que debe minimizarse.
2. Algoritmo TOSFIT: Introducen un algoritmo práctico que combina el conocimiento de pre-entrenamiento de LLMs con la adaptación bayesiana rigurosa, resolviendo el problema de la maximización intratable en espacios discretos.
3. Validación Empírica: Demuestran que el enfoque es superior a métodos de búsqueda evolutiva, aprendizaje por refuerzo (Actor-Critic) y BO clásica en contextos de in-context learning.

4. Resultados Experimentales

El método se evaluó en tres tareas diversas con modelos LLM de 0.6B a 8B parámetros:

1. Refinamiento de Respuestas FAQ: Optimización de texto para alinear semánticamente con una respuesta "ground truth".
2. Búsqueda de Proteínas: Diseño de secuencias de aminoácidos con alta estabilidad térmica (espacio de búsqueda astronómico).
3. Diseño de Circuitos Cuánticos: Generación de código Qiskit válido para preparar estados de baja energía.

Hallazgos Principales:

• Eficiencia de Muestra (Sample Efficiency): TOSFIT supera consistentemente a todos los baselines (incluyendo FIBO, búsqueda evolutiva y Actor-Critic) en la velocidad de convergencia hacia la mejor recompensa observada.
• Eficiencia Computacional: A pesar del costo adicional del fine-tuning, TOSFIT es más eficiente computacionalmente en términos de tiempo total para alcanzar un objetivo, debido a su alta eficiencia de muestra.
• Importancia del Contexto y la Adaptación Cuidadosa:
  • Los experimentos muestran que inicializar con un LLM pre-entrenado y bien condicionado es crucial.
  • Una tasa de aprendizaje demasiado alta o una adaptación descuidada lleva al "olvido" del prior y al estancamiento.
  • La optimización en batch (múltiples candidatos en paralelo) mejora la eficiencia de iteración sin sacrificar significativamente la eficiencia de muestra.

5. Significado e Impacto

• Superación de Limitaciones de la BO: TOSFIT resuelve el cuello de botella histórico de la Optimización Bayesiana en espacios discretos masivos, permitiendo su aplicación en dominios científicos complejos donde antes era inviable.
• Sinergia LLM + Optimización Bayesiana: El trabajo demuestra que los LLMs no son solo generadores de texto, sino políticas de muestreo potentes que pueden ser ajustadas rigurosamente mediante principios bayesianos para la toma de decisiones bajo incertidumbre.
• Aplicabilidad Científica: Ofrece una herramienta viable para la aceleración de descubrimientos en biotecnología (diseño de proteínas) y computación cuántica, donde el espacio de búsqueda es demasiado grande para la fuerza bruta y la falta de gradientes impide métodos tradicionales.
• Fundamento Teórico: Proporciona una justificación teórica sólida para el uso de fine-tuning en la optimización, vinculando la divergencia de la política con el regret, lo que guía el diseño de algoritmos futuros.

En conclusión, TOSFIT representa un avance significativo al transformar el muestreo de Thompson en un proceso de optimización de políticas mediante fine-tuning, logrando un equilibrio óptimo entre la explotación de conocimiento previo y la exploración guiada por la incertidumbre en espacios de búsqueda combinatorios masivos.