LatSearch: Latent Reward-Guided Search for Faster Inference-Time Scaling in Video Diffusion

1. El Problema

La generación de video basada en modelos de difusión ha avanzado rápidamente, pero enfrenta desafíos significativos al intentar escalar la inferencia para mejorar la calidad, similar a lo logrado en los Modelos de Lenguaje Grandes (LLMs).

• Limitaciones de los métodos actuales:
  • Optimización de ruido inicial: Métodos anteriores intentan encontrar "ruido dorado" (golden noise) inicial, pero a menudo dependen de priores fijos o no pueden corregir errores una vez iniciada la trayectoria de denoising.
  • Búsqueda basada en salida final: Las estrategias de búsqueda (como Best-of-N o búsqueda evolutiva) generan múltiples candidatos y los evalúan solo después de decodificarlos completamente a video. Esto provoca:
    • Retraso en la recompensa: La señal de calidad llega demasiado tarde en el proceso.
    • Acumulación de errores: Los errores iniciales se propagan a lo largo de toda la trayectoria sin corrección intermedia.
    • Alto costo computacional: Decodificar múltiples videos completos es extremadamente costoso, lo que impide el uso de algoritmos de búsqueda más potentes.
• Falta de evaluación intermedia: No existen mecanismos eficientes para evaluar latentes parcialmente denoizados durante el proceso, lo que impide estrategias flexibles como la parada temprana o la poda de candidatos de baja calidad.

2. Metodología: LatSearch

Los autores proponen LatSearch, un marco de escalado en tiempo de inferencia que integra un modelo de recompensa latente con un mecanismo de búsqueda guiada por recompensas.

A. Modelo de Recompensa Latente

A diferencia de los verificadores tradicionales que operan en el espacio de video decodificado, este modelo evalúa latentes parcialmente denoizados en cualquier paso de tiempo ($t$).

• Entradas: Representación latente ($z_t$), prompt de texto ($p$) y el paso de tiempo ($t$).
• Salidas: Tres dimensiones de puntuación:
  1. Calidad Visual (VQ).
  2. Calidad de Movimiento (MQ).
  3. Alineación Texto-Video (TA).
• Construcción de Datos y Entrenamiento:
  • Como no existen etiquetas de "verdad fundamental" a nivel de latente, se utiliza una estrategia de asignación de crédito basada en similitud. Se mide la similitud coseno entre un latente intermedio $z_t$ y el latente final limpio $z_0$.
  • Las recompensas a nivel de video se asignan a los latentes intermedios proporcionalmente a esta similitud.
  • Función de Pérdida: Se entrena combinando una pérdida de regresión (para ajustar la magnitud absoluta) y una pérdida de preferencia (inspirada en RLHF, para aprender el ordenamiento relativo entre candidatos), lo que mejora la robustez y la precisión.

B. Mecanismo de Búsqueda: RGRP (Resampling y Poda Guiados por Recompensa)

LatSearch mantiene múltiples trayectorias de candidatos durante el proceso de difusión e implementa dos fases clave:

1. Resampling (Muestreo) Guiado por Recompensa:
   • En pasos de tiempo seleccionados, los candidatos se evalúan con el modelo de recompensa latente.
   • En lugar de seleccionar determinísticamente el mejor, se realiza un muestreo probabilístico basado en las probabilidades normalizadas por la recompensa. Esto evita la sobreconfianza en un modelo de recompensa imperfecto y fomenta la exploración.
   • Se aplica un operador de unicidad para eliminar candidatos duplicados (mismas semillas iniciales), ahorrando cómputo.
2. Poda (Pruning):
   • Al final del programa de búsqueda planificado, se calcula una recompensa acumulada para cada trayectoria.
   • Se retiene únicamente el candidato con la mayor recompensa acumulada para continuar hasta la decodificación final, eliminando redundancias y reduciendo drásticamente el costo computacional restante.

3. Contribuciones Clave

1. Evaluación Intermedia en Espacio Latente: Introducción de un modelo de recompensa capaz de evaluar latentes parciales en cualquier paso de tiempo, proporcionando supervisión a nivel de proceso en lugar de solo al final.
2. Estrategia de Asignación de Crédito: Desarrollo de un método basado en similitud coseno para mapear recompensas de video final a latentes intermedios, haciendo factible el aprendizaje de recompensas a nivel de latente.
3. Algoritmo RGRP: Diseño de un algoritmo de búsqueda que combina muestreo probabilístico (para diversidad) y poda acumulativa (para eficiencia), equilibrando exploración y explotación sin el costo de decodificar videos completos repetidamente.
4. Eficiencia Computacional: Lograr mejoras de calidad significativas con un aumento mínimo en el tiempo de inferencia en comparación con métodos de búsqueda existentes.

4. Resultados Experimentales

El método se evaluó utilizando el benchmark VBench-2.0 sobre el modelo base Wan2.1-1.3B y Wan2.1-14B.

• Calidad vs. Costo:
  • LatSearch mejora consistentemente la calidad del video en todas las dimensiones (Creatividad, Sentido Común, Controlabilidad, Fidelidad Humana y Física) en comparación con el modelo base.
  • Comparación con el Estado del Arte (SOTA):
    • Logra una calidad comparable o superior a métodos como EvoSearch y VideoReward.
    • Reducción de tiempo: Es hasta un 79% más rápido que los métodos de búsqueda más potentes (ej. EvoSearch tarda ~790s, LatSearch ~164s para resultados similares).
    • En términos de eficiencia, LatSearch logra un aumento de calidad del 3.35% sobre la línea base con solo un costo computacional de 2.13x, mientras que EvoSearch requiere un costo de 10.15x para un margen de mejora similar.
• Análisis de Componentes:
  • La pérdida de preferencia mejora la precisión de la predicción de recompensas en un 3-4% en pasos intermedios.
  • La estrategia de similitud coseno para la asignación de crédito superó a otras estrategias (uniforme, exponencial, error L2).
  • El presupuesto de búsqueda (N=6) ofrece el mejor equilibrio entre calidad y latencia.

5. Significado e Impacto

LatSearch representa un avance significativo en la generación de video al demostrar que el escalado en tiempo de inferencia no requiere un costo prohibitivo.

• Paradigma de Eficiencia: Cambia el enfoque de "decodificar todo para evaluar" a "evaluar en el espacio latente", lo que permite algoritmos de búsqueda más sofisticados que antes eran computacionalmente inviables.
• Control y Calidad: Al permitir la poda temprana de trayectorias de baja calidad, mitiga la acumulación de errores, resultando en videos más coherentes y fieles al prompt.
• Escalabilidad: Al ser agnóstico a la arquitectura del modelo (funciona en Wan 1.3B y 14B), el método es aplicable a futuros modelos de difusión más grandes, facilitando la generación de video de alta fidelidad con recursos limitados.

En resumen, LatSearch cierra la brecha entre la teoría de escalado en tiempo de inferencia y la práctica eficiente en video, ofreciendo una solución viable para generar videos de alta calidad con una fracción del tiempo de cómputo de las técnicas actuales.