FA-Seg: A Fast and Accurate Diffusion-Based Method for Open-Vocabulary Segmentation

1. Planteamiento del Problema

La Segmentación Semántica de Vocabulario Abierto (OVSS) tiene como objetivo segmentar objetos en imágenes basándose en categorías de texto arbitrarias, sin requerir anotaciones densas a nivel de píxel para cada clase posible.

• Limitaciones de los Métodos Existentes:
  • Modelos Visión-Lenguaje (ej. CLIP): Aunque son efectivos para el reconocimiento zero-shot, dependen de embebidos globales de imagen-texto, lo que a menudo resulta en una precisión espacial deficiente y bordes borrosos a nivel de píxel.
  • Modelos Basados en Difusión: Aunque codifican naturalmente características espaciales de alta granularidad mediante mecanismos de atención, los métodos existentes suelen sufrir altos costos computacionales. Muchos requieren múltiples pasos de inferencia por clase, optimización extensa o generación de datos sintéticos, creando una compensación entre la calidad de la segmentación y la eficiencia de la inferencia.
• La Brecha: Existe una necesidad crítica de un método OVSS sin entrenamiento que logre una alta precisión de segmentación (precisión espacial de alta granularidad) manteniendo al mismo tiempo una eficiencia computacional adecuada para su implementación en el mundo real.

2. Metodología: FA-Seg

FA-Seg es un marco de trabajo sin entrenamiento que aprovecha modelos de difusión pre-entrenados de texto a imagen (específicamente Stable Diffusion) para realizar la segmentación en una sola ejecución de inferencia. El pipeline consta de tres etapas principales:

A. Construcción de Doble Prompt

Para permitir tanto la reconstrucción fiel de la imagen como la extracción de atención específica por clase, FA-Seg construye dos prompts complementarios:

1. Prompt de Descripción ($\mathcal{P}$): Generado por un modelo de descripción de imágenes (ej. BLIP) para guiar la inversión DDIM hacia la reconstrucción de la imagen de entrada.
2. Prompt de Clase ($\mathcal{P}'$): Una lista de categorías de objetos candidatas (generadas mediante TagCLIP + BLIP) añadida a la descripción.

• Mecanismo: El modelo de difusión utiliza $\mathcal{P}$ para la reconstrucción, pero aprovecha $\mathcal{P}'$ para generar mapas de atención cruzada específicamente para las clases candidatas. Esto permite que el modelo se centre en conceptos semánticos específicos sin necesidad de re-inferencia iterativa.

B. Inversión Eficiente y Extracción de Atención

• Inversión Rápida: En lugar de la inversión DDIM estándar de múltiples pasos, FA-Seg emplea Flujo Rectificado 2 (2-RF). Este modelo destilado permite una reconstrucción precisa de la imagen en un proceso de (1+1) pasos (un paso hacia adelante, un paso hacia atrás), reduciendo drásticamente el tiempo de inferencia.
• Extracción de Atención: Durante el proceso de eliminación de ruido, el método extrae:
  • Mapas de Atención Cruzada: Que vinculan los tokens de texto (del prompt de clase) con las regiones de la imagen.
  • Mapas de Atención Auto: Que capturan la estructura espacial interna y las similitudes por pares de la imagen.
  • Estos se extraen a través de múltiples resoluciones ($r \in \{8, 16, 32, 64\}$).

C. Método de Refinamiento Jerárquico de Atención (HARD)

Para producir máscaras de alta calidad, FA-Seg fusiona mapas de atención de diferentes resoluciones:

1. Fusión de Atención Cruzada: Los mapas de atención cruzada de diferentes resoluciones se aumentan de tamaño y se agregan utilizando una suma ponderada. Las resoluciones más bajas (ej. 16) se ponderan más alto para el contexto semántico, mientras que las resoluciones más altas se ponderan menos para reducir el ruido.
2. Refinamiento de Atención Auto: El mapa de atención cruzada fusionado se refina adicionalmente utilizando matrices de afinidad de atención auto. Los mapas de atención auto proporcionan detalles estructurales (bordes y formas) que carece la atención cruzada. El método transforma los tensores de atención auto para que coincidan con la resolución del mapa de atención cruzada y los aplica como pesos multiplicativos para afilar los bordes de los objetos.
3. Volteo en Tiempo de Prueba (TTF): Para mejorar la consistencia espacial, el proceso se repite en una versión horizontalmente volteada de la imagen de entrada. Los mapas de atención resultantes se realinean y promedian con los mapas originales para reducir el ruido y mejorar la robustez.

3. Contribuciones Clave

1. Inferencia Unificada en Una sola Ejecución: A diferencia de los métodos anteriores que requieren ejecuciones de inferencia separadas para cada clase, FA-Seg genera máscaras de segmentación para todas las clases candidatas simultáneamente en una sola inferencia de (1+1) pasos.
2. Mecanismo de Doble Prompt: Una estrategia novedosa que separa el prompt de reconstrucción del prompt consciente de la clase, permitiendo que el modelo de difusión extraiga mapas de atención discriminativos para categorías específicas sin perder la fidelidad de la imagen.
3. Refinamiento Jerárquico de Atención (HARD): Una estrategia de fusión multi-escala que combina la atención cruzada (para la localización semántica) y la atención auto (para el refinamiento estructural) a través de resoluciones para lograr bordes de objetos precisos.
4. Eficiencia sin Entrenamiento: El método no requiere ajuste fino, ni generación de datos sintéticos, ni recursos computacionales mínimos, logrando un rendimiento de vanguardia con menos de 1 segundo de tiempo de inferencia.

4. Resultados Experimentales

Los autores evaluaron FA-Seg en tres benchmarks estándar: PASCAL VOC, PASCAL Context y COCO Object.

• Rendimiento (mIoU):
  • FA-Seg logró un mIoU promedio del 43.8% en los tres conjuntos de datos.
  • Estableció un nuevo estado del arte (SOTA) para métodos basados en difusión sin entrenamiento, superando a competidores como InvSeg (42.4%), DiffSegmentor (41.8%) y varios enfoques basados en CLIP.
  • Puntuaciones específicas: 64.0% (PASCAL VOC), 31.3% (PASCAL Context), 36.2% (COCO).
• Eficiencia:
  • Tiempo de Inferencia: 0.36 segundos por imagen en una GPU RTX 4090. Esto es significativamente más rápido que InvSeg (7.9s en H100) y DiffSegmentor (el tiempo escala con el número de clases).
  • Memoria: Consume solo 13.4 GB de VRAM, en comparación con los 32.4 GB de InvSeg.
• Estudios de Ablación:
  • Eliminar el refinamiento de atención auto causó la mayor caída de rendimiento (-15.1%), demostrando su papel crítico en la precisión espacial.
  • La fusión ponderada de mapas multi-resolución superó a la fusión por media simple.
  • La combinación TagCLIP+BLIP para la generación de clases candidatas demostró ser la más robusta.

5. Significado e Impacto

• Cerrando la Brecha: FA-Seg cierra con éxito la brecha entre la calidad de la segmentación y la eficiencia de la inferencia, demostrando que la segmentación de vocabulario abierto de alta precisión no requiere un entrenamiento pesado ni optimización multi-paso.
• Implementación Práctica: Al reducir el tiempo de inferencia a menos de un segundo y el uso de memoria a ~13 GB, FA-Seg hace viable la segmentación de vocabulario abierto en tiempo real para aplicaciones del mundo real (ej. robótica, conducción autónoma, herramientas interactivas).
• Base para Trabajos Futuros: El artículo establece una línea base sólida para la segmentación por difusión sin entrenamiento, sugiriendo que las mejoras futuras podrían centrarse en la generación adaptativa de candidatos y la segmentación a nivel de instancia.

En conclusión, FA-Seg representa un avance significativo en el campo de la segmentación de vocabulario abierto al aprovechar las capacidades espaciales inherentes de los modelos de difusión a través de un marco altamente optimizado, sin entrenamiento y computacionalmente eficiente.