YOSE: You Only Select Essential Tokens for Efficient DiT-based Video Object Removal

1. Planteamiento del Problema

La Eliminación de Objetos en Video (VOR) busca borrar objetos no deseados de los videos manteniendo el realismo espacial y la consistencia temporal. Los métodos recientes de última generación (SOTA), como MiniMax Remover, aprovechan los Transformadores de Difusión (DiT) para lograr resultados de alta calidad. Sin embargo, estos métodos sufren de una latencia de inferencia significativa e ineficiencia computacional.

• El Problema Central: Los modelos VOR basados en DiT existentes realizan cálculos densos sobre todo el espacio de tokens espacio-temporal (todos los fotogramas y píxeles), incluso cuando solo una región enmascarada pequeña requiere procesamiento.
• La Consecuencia: La complejidad computacional permanece constante independientemente del tamaño de la máscara. Dado que las tareas VOR del mundo real suelen involucrar ratios de máscara pequeños (por ejemplo, ~70% de los casos tienen máscaras <20%), esto resulta en una redundancia masiva. Por ejemplo, MiniMax Remover opera a solo ~10 FPS, limitando su aplicación práctica.

2. Metodología: Marco YOSE

Los autores proponen YOSE (You Only Select Essential Tokens), un marco de ajuste fino ligero y tipo "plug-in" diseñado para acelerar la VOR basada en DiT sin reestructurar el modelo subyacente. YOSE introduce dos componentes principales:

A. Indexación de Longitud Variable por Lotes (BVI)

BVI es un operador de indexación dinámica diferenciable que permite al modelo procesar solo los tokens dentro de las regiones enmascaradas (tokens esenciales).

• Mecanismo: En lugar de una indexación rígida (que bloquea el flujo de gradientes), BVI utiliza una selección basada en interpolación continua (mediante grid_sample) para mapear los tokens enmascarados a una secuencia de longitud variable.
• Índice de Avance ($Ind_F$): Selecciona los tokens esenciales del área enmascarada para su procesamiento.
• Índice de Retroceso ($Ind_B$): Restaura los resultados procesados a sus posiciones espacio-temporales originales.
• Beneficio: Soporta el procesamiento de tokens de longitud variable a través de un lote, permitiendo que el modelo maneje diferentes tamaños de máscara de manera eficiente mientras mantiene el flujo de gradientes para el entrenamiento de extremo a extremo. Esto reduce la complejidad computacional de $O(N_{total})$ a $O(N_{masked})$.

B. Módulo Simulador de Proceso de Difusión (DiffSim)

Procesar solo los tokens enmascarados corre el riesgo de romper la continuidad semántica con las regiones no enmascaradas (fondo). DiffSim aborda esto simulando la influencia de las regiones no enmascaradas sin procesarlas.

• Mecanismo:
  1. Construye un Latente Residual ($Res_{Nis}$) que representa la diferencia entre el latente de ruido y el latente de video enmascarado.
  2. Genera características Key & Value (KV) proxy para las regiones no enmascaradas utilizando parámetros aprendibles: un parámetro de combinación ($G$), un parámetro de escala ($S$) y un sesgo ($Bias$).
  3. La fórmula para la KV simulada es: $KV = G \cdot Lat_{mask} + (1-G) \cdot Res_{Nis}$, seguida de modulación de escala y sesgo.
• Beneficio: Estas características simuladas se inyectan en el mecanismo de auto-atención del DiT. Esto permite que los tokens "internos" (enmascarados) atiendan al contexto "externo" (no enmascarado), preservando la consistencia semántica y temporal a través de los límites sin el costo del cálculo de tokens completos.

C. Estrategia de Fusión

Para eliminar artefactos de borde causados por desajustes estadísticos entre las regiones procesadas y no procesadas:

• La máscara se dilata para crear una región superpuesta.
• La media y la varianza de la región superpuesta generada se alinean con la región no enmascarada original.
• Se aplica una estrategia de mezcla ponderada para garantizar una transición suave.

3. Contribuciones Clave

1. Marco YOSE: Un marco de ajuste fino consciente de la máscara que acelera la VOR basada en DiT (por ejemplo, MiniMax Remover) procesando solo tokens esenciales, logrando una complejidad computacional lineal en relación con el ratio de máscara.
2. Indexación de Longitud Variable por Lotes (BVI): Un operador diferenciable que permite la selección y el procesamiento eficiente de tokens de longitud variable, resolviendo los problemas de flujo de gradientes y paralelización de lotes asociados con el enmascaramiento dinámico.
3. Simulador de Proceso de Difusión (DiffSim): Un módulo novedoso que aproxima la influencia de difusión de las regiones no enmascaradas dentro de los bloques DiT, asegurando la consistencia semántica global sin cálculos redundantes.
4. Rendimiento: El método logra una aceleración de hasta 2.5× (alcanzando ~25 FPS para ratios de máscara del 20%) manteniendo una calidad visual comparable o ligeramente superior a la línea base.

4. Resultados Experimentales

Los autores evaluaron YOSE en MiniMax Remover y VACE utilizando los conjuntos de datos DAVIS y YouTube-VOS.

• Velocidad y Eficiencia:
  • FLOPs y Latencia: YOSE exhibe una relación lineal entre el tiempo de inferencia y el ratio de máscara.
  • Aceleración: Logra una aceleración de 3.3× en un ratio de máscara del 5% y una aceleración de 2.5× en un ratio de máscara del 20% (escenarios típicos del mundo real).
  • FPS: Aumenta de 10 FPS (MiniMax) a **25 FPS** (YOSE) a una resolución de 480p.
• Calidad Visual:
  • Métricas: YOSE mantiene puntuaciones comparables o ligeramente mejoradas en las métricas VBench (Suavidad del Movimiento, Calidad Estética) y métricas de fondo (PSNR, SSIM, LPIPS).
  • Preservación del Fondo: Notablemente, YOSE (VACE) mejoró el PSNR del fondo en 5.47 dB y el SSIM en 0.07 en comparación con el VACE original, evitando efectivamente la "alucinación" de artefactos con forma de máscara en el fondo.
• Estudios de Ablación:
  • Eliminar DiffSim provoca una caída tanto en la suavidad del movimiento como en la calidad de reconstrucción del fondo.
  • Eliminar la Estrategia de Fusión resulta en artefactos de borde visibles.
  • Eficiencia de Entrenamiento: BVI permite el entrenamiento con un tamaño de lote de 4 en 8 GPUs en ~4 horas, en comparación con ~11 horas para enfoques convencionales.

5. Significado

YOSE representa un paso significativo hacia la edición de video práctica y en tiempo real. Al desacoplar el costo computacional del tamaño total del video y vincularlo estrictamente al área de interés (la máscara), resuelve el principal cuello de botella de los modelos actuales de video generativo basados en DiT.

• Escalabilidad: Hace viable la eliminación de objetos en video de alta fidelidad en hardware de consumo o para aplicaciones en tiempo real donde la latencia es crítica.
• Generalización: El marco es agnóstico al modelo (plug-and-play) y ha sido validado exitosamente en múltiples arquitecturas basadas en DiT (MiniMax Remover, VACE).
• Calidad vs. Velocidad: Demuestra que una aceleración agresiva no necesariamente requiere un compromiso en la fidelidad visual; en algunos casos, el procesamiento selectivo incluso mejora la estabilidad del fondo al evitar que el modelo sobreinterprete las regiones no enmascaradas.