Resumen Técnico: Guía de Recompensa por Muestreo de Mirada Adelante (LiDAR) para el Escalamiento en Tiempo de Prueba de Modelos de Difusión

Planteamiento del Problema

Los modelos de difusión han logrado un rendimiento generativo notable en imágenes, vídeos y lenguaje; sin embargo, las muestras generadas frecuentemente no logran alinearse plenamente con la intención humana. Aunque existen funciones de recompensa que capturan las preferencias humanas, integrar estas en el proceso de muestreo presenta desafíos significativos. El objetivo central del escalamiento en tiempo de prueba es muestrear de una distribución sesgada por la recompensa $p_\theta^r(x_0|c) \propto p_\theta(x_0|c) \exp(\lambda \cdot r(x_0, c))$.

Lograr esto requiere computar la Recompensa Futura Esperada (EFR), $r_t^\lambda(x_t, c)$, para partículas intermedias $x_t$ con el fin de guiar el proceso de difusión inversa. Los métodos existentes enfrentan limitaciones críticas:

1. Rollout hacia atrás (Backward Rollout): Muestrear directamente desde $x_t$ hacia $x_0$ para estimar la EFR es computacionalmente prohibitivo debido a la necesidad de un desruido iterativo en cada paso de tiempo.
2. Aproximaciones basadas en Tweedie: Los métodos como la Guía de Gradiente (Gradient Guidance) y el Método de Monte Carlo Secuencial (SMC) dependen de aproximaciones de Taylor de primer orden (fórmula de Tweedie) para estimar la EFR. Estos introducen sesgo, particularmente cuando la fuerza de la señal de recompensa ($\lambda$) es alta, y a menudo requieren retropropagación a través de redes neuronales (Gradient Guidance) o sufren de colapso de muestras en espacios de alta dimensión (SMC).

Metodología: Muestreo LiDAR

Los autores proponen LiDAR (Lookahead Sample Reward Guidance), un marco que desacopla la dependencia neuronal entre la partícula intermedia $x_t$ y la EFR, permitiendo una guía de forma cerrada y libre de derivadas.

1. Reformulación de la Recompensa Futura Esperada

El artículo deriva una nueva formulación para la EFR (Teorema 3.1) que expresa la expectativa de recompensa utilizando únicamente muestras marginales de un modelo de difusión preentrenado y el núcleo de perturbación hacia adelante, en lugar de condicionarse a la partícula actual $x_t$ mediante redes neuronales:
$$r_t^\lambda(x_t, c) = \log \mathbb{E}_{p_\theta(x_0|c)} \left[ \frac{p(x_t|x_0)}{\mathbb{E}_{p_\theta(x_0|c)}[p(x_t|x_0)]} \exp(\lambda \cdot r(x_0, c)) \right]$$
Este enfoque de "rollout hacia adelante" permite que la EFR se compute para cualquier $x_t$ utilizando un conjunto pregenerado de muestras marginales $\{\hat{x}_0^i\}$, eliminando la necesidad de retropropagación a través de la red de puntuación $s_\theta$ o el modelo de recompensa durante el paso de guía.

2. Muestreo de Mirada Adelante Eficiente (Efficient Lookahead Sampling)

Para mitigar el costo de generar muestras marginales para cada prompt, los autores introducen una estrategia de muestreo de mirada adelante eficiente. En lugar de utilizar el proceso de difusión completo y lento, emplean un generador rápido $q(x_0|c)$, como un resolvedor de ODE de pocos pasos (por ejemplo, DPM-Solver) o un modelo destilado, para generar un conjunto de $n$ muestras de mirada adelante $\{\hat{x}_0^i\}_{i=1}^n$. Estas muestras son anotadas con valores de recompensa $r(\hat{x}_0^i, c)$.

3. Guía de Forma Cerrada Libre de Derivadas

Utilizando el conjunto finito de muestras de mirada adelante, los autores derivan una expresión de forma cerrada para el gradiente de la recompensa empírica de mirada adelante (Teorema 3.3). El término de guía dirige la partícula actual $x_t$ hacia muestras de mirada adelante con alta recompensa y repele aquellas con baja recompensa:
$$\nabla_{x_t} \hat{r}_t^\lambda(x_t, c) = \sum_{i=1}^n (w_i^r - w_i) \frac{\hat{x}_0^i}{\sigma_t^2}$$
donde $w_i^r$ y $w_i$ son pesos softmax basados en los valores de recompensa y la distancia entre $x_t$ y las muestras de mirada adelante $\hat{x}_0^i$. Esta formulación no requiere diferenciación neuronal.

4. Leyes de Escalamiento Teóricas

El artículo establece leyes de escalamiento para el método:

• Precisión de la Mirada Adelante ($\delta$): La distancia de variación total entre la distribución objetivo y la distribución generada por LiDAR disminuye como $O(1/\sqrt{\delta})$, donde $\delta$ es el número de pasos de discretización en el resolvedor de mirada adelante.
• Tamaño de la Muestra ($n$): La varianza de la estimación de la recompensa empírica disminuye como $O(1/\sqrt{n})$, lo que implica que aumentar el número de muestras de mirada adelante mejora la precisión de la aproximación.

Resultados Clave

Los experimentos se realizaron en SD v1.5, SDXL, FLUX (flow matching) y UDLM (difusión discreta) utilizando ImageReward, CLIP, HPS y GenEval como métricas.

• Rendimiento vs. Eficiencia: En SDXL, LiDAR logra un rendimiento en GenEval comparable al de los métodos más recientes de guía de gradiente (DATE), pero con una aceleración de 9.5× en el tiempo de inferencia.
• Eficiencia de Memoria: A diferencia de los métodos de guía de gradiente que requieren aproximadamente 3× más memoria debido a la retropropagación, el uso de memoria de LiDAR es casi idéntico al del muestreo vainilla, ya que la etapa de muestreo objetivo (Algoritmo 2) no requiere el cálculo de gradientes.
• Robustez al "Reward Hacking": Mientras que los métodos de guía de gradiente a menudo degradan en métricas como CLIP y HPS al optimizar para ImageReward (reward hacking), LiDAR mejora todas las métricas simultáneamente. Esto se atribuye a la correlación entre ImageReward y otras métricas en las muestras de mirada adelante, lo que guía naturalmente al muestreador en una dirección equilibrada.
• Mejoras Ortogonales: LiDAR proporciona ganancias de rendimiento adicionales cuando se aplica a modelos que ya han sido ajustados con Optimización de Preferencia Directa (DPO), demostrando su ortogonalidad respecto al alineamiento basado en entrenamiento.
• Escalabilidad: El método escala eficazmente con el número de muestras de mirada adelante ($n$) y los pasos del resolvedor ($\delta$), y puede combinarse con otras técnicas como SMC y Best-of-N.

Significado y Reivindicaciones

El artículo afirma que LiDAR aborda sistemáticamente las limitaciones de eficiencia y de reward hacking de los métodos de escalamiento en tiempo de prueba existentes. Al reformular el cálculo de la EFR para que dependa del rollout hacia adelante y de muestras marginales, los autores permiten una guía libre de derivadas que es tanto computacionalmente eficiente como teóricamente fundamentada.

Los autores enfatizan que LiDAR logra un compromiso entre eficiencia y rendimiento superior al de las líneas base existentes. Específicamente, evita la sobrecarga computacional de la retropropagación neuronal y los problemas de colapso de muestras de SMC, proporcionando al mismo tiempo estimaciones de EFR más precisas que los métodos basados en aproximaciones de Taylor. El trabajo sugiere que el desacoplamiento de la señal de recompensa de la dependencia neuronal de la partícula intermedia es un camino viable hacia un modelado generativo escalable y alineado con los humanos.