Gradient-Free Noise Optimization for Reward Alignment in Generative Models

Resumen Técnico: ZeNO (Optimización de Ruido de Orden Cero)

Enunciado del Problema
Los métodos existentes de alineación de recompensas para modelos generativos (por ejemplo, modelos de difusión y de flujo) suelen depender de la optimización de trayectorias estocásticas de múltiples pasos. Aunque efectivos en estos contextos, tales enfoques son difíciles de extender a generadores deterministas o de un solo paso (por ejemplo, modelos de consistencia, modelos de difusión destilados), los cuales no exponen naturalmente la misma estructura de trayectoria. Una alternativa natural es optimizar directamente el ruido de entrada. Sin embargo, la optimización de ruido basada en gradientes estándar requiere retropropagación a través tanto del generador como de la función de recompensa. Esto limita la aplicabilidad a pipelines diferenciables y puede ser poco fiable cuando los gradientes reflejan sensibilidades locales espurias en lugar de mejoras semánticas significativas. Además, muchas funciones de recompensa prácticas (por ejemplo, diseñabilidad de proteínas, alineación composicional texto-imagen) son no diferenciables o cajas negras, lo que hace inviables los métodos basados en gradientes.

Metodología: ZeNO
Los autores proponen ZeNO (Optimización de Ruido de Orden Cero), un marco libre de gradientes que formula la optimización de ruido como un problema de control de integral de camino (PIC). La idea central es tratar la evolución del ruido de entrada como un proceso estocástico controlable, optimizándolo utilizando únicamente evaluaciones de recompensa de orden cero.

1. Formulación de Control de Integral de Camino:
   El método modela la optimización de ruido como una Ecuación Diferencial Estocástica (EDE) controlada. El objetivo es dirigir la trayectoria del ruido hacia regiones que produzcan estados terminales de alta recompensa, mientras se regulariza el esfuerzo de control. El control óptimo $u^*(z, t)$ se deriva en una forma de integral de camino:
   $$u^*(z, t) = \frac{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda) d w_t]}{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda)]}$$
   Crucialmente, este estimador depende únicamente de evaluaciones de costo terminal (recompensas) y no requiere diferenciación a través del generador $G_\theta$ ni de la función de recompensa $r$.

2. Dinámicas de Referencia (Ornstein–Uhlenbeck):
   Para asegurar que el ruido actualizado permanezca compatible con los generadores preentrenados, ZeNO adopta el proceso de Ornstein–Uhlenbeck (OU) como dinámicas de referencia. A diferencia de las EDEs con explosión de varianza, el proceso OU preserva la distribución estacionaria gaussiana $N(0, I)$, que coincide con el prior de la mayoría de los generadores preentrenados. Esta elección conecta la actualización del ruido con la dinámica de Langevin que apunta implícitamente a una distribución inclinada por la recompensa $p^*(z) \propto p(z) \exp(r(G_\theta(z))/\lambda)$.

3. Estimador Linealizado y Reducción de Varianza:
   El estimador estándar de integral de camino sufre de alta varianza debido al ponderado exponencial. Los autores introducen una aproximación linealizada (Proposición 1) restando la recompensa media $\bar{r}$ como línea base y linealizando el término exponencial bajo el supuesto de una varianza de recompensa pequeña:
   $$u^*(z_t, t) \approx \mathbb{E}[(r(G_\theta(z_T)) - \bar{r}) d w_t] / \lambda$$
   Este estimador coincide con el estimador de la función de puntuación utilizado en Estrategias de Evolución Natural y supera empíricamente a las variantes exponenciales.

4. Horizonte Recedente de Un Paso:
   Para evitar el costo computacional de simulaciones de largo horizonte, ZeNO emplea una aproximación de horizonte recedente de un paso ($H=1$). En cada iteración, se muestrean $N$ candidatos de ruido perturbados desde las dinámicas de referencia, se evalúan mediante el modelo de recompensa y se utiliza la estimación de control linealizada para actualizar el estado actual del ruido. Esto permite la escalabilidad en tiempo de inferencia aumentando el número de partículas ($N$) o iteraciones ($M$) sin modificar el generador.

Contribuciones Clave

• Alineación Libre de Gradientes: ZeNO habilita la alineación de recompensas para generadores deterministas y de un solo paso sin requerir retropropagación a través del generador o del modelo de recompensa.
• Compatibilidad con Cajas Negras: El método es aplicable a funciones de recompensa no diferenciables y de caja negra, un escenario donde los métodos existentes de optimización de ruido basados en gradientes (por ejemplo, ORIGEN, ReNO) fallan.
• Conexión Teórica: El marco establece un vínculo teórico entre la optimización de ruido de orden cero y el muestreo de Langevin dirigido a una distribución inclinada por la recompensa, proporcionando una justificación principista para la regla de actualización.
• Escalabilidad en Tiempo de Inferencia: ZeNO introduce un mecanismo flexible para intercambiar el presupuesto computacional (vía $N$ y $M$) por un mejor rendimiento en recompensas.

Resultados Experimentales
Los autores evalúan ZeNO en diversos generadores (SDXL-Turbo, DMD2, LCM, DFGAN) y funciones de recompensa (puntuación estética, PickScore, GenEval y diseñabilidad de proteínas).

• Generación de Imágenes: ZeNO mejora consistentemente las puntuaciones de recompensa en todos los generadores y tipos de recompensa probados. Logra un rendimiento competitivo o superior en comparación con las líneas base basadas en gradientes (ORIGEN, $\Psi$-Sampler) y el muestreo Best-of-N. Cabe destacar que, en el generador LCM, los métodos basados en gradientes degradaron el rendimiento por debajo de la línea base no controlada, mientras que ZeNO mantuvo la mejora, destacando la robustez de los métodos de orden cero frente a las sensibilidades locales de los gradientes.
• Recompensas No Diferenciables: En la tarea GenEval (conteo y posicionamiento de objetos), donde las recompensas son no diferenciables, ZeNO superó a Best-of-N en tareas de conteo, demostrando que la optimización iterativa de ruido extrae más valor del mismo presupuesto computacional que el muestreo independiente.
• Generación de Estructuras de Proteínas: Aplicado a un generador de esqueleto de proteínas de un solo paso basado en Riemannian MeanFlow (RMF), ZeNO mejoró significativamente la diseñabilidad (medida por la RMSD de autoconsistencia). Alcanzó niveles de diseñabilidad comparables a una línea base de RMF de 5 pasos utilizando únicamente una evaluación de flujo, demostrando su capacidad para compensar la reducción de pasos de muestreo mediante la optimización de ruido.
• Diversidad: El análisis mediante la Puntuación Vendi indica que ZeNO mejora la recompensa preservando la diversidad de las muestras, evitando el colapso a un único modo observado a menudo en enfoques basados en gradientes.

Significado y Afirmaciones
El artículo afirma que ZeNO proporciona una solución práctica y teóricamente fundamentada para alinear la creciente clase de modelos generativos eficientes, deterministas y de un solo paso con funciones de recompensa. Su principal significado radica en su capacidad para operar en escenarios donde los gradientes no están disponibles o son poco fiables, como en modelos de recompensa de caja negra o tareas complejas de diseño biológico. Al formular la optimización de ruido como un problema de control de integral de camino, ZeNO ofrece un marco escalable y agnóstico al generador para la alineación en tiempo de inferencia que no requiere reentrenamiento ni modificación del modelo generativo subyacente. Los autores señalan una limitación: el método requiere múltiples evaluaciones de recompensa por paso de actualización, lo cual puede ser computacionalmente exigente si el propio modelo de recompensa es costoso.