Resumen Técnico: Ajuste Adjoint Eficiente para el Ajuste Fino de Modelos de Difusión

Enunciado del Problema

El ajuste fino basado en recompensas se ha convertido en el enfoque estándar para alinear modelos de difusión y flujo preentrenados con preferencias humanas en la generación de texto a imagen (T2I). Entre estos, Ajuste Adjoint (AM) ofrece una formulación fundamentada al plantear el ajuste fino por recompensas como un problema de Control Óptico Estocástico (SOC). AM busca transformar un modelo preentrenado en una distribución objetivo $p^\star_1(x) \propto e^{\beta r(x)}p_{data}(x)$, donde $r(x)$ es una función de recompensa.

Sin embargo, AM adolece de ineficiencias computacionales significativas que obstaculizan su aplicación práctica:

1. Cuello de botella de la simulación hacia adelante: AM requiere simular trayectorias generativas completas utilizando un solucionador de ecuaciones diferenciales estocásticas (SDE) bajo dinámicas "sin memoria". Esto requiere un gran número de evaluaciones de funciones (NFEs) y exige almacenar toda la trayectoria en memoria.
2. Cuello de botella de la simulación hacia atrás: Para calcular la señal de control, AM requiere una simulación hacia atrás de Ecuación Diferencial Ordinaria (ODE) del estado adjunto a lo largo de la trayectoria muestreada. Esto implica calcular Productos Jacobiano-Vectores (JVPs) a través de la red preentrenada en cada paso de integración, lo cual es computacionalmente costoso.

Los autores observan que estos cuellos de botella están intrínsecamente ligados a la deriva base no trivial heredada del modelo preentrenado en las formulaciones estándar de AM.

Metodología: Ajuste Adjoint Eficiente (EAM)

El artículo propone Ajuste Adjoint Eficiente (EAM), que reformula el problema SOC para eliminar ambos cuellos de botella computacionales, preservando al mismo tiempo la garantía teórica de alcanzar la distribución objetivo inclinada por la recompensa.

1. Rediseño de las Dinámicas Base

EAM introduce dos condiciones estructurales para las dinámicas base:

• Linealidad: La deriva base $b(x,t)$ se rediseña para ser lineal en $x$ (es decir, $b(x,t) = D(t)x$).
• Ausencia de memoria: La dinámica base permanece sin memoria, permitiendo que los estados intermedios se muestreen directamente desde el punto final.

Al fijar el coeficiente de difusión $\sigma(t)$ para que coincida con la configuración estándar, los autores derivan una familia específica de derivas lineales admisibles (Ecuación 15) que satisfacen estas condiciones. Esta deriva lineal permite una solución de forma cerrada para el estado adjunto, eliminando la necesidad de simulación ODE hacia atrás y evaluaciones de JVP.

2. Corrección del Costo Terminal

Cambiar la deriva base altera la distribución terminal del sistema no controlado. Para asegurar que la distribución controlada óptima coincida aún con la distribución objetivo inclinada por la recompensa deseada $p^\star_1$, EAM deriva un costo terminal corregido $g(x)$. El gradiente de este costo, $\nabla g(x)$, incluye términos para la puntuación preentrenada, la nueva puntuación terminal base y el gradiente de la recompensa (Ecuación 16).

3. Algoritmo de Entrenamiento

El bucle de entrenamiento de EAM (Algoritmo 1) procede de la siguiente manera:

1. Generación del punto final: En lugar de un despliegue completo de SDE, la imagen del punto final $X_1$ se genera utilizando un solucionador ODE determinista de pocos pasos (por ejemplo, 10 NFEs).
2. Construcción del estado intermedio: Los estados intermedios $X_t$ se muestrean directamente desde el kernel de ruido original $q_t(\cdot | X_1)$ condicionado al punto final generado, evitando la necesidad de simulación SDE hacia adelante.
3. Adjoint de forma cerrada: El estado adjunto se calcula mediante una escala de forma cerrada del gradiente del punto final $\nabla g(X_1)$, eliminando por completo la simulación hacia atrás.
4. Parametrización del control: El control $u(x,t)$ se parametriza mediante un modelo de velocidad entrenable (por ejemplo, LoRA) inicializado desde el modelo preentrenado, permitiendo que el ajuste fino comience desde las dinámicas preentrenadas.

Contribuciones Clave

• Propuesta de algoritmo: La introducción de EAM, un algoritmo de ajuste fino basado en gradientes de recompensa que logra los mismos objetivos teóricos que AM pero con un costo computacional significativamente reducido.
• Caracterización teórica: Una caracterización de derivas base lineales que satisfacen los requisitos para la construcción eficiente de trayectorias (ausencia de memoria) y el cálculo de adjuntos de forma cerrada.
• Validación empírica: Demostración de que EAM converge hasta 4$\times$ más rápido que AM en benchmarks estándar T2I, mientras iguala o supera su rendimiento en múltiples métricas.

Resultados

Los autores evaluaron EAM en Stable Diffusion 3.5-Medium (SD3.5-M) utilizando el conjunto de datos Pick-a-Pic, comparándolo con AM y la línea base SD3.5-M.

• Eficiencia: EAM reduce el costo de entrenamiento por iteración hasta en 4$\times$. Esto se logra reemplazando simulaciones SDE de 40 pasos con muestreo de puntos finales ODE de 10 pasos y eliminando las simulaciones de adjuntos hacia atrás.
• Rendimiento: En benchmarks que incluyen PickScore, ImageReward, HPSv2.1, CLIPScore y Aesthetics, EAM coincide consistentemente o supera a AM.
  • En configuraciones de recompensa única (PickScore), EAM logró un PickScore de 23.16 en comparación con 23.00 de AM.
  • En configuraciones de múltiples recompensas (PickScore + HPSv2.1 + Aesthetics), EAM logró un PickScore de 22.93 y una puntuación de Aesthetics de 6.56, superando a AM (6.21).
• Mejoras cualitativas: Las comparaciones visuales indican que EAM preserva mejor los detalles finos y las composiciones complejas en comparación con AM, que a veces suaviza los detalles. Esto se atribuye a la consistencia de EAM entre las trayectorias de entrenamiento (ruido condicionado al punto final) y las trayectorias de inferencia.

Significado y Afirmaciones

El artículo afirma que las limitaciones principales de Ajuste Adjoint no son inherentes a la propia formulación SOC, sino que son artefactos de la elección específica de la deriva base. Al rediseñar la deriva base para que sea lineal y sin memoria, y ajustar correspondientemente el costo terminal, EAM elimina los dos grandes cuellos de botella computacionales (simulación SDE hacia adelante y simulación de adjunto hacia atrás) sin sacrificar las garantías teóricas de alineación.

Los autores posicionan a EAM como una solución práctica y escalable para el ajuste fino por recompensas que conserva la eficiencia de muestreo de los métodos basados en gradientes mientras reduce drásticamente la sobrecarga computacional, haciendo más accesible la alineación de alta calidad de modelos de difusión. El trabajo no afirma resolver todos los desafíos de optimización (por ejemplo, se señala que la reproducción de búfer o los programas adaptativos son trabajos futuros), pero establece una nueva línea base para la eficiencia en este dominio.