🤖 machine learning

$f$ -Trajectory Balance: A Loss Family for Tuning GFlowNets, Generative Models, and LLMs with Off- and On-Policy Data

Este artículo introduce el Equilibrio de Trayectoria $f$ , una familia de funciones de pérdida que extiende el enfoque del error cuadrático medio a todas las $f$ -divergencias, permitiendo el entrenamiento de modelos generativos y LLMs con datos off-policy mientras se preservan las propiedades de optimización específicas (como la cobertura de modos) de sus correspondientes gradientes de $f$ -divergencia on-policy.

Autores originales: Jake Fawkes, Jason Hartford

Publicado 2026-05-18

📖 5 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Jake Fawkes, Jason Hartford

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que estás enseñando a un robot a explorar un vasto sistema de cuevas oscuras. La cueva tiene muchas cámaras ocultas (llamadas "modos"), algunas de las cuales contienen oro (recompensas altas), y otras están simplemente vacías o son peligrosas. Tu objetivo es enseñar al robot a encontrar tantas cámaras de oro diferentes como sea posible, en lugar de solo encontrar una gran mina de oro e ignorar todo lo demás.

Este artículo presenta un nuevo conjunto de "herramientas de enseñanza" (funciones de pérdida) para ayudar a entrenar a estos robots, ya sea que estén generando moléculas para nuevos fármacos, creando arte o escribiendo código.

Aquí tienes el desglose de las ideas del artículo utilizando analogías simples:

1. El Problema: La "Fiebre del Oro" vs. El "Explorador"

En el pasado, la forma estándar de enseñar a estos robots era como una Fiebre del Oro. El robot encontraría un lugar con un poco de oro, se emocionaría y luego pasaría todo su tiempo cavando allí. Ignoraría todas las demás minas de oro en la cueva.

En términos técnicos: Esto se llama "búsqueda de modos". El modelo colapsa sobre unas pocas respuestas de alta probabilidad e ignora el resto de la diversidad.
El objetivo del artículo: Queremos "cobertura de modos". Queremos que el robot se expanda y explore todas las minas de oro, incluso las más pequeñas, para obtener un mapa completo de la cueva.

2. La Vieja Herramienta: La Pérdida "KL al Cuadrado"

Recientemente, los investigadores encontraron un truco ingenioso para enseñar a los robots utilizando un método de "error cuadrático" (midiendo la distancia entre lo que el robot piensa y lo que debería pensar).

La Analogía: Imagina a un profesor calificando a un estudiante. Si el estudiante se equivoca en la respuesta, el profesor no solo dice "Incorrecto". Calculan el cuadrado del error.
El Beneficio: Este método es muy estable. Funciona incluso si el estudiante está aprendiendo de notas antiguas (datos fuera de política) en lugar de lecciones en vivo.
El Defecto: Aunque es estable, este método específico "cuadrado" aún actúa como una Fiebre del Oro. Empuja naturalmente al robot a centrarse en solo unas pocas respuestas, perdiendo la diversidad de la cueva.

3. La Nueva Solución: La Familia "f-Balance de Trayectoria"

Los autores se dieron cuenta de que el método "cuadrado" era solo un sabor específico de una familia mucho más grande de herramientas de enseñanza. Ellos llaman a esta familia f-Balance de Trayectoria.

Piensa en esta familia como un diafragma o un botón de volumen en una radio.

Girando el dial en una dirección (Números bajos): Obtienes herramientas que actúan como un Super Explorador. Estas herramientas fuerzan al robot a expandirse y encontrar cada mina de oro posible, incluso las de difícil acceso. Esto es excelente para el descubrimiento de fármacos, donde quieres encontrar cualquier molécula que pueda funcionar, no solo la más obvia.
Girando el dial en la otra dirección (Números altos): Obtienes herramientas que actúan como un Buscador de Oro. Estas herramientas le dicen al robot que ignore las minas pequeñas y se concentre intensamente en el montón de oro más grande y obvio. Esto es útil si solo quieres la mejor respuesta individual.
El Terreno Intermedio: Puedes ajustar el dial en cualquier punto intermedio para obtener una mezcla de exploración y enfoque.

4. Por Qué Esto Es Importante

El artículo demuestra dos cosas principales:

El "Interruptor Mágico": Mostraron que puedes cambiar la herramienta estándar "cuadrada" por cualquiera de estas nuevas herramientas de "dial", y las matemáticas siguen funcionando perfectamente. El robot aprende igual de bien, pero con una personalidad diferente (más exploradora o más enfocada).
Estabilidad: Al igual que la vieja herramienta, estas nuevas herramientas funcionan incluso si el robot está aprendiendo de datos antiguos (fuera de política). Esto es crucial para aplicaciones del mundo real como los Modelos de Lenguaje Grandes (LLM), donde el entrenamiento a menudo ocurre de forma asíncrona (el robot aprende de datos generados hace un tiempo).

5. Pruebas del Mundo Real (Los Mapas de la Cueva)

Los autores probaron estas herramientas en tres "cuevas" diferentes:

El Juego de la Rejilla: Un laberinto informático simple con cuatro esquinas llenas de oro. La herramienta estándar solo encontró una esquina. Las nuevas herramientas de "Explorador" encontraron las cuatro esquinas rápidamente.
Descubrimiento de Moléculas (SynFlowNet): Intentaron generar nuevas moléculas químicas. Al girar el dial hacia "Explorador", generaron una variedad mucho más amplia de moléculas únicas que podrían ser potencialmente fármacos, en lugar de solo variaciones de los mismos pocos químicos.
Modelos de Lenguaje (LLM): Ajustaron modelos de IA para resolver problemas matemáticos. Las nuevas herramientas permitieron a la IA explorar diferentes formas de resolver problemas sin quedarse atrapada en un bucle repitiendo la misma respuesta, incluso cuando los datos de entrenamiento estaban retrasados (asíncronos).

Resumen

Este artículo no inventa un nuevo tipo de robot. En cambio, inventa un nuevo conjunto de instrucciones sobre cómo entrenarlos.

Antigua forma: "Encuentra la mina de oro más grande y cava allí." (Estable, pero aburrida).
Nueva forma: "Aquí tienes un dial. Puedes elegir ser un Buscador de Oro, un Super Explorador o cualquier cosa intermedia. Y no importa cuál elijas, el entrenamiento será estable y funcionará con datos antiguos."

Esto le da a los ingenieros un botón simple para controlar qué tan creativos o enfocados deben ser sus modelos de IA, sin romper el proceso de entrenamiento.

Resumen Técnico: Equilibrio de Trayectoria f

Enunciado del Problema

El alineamiento y el ajuste fino de modelos generativos, incluidos los Modelos de Lenguaje Grandes (LLM) y las Redes de Flujo Generativo (GFlowNets), suelen depender de minimizar la discrepancia entre la distribución de un modelo y una distribución objetivo. En el ajuste fino mediante Aprendizaje por Refuerzo (RL), esto se formula a menudo como la maximización de la recompensa esperada sujeta a una penalización de divergencia Kullback-Leibler (KL).

Un desafío significativo en este dominio es la compensación entre los comportamientos de búsqueda de modos y de cobertura de modos. Los objetivos estándar, como la divergencia KL inversa (a menudo implementada mediante Equilibrio de Trayectoria o pérdida KL al cuadrado), tienden a ser de búsqueda de modos, lo que provoca que los modelos colapsen sobre unos pocos modos de alta probabilidad. Esto es subóptimo para tareas que requieren exploración, como el descubrimiento de fármacos o la generación de candidatos diversos, donde se prefiere la "cobertura de modos" (capturar la diversidad completa de la distribución objetivo).

Además, los métodos existentes a menudo luchan con la validez fuera de política. Aunque estimadores como REINFORCE o PPO son comunes, sufren de alta varianza. Los enfoques recientes que utilizan la pérdida KL al cuadrado ( $KL_{sq}$ ) ofrecen gradientes de baja varianza en política y validez fuera de política, pero están intrínsecamente ligados a las propiedades de la divergencia KL inversa, lo que limita la capacidad de controlar la compensación entre exploración y explotación.

Metodología

Los autores proponen el Equilibrio de Trayectoria f, una familia de funciones de pérdida sustitutas que generaliza la pérdida KL al cuadrado a toda la familia de divergencias f.

Fundamento Teórico

La contribución teórica central es establecer una correspondencia uno a uno entre las funciones de pérdida invariantes a la traslación en log-probabilidades y las divergencias f.

De Divergencia a Pérdida: Para cualquier función convexa $f$ que define una divergencia f $D_f(p_\theta \| p^*)$ , los autores derivan una función de pérdida $L_f$ definida sobre la diferencia de log-probabilidades $\Delta_\theta(y) = \log p_\theta(y) - \log p^*(y)$ :
$L_f(\Delta_\theta(y)) = \int_0^{\Delta_\theta(y)} (f'(e^t) - f'(1)) dt$
Esta pérdida tiene dos propiedades críticas:
- Equivalencia de Gradiente en Política: Cuando las muestras se extraen de la política del modelo ( $\mu = p_\theta$ ), los gradientes auto-diferenciados esperados de $L_f$ coinciden con los gradientes de la divergencia f correspondiente $D_f$ .
- Validez Fuera de Política: Cuando se evalúa en datos fuera de política (muestreados desde una distribución de comportamiento $\mu$ ), la pérdida sigue siendo un objetivo válido con el mismo minimizador global ( $p_\theta = p^*$ ), siempre que el soporte de la objetivo esté contenido en el soporte de la política de comportamiento.
De Pérdida a Divergencia: Por el contrario, cualquier pérdida $\ell$ estrictamente convexa e invariante a la traslación en log-probabilidades corresponde a minimizar una divergencia f específica. El artículo demuestra que estas mapeos son inversos entre sí.

Implementaciones Prácticas

Para abordar los desafíos prácticos en el entrenamiento, los autores introducen varias extensiones:

Pérdida DevGrad (Normalización por Lote): Para manejar funciones de partición intratables (comunes en GFlowNets y objetivos no normalizados), los autores generalizan el estimador VarGrad. En lugar de usar la varianza del lote, utilizan la desviación generalizada por lote de las diferencias de log-probabilidades. Esto implica estimar la constante de normalización $\log Z$ minimizando la pérdida sobre el lote, centrando efectivamente los coeficientes de la función de puntuación y reduciendo la varianza.
Pérdida Templada: Para prevenir el desbordamiento numérico al tratar con parámetros de temperatura pequeños $\beta$ (donde la distribución objetivo es $p^* \propto \pi_{ref} \exp(r/\beta)$ ), los autores proponen una pérdida templada. Esto escala la pérdida por $1/\beta$ y la aplica a la distribución templada, asegurando que los gradientes permanezcan acotados y estables.
Equilibrio de Trayectoria f para GFlowNets: El marco se aplica a GFlowNets imponiendo restricciones de equilibrio de trayectoria utilizando la pérdida $L_f$ . Esto generaliza el Equilibrio de Trayectoria estándar (que corresponde a la KL inversa) a divergencias f arbitrarias.

Contribuciones Clave

Familia de Pérdidas Generalizada: Derivación de una familia general de pérdidas sustitutas invariantes a la traslación ( $L_f$ ) cuyos gradientes en política coinciden con los de la divergencia f correspondiente, manteniendo al mismo tiempo la validez fuera de política.
Equivalencia Teórica: Demostración de la relación inversa entre las pérdidas de log-probabilidad invariantes a la traslación y las divergencias f, permitiendo el diseño sistemático de nuevas funciones de pérdida con propiedades deseadas (por ejemplo, cobertura de modos).
Extensión a Objetivos No Normalizados: Generalización de la técnica VarGrad a divergencias f arbitrarias mediante la pérdida DevGrad, abordando el problema de la función de partición intratable en el ajuste fino de RL.
Control de la Exploración: Demostración de que la elección de la divergencia f (parametrizada por $\alpha$ en las $\alpha$ -divergencias) controla directamente el comportamiento de búsqueda de modos frente a cobertura de modos. Valores más bajos de $\alpha$ (por ejemplo, KL Directa, Hellinger) promueven la cobertura de modos, mientras que valores más altos (por ejemplo, $\chi^2$ de Pearson) promueven la búsqueda de modos.

Resultados Experimentales

Los autores evalúan la familia de pérdidas en cuatro tareas distintas:

Cuadrícula Sintética (GFlowNets): En una tarea de cuadrícula 2D diseñada para probar la cobertura de modos, las pérdidas con $\alpha$ más bajo (KL Directa, Hellinger) descubrieron con éxito los cuatro modos de la distribución objetivo significativamente más rápido y de manera más fiable que el Equilibrio de Trayectoria estándar (KL Inversa). Por el contrario, el $\chi^2$ de Pearson ( $\alpha=2$ ) exhibió un fuerte comportamiento de búsqueda de modos, quedando atrapado en el primer modo encontrado.
Descubrimiento de Moléculas (SynFlowNet): Aplicado al marco SynFlowNet para generar moléculas sintéticamente accesibles. Los autores encontraron que ajustar $\alpha$ permitía controlar la diversidad de las moléculas generadas. Un programa de recocido (comenzando con $\alpha$ más bajo para exploración y aumentando a $\alpha=1$ para explotación) produjo el mejor compromiso, generando moléculas de mayor recompensa con mayor diversidad en comparación con el Equilibrio de Trayectoria estándar.
Modelos de Difusión: En el muestreo condicional de dígitos MNIST (impares vs. pares), el Equilibrio de Trayectoria estándar falló en muestrear uniformemente a través de los modos objetivo (sobremuestreando dígitos específicos como 0 y 6). Las pérdidas del Equilibrio de Trayectoria f, particularmente aquellas con $\alpha$ más bajo, lograron un muestreo más uniforme a través de los modos objetivo.
Ajuste Fino Asíncrono de LLM: Las pérdidas se aplicaron para ajustar fino LLMs (familias Qwen y OLMo) en problemas de matemáticas (GSM8k, MATH) con un retraso asíncrono de 50 pasos. Las pérdidas de divergencia f demostraron un entrenamiento estable y compensaciones consistentes entre entropía y recompensa, mientras que las implementaciones estándar de PPO mostraron inestabilidad bajo las mismas condiciones asíncronas.

Significado y Afirmaciones

El artículo afirma que la efectividad de la pérdida KL al cuadrado no es un fenómeno aislado, sino un caso específico de una familia más amplia de pérdidas. Al extender la propiedad de "coincidencia de gradientes en política" y "validez fuera de política" a toda la familia de divergencias f, los autores proporcionan un marco unificado para el ajuste fino de modelos generativos.

El significado principal radica en la capacidad de controlar explícitamente la compensación entre exploración y explotación simplemente seleccionando una función de pérdida sustituta diferente (o ajustando un parámetro $\alpha$ ), sin cambiar el algoritmo de entrenamiento subyacente. Esto permite a los profesionales elegir entre comportamientos de búsqueda de modos (para maximización de recompensa) y cobertura de modos (para diversidad y exploración), manteniendo al mismo tiempo la estabilidad y la compatibilidad fuera de política requeridas para el entrenamiento asíncrono a gran escala moderno de LLMs y GFlowNets. El trabajo sugiere que este marco es particularmente valioso para aplicaciones como el descubrimiento de fármacos y agentes exploratorios donde cubrir el espacio de candidatos de alta recompensa es crítico.