Escalado de la Inferencia en Modelos de Difusión Un Nuevo Paradigma

Introducción

Recientes avances en Modelos de Lenguaje Grandes (LLMs) han demostrado la efectividad del escalado durante la inferencia. Modelos como o1, o3, DeepSeek R1, QwQ y Step Reasoner mini han mostrado que un incremento en la computación durante la inferencia puede mejorar significativamente el rendimiento. Esto plantea la pregunta: ¿se puede aplicar este principio también a los modelos de difusión?

Un equipo liderado por Xie Saining en la Universidad de Nueva York ha explorado esta misma pregunta. A través de una investigación sistemática utilizando un marco de búsqueda general, encontraron que el escalado en tiempo de inferencia es realmente efectivo para los modelos de difusión. Aumentar la computación durante la inferencia lleva a una mejora significativa en la calidad de las muestras generadas. Además, la complejidad de las imágenes permite diferentes combinaciones de componentes dentro del marco, adaptadas a varios escenarios de aplicación.

Hallazgos Clave

• El escalado en tiempo de inferencia es efectivo para los modelos de difusión: Asignar más recursos computacionales durante la inferencia resulta en muestras de mayor calidad.
• Flexibilidad en combinaciones de componentes: El marco permite diferentes configuraciones de componentes, atendiendo a diversas aplicaciones.
• Más allá de los pasos de eliminación de ruido: La investigación sugiere que buscar un mejor ruido durante el muestreo es otra dimensión para escalar NFE, más allá de simplemente aumentar los pasos de eliminación de ruido.
• Dos Ejes de Diseño: El marco se centra en dos ejes de diseño clave:
  • Verificadores: Proporcionar retroalimentación durante el proceso de búsqueda.
  • Algoritmos: Encontrar mejores candidatos de ruido.

Metodología de Investigación

El equipo exploró tres escenarios diferentes para los verificadores, simulando varios casos de uso:

1. Escenarios donde se dispone de información privilegiada sobre la evaluación final.
2. Escenarios donde se dispone de información condicional para guiar la generación.
3. Escenarios sin información adicional disponible.

Para los algoritmos, investigaron:

1. Búsqueda Aleatoria: Seleccionar el mejor de un conjunto fijo de candidatos.
2. Búsqueda de Orden Cero: Mejorar iterativamente los candidatos de ruido utilizando la retroalimentación del verificador.
3. Búsqueda de Trayectoria: Mejorar iterativamente las trayectorias de muestreo de difusión utilizando la retroalimentación del verificador.

Inicialmente, el estudio exploró estos diseños en una configuración relativamente simple de generación condicional de clase ImageNet. Posteriormente, aplicaron estos diseños a una generación condicional de texto a mayor escala y evaluaron su marco propuesto.

Escalado del Tiempo de Inferencia

El artículo propone un marco para escalar el tiempo de inferencia en modelos de difusión, enmarcando el desafío como una búsqueda del ruido de muestreo óptimo. El proceso involucra dos componentes centrales:

• Verificadores: Estos son modelos pre-entrenados que evalúan la calidad de las muestras generadas. Toman muestras generadas, junto con condiciones opcionales, y producen una puntuación escalar.
• Algoritmos: Estos algoritmos utilizan las puntuaciones del verificador para encontrar mejores muestras candidatas. La función toma un verificador (V), un modelo de difusión pre-entrenado (_θ) y un conjunto de muestras generadas y condiciones, produciendo el mejor ruido inicial.

El presupuesto total de inferencia se mide por el número total de evaluaciones de función (NFE), incluyendo tanto los pasos de eliminación de ruido como los costos de búsqueda.

Verificadores de Búsqueda

Los investigadores comenzaron con un verificador Oracle, que tiene información completa sobre la evaluación final de las muestras seleccionadas. Para ImageNet, esto incluyó métricas como FID e IS. Luego exploraron modelos pre-entrenados más accesibles como verificadores supervisados, como CLIP y DINO. Estos modelos se utilizaron para clasificar muestras, seleccionando la muestra con el logit más alto correspondiente a la etiqueta de clase.

Sin embargo, observaron que estos clasificadores, operando puntualmente, solo se alinean parcialmente con los objetivos de la puntuación FID. Esto llevó a una reducción en la varianza de la muestra y al colapso de modo a medida que aumentaba la computación. Este fenómeno, denominado "verifier hacking", fue acelerado por el espacio de búsqueda sin restricciones del algoritmo de búsqueda aleatoria.

Curiosamente, el estudio encontró que los verificadores no necesariamente necesitan información condicional para guiar eficazmente la búsqueda. Observaron una fuerte correlación entre los logits de los clasificadores DINO/CLIP y la similitud coseno del espacio de características entre la predicción x a un nivel de ruido bajo y la muestra limpia final. Esto llevó al uso de verificadores auto-supervisados, que no requieren información condicional adicional y aún así exhibieron un comportamiento de escalado efectivo.

Algoritmos de Búsqueda

Para mitigar el "verifier hacking", los investigadores exploraron algoritmos de búsqueda más refinados que optimizan gradualmente las muestras candidatas. Esto incluyó un método de búsqueda de orden cero:

1. Comenzar con un ruido gaussiano aleatorio como punto de pivote.
2. Encontrar N candidatos en el vecindario del punto de pivote.
3. Ejecutar los candidatos a través del solucionador ODE para obtener muestras y puntuaciones de verificador.
4. Actualizar el punto de pivote con el mejor candidato y repetir los pasos 1-3.

También investigaron un algoritmo de búsqueda de trayectoria, que explora la posibilidad de buscar a lo largo de la trayectoria de muestreo:

1. Muestrear N muestras de ruido inicial y ejecutar el solucionador ODE hasta un nivel de ruido σ.
2. Agregar ruido a cada muestra y simular un proceso de ruido hacia adelante.
3. Ejecutar un solucionador ODE en cada muestra ruidosa y mantener los N candidatos principales basados en las puntuaciones del verificador, repitiendo hasta que el solucionador ODE alcance σ=0.
4. Buscar aleatoriamente las N muestras restantes y mantener la mejor.

Tanto los algoritmos de búsqueda de orden cero como de trayectoria mantienen una fuerte localidad en comparación con la búsqueda aleatoria.

Escalado en Escenarios de Texto a Imagen

El equipo examinó las capacidades de escalado del marco de búsqueda en tareas de texto a imagen a mayor escala. Utilizaron los conjuntos de datos DrawBench y T2I-CompBench para la evaluación, con el modelo FLUX.1-dev como base. También ampliaron la selección de verificadores supervisados, incluyendo Aesthetic Score Predictor, CLIPScore e ImageReward. Además, crearon un Verifier Ensemble combinando estos tres verificadores.

Análisis: Alineación Verificador-Tarea

El estudio comparó los resultados de varias combinaciones de verificador-algoritmo en diferentes conjuntos de datos. En DrawBench, encontraron que el uso de todos los verificadores generalmente mejoraba la calidad de la muestra. Sin embargo, observaron que el uso de los verificadores Aesthetic y CLIP de forma aislada podría llevar a un sobreajuste de sus sesgos, resultando en impactos negativos entre sí. Esto se debe a una falta de coincidencia en su enfoque de evaluación: Aesthetic Score se centra en la calidad visual, a menudo favoreciendo imágenes muy estilizadas, mientras que CLIP prioriza la alineación visual-texto, a veces sacrificando la calidad visual.

Notaron que algunos verificadores son más adecuados para tareas específicas, y la efectividad de un verificador depende de su alineación con los requisitos de la tarea.

Rendimiento del Algoritmo

Los tres algoritmos de búsqueda (Aleatorio, Orden Cero y Trayectoria) mejoraron eficazmente la calidad del muestreo en DrawBench. Sin embargo, la Búsqueda Aleatoria superó en algunos aspectos debido a la naturaleza local de los otros dos métodos. La búsqueda aleatoria convergió más rápidamente al sesgo del verificador, mientras que los otros dos algoritmos requieren una mejora en candidatos menos que óptimos.

Compatibilidad con el Ajuste Fino

El equipo investigó la compatibilidad de su método de búsqueda con modelos ajustados. Utilizaron un modelo Stable Diffusion XL ajustado con DPO y encontraron que el método de búsqueda podría generalizarse a diferentes modelos y mejorar el rendimiento de los modelos ya alineados.

Efectos de Diferentes Dimensiones de la Computación de Inferencia

El estudio exploró cómo diferentes aspectos de la computación de inferencia afectan los resultados:

• Número de Iteraciones de Búsqueda: Aumentar las iteraciones acerca el ruido al óptimo.
• Computación por Iteración de Búsqueda: Ajustar el número de pasos de eliminación de ruido por iteración revela diferentes regiones computacionalmente óptimas.
• Computación de Generación Final: El equipo utilizó configuraciones óptimas para los pasos finales de eliminación de ruido para garantizar la máxima calidad de la muestra final.

Efectividad de la Inversión en Computación

Los investigadores exploraron la efectividad del escalado en tiempo de inferencia en modelos de difusión más pequeños. Encontraron que, para ImageNet, escalar modelos más pequeños puede ser muy eficiente. En ciertos casos, la búsqueda en un modelo más pequeño puede superar a los modelos más grandes sin búsqueda. Sin embargo, la efectividad depende del rendimiento de referencia del modelo más pequeño.

En configuraciones basadas en texto, PixArt-Σ, utilizando solo una fracción de la computación, superó a FLUX-1.dev. Estos resultados demuestran que los recursos computacionales significativos gastados durante el entrenamiento pueden compensarse con cantidades más pequeñas de computación durante la generación, lo que resulta en muestras de mayor calidad de manera más eficiente.