Resumen Técnico: JAGG (Jacobian-Aggregated Group Gradient)

Planteamiento del Problema

La Optimización de Política de Grupo Relativo (GRPO, por sus siglas en inglés) ha emergido como un algoritmo de aprendizaje por refuerzo (RL) líder para alinear modelos generativos con las preferencias humanas, eliminando con éxito la necesidad de una red de valor separada en los Modelos de Lenguaje Extensos (LLMs). Sin embargo, extender GRPO a modelos de difusión y de ajuste de flujo (específicamente, Diffusion Transformers o DiTs) introduce un cuello de botella computacional severo.

En el GRPO basado en difusión, el proceso de generación se trata como un problema de decisión de múltiples pasos donde la política se optimiza sobre una trayectoria de muestreo de $W$ pasos temporales. El entrenamiento estándar requiere un paso completo de ida (forward) y vuelta (backward) a través de la arquitectura de alta capacidad del DiT para cada uno de los pasos en la trayectoria para calcular el gradiente de la política. Para la generación de texto a imagen (T2I) de alta resolución, donde las trayectorias suelen abarcar docenas de pasos, esto resulta en costos de entrenamiento prohibitivamente caros. A diferencia de los modelos de lenguaje donde los pasos de tokens son ligeros, cada paso backward a través de un DiT moderno incurre en el costo total de un paso forward de un transformer, lo que hace que el requisito de backward por paso sea el gasto dominante en el entrenamiento.

Aunque los métodos de aceleración de inferencia libres de entrenamiento (por ejemplo, $\Delta$-DiT, ScalingCache) han explotado con éxito la observación de que los estados ocultos y las predicciones de velocidad de un DiT varían de forma suave y casi lineal a lo largo de la trayectoria de eliminación de ruido para omitir evaluaciones intermedias, seguía siendo una pregunta abierta si esta misma linealidad podía aprovecharse para reducir el costo del paso backward en el entrenamiento de RL.

Metodología: JAGG

Los autores proponen JAGG (Jacobian-Aggregated Group Gradient), un método de aceleración con principios matemáticos que reduce el número de pasos backward de la red de $W$ a 2 para un grupo de $W$ pasos de eliminación de ruido consecutivos.

Idea Central

El método se basa en la idea de que el Jacobiano del campo de velocidad de eliminación de ruido $v_\theta(z, t)$ con respecto a los parámetros del modelo $\theta$ exhibe una linealidad aproximada en los inputs $(z, t)$, particularmente durante el régimen de alto ruido del proceso de eliminación de ruido. Esto permite que los Jacobianos de los pasos intermedios se aproximen mediante interpolación lineal utilizando únicamente los Jacobianos de los extremos ($J_0$ y $J_{W-1}$).

Flujo de Trabajo Algorítmico

JAGG procesa un grupo de $W$ pasos temporales a través de cuatro fases:

1. Pasos Forward: El modelo realiza pasos forward para los $W$ pasos. Sin embargo, el seguimiento de gradientes se habilita solo para los dos extremos (paso 0 y paso $W-1$). Los $W-2$ pasos intermedios se ejecutan usando torch.no_grad() para ahorrar memoria y computación.
2. Cómputo del Gradiente Upstream: Para cada paso $j$, se crea un proxy desvinculado (detached) de la salida del modelo. Se calcula la pérdida por paso $L_j$ (derivada del objetivo tipo PPO) y el gradiente upstream $s_j = \partial L_j / \partial \hat{v}_j$ mediante diferenciación automática. Crucialmente, este paso no involucra un paso backward del transformer; el gradiente fluye solo a través del paso SDE ligero y el objetivo de PPO.
3. Agregación Ponderada por Paso Temporal: Los gradientes upstream $s_j$ se agregan en dos señales compuestas, $s_{base}$ y $s_{corr}$, utilizando pesos de interpolación $\alpha_j$ derivados del paso temporal $t_j$:
   $$\alpha_j = \frac{t_0 - t_j}{t_0 - t_{W-1}}$$
   $$s_{base} = \sum_{j=0}^{W-1} (1 - \alpha_j) s_j, \quad s_{corr} = \sum_{j=0}^{W-1} \alpha_j s_j$$
   Esta agregación reconstruye efectivamente la señal de gradiente completa utilizando la suposición de interpolación lineal.
4. Backward Conjunto: Se emite una única llamada a torch.autograd.backward sobre las dos salidas de los extremos $[\hat{v}_0, \hat{v}_{W-1}]$ con los gradientes agregados $[s_{base}, s_{corr}]$. Esto activa exactamente un FSDP all-gather por módulo, reemplazando los $W$ pasos backward separados requeridos por el entrenamiento estándar.

Garantía Teórica

Los autores proporcionan una prueba formal (Teorema 1) que establece que, si el campo de velocidad $v_\theta(z, t)$ es estrictamente lineal en $(z, t)$, la interpolación del Jacobiano ponderado por $t$ es exacta con cero error de aproximación. Bajo estas condiciones, el paso backward de los dos extremos recupera el gradiente de grupo exacto.

Enrutamiento Adaptativo ($jagg\_frac$)

Reconociendo que la suposición de linealidad se degrada en las etapas finales (de bajo ruido) de la eliminación de ruido, los autores introducen un parámetro de enrutamiento $jagg\_frac$. Este parámetro aplica JAGG solo a la primera fracción de los grupos de eliminación de ruido (donde la suposición de linealidad de alto ruido se cumple, verificado empíricamente mediante la similitud de coseno de los gradientes) y utiliza pasos backward exactos por paso para los grupos restantes. Para los modelos T2I probados, $jagg\_frac$ se estableció en 0.6.

Contribuciones Clave

1. Identificación del Cuello de Botella: El artículo identifica y caracteriza formalmente el cuello de botella computacional en el entrenamiento de difusión GRPO causado por los pasos backward por cada paso temporal a lo largo de la trayectoria de muestreo.
2. Marco de Trabajo JAGG: Propone un marco de agregación de Jacobianos que reduce los pasos backward de $W$ a 2 por grupo, fundamentado en un análisis teórico riguroso de la linealidad del Jacobiano de la velocidad.
3. Prueba Teórica: Los autores demuestran que, bajo condiciones de velocidad lineal, la interpolación propuesta es exacta, proporcionando una base teórica sólida para la aproximación.
4. Validación Empírica: Experimentos extensos en múltiples backbones de difusión (Flux, QwenImage) y algoritmos de RL (DanceGRPO, Flow-GRPO-Fast) demuestran la generalidad y efectividad del método.

Resultados Experimentales

Los autores evaluaron JAGG en benchmarks de T2I utilizando HPSv2.1, PickScore y CLIPScore como métricas.

• Eficiencia: JAGG logró una aceleración constante de tiempo de ejecución (wall-clock speedup) del 17.3% al 18.2% por paso de entrenamiento en configuraciones estándar de DanceGRPO (reduciendo el tiempo de ~835s a ~691s en QwenImage y de ~279s a ~228s en Flux). Al aplicarse a Flow-GRPO-Fast, la etapa de actualización sola vio aceleraciones del 29.9% al 30.0%.
• Calidad: El método entregó estas aceleraciones con una degradación insignificante en la calidad de la generación. En muchos casos, JAGG siguió de cerca las curvas de recompensa de la línea base y igualó o superó ligeramente a la línea base al alcanzar la convergencia.
• Ablación: Una variante de ablación que realizaba pasos backward en los extremos pero sin la interpolación del Jacobiano (asignando la pérdida completa a los extremos de forma independiente) se desempeñó consistentemente por debajo de la propuesta. Esto confirma que la agregación de la información del gradiente intermedio mediante la interpolación del Jacobiano es crítica para un entrenamiento estable.
• Calidad del Gradiente: El análisis empírico mostró que en los grupos de alto ruido (0–2), la similitud de coseno entre el gradiente agregado por JAGG y el gradiente exacto fue alta (0.74–0.85), validando la suposición de linealidad. La similitud disminuyó en los grupos posteriores, justificando la estrategia de enrutamiento $jagg\_frac$.

Significado y Reivindicaciones

El artículo afirma que JAGG representa un paso significativo hacia la viabilidad del entrenamiento de difusión GRPO a gran escala. Al explotar la casi-linealidad de los Jacobianos de velocidad en el régimen de alto ruido, el método reduce la carga computacional del entrenamiento de RL sin sacrificar la calidad de la generación.

Los autores enfatizan que JAGG no es simplemente un atajo heurístico, sino que está fundamentado en las propiedades matemáticas de la trayectoria de eliminación de ruido. El método se presenta como ortogonal a las técnicas existentes de aceleración de inferencia, apuntando específicamente al cuello de botella del paso backward único de la difusión. Los resultados sugieren que JAGG permite una adopción más amplia del alineamiento mediante RL para modelos generativos visuales al reducir significativamente la barrera del costo de entrenamiento.