Reference-guided Policy Optimization for Molecular Optimization via LLM Reasoning

El artículo presenta RePO, un nuevo enfoque de optimización que combina aprendizaje por refuerzo con recompensas verificables y guía supervisada basada en referencias para superar las limitaciones de los modelos de lenguaje actuales en la optimización de moléculas, logrando un mejor equilibrio entre la exploración de nuevas estructuras y la explotación de referencias existentes sin necesidad de datos de trayectoria.

Lo esencial

¡Hola! Imagina que tienes un chef de cocina muy inteligente (el modelo de lenguaje o LLM) al que le pides que cree una receta nueva basada en una que ya existe.

El objetivo es: Mejorar el sabor de la receta (optimizar una propiedad química), pero sin cambiar tanto los ingredientes que ya no se parezca a la original (mantener la similitud estructural).

El problema es que el chef, aunque es muy listo, a veces se pierde:

1. Si solo le das la receta final perfecta sin explicarle cómo la hizo, el chef se vuelve perezoso y deja de pensar en los pasos intermedios (solo copia la respuesta).
2. Si le dices "prueba mil cosas hasta que salga algo bueno" sin darle una referencia, el chef se vuelve demasiado conservador y solo hace cambios minúsculos que no mejoran nada, o se pierde en un laberinto gigante de ingredientes posibles.

Aquí es donde entra el RePO (Optimización de Política Guiada por Referencia), la solución que proponen en este paper.

🍳 La Analogía del Chef y el "Chef Maestro"

Imagina que quieres mejorar un pastel. Tienes un Pastel Original y un Pastel de Referencia (el ejemplo de éxito que te dan en el dataset).

• El problema de los métodos antiguos (SFT y RLVR):

  • Método A (Solo copiar): Le dices al chef: "Haz exactamente este pastel de referencia". El chef lo copia, pero deja de pensar por qué es bueno. Si le pides que lo mejore, no sabe cómo empezar porque solo sabe copiar.
  • Método B (Solo probar a ciegas): Le dices: "Prueba mil combinaciones y si el pastel sabe mejor, repítelo". Como hay millones de combinaciones posibles, el chef se asusta. Solo prueba cambios muy pequeños (como cambiar una pizca de sal) por miedo a arruinarlo, y nunca logra un pastel espectacular.

• La solución RePO (El Chef con un Asistente):
  RePO es como tener un sistema de entrenamiento híbrido para el chef:

  1. La Exploración (El RL): Le permitimos al chef probar sus propias ideas y crear nuevos pasteles. Si un pastel sale delicioso (mejora la propiedad), le damos una medalla (recompensa). Esto le anima a explorar y ser creativo.
  2. La Guía (La Referencia): Pero, ¡ojo! Mientras el chef piensa en los pasos (el "razonamiento"), le mostramos el Pastel de Referencia al final. Le decimos: "Piensa todo lo que quieras, pero asegúrate de que tu resultado final se parezca a este ejemplo de éxito".

¿Cómo funciona mágicamente?

El truco de RePO es que no le dice al chef cómo pensar paso a paso, solo le dice qué resultado final buscar.

• Sin atascar el pensamiento: El chef sigue usando su propia lógica para decidir qué ingredientes cambiar (el razonamiento). No le copian la receta mental, solo le dan un objetivo claro.
• Evitar el caos: Como tiene el ejemplo de referencia como "ancla", no se pierde en combinaciones imposibles o tóxicas. Sabe que debe mantener la estructura base.

🏆 ¿Qué lograron?

En sus pruebas (como en un concurso de química llamado TOMG-Bench), este nuevo método (RePO) fue mucho mejor que los anteriores:

• Encontró mejores recetas: Logró mejorar más las propiedades de las moléculas (como hacerlas más efectivas como medicamentos).
• Mantuvo la esencia: Las moléculas nuevas seguían pareciéndose a las originales (no se rompieron las reglas químicas).
• Aprendió a pensar: A diferencia de los otros métodos que dejaron de razonar, RePO mantuvo la capacidad del modelo de explicar por qué hizo los cambios (ej: "Cambiamos el bromo por cloro para reducir el peso").

En resumen

Imagina que estás aprendiendo a conducir un coche de carreras.

• Los métodos viejos te decían: "Copia exactamente el trayecto del campeón" (y dejabas de pensar) O "Maneja a lo loco hasta que ganes" (y te estrellabas).
• RePO te dice: "Maneja libremente, explora la pista y prueba nuevas rutas, pero ten siempre en la mira la meta del campeón para asegurarte de que llegas bien".

Gracias a esto, la Inteligencia Artificial puede ahora ayudar a los científicos a diseñar nuevos medicamentos y materiales de forma más rápida, segura y creativa, entendiendo mejor la química detrás de sus decisiones.

Resumen técnico

Aquí tienes un resumen técnico detallado del artículo "Reference-Guided Policy Optimization for Molecular Optimization via LLM Reasoning" (Optimización de Políticas Guiada por Referencia para Optimización Molecular mediante Razonamiento de LLM), publicado en ICLR 2026.

1. Planteamiento del Problema

El artículo aborda el desafío de la optimización molecular basada en instrucciones utilizando Modelos de Lenguaje Grandes (LLM). En esta tarea, el modelo debe modificar una molécula de entrada ($m_0$) para mejorar una propiedad objetivo (como la solubilidad LogP o la similitud con fármacos QED), manteniendo al mismo tiempo una alta similitud estructural con la molécula original.

El conflicto central (Desajuste de Supervisión):

• Objetivos competitivos: Mejorar una propiedad química a menudo requiere ediciones estructurales no triviales (ej. añadir grupos funcionales), lo que reduce la similitud estructural y puede romper la validez química.
• Limitación de los datos: Los conjuntos de datos existentes suelen proporcionar solo una molécula de referencia optimizada ($m_{ref}$) por instrucción, sin incluir la trayectoria de razonamiento paso a paso (cómo se llegó a esa solución).
• Fallo de los métodos actuales:
  • SFT (Fine-Tuning Supervisado) estándar: Al entrenar solo con la respuesta final (molécula), el modelo colapsa el razonamiento, generando salidas directas y cortas sin explorar el espacio químico. Esto limita la capacidad de encontrar soluciones óptimas.
  • RLVR (Aprendizaje por Refuerzo con Recompensas Verificables, ej. GRPO): Al iniciar desde un modelo base o uno con SFT, el RL sufre de recompensas dispersas. Es difícil encontrar moléculas que simultáneamente mejoren la propiedad y mantengan la similitud. Sin una guía adecuada, el modelo se vuelve excesivamente conservador (edits mínimos) o falla en generar trayectorias de razonamiento útiles.

2. Metodología: RePO (Reference-Guided Policy Optimization)

Los autores proponen RePO, un enfoque de optimización que combina la exploración guiada por recompensas con la explotación de referencias, sin necesidad de datos de trayectorias intermedias etiquetadas.

Objetivo de la Función de Pérdida (Ecuación 4):
RePO actualiza la política $\pi_\theta$ maximizando una función de objetivo compuesta por tres términos:

1. Término de Exploración (RLVR/GRPO):

   • Utiliza un mecanismo similar a GRPO donde se muestrean múltiples candidatos ($G$) para una instrucción dada.
   • Calcula ventajas relativas a nivel de grupo basadas en una recompensa escalar $r(o_i, q)$ que evalúa la satisfacción de la propiedad objetivo y la restricción de similitud.
   • Este término fomenta la exploración de nuevas estructuras químicas válidas.

2. Término de Guía de Referencia (Answer-Level Guidance):

   • Aumenta la probabilidad de la molécula de referencia ($m_{ref}$) condicionada al prefijo de razonamiento ($t_i$) generado por el modelo.
   • Clave técnica: La guía se aplica solo a los tokens de la respuesta final (la molécula), no a los tokens de razonamiento intermedios.
   • Esto actúa como un "ancla" que reduce la dispersión de la recompensa al inicio del entrenamiento, guiando al modelo hacia soluciones prometedoras sin forzarle a imitar el proceso de pensamiento específico de la referencia, preservando así la diversidad de rutas de razonamiento.

3. Regularización KL:

   • Mantiene la política actualizada cerca de una política de referencia ($\pi_{ref}$) para evitar el colapso de la diversidad y estabilizar el entrenamiento.

Diseño de Recompensas:

• Similitud Estructural ($r_{struct}$): Se utiliza la similitud de Tanimoto basada en huellas dactilares moleculares (ECFP4).
• Propiedad Objetivo ($r_{prop}$): Se define como una recompensa binaria (1 si mejora la propiedad según la instrucción, 0 en caso contrario), lo que resulta más estable que las recompensas continuas en este contexto.

3. Contribuciones Clave

1. Identificación del Desajuste: Demostraron empíricamente que el SFT "solo respuesta" colapsa el razonamiento de múltiples pasos y que el RLVR puro sufre de señales de aprendizaje escasas en el espacio químico restringido.
2. Propuesta de RePO: Introdujeron un marco que desacopla la guía de la respuesta (para estabilizar el aprendizaje) de la exploración de la trayectoria (para permitir la innovación), resolviendo el problema de la falta de datos de trayectorias.
3. Validación Empírica: Demostraron que RePO supera consistentemente a las líneas base (SFT, GRPO, GRPO con inicialización SFT) en tareas de un solo objetivo y múltiples objetivos, manteniendo un equilibrio superior entre éxito de optimización y similitud estructural.

4. Resultados Experimentales

Los experimentos se realizaron en los benchmarks TOMG-Bench (optimización de un solo objetivo) y MuMOInstruct (optimización multi-objetivo).

• Rendimiento General: RePO logró el mejor puntaje en la métrica compuesta Tasa de Éxito × Similitud (SR × Sim) en 4 de 6 tareas de un solo objetivo.
  • En tareas de optimización de propiedades (QED, LogP, MR), RePO superó a GRPO y SFT, mostrando mejoras de hasta un 17.4% en la tasa de éxito sobre GRPO.
  • En tareas multi-objetivo (MuMOInstruct), RePO demostró una mejor capacidad para equilibrar objetivos conflictivos (ej. aumentar permeabilidad BBB y reducir toxicidad) en comparación con SFT (que a menudo sacrifica la similitud) y GRPO (que a menudo es demasiado conservador).
• Generalización: RePO mantuvo su ventaja en instrucciones no vistas (unseen instruction styles), indicando una mejor capacidad de generalización.
• Análisis de Mecanismos:
  • Calidad del Razonamiento: Evaluaciones con "LLM-as-a-judge" mostraron que RePO genera explicaciones químicas más coherentes y válidas que GRPO (que a menudo produce razonamientos inválidos o químicamente incorrectos).
  • Escalado de Inferencia: RePO se beneficia significativamente de aumentar el presupuesto de cómputo en inferencia (muestreo best-of-k), mejorando tanto la tasa de éxito como la similitud.
  • Robustez: El método es robusto ante la corrupción de las referencias de demostración y funciona bien con diferentes arquitecturas de modelos (Qwen, Llama).

5. Significado e Impacto

El trabajo de RePO es significativo porque:

• Cierra la brecha entre SFT y RL: Proporciona una solución práctica para aplicar RL en dominios científicos donde los datos de trayectorias de razonamiento son inexistentes o costosos de generar.
• Mejora la confiabilidad en Química: Al combinar la exploración con una guía de respuesta, RePO reduce la generación de moléculas inválidas o químicamente ilógicas, un problema común en los LLM aplicados a la ciencia.
• Marco Generalizable: Aunque se centra en la optimización molecular, el enfoque de "guía de respuesta + exploración de trayectoria" es aplicable a otros problemas de razonamiento científico complejo con espacios de búsqueda vastos y restricciones estrictas.

En resumen, RePO representa un avance crucial en el uso de LLMs para el descubrimiento de fármacos, permitiendo que los modelos no solo "adivinen" la molécula correcta, sino que aprendan a razonar sobre las modificaciones químicas de manera eficiente y guiada.