Rethinking Test Time Scaling for Flow-Matching Generative Models

Resumen Técnico: Rethinking Test-Time Scaling for Generative Flow Matching Models

1. El Problema

La escalabilidad en tiempo de prueba (Test-Time Scaling o TTS) busca mejorar la calidad de salida de modelos generativos preentrenados asignando más capacidad computacional durante la inferencia, sin modificar los pesos del modelo. Mientras que en los modelos de difusión estocásticos (como SD1.5/SDXL) existen métodos exitosos de búsqueda local y exploración de trayectorias, estos enfoques enfrentan desafíos críticos al aplicarse a los modelos modernos de ajuste de flujo (Flow Matching), como Flux.1-dev:

• Incompatibilidad con la determinismo: Los modelos de flujo modernos utilizan solucionadores de Ecuaciones Diferenciales Ordinarias (ODE) deterministas para una generación eficiente. Esto elimina la estocasticidad natural necesaria para la "bifurcación" de trayectorias locales (crear variaciones a partir de un punto intermedio).
• Costo computacional excesivo: Intentar simular estocasticidad en modelos de flujo (usando Ecuaciones Diferenciales Estocásticas - SDE o mecanismos de ruido-re-ruido) requiere muchos más pasos de inferencia, anulando las ganancias de eficiencia de los ODE.
• Diversidad limitada y sesgo de recompensa:
  • Los modelos de flujo tienden a converger en regiones de alta densidad del manifold de datos, produciendo salidas muy similares para un mismo prompt, lo que reduce el espacio de búsqueda efectivo.
  • Los modelos de recompensa (reward models) estándar están entrenados solo en imágenes limpias. Al evaluar imágenes intermedias (ruidosas) durante la inferencia, sufren un "gap de dominio", favoreciendo incorrectamente imágenes con detalles de alta frecuencia (estilo caricaturesco) y descartando candidatos prometedores que aún no han desarrollado detalles finos.

2. Metodología Propuesta: DOG-Trim

Los autores proponen un pipeline llamado DOG-Trim (Diversity enhanced Order aligned Global flow Trimming), que abandona la búsqueda local costosa a favor de una estrategia de poda global optimizada. El sistema consta de tres componentes principales:

A. Estrategia de Búsqueda Global con Poda (Global Pruning)
En lugar de refinar localmente unas pocas trayectorias, el método genera un gran número de candidatos iniciales ($N$) en paralelo. Estos candidatos se desruidan en paralelo hasta un paso intermedio predefinido. En estos puntos de verificación:

1. Se estiman las imágenes limpias (Euler estimates) a partir del estado ruidoso actual.
2. Se decodifican a espacio de imagen y se evalúan con un modelo de recompensa.
3. Se eliminan (poda) los candidatos con puntuación más baja, conservando solo una fracción $\gamma$ para continuar la inferencia.
   Este enfoque maximiza la exploración del espacio de búsqueda bajo un presupuesto computacional fijo.

B. Repel (Mecanismo de Repulsión de Tokens)
Para abordar la falta de diversidad en los modelos de flujo, se introduce Repel.

• Funcionamiento: Es un mecanismo a nivel de token que actúa durante los primeros pasos de muestreo (cuando se establece la estructura semántica global).
• Mecanismo: Calcula la similitud coseno entre los tokens de imagen de diferentes trayectorias en el mismo lote. Si dos tokens de diferentes imágenes son demasiado similares, el mecanismo "empuja" el token hacia una dirección opuesta, forzando a las trayectorias a divergir.
• Ventaja: Aumenta la diversidad de la muestra sin un costo computacional significativo y evita que la búsqueda global colapse en soluciones idénticas.

C. NARF (Ajuste Fino de Recompensa Consciente del Ruido)
Para solucionar el sesgo de los modelos de recompensa en etapas tempranas (imágenes ruidosas), se propone NARF (Noise-Aware Reward Fine-tuning).

• Auto-distilación: Se entrena un modelo de recompensa para predecir la puntuación de la imagen final limpia basándose en estados intermedios ruidosos.
• Proceso de Entrenamiento:
  1. Se generan trayectorias completas desde el modelo de flujo.
  2. La "verdad fundamental" (pseudo-label) para un estado intermedio es la puntuación que el modelo de recompensa original asigna a la imagen final limpia de esa misma trayectoria.
  3. Se utiliza un aprendizaje curricular: el entrenamiento comienza con imágenes casi limpias (paso tardío) y progresa gradualmente hacia estados más ruidosos (paso temprano), asegurando una convergencia estable.
• Resultado: El modelo de recompensa ajustado puede "mirar hacia adelante" y predecir correctamente la calidad final de una imagen ruidosa, evitando podar candidatos prometedores prematuramente.

3. Contribuciones Clave

1. Estrategia de Búsqueda Global: Demostración de que, para modelos de flujo basados en ODE, la búsqueda global con poda es superior a las estrategias locales o híbridas (como las basadas en SDE) debido a la ineficiencia de la exploración local en estos modelos.
2. Mecanismo Repel: Un método eficiente a nivel de token para aumentar la diversidad de muestreo en trayectorias paralelas, expandiendo el espacio de búsqueda efectivo.
3. NARF: Una estrategia de auto-distilación eficiente para adaptar los modelos de recompensa al dominio de imágenes ruidosas, permitiendo una poda temprana y precisa.

4. Resultados Experimentales

Los experimentos se realizaron principalmente sobre el modelo Flux.1-dev comparando con el estado del arte (SOTA) bajo el mismo presupuesto computacional (equivalente a generar 6 imágenes completas).

• Rendimiento Superior: El método DOG-Trim supera significativamente a los métodos existentes (como Best-of-N, Greedy Search, FK-Steering e iSMC).
  • En métricas de recompensa (HPSv2.1, PickScore, ImageReward), DOG-Trim logra aproximadamente el doble de mejora en el rendimiento relativo al baseline sin escalabilidad en comparación con el mejor método existente.
  • Por ejemplo, en HPSv2.1, DOG-Trim alcanza 0.3314 frente a 0.3225 del mejor competidor (iSMC-ODE).
• Análisis de Diversidad: La introducción de Repel aumenta la varianza en las puntuaciones de recompensa, permitiendo encontrar mejores candidatos que de otro modo no se habrían explorado.
• Validación de Poda Temprana: El uso de NARF mejora drásticamente la concordancia de ranking (medida por el coeficiente de Kendall) entre las puntuaciones de etapas tempranas y finales, reduciendo la tasa de poda incorrecta de los mejores candidatos.
• Generalización: El modelo de recompensa ajustado con NARF (entrenado en Flux) generaliza bien a otros modelos como SDXL y SD3, demostrando que la técnica no está fuertemente acoplada a la arquitectura específica.

5. Significado e Impacto

Este trabajo es fundamental para el avance de la generación de imágenes de alta calidad con modelos de flujo modernos (como Flux):

• Cambio de Paradigma: Establece que la escalabilidad en tiempo de prueba para modelos deterministas (ODE) no debe intentar imitar la estocasticidad de los modelos de difusión antiguos (SDE), sino aprovechar la eficiencia de los ODE mediante una exploración global masiva y una poda inteligente.
• Eficiencia Computacional: Proporciona una ruta para obtener imágenes de mayor calidad sin aumentar el tiempo de inferencia por imagen individual, sino optimizando la asignación de recursos hacia la exploración de más candidatos.
• Solución al Sesgo de Recompensa: Resuelve un problema técnico crítico (el sesgo de los reward models en estados ruidosos) que había limitado la efectividad de la poda temprana en la literatura anterior.

En conclusión, DOG-Trim demuestra que, con las herramientas correctas (diversidad forzada y recompensas conscientes del ruido), es posible escalar la inferencia de modelos de flujo de manera mucho más efectiva que con los enfoques tradicionales de búsqueda local.