Resumen Técnico: Contexto Disperso – Generación Eficiente Condicionada por Referencia mediante el Descarte de Tokens

Planteamiento del Problema

Los modelos de difusión basados en referencias se han vuelto fundamentales para la generación de imágenes controlable, permitiendo tareas como la edición de imágenes y la síntesis dirigida por sujetos al condicionar la generación mediante imágenes de entrada. Sin embargo, estos modelos enfrentan cuellos de botella computacionales significativos. En los modelos de difusión modernos basados en Transformers (DiTs), las imágenes de referencia se codifican típicamente en rejillas espaciales densas de tokens mediante un VAE y se concatenan con los tokens de entrada ruidosos. Dado que las operaciones de auto-atención escalan cuadráticamente con el número de tokens, la inclusión de una o más imágenes de referencia de alta resolución aumenta drásticamente el tiempo de inferencia y el uso de memoria. Este costo resulta prohibitivo en aplicaciones del mundo real que involucran múltiples referencias. Aunque existen técnicas de eficiencia para la generación impulsada por prompts, estas siguen siendo poco exploradas para los modelos condicionados por referencia, los cuales presentan desafíos únicos debido a la representación densa de las entradas visuales.

Metodología

El artículo introduce Sparse Context (Contexto Disperso), un marco diseñado para construir representaciones de referencia dispersas reteniendo solo un subconjunto reducido de tokens de referencia. La metodología aborda el problema a través de tres componentes principales:

1. Observación de la Redundancia

Los autores demuestran primero que los tokens de referencia contienen una redundancia significativa. A través de un experimento de "Poda de Contexto Naive" (Naïve Context Pruning) en un modelo FLUX.2-Klien preentrenado, muestran que el descarte aleatorio de hasta el 80% de los tokens de referencia durante la inferencia preserva en gran medida la disposición general (layout) de la escena, a pesar de una pérdida en los detalles finos. Esto sugiere que los modelos condicionados por referencia son inherentemente robustos a entradas dispersas, pero la brecha de rendimiento entre referencias completas y parciales se ve exacerbada por un desajuste entre la distribución de entrenamiento y prueba.

2. Estrategia de Entrenamiento: Descarte Aleatorio de Tokens

Para cerrar la brecha de distribución y permitir una generación dispersa robusta, los autores ajustan (fine-tune) un generador condicionado por referencia de última generación utilizando el descarte aleatorio de tokens durante el entrenamiento.

• Proceso: Durante el entrenamiento, se descarta una fracción aleatoria de los tokens de referencia (manteniendo una fracción $f \in [0.05, 0.25]$) antes de la concatenación con los tokens objetivo ruidosos.
• Objetivo: Esto incentiva al modelo a aprender representaciones robustas que funcionen eficazmente con información de referencia parcial, desacoplando el modelo de cualquier regla específica de selección de tokens.
• Implementación: Se aplica un LoRA (Adaptación de Bajo Rango) ligero al modelo base FLUX-2-Klein-9B sobre un conjunto de datos curado que cubre la edición impulsada por instrucciones y la personalización.

3. Estrategia de Inferencia: Selección de Tokens Basada en Tareas

Mientras que el entrenamiento utiliza el descarte aleatorio para asegurar la robustez, la inferencia emplea estrategias de selección de tokens basadas en tareas para maximizar la calidad dentro de un presupuesto de cómputo fijo:

• Edición de Imágenes Impulsada por Instrucciones: Para preservar la estructura espacial, el método utiliza una estrategia de muestreo sensible a los bordes. Se extrae un mapa de bordes Canny de la imagen de referencia, y los tokens se muestrean basados en esta distribución de probabilidad, priorizando los límites estructurales sobre las regiones uniformes.
• Personalización: Para preservar la identidad del sujeto, el método utiliza un muestreo basado en la prominencia (saliency). Un mapa de prominencia aproxima la máscara del sujeto, permitiendo que el modelo concentre los tokens en las regiones más informativas de la imagen de referencia.

Contribuciones Clave

1. Marco de Contexto Disperso (Sparse Context): Un enfoque novedoso que construye representaciones de referencia dispersas mediante el descarte de tokens, reduciendo significativamente el costo cuadrático de la auto-atención sin modificar la arquitectura subyacente del modelo.
2. Estrategia de Ajuste Fino Robusta: Un régimen de entrenamiento utilizando el descarte aleatorio de tokens que permite al modelo generalizar a través de diversos niveles de dispersión, permitiendo un control flexible sobre el presupuesto de cómputo en tiempo de inferencia.
3. Inferencia Basada en Tareas: La propuesta de heurísticas específicas de selección de tokens (basadas en bordes para edición, basadas en prominencia para personalización) que superan al descarte aleatorio al priorizar la información crítica para la tarea.
4. Ortogonalidad con Métodos Existentes: Demostración de que Sparse Context es compatible y complementario con otras técnicas de eficiencia, como la Fusión de Tokens (Token Merging - ToMe) y el almacenamiento en caché de KV (KV-caching), permitiendo ganancias de aceleración compuestas.

Resultados Experimentales

El método fue evaluado en edición de imágenes impulsada por instrucciones (PIE-Bench) y personalización basada en referencias (Subject-200K, CustomDiffusion-105) utilizando el modelo FLUX-2-Klein-9B.

• Eficiencia:
  • Referencia Única: Logra una aceleración de inferencia de hasta 2× (por ejemplo, 1.83× con retención del 10% de los tokens).
  • Multi-Referencia: Logra una aceleración de hasta 4× para escenarios con múltiples referencias (por ejemplo, 6 referencias), ya que los tokens de referencia dominan el cómputo de la atención.
  • Eficiencia Combinada: Al integrarse con el almacenamiento en caché de KV para la generación de múltiples referencias, el método logra aceleraciones de hasta 6.72×.
• Calidad:
  • Edición de Imágenes: El método preserva la estructura de la escena e la identidad significativamente mejor que la línea base de Poda de Contexto Naive. Las métricas cuantitativas (LPIPS, DINO, CLIP-I) muestran que Sparse Context mantiene una alta fidelidad incluso con fracciones de tokens bajas (por ejemplo, $f=0.1$), mientras que la línea base sufre de corrupción de la identidad.
  • Personalización: El muestreo basado en prominencia logra preservar detalles finos del sujeto (por ejemplo, colores específicos, patrones) que se pierden con el descarte aleatorio.
• Integración: La combinación de Sparse Context con Token Merging mejora aún más el rendimiento, demostrando que el modelo ajustado aprovecha eficazmente otras técnicas de reducción de tokens.

Significado y Reivindicaciones

El artículo afirma que Sparse Context desbloquea el potencial de eficiencia no aprovechado inherente a los modelos de difusión basados en referencias. Al reconocer que las rejillas densas de tokens suelen ser innecesarias, los autores demuestran que la generación basada en referencias puede hacerse significativamente más eficiente sin comprometer la controlabilidad o la fidelidad visual.

Los autores enfatizan que su enfoque es simple y modular, requiriendo solo un ajuste fino ligero y ofreciendo un control flexible en el tiempo de inferencia. Afirman que este método aborda un cuello de botella crítico en el despliegue de modelos generativos de alta resolución y múltiples referencias, haciéndolos más prácticos para aplicaciones del mundo real. El trabajo destaca que la redundancia en las representaciones de referencia es una propiedad fundamental que puede explotarse para reducir los costos computacionales, un hallazgo que se extiende más allá de los modelos específicos probados hacia la clase más amplia de sistemas de generación de imagen y video basados en DiT.