Resumen Técnico: Más allá del Proxy: Guía Destilada por Trayectorias para el Entrenamiento de GFlowNet Offline

Enunciado del Problema

Las Redes de Flujo Generativo (GFlowNets) son modelos potentes para muestrear objetos composicionales proporcionales a sus recompensas. Sin embargo, en muchas aplicaciones prácticas, como el descubrimiento de moléculas, el diseño de secuencias biológicas y los sistemas de recomendación, la interacción activa con el entorno para consultar recompensas es inviable debido a los altos costos experimentales o a la necesidad de evaluación humana. En consecuencia, las GFlowNets deben entrenarse utilizando conjuntos de datos estáticos offline.

Los paradigmas actuales de entrenamiento offline suelen depender del entrenamiento de un modelo de recompensa proxy para aproximar las señales de recompensa a partir del conjunto de datos, al cual la GFlowNet consulta para la optimización de la política. Este enfoque enfrenta desafíos significativos:

1. Escasez de Datos y Costos: Construir un modelo proxy confiable a menudo requiere grandes conjuntos de datos diversos y una experiencia significativa en el dominio.
2. Propagación de Errores: Las consultas de recompensa para muestras o estados fuera de distribución (OOD) estimados incorrectamente por el proxy pueden propagar errores a través de los gradientes, comprometiendo la política aprendida.
3. Limitaciones de los Métodos sin Proxy: Los enfoques existentes sin proxy (por ejemplo, RO-GFlowNets, COFlowNet) a menudo imponen restricciones gruesas para alinear la política con el conjunto de datos. Estas restricciones pueden limitar la generalización, inhibir una exploración efectiva y aumentar la sensibilidad a la calidad de los datos.

La pregunta central abordada es: ¿Cómo podemos aprovechar eficazmente las trayectorias offline para entrenar GFlowNets cuando las recompensas terminales son insuficientes para construir un modelo proxy confiable?

Metodología: GFlowNet Destilado por Trayectorias (TD-GFN)

Los autores proponen TD-GFN, un nuevo marco sin proxy que extrae información estructural rica del Grafo Acíclico Dirigido (DAG) del entorno para guiar el entrenamiento de la política sin depender de un modelo de recompensa proxy. La metodología consta de tres fases principales:

1. Extracción de Recompensas de Arista mediante Aprendizaje por Refuerzo Inverso (IRL)

En lugar de predecir recompensas terminales, TD-GFN utiliza IRL para destilar recompensas de arista a nivel de transición ($R_E$) a partir de trayectorias offline.

• Reequilibrio: Para abordar el problema de que los datos offline pueden no representar el comportamiento experto, los autores introducen una estrategia de reequilibrio. Las trayectorias se muestrean con una probabilidad proporcional a sus recompensas terminales ($P(\tau) \propto R(s_T)$), creando un conjunto de datos pseudo-experto.
• Entrenamiento Adversarial: Utilizando un marco modificado de GAIL (Aprendizaje de Imitación Adversarial Generativo), un discriminador $D_\phi$ distingue entre transiciones en el conjunto de datos reequilibrado y aquellas generadas por una política $\pi_\psi$.
• Derivación de Recompensa: Se extrae una función de recompensa de nivel de arista sin sesgo del discriminador:
  $$R_E(s, s') = \log D_\phi(s, s') - \log(1 - D_\phi(s, s'))$$
  Teóricamente, esta recompensa aproxima la puntuación de flujo entrante condicional en logaritmo ($\log P_B(s|s')$), cuantificando la importancia de transiciones específicas para aprender una política efectiva. Crucialmente, estas recompensas proporcionan una guía intermedia densa en lugar de predicciones terminales.

2. Poda del DAG Guiada por Recompensas

Para garantizar la robustez y prevenir la propagación de errores derivada de estimaciones de IRL potencialmente inexactas, TD-GFN no utiliza las recompensas de arista para dar forma directamente a los gradientes de la política. En su lugar, las utiliza para podar el DAG del entorno.

• Umbralización: Las aristas con recompensas por debajo de un umbral estadístico (media menos $K$ veces la desviación estándar de la distribución de recompensas) se eliminan.
• Garantía Teórica: Los autores demuestran que esta estrategia de poda preserva todos los modos de "cuello de botella alto" (estados terminales alcanzables mediante caminos con puntuaciones de arista altas) mientras elimina transiciones de baja utilidad. Esto crea un espacio de acciones compacto y expresivo que prioriza trayectorias de alto valor, al tiempo que aísla la política de errores de aproximación del IRL.

3. Muestreo Backward Priorizado

El entrenamiento se realiza en el subgrafo podado utilizando un mecanismo de muestreo backward priorizado.

• Estrategia de Muestreo: Los estados terminales se muestrean a partir de la intersección del conjunto de datos y el grafo podado con una probabilidad proporcional a sus recompensas de verdad terrenal (ground-truth).
• Construcción de Trayectoria Backward: Desde cada terminal muestreado, se construyen trayectorias en reversa hacia la raíz utilizando una política backward aprendida $P_B$, definida por las recompensas de arista aprendidas:
  $$P_B(s_t | s_{t+1}) = \frac{\exp\{R_E(s_t, s_{t+1})\}}{\sum \exp\{R_E(s, s_{t+1})\}}$$
• Actualizaciones de Gradiente: La política de la GFlowNet se optimiza utilizando estas trayectorias muestreadas y sus recompensas terminales de verdad terrenal. Al depender exclusivamente de las recompensas terminales registradas para las actualizaciones de gradiente, el método evita los problemas de propagación de errores inherentes a los métodos basados en proxy.

Contribuciones Clave

1. Un Nuevo Paradigma sin Proxy: TD-GFN elimina la dependencia de modelos proxy y consultas de recompensa OOD. Emplea IRL para destilar recompensas de arista a nivel de transición a partir de conjuntos de datos reequilibrados, proporcionando una guía estructurada para el entrenamiento y la generalización de la política.
2. Mecanismo de Guía Robusto: El marco introduce dos mecanismos de guía indirecta: poda del DAG y muestreo backward priorizado. Este diseño asegura que las actualizaciones de gradiente dependan exclusivamente de recompensas de verdad terrenal, previniendo la propagación de errores mientras se aprovecha una guía estructural densa para una exploración eficiente.
3. Rendimiento de Vanguardia: Experimentos extensivos demuestran que TD-GFN logra una velocidad de convergencia superior y una calidad de muestra en comparación con una amplia gama de líneas base, incluidos métodos basados en proxy (Proxy-GFN), algoritmos de RL offline (CQL, IQL) y otras variantes de GFlowNet sin proxy (COFlowNet, Dataset-GFN).

Resultados Experimentales

Los autores evaluaron TD-GFN en tres puntos de referencia:

• Hypergrid (Sintético): En una cuadrícula de 84 dimensiones con recompensas dispersas y multimodales, TD-GFN identificó todos los 16 modos de recompensa con menos de 5,000 visitas a estados (una mejora de seis veces sobre las líneas base) y convergió al Error L1 Empírico más bajo dos veces más rápido que el competidor mejor (COFlowNet). Demostró una fuerte robustez bajo datos ruidosos (conjunto de datos Mixed) y escasez de datos (10% de los datos).
• Diseño de Secuencias Biológicas (Péptidos Antimicrobianos): TD-GFN generó secuencias con las recompensas más altas y la mayor diversidad entre las 100 primeras secuencias, superando tanto a los métodos basados en proxy como a otros métodos offline.
• Diseño de Moléculas: TD-GFN descubrió más moléculas de alta recompensa y modos diversos que las líneas base, mientras requería significativamente menos trayectorias para converger (2.75 × 10^4 vs. 43.7 × 10^4 para Proxy-GFN). Igualó el rendimiento de un Oracle-GFN (que utiliza recompensas de verdad terrenal directamente) en la identificación de moléculas de alta recompensa.

Significado y Afirmaciones

El artículo afirma que TD-GFN establece un paradigma más robusto, estable y eficiente para el entrenamiento de GFlowNet offline. Al aprovechar la información estructural inherente en las trayectorias offline a través de IRL, el método cierra la brecha entre los conjuntos de datos offline limitados y los requisitos complejos de exploración del modelado generativo.

Los autores enfatizan que TD-GFN:

• Desbloquea el Potencial: Permite el despliegue de GFlowNets en dominios complejos donde las consultas activas de recompensa son prohibitivas.
• Previene la Propagación de Errores: Al desacoplar el mecanismo de guía (recompensas de arista) del mecanismo de actualización de gradiente (recompensas terminales), evita los peligros de los modelos proxy.
• Mejora la Exploración: La guía estructural densa permite que el modelo explore eficazmente regiones de alta recompensa, incluso cuando están escasamente representadas o no observadas en los datos de entrenamiento.

El trabajo sugiere que la información a nivel de trayectoria, a menudo pasada por alto en los métodos basados en proxy, es una piedra angular para el entrenamiento efectivo de GFlowNet offline, ofreciendo una solución escalable para aplicaciones del mundo real en el descubrimiento de fármacos, el diseño biológico y la optimización combinatoria.